[TensorRT] Copy tensorrt.host_tensor attribute in outline pass

Author: yizhuoz004

mlir-tensorrt/compiler/lib/Compiler/TensorRTToExecutable/Passes.cpp

@@ -280,16 +280,22 @@ 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();
 
   SmallVector<Attribute> profileAttrsPerInput;
   SmallVector<Attribute> dimensionNamesAttrsPerInput;
   for (Value v : inputs) {
     auto rtt = dyn_cast<RankedTensorType>(v.getType());
-    if (!rtt || rtt.hasStaticShape()) {
+    if (!rtt) {
       profileAttrsPerInput.push_back(Attribute{});
       dimensionNamesAttrsPerInput.push_back(Attribute{});
       continue;
     }
+    if (rtt.hasStaticShape()) {
+      dimensionNamesAttrsPerInput.push_back(Attribute{});
+    }
 
     auto blockArg = dyn_cast<BlockArgument>(v);
     if (!blockArg || blockArg.getOwner()->getParentOp() != parentFunc) {
@@ -299,18 +305,30 @@ outlineOp(RewriterBase &rewriter, tensorrt::TensorRTModuleOp trtModule,
     }
 
     int64_t argIndex = blockArg.getArgNumber();
-    profileAttrsPerInput.push_back(
+    // Get shape profile and dynamision name attributes of the input
+    if (rtt.hasStaticShape()) {
+      // static-shaped argument can only have value bound attr (shape input)
+      profileAttrsPerInput.push_back(
+        parentFunc.getArgAttrOfType<tensorrt::ShapeProfileAttr>(
+            argIndex, tensorrtValueBoundsAttrName));
+      dimensionNamesAttrsPerInput.push_back(Attribute{});
+      // Get host tensor attribute of the input
+      auto hostTensorAttr = parentFunc.getArgAttr(argIndex, hostTensorAttrName);
+      if (hostTensorAttr) {
+        func->setArgAttr(argIndex, hostTensorAttrName, hostTensorAttr);
+      }
+    } else {
+      profileAttrsPerInput.push_back(
         parentFunc.getArgAttrOfType<tensorrt::ShapeProfileAttr>(
             argIndex, tensorrtShapeBoundsAttrName));
-
-    dimensionNamesAttrsPerInput.push_back(
+      dimensionNamesAttrsPerInput.push_back(
         parentFunc.getArgAttrOfType<DictionaryAttr>(
             argIndex, tensorrtDimensionNamesAttrName));
-
-    if (!profileAttrsPerInput.back()) {
-      return emitError(blockArg.getLoc())
-             << "Profile attribute (" << tensorrtShapeBoundsAttrName
-             << ") of argument " << argIndex << " is not set";
+      if (!profileAttrsPerInput.back()) {
+        return emitError(blockArg.getLoc())
+                << "Profile attribute (" << tensorrtShapeBoundsAttrName
+                << ") of argument " << argIndex << " is not set";
+      }
     }
   }