Add type canonicalizer for tensorrt ops

Author: yizhuoz004

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/RefineTypes.cpp

@@ -54,6 +54,7 @@ static void updateTypeInPlaceAndMaybeInsertCast(RewriterBase &rewriter,
 
   // If all the users are StableHLO ops or plugins, then they all allow in-place
   // update of operand types.
+  // Need to make exception for tensorrt dialect?
   auto isOpaquePlugin = [](Operation *op) {
     return llvm::isa<tensorrt::OpaquePluginOp>(op);
   };
@@ -294,6 +295,30 @@ struct StableHloRefineTypeFromWithShapeGeneric
   }
 };
 
+struct TensorRTResizeAbsorbCastPattern: public OpRewritePattern<tensor::CastOp> {
+  using OpRewritePattern<tensor::CastOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(tensor::CastOp castOp,
+                                PatternRewriter &rewriter) const override {
+    auto resizeOp = castOp.getOperand().getDefiningOp();
+
+    // tensor.cast is the only user of the tensorrt op
+    // tensor.cast must be a refining op
+    if (!resizeOp || !resizeOp->hasOneUse() ||
+        !isa<tensorrt::TensorRTDialect>(resizeOp->getDialect()) || 
+        !tensorrt::isTargetRefinementOfSource(
+            castOp.getSource().getType().getShape(),
+            castOp.getType().getShape()))
+      return failure();
+    
+    // how to rewrite?
+    rewriter.modifyOpInPlace(resizeOp, [&]() {
+        resizeOp->getResult(0).setType(
+            cast<RankedTensorType>(castOp.getType()));
+      });
+    return success();
+  }
+};
+
 class PlanRefineTypesPass
     : public plan::impl::PlanRefineTypesPassBase<PlanRefineTypesPass> {
   using Base::Base;
@@ -315,7 +340,8 @@ class PlanRefineTypesPass
         RefineDynamicIota,
         SimplifyIdentityDynamicBroadcast,
         StableHloRefineTypeFromWithShapeGeneric,
-        WithShapeAbsorbCastPattern
+        WithShapeAbsorbCastPattern,
+        TensorRTResizeAbsorbCastPattern
       >(ctx);
     // clang-format on
     stablehlo::populateStablehloRefineShapesPatterns(&patterns, ctx);

mlir-tensorrt/test/Dialect/Plan/refine-types.mlir

@@ -130,4 +130,21 @@ func.func @tensorrt_opaque_plugin_no_cast() -> tensor<?xf32> {
 // CHECK: %[[v1:.*]] = stablehlo.dynamic_reshape %{{.*}}, %{{.*}} : (tensor<64xf32>, tensor<1xi32>) -> tensor<64xf32>
 // CHECK: %[[v2:.*]] = tensorrt.opaque_plugin
 // CHECK-SAME: (%[[v1]]) : (tensor<64xf32>) -> tensor<?xf32>
-// CHECK: return %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
+
+// -----
+
+func.func @tensorrt_resize_no_cast() -> tensor<1x1x8x8xf32> {
+  %cst = stablehlo.constant dense<1.000000e+00> : tensor<1x1x4x4xf32>
+  %c = stablehlo.constant dense<[1, 1, 8, 8]> : tensor<4xi32>
+  %result = tensorrt.resize_linear {
+    coordinateTransformation = #tensorrt.resize_coordinate_transformation<kALIGN_CORNERS>,
+    selectorForSinglePixel = #tensorrt.resize_selector<kUPPER>
+  } %cst, %c : (tensor<1x1x4x4xf32>, tensor<4xi32>) -> tensor<?x?x?x?xf32>
+  %cast = tensor.cast %result : tensor<?x?x?x?xf32> to tensor<1x1x8x8xf32>
+  return %cast : tensor<1x1x8x8xf32>
+}
+// CHECK-LABEL: func.func @tensorrt_resize_no_cast
+// CHECK-SAME: -> tensor<1x1x8x8xf32>
+// CHECK: %[[v1:.*]] = tensorrt.resize_linear {coordinateTransformation = #tensorrt.resize_coordinate_transformation<kALIGN_CORNERS>, selectorForSinglePixel = #tensorrt.resize_selector<kUPPER>} %cst, %c : (tensor<1x1x4x4xf32>, tensor<4xi32>) -> tensor<1x1x8x8xf32>
+// CHECK: return %[[v1]] : tensor<1x1x8x8xf32>