Add in plan-refine-types pass, fix unit test

Author: yizhuoz004

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

@@ -54,10 +54,10 @@ static void updateTypeInPlaceAndMaybeInsertCast(RewriterBase &rewriter,
 
   // If all the users are StableHLO ops or plugins, then they all allow in-place
   // update of operand types.
-  auto isOpaquePlugin = [](Operation *op) {
+  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);
@@ -377,6 +377,7 @@ class PlanRefineTypesPass
       >(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,11 +126,10 @@ 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>
 
 // -----