Execute ShapeRefinementPipeline and canonicalizer multiple times to reason out more static shapes

Author: Vremold

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -213,6 +213,7 @@ class LowerToBackendContractPass
 
     OpPassManager pm(module.getOperationName());
     TorchLoweringPipelineOptions options;
+    options.maxIterations = maxIterations;
     options.decompose = decompose;
     options.backendLegalOps = backendLegalOps;
     createTorchSimplificationPipeline(pm, options);

lib/Dialect/Torch/Transforms/Passes.cpp

@@ -124,12 +124,14 @@ void mlir::torch::Torch::createTorchSimplificationPipeline(
   pm.addNestedPass<func::FuncOp>(Torch::createMaximizeValueSemanticsPass());
   // Update the return op to return value tensors.
   pm.addPass(Torch::createRefinePublicReturnPass());
-  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
-  // Do shape refinement.
-  // This should be run before RefineTypes (which primarily does dtype
-  // inference), because Torch type promotion rules actually depend on the shape
-  // of the operand.
-  createTorchShapeRefinementPipeline(pm);
+  for (int i = 0; i < options.maxIterations; ++i) {
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    // Do shape refinement.
+    // This should be run before RefineTypes (which primarily does dtype
+    // inference), because Torch type promotion rules actually depend on the shape
+    // of the operand.
+    createTorchShapeRefinementPipeline(pm);
+  }
   // Refine types in the program, which mainly means inferring dtypes of ops.
   pm.addNestedPass<func::FuncOp>(Torch::createRefineTypesPass());
   // Propagate to ABI return types the shape/dtype information discovered by