diff --git a/.github/workflows/mlir-tensorrt-ci.yml b/.github/workflows/mlir-tensorrt-ci.yml index 9e6f6710c..1008c5ee9 100644 --- a/.github/workflows/mlir-tensorrt-ci.yml +++ b/.github/workflows/mlir-tensorrt-ci.yml @@ -36,6 +36,48 @@ jobs: sudo apt-get autoremove -y sudo apt-get autoclean -y + + # Value of `github.workspace` is /home/runner/work/{repo_name}/{repo-name} + # i.e. /home/runner/work/TensorRT-Incubator/TensorRT-Incubator in our case. + # After this action, repo is cloned inside above path. + - uses: actions/checkout@v4 + with: + fetch-depth: 5 + + - name: Validate commit message + if: ${{ github.event_name == 'pull_request' }} + env: + PR_HEAD_COMMIT_SHA: ${{ github.event.pull_request.head.sha }} + run: | + cat > commit_message_checker.py <") + print("") + print("NOTE: Body should start on new line. '2 spaces + enter' for new line!") + print("NOTE: Body should be at least two lines.") + sys.exit(1) + EOF + + python3 commit_message_checker.py ${PR_HEAD_COMMIT_SHA} # Run initial format check - name: Run python format and clang check diff --git a/mlir-tensorrt/compiler/lib/Compiler/TensorRTToExecutable/Passes.cpp b/mlir-tensorrt/compiler/lib/Compiler/TensorRTToExecutable/Passes.cpp index 01a4a6a31..8dc12a0ed 100644 --- a/mlir-tensorrt/compiler/lib/Compiler/TensorRTToExecutable/Passes.cpp +++ b/mlir-tensorrt/compiler/lib/Compiler/TensorRTToExecutable/Passes.cpp @@ -280,14 +280,15 @@ outlineOp(RewriterBase &rewriter, tensorrt::TensorRTModuleOp trtModule, mlir::tensorrt::TensorRTDialect::getShapeProfileArgAttrName(); StringRef tensorrtDimensionNamesAttrName = mlir::tensorrt::TensorRTDialect::getDimensionNamesArgAttrName(); + StringRef tensorrtValueBoundsAttrName = + mlir::tensorrt::TensorRTDialect::getShapeTensorValueBoundsArgAttrName(); + StringRef hostTensorAttrName = mlir::getHostTensorArgAttrName(); + StringRef memorySpaceAttrName = + plan::PlanDialect::getMemorySpaceConstraintAttrName(); - SmallVector profileAttrsPerInput; - SmallVector dimensionNamesAttrsPerInput; for (Value v : inputs) { auto rtt = dyn_cast(v.getType()); - if (!rtt || rtt.hasStaticShape()) { - profileAttrsPerInput.push_back(Attribute{}); - dimensionNamesAttrsPerInput.push_back(Attribute{}); + if (!rtt) { continue; } @@ -299,30 +300,42 @@ outlineOp(RewriterBase &rewriter, tensorrt::TensorRTModuleOp trtModule, } int64_t argIndex = blockArg.getArgNumber(); - profileAttrsPerInput.push_back( - parentFunc.getArgAttrOfType( - argIndex, tensorrtShapeBoundsAttrName)); - - dimensionNamesAttrsPerInput.push_back( - parentFunc.getArgAttrOfType( - argIndex, tensorrtDimensionNamesAttrName)); - - if (!profileAttrsPerInput.back()) { - return emitError(blockArg.getLoc()) - << "Profile attribute (" << tensorrtShapeBoundsAttrName - << ") of argument " << argIndex << " is not set"; + // Get shape profile and dynamision name attributes of the input + if (rtt.hasStaticShape()) { + // static-shaped argument can only have value bound attr (shape input) + auto valueBoundAttr = + parentFunc.getArgAttrOfType( + argIndex, tensorrtValueBoundsAttrName); + if (valueBoundAttr) { + func->setArgAttr(argIndex, tensorrtValueBoundsAttrName, valueBoundAttr); + } + // Get memory space attribute of the input + auto memorySpaceAttr = + parentFunc.getArgAttr(argIndex, memorySpaceAttrName); + if (memorySpaceAttr) { + func->setArgAttr(argIndex, memorySpaceAttrName, memorySpaceAttr); + // Add tensorrt.host_tensor attr, it is needed by NetworkEncoder for now + func->setArgAttr(argIndex, hostTensorAttrName, rewriter.getUnitAttr()); + } + } else { + auto shapeBoundAttr = + parentFunc.getArgAttrOfType( + argIndex, tensorrtShapeBoundsAttrName); + if (!shapeBoundAttr) { + return emitError(blockArg.getLoc()) + << "Profile attribute (" << tensorrtShapeBoundsAttrName + << ") of argument " << argIndex << " is not set"; + } + func->setArgAttr(argIndex, tensorrtShapeBoundsAttrName, shapeBoundAttr); + auto dimensionNameAttr = parentFunc.getArgAttrOfType( + argIndex, tensorrtDimensionNamesAttrName); + if (dimensionNameAttr) { + func->setArgAttr(argIndex, tensorrtDimensionNamesAttrName, + dimensionNameAttr); + } } } - for (unsigned idx = 0; idx < func->getNumArguments(); idx++) { - if (profileAttrsPerInput[idx]) - func->setArgAttr(idx, tensorrtShapeBoundsAttrName, - profileAttrsPerInput[idx]); - if (dimensionNamesAttrsPerInput[idx]) - func->setArgAttr(idx, tensorrtDimensionNamesAttrName, - dimensionNamesAttrsPerInput[idx]); - } - rewriter.setInsertionPoint(inlineGroupOp); auto callOp = rewriter.create( inlineGroupOp.getLoc(), inlineGroupOp.getResultTypes(), inputs, diff --git a/mlir-tensorrt/tensorrt/test/Target/TensorRT/translate-to-tensorrt.mlir b/mlir-tensorrt/tensorrt/test/Target/TensorRT/translate-to-tensorrt.mlir index 45aaac551..6d7d99488 100644 --- a/mlir-tensorrt/tensorrt/test/Target/TensorRT/translate-to-tensorrt.mlir +++ b/mlir-tensorrt/tensorrt/test/Target/TensorRT/translate-to-tensorrt.mlir @@ -65,3 +65,27 @@ func.func @trt_dim_names( %0 = tensorrt.identity %arg0 : tensor to tensor return %0 : tensor } + +// CHECK-LABEL: @trt_host_input +// CHECK-SAME: tensorrt.engine +func.func @trt_host_input(%arg0: tensor {tensorrt.dimension_names = {}, tensorrt.shape_profile = #tensorrt.shape_profile}, %arg1: tensor {plan.memory_space = #plan.memory_space, tensorrt.value_bounds = #tensorrt.shape_profile}) -> tensor { + %0 = tensorrt.element_wise (%arg0, %arg0 : tensor, tensor) -> tensor + %1 = tensorrt.shape %0 : tensor -> tensor<2xi32> + %2 = tensorrt.slice %1[0][1][1] : tensor<2xi32> to tensor<1xi32> + %3 = tensorrt.collapse_rank %2 : tensor<1xi32> to tensor + %cst_i32 = tensorrt.constant dense<1> : tensor + %4 = tensorrt.element_wise (%3, %cst_i32 : tensor, tensor) -> tensor + %5 = tensorrt.slice %1[1][1][1] : tensor<2xi32> to tensor<1xi32> + %6 = tensorrt.collapse_rank %5 : tensor<1xi32> to tensor + %7 = tensorrt.element_wise (%4, %6 : tensor, tensor) -> tensor + %cst_i32_0 = tensorrt.constant dense<1> : tensor + %8 = tensorrt.element_wise (%arg1, %cst_i32_0 : tensor, tensor) -> tensor + %9 = tensorrt.element_wise (%7, %8 : tensor, tensor) -> tensor + %cst_i32_1 = tensorrt.constant dense<1> : tensor<1xi32> + %10 = tensorrt.reshape %9 shape(%cst_i32_1: tensor<1xi32>) : tensor to tensor + %cst_i32_2 = tensorrt.constant dense<1> : tensor<1xi32> + %11 = tensorrt.reshape %arg1 shape(%cst_i32_2: tensor<1xi32>) : tensor to tensor + %12 = tensorrt.concatenation {axis = 0 : i32} ins(%10, %11 : tensor, tensor) -> tensor<2xi32> + %13 = tensorrt.reshape %0 shape(%12: tensor<2xi32>) : tensor to tensor + return %13 : tensor +}