[tensorrt] Refine types for tensorrt ops with plan.with_shape

Adds `TensorRTRefineTypeFromWithShapeGeneric` pattern to refines
types of tensorrt ops in `PlanRefineTypesPass`.

Author: yizhuoz004

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

@@ -52,12 +52,12 @@ static void updateTypeInPlaceAndMaybeInsertCast(RewriterBase &rewriter,
   rewriter.modifyOpInPlace(toUpdate.getDefiningOp(),
                            [&]() { toUpdate.setType(newType); });
 
-  // If all the users are StableHLO ops or plugins, then they all allow in-place
-  // update of operand types.
-  auto isOpaquePlugin = [](Operation *op) {
-    return llvm::isa<tensorrt::OpaquePluginOp>(op);
+  // If all the users are StableHLO or TensorRT ops, then they all allow
+  // in-place update of operand types.
+  auto isTensorRTOp = [](Operation *op) {
+    return llvm::isa<tensorrt::TensorRTDialect>(op->getDialect());
   };
-  if (stablehlo::canUpdateTypeWithoutCast(toUpdate, isOpaquePlugin))
+  if (stablehlo::canUpdateTypeWithoutCast(toUpdate, isTensorRTOp))
     return;
 
   OpBuilder::InsertionGuard g(rewriter);
@@ -294,6 +294,38 @@ struct StableHloRefineTypeFromWithShapeGeneric
   }
 };
 
+/// Given a pattern `plan.with_shape(tensorrt_op, dims...)`, if inspection of
+/// `dims` yields an opportunity to refine the type of `with_shape`, then
+/// `tensorrt_op` can also be refined. The refinements are made (and casts are
+/// inserted if required).
+struct TensorRTRefineTypeFromWithShapeGeneric
+    : public OpRewritePattern<WithShapeOp> {
+  using OpRewritePattern<WithShapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(WithShapeOp withOp,
+                                PatternRewriter &rewriter) const override {
+    auto producer = withOp.getOperand().getDefiningOp();
+    if (!producer || !producer->hasOneUse() ||
+        !isa<tensorrt::TensorRTDialect>(producer->getDialect()))
+      return failure();
+
+    // Create a new shape and try to refine it.
+    std::optional<SmallVector<int64_t>> newShape =
+        getRefinedShape(withOp.getShape(), withOp.getOperand().getType());
+    if (!newShape)
+      return failure();
+
+    // Update type of the producer.
+    updateTypeInPlaceAndMaybeInsertCast(
+        rewriter, withOp.getOperand(),
+        withOp.getOperand().getType().clone(*newShape));
+
+    // Update type of the WithShapeOp.
+    updateTypeInPlaceAndMaybeInsertCast(rewriter, withOp.getResult(),
+                                        withOp.getType().clone(*newShape));
+    return success();
+  }
+};
+
 class PlanRefineTypesPass
     : public plan::impl::PlanRefineTypesPassBase<PlanRefineTypesPass> {
   using Base::Base;
@@ -315,10 +347,12 @@ class PlanRefineTypesPass
         RefineDynamicIota,
         SimplifyIdentityDynamicBroadcast,
         StableHloRefineTypeFromWithShapeGeneric,
-        WithShapeAbsorbCastPattern
+        WithShapeAbsorbCastPattern,
+        TensorRTRefineTypeFromWithShapeGeneric
       >(ctx);
     // clang-format on
     stablehlo::populateStablehloRefineShapesPatterns(&patterns, ctx);
+    stablehlo::populateStablehloCanonicalizationPatterns(ctx, &patterns);
     if (failed(applyPatternsAndFoldGreedily(funcTarget, std::move(patterns),
                                             config))) {
       emitError(funcTarget.getLoc())

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

@@ -126,8 +126,36 @@ func.func @tensorrt_opaque_plugin_no_cast() -> tensor<?xf32> {
   return %3 : tensor<?xf32>
 }
 // CHECK-LABEL: func.func @tensorrt_opaque_plugin_no_cast
-// CHECK-SAME: () -> 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-SAME: () -> tensor<1xf32>
+// CHECK: %[[cst:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<64xf32>
+// CHECK: %[[v0:.*]] = tensorrt.opaque_plugin
+// CHECK: return %{{.*}} : tensor<1xf32>
+
+// -----
+
+func.func @refine_tensorrt_resize_with_shape() -> tensor<?x?x?x?xf32> {
+  %c3 = arith.constant 3 : index
+  %c2 = arith.constant 2 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %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>
+  %dim_i32_0 = tensor.extract %c[%c0] : tensor<4xi32>
+  %dim_i32_1 = tensor.extract %c[%c1] : tensor<4xi32>
+  %dim_i32_2 = tensor.extract %c[%c2] : tensor<4xi32>
+  %dim_i32_3 = tensor.extract %c[%c3] : tensor<4xi32>
+  %dim_0 = arith.index_cast %dim_i32_0 : i32 to index
+  %dim_1 = arith.index_cast %dim_i32_1 : i32 to index
+  %dim_2 = arith.index_cast %dim_i32_2 : i32 to index
+  %dim_3 = arith.index_cast %dim_i32_3 : i32 to index
+  %1 = plan.with_shape %result(%dim_0, %dim_1, %dim_2, %dim_3) : (tensor<?x?x?x?xf32>, index, index, index, index) -> tensor<?x?x?x?xf32>
+  return %1 : tensor<?x?x?x?xf32>
+}
+// CHECK-LABEL: func.func @refine_tensorrt_resize_with_shape
+// CHECK-SAME: -> tensor<1x1x8x8xf32>
+// CHECK: %[[v0:.*]] = 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 %[[v0]] : tensor<1x1x8x8xf32>