[mlir] Add `test-convergence` option to Canonicalizer tests

[springerm]

Summary

This new option is set to false by default. It should be set only in Canonicalizer tests to detect faulty canonicalization patterns. I.e., patterns that prevent the canonicalizer from converging. The canonicalizer should always convergence on such small unit tests that we have in canonicalize.mlir.

Two faulty canonicalization patterns were detected and fixed with this change.

Depends On: D140729

[mehdi_amini]

mlir/test/Transforms/canonicalize.mlir

1

Can you keep the -pass-pipeline form? In general I would prefer we stick to it: it is less ambiguous in the intent.

[springerm]

mlir/test/Transforms/canonicalize.mlir

1

I've seen mixed usage across all canonicalize.mlir. Why is -pass-pipeline better here? It's just a single pass in these tests.

mlir/include/mlir/Transforms/Passes.td

Option<"topDownProcessingEnabled", "top-down", "bool",
            "Seed the worklist in general top-down order">,
     Option<"enableRegionSimplification", "region-simplify", "bool",
            /*default=*/"true",
            "Perform control flow optimizations to the region tree">,
     Option<"maxIterations", "max-iterations", "int64_t",
            /*default=*/"10",
            "Max. iterations between applying patterns / simplifying regions">,
     Option<"maxNumRewrites", "max-num-rewrites", "int64_t", /*default=*/"-1",
-           "Max. number of pattern rewrites within an iteration">
+           "Max. number of pattern rewrites within an iteration">,
+    Option<"testConvergence", "test-convergence", "bool", /*default=*/"false",
+           "Test only: Fail pass on non-convergence to detect cyclic pattern">
   ] # RewritePassUtils.options;
 }
 
 def ControlFlowSink : Pass<"control-flow-sink"> {
   let summary = "Sink operations into conditional blocks";
   let description = [{
     This pass implements control-flow sink on operations that implement
     `RegionBranchOpInterface` by moving dominating operations whose only uses

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

LogicalResult matchAndRewrite(LaunchOp op,
     // If the range implies a single value for `id`, replace `id`'s uses by
     // zero.
     Value zero;
     bool simplified = false;
     auto constPropIdUses = [&](Value id, Value size) {
       // Check if size is trivially one.
       if (!matchPattern(size, m_One()))
         return;
+      if (id.getUses().empty())
+        return;
       if (!simplified) {
         // Create a zero value the first time.
         OpBuilder::InsertionGuard guard(rewriter);
         rewriter.setInsertionPointToStart(&op.getBody().front());
         zero =
             rewriter.create<arith::ConstantIndexOp>(op.getLoc(), /*value=*/0);
       }
-      id.replaceAllUsesWith(zero);
+      rewriter.replaceAllUsesWith(id, zero);
       simplified = true;
     };
     constPropIdUses(op.getBlockIds().x, op.getGridSizeX());
     constPropIdUses(op.getBlockIds().y, op.getGridSizeY());
     constPropIdUses(op.getBlockIds().z, op.getGridSizeZ());
     constPropIdUses(op.getThreadIds().x, op.getBlockSizeX());
     constPropIdUses(op.getThreadIds().y, op.getBlockSizeY());
     constPropIdUses(op.getThreadIds().z, op.getBlockSizeZ());

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

 struct RemoveConstantIfCondition : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
     // Early return if there is no condition.
     Value ifCond = op.getIfCond();
     if (!ifCond)
-      return success();
+      return failure();
 
     IntegerAttr constAttr;
-    if (matchPattern(ifCond, m_Constant(&constAttr))) {
-      if (constAttr.getInt())
-        rewriter.updateRootInPlace(op,
-                                   [&]() { op.getIfCondMutable().erase(0); });
-      else
-        rewriter.eraseOp(op);
-    }
+    if (!matchPattern(ifCond, m_Constant(&constAttr)))
+      return failure();
+      rewriter.updateRootInPlace(op, [&]() { op.getIfCondMutable().erase(0); });
 
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ParallelOp

mlir/lib/Transforms/Canonicalizer.cpp

LogicalResult initialize(MLIRContext *context) override {
     return success();
   }
   void runOnOperation() override {
     GreedyRewriteConfig config;
     config.useTopDownTraversal = topDownProcessingEnabled;
     config.enableRegionSimplification = enableRegionSimplification;
     config.maxIterations = maxIterations;
     config.maxNumRewrites = maxNumRewrites;
+    LogicalResult converged =
+        applyPatternsAndFoldGreedily(getOperation(), patterns, config);
     // Canonicalization is best-effort. Non-convergence is not a pass failure.
-    (void)applyPatternsAndFoldGreedily(getOperation(), patterns, config);
+    if (testConvergence && failed(converged))
+      signalPassFailure();
   }
 
   FrozenRewritePatternSet patterns;
 };
 } // namespace
 
 /// Create a Canonicalizer pass.
 std::unique_ptr<Pass> mlir::createCanonicalizerPass() {

mlir/test/Dialect/AMDGPU/canonicalize.mlir

-// RUN: mlir-opt %s -split-input-file -canonicalize  | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
 
 // CHECK-LABEL: func @known_oob_load
 func.func @known_oob_load(%arg0: memref<4xf32>) -> f32 {
   // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f32
   // CHECK: return %[[zero]]
   %c4_i32 = arith.constant 4 : i32
   %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c4_i32] : memref<4xf32>, i32 -> f32
   func.return %0 : f32

mlir/test/Dialect/Affine/canonicalize.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="top-down=0" | FileCheck %s --check-prefix=CHECK-BOTTOM-UP
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence top-down=0" | FileCheck %s --check-prefix=CHECK-BOTTOM-UP
 
 // -----
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 + 1)>
 
 // CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() {
 func.func @compose_affine_maps_1dto2d_no_symbols() {

mlir/test/Dialect/Arith/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize --split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file | FileCheck %s
 
 // CHECK-LABEL: @select_same_val
 //       CHECK:   return %arg1
 func.func @select_same_val(%arg0: i1, %arg1: i64) -> i64 {
   %0 = arith.select %arg0, %arg1, %arg1 : i64
   return %0 : i64
 }
 

mlir/test/Dialect/Bufferization/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize --split-input-file \
-// RUN:   -allow-unregistered-dialect |\
+// RUN: mlir-opt %s \
+// RUN:   -canonicalize="test-convergence" \
+// RUN:   --split-input-file -allow-unregistered-dialect | \
 // RUN: FileCheck %s
 
 // Basic folding of to_tensor(to_memref(t)) -> t
 // CHECK-LABEL: func @tensor_load_of_buffer_cast(
 func.func @tensor_load_of_buffer_cast(%arg0: tensor<?xf32>) -> tensor<?xf32> {
   %0 = bufferization.to_memref %arg0 : memref<?xf32>
   %1 = bufferization.to_tensor %0 : memref<?xf32>
   return %1 : tensor<?xf32>

mlir/test/Dialect/Builtin/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // UnrealizedConversionCastOp
 //===----------------------------------------------------------------------===//
 
 // Test folding conversion casts feeding into other casts.
 // CHECK-LABEL: func @multiple_conversion_casts
 // CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]:

mlir/test/Dialect/Complex/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" | FileCheck %s
 
 // CHECK-LABEL: func @create_of_real_and_imag
 // CHECK-SAME: (%[[CPLX:.*]]: complex<f32>)
 func.func @create_of_real_and_imag(%cplx: complex<f32>) -> complex<f32> {
   // CHECK-NEXT: return %[[CPLX]] : complex<f32>
   %real = complex.re %cplx : complex<f32>
   %imag = complex.im %cplx : complex<f32>
   %complex = complex.create %real, %imag : complex<f32>

 // CHECK-LABEL: func @complex_sub_zero
 func.func @complex_sub_zero() -> complex<f32> {
   %complex1 = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
   %complex2 = complex.constant [0.0 : f32, 0.0 : f32] : complex<f32>
   // CHECK: %[[CPLX:.*]] = complex.constant [1.000000e+00 : f32, 0.000000e+00 : f32] : complex<f32>
   // CHECK-NEXT: return %[[CPLX:.*]] : complex<f32>
   %sub = complex.sub %complex1, %complex2 : complex<f32>
   return %sub : complex<f32>
-}
+}
 

mlir/test/Dialect/ControlFlow/canonicalize.mlir

-// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='builtin.module(func.func(canonicalize))' -split-input-file | FileCheck --dump-input-context 20 %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck --dump-input-context 20 %s
 
 /// Test the folding of BranchOp.
 
 // CHECK-LABEL: func @br_folding(
 func.func @br_folding() -> i32 {
   // CHECK-NEXT: %[[CST:.*]] = arith.constant 0 : i32
   // CHECK-NEXT: return %[[CST]] : i32
   %c0_i32 = arith.constant 0 : i32

mlir/test/Dialect/GPU/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // Fold all the gpu.wait ops as they are redundant.
 // CHECK-LABEL: func @fold_wait_op_test1
 func.func @fold_wait_op_test1() {
   %1 = gpu.wait async
   gpu.wait []
   %3 = gpu.wait async
   gpu.wait [%3]

mlir/test/Dialect/LLVMIR/canonicalize.mlir

-// RUN: mlir-opt -canonicalize %s -split-input-file | FileCheck %s
+// RUN: mlir-opt -canonicalize="test-convergence" %s -split-input-file | FileCheck %s
 
 // CHECK-LABEL: fold_extractvalue
 llvm.func @fold_extractvalue() -> i32 {
   //  CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
   %c0 = arith.constant 0 : i32
   //  CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
   %c1 = arith.constant 1 : i32
 

mlir/test/Dialect/Linalg/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @memref_cast(
 func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c8 = arith.constant 8 : index
   %c16 = arith.constant 16 : index
   %1 = memref.alloc (%b) : memref<?xi8>

mlir/test/Dialect/Math/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" | FileCheck %s
 
 // CHECK-LABEL: @ceil_fold
 // CHECK: %[[cst:.+]] = arith.constant 1.000000e+00 : f32
 // CHECK: return %[[cst]]
 func.func @ceil_fold() -> f32 {
   %c = arith.constant 0.3 : f32
   %r = math.ceil %c : f32
   return %r : f32

mlir/test/Dialect/MemRef/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: func @subview_of_size_memcast
 //  CHECK-SAME:   %[[ARG0:.[a-z0-9A-Z_]+]]: memref<4x6x16x32xi8>
 //       CHECK:   %[[S:.+]] = memref.subview %[[ARG0]][0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : memref<4x6x16x32xi8> to memref<16x32xi8, strided{{.*}}>
 //       CHECK:   %[[M:.+]] = memref.cast %[[S]] : memref<16x32xi8, {{.*}}> to memref<16x32xi8, strided{{.*}}>
 //       CHECK:   return %[[M]] : memref<16x32xi8, strided{{.*}}>
 func.func @subview_of_size_memcast(%arg : memref<4x6x16x32xi8>) ->
   memref<16x32xi8, strided<[32, 1], offset: ?>>{

mlir/test/Dialect/OpenACC/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
 
 func.func @testenterdataop(%a: memref<10xf32>) -> () {
   %ifCond = arith.constant true
   acc.enter_data if(%ifCond) create(%a: memref<10xf32>)
   return
 }
 
 // CHECK: acc.enter_data create(%{{.*}} : memref<10xf32>)

mlir/test/Dialect/OpenMP/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
 
 func.func @update_no_op(%x : memref<i32>) {
   omp.atomic.update %x : memref<i32> {
   ^bb0(%xval : i32):
     omp.yield(%xval : i32)
   }
   return
 }

mlir/test/Dialect/PDL/canonicalize.mlir

-// RUN: mlir-opt -canonicalize %s | FileCheck %s
+// RUN: mlir-opt -canonicalize="test-convergence" %s | FileCheck %s
 
 pdl.pattern @operation_op : benefit(1) {
   %root = operation "foo.op"
   rewrite %root {
     // CHECK: operation "bar.unused"
     %unused_rewrite = operation "bar.unused"
     erase %root
   }
 }

mlir/test/Dialect/Quant/canonicalize.mlir

-// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(func.func(canonicalize))' | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' | FileCheck %s
 
 // -----
 // CHECK-LABEL: redundant_scast
 func.func @redundant_scast() -> tensor<4xi8> {
   // CHECK-NEXT: arith.constant dense<10> : tensor<4xi8>
   // CHECK-NEXT: return
   %cst = arith.constant dense<5> : tensor<4xi8>
   %1 = "quant.scast"(%cst) : (tensor<4xi8>) -> tensor<4x!quant.uniform<u8:f32, 7.812500e-03:128>>

mlir/test/Dialect/SCF/canonicalize.mlir

-// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize))' -split-input-file | FileCheck %s
-
-
-// -----
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck %s
 
 func.func @single_iteration_some(%A: memref<?x?x?xi32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   %c6 = arith.constant 6 : index
   %c7 = arith.constant 7 : index

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

-// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(func.func(canonicalize))' | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // spirv.AccessChain
 //===----------------------------------------------------------------------===//
 
 func.func @combine_full_access_chain() -> f32 {
   // CHECK: %[[INDEX:.*]] = spirv.Constant 0
   // CHECK-NEXT: %[[VAR:.*]] = spirv.Variable

mlir/test/Dialect/Shape/canonicalize.mlir

-// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize="test-convergence" %s | FileCheck %s
 
 // CHECK-LABEL: func @f
 func.func @f(%arg0: tensor<2x3x4xf32>) -> tensor<3xindex> {
   // CHECK: shape.const_shape [2, 3, 4] : tensor<3xindex>
   %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<3xindex>
   return %0 : tensor<3xindex>
 }
 

mlir/test/Dialect/Tensor/canonicalize.mlir

-// RUN: mlir-opt %s -split-input-file -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
 
 // Checks that NOP casts are removed.
 // CHECK-LABEL: cast_values
 func.func @cast_values(%arg0: tensor<*xi32>) -> tensor<2xi32> {
   // NOP cast
   %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32>
   // CHECK-NEXT: %[[RET:.*]] = tensor.cast %arg0 : tensor<*xi32> to tensor<2xi32>
   %2 = tensor.cast %0 : tensor<*xi32> to tensor<2xi32>

mlir/test/Dialect/Tosa/canonicalize.mlir

-// RUN: mlir-opt --canonicalize %s | FileCheck %s
+// RUN: mlir-opt -canonicalize="test-convergence" %s | FileCheck %s
 
 // CHECK-LABEL: @argmax_nofold
 func.func @argmax_nofold(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
   // CHECK: "tosa.argmax"
   %0 = "tosa.argmax"(%arg0) {axis = 0 : i64}: (tensor<?x1xf32>) -> tensor<?x1xf32>
   return %0 : tensor<?x1xf32>
 }
 

mlir/test/Dialect/Vector/canonicalize.mlir

-// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize))' -split-input-file -allow-unregistered-dialect | FileCheck %s
-
-// -----
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: create_vector_mask_to_constant_mask
 func.func @create_vector_mask_to_constant_mask() -> (vector<4x3xi1>) {
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   // CHECK: vector.constant_mask [3, 2] : vector<4x3xi1>
   %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
   return %0 : vector<4x3xi1>

mlir/test/Pass/run-reproducer.mlir

 // RUN: mlir-opt %s -dump-pass-pipeline 2>&1 | FileCheck %s
 // RUN: mlir-opt %s -mlir-print-ir-before=cse 2>&1 | FileCheck -check-prefix=BEFORE %s
 
 func.func @foo() {
   %0 = arith.constant 0 : i32
   return
 }
 
 func.func @bar() {
   return
 }
 
 {-#
   external_resources: {
     mlir_reproducer: {
       verify_each: true,
-      // CHECK:  builtin.module(func.func(cse,canonicalize{ max-iterations=1 max-num-rewrites=-1 region-simplify=false top-down=false}))
+      // CHECK:  builtin.module(func.func(cse,canonicalize{ max-iterations=1 max-num-rewrites=-1 region-simplify=false test-convergence=false top-down=false}))
       pipeline: "builtin.module(func.func(cse,canonicalize{max-iterations=1 max-num-rewrites=-1 region-simplify=false top-down=false}))",
       disable_threading: true
     }
   }
 #-}
 
 // BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @foo()
 // BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @bar()
 // BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
 // BEFORE-NOT: // -----// IR Dump After

mlir/test/Transforms/canonicalize.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(canonicalize))' -split-input-file | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @test_subi_zero
 func.func @test_subi_zero(%arg0: i32) -> i32 {
   // CHECK-NEXT: %c0_i32 = arith.constant 0 : i32
   // CHECK-NEXT: return %c0
   %y = arith.subi %arg0, %arg0 : i32
   return %y: i32
 }