[mlir][Transforms] Simplify region before simplifying operation in CSE.

[mravishankar]

Summary

This covers more options for CSE. It also ensures that two operations
that have same operands but different regions to begin with, but same
regions after simplifyRegions, don't get both added to the list of
knownValues.

Fixes #59135

[mravishankar]

mlir/test/Transforms/cse.mlir

465

Damn, forgot to write the actual CHECKs. Will update shortly.

[rriddle]

Makes sense to me!

mlir/lib/Transforms/CSE.cpp

319

This will miss the possible simplification of the region op?

[mravishankar]

Add missing CHECK lines.

[mravishankar]

Refactor to account for invalid continue.

mlir/lib/Transforms/CSE.cpp

319

Good catch! Hopefully this is fixed now...

[Mogball]

Thanks for the fix

[mravishankar]

Rebase.

[rriddle]

Forgot to hit approve, but LGTM of course.

mlir/lib/Transforms/CSE.cpp

LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
   // Otherwise, we add this operation to the known values map.
   knownValues.insert(op, op);
   return failure();
 }
 
 void CSE::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
                         bool hasSSADominance) {
   for (auto &op : *bb) {
-    // If the operation is simplified, we don't process any held regions.
-    if (succeeded(simplifyOperation(knownValues, &op, hasSSADominance)))
-      continue;
-
     // Most operations don't have regions, so fast path that case.
-    if (op.getNumRegions() == 0)
-      continue;
-
-    // If this operation is isolated above, we can't process nested regions with
-    // the given 'knownValues' map. This would cause the insertion of implicit
-    // captures in explicit capture only regions.
+    if (op.getNumRegions() != 0) {
+      // If this operation is isolated above, we can't process nested regions
+      // with the given 'knownValues' map. This would cause the insertion of
+      // implicit captures in explicit capture only regions.
       if (op.mightHaveTrait<OpTrait::IsIsolatedFromAbove>()) {
         ScopedMapTy nestedKnownValues;
         for (auto &region : op.getRegions())
           simplifyRegion(nestedKnownValues, region);
-      continue;
-    }
-
+      } else {
         // Otherwise, process nested regions normally.
         for (auto &region : op.getRegions())
           simplifyRegion(knownValues, region);
       }
+    }
+
+    // If the operation is simplified, we don't process any held regions.
+    if (succeeded(simplifyOperation(knownValues, &op, hasSSADominance)))
+      continue;
+  }
   // Clear the MemoryEffects cache since its usage is by block only.
   memEffectsCache.clear();
 }
 
 void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;

mlir/test/Transforms/cse.mlir

%1 = test.cse_of_single_block_op inputs(%a, %b) {
     test.region_yield %2 : f32
   } : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
   return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
 }
 // CHECK-LABEL: func @cse_single_block_with_commutative_ops
 //       CHECK:   %[[OP:.+]] = test.cse_of_single_block_op
 //   CHECK-NOT:   test.cse_of_single_block_op
 //       CHECK:   return %[[OP]], %[[OP]]
+
+func.func @failing_issue_59135(%arg0: tensor<2x2xi1>, %arg1: f32, %arg2 : tensor<2xi1>) -> (tensor<2xi1>, tensor<2xi1>) {
+  %false_2 = arith.constant false 
+  %true_5 = arith.constant true 
+  %9 = test.cse_of_single_block_op inputs(%arg2) {
+  ^bb0(%out: i1):
+    %true_144 = arith.constant true
+    test.region_yield %true_144 : i1
+  } : tensor<2xi1> -> tensor<2xi1>
+  %15 = test.cse_of_single_block_op inputs(%arg2) {
+  ^bb0(%out: i1):
+    %true_144 = arith.constant true
+    test.region_yield %true_144 : i1
+  } : tensor<2xi1> -> tensor<2xi1>
+  %93 = arith.maxsi %false_2, %true_5 : i1
+  return %9, %15 : tensor<2xi1>, tensor<2xi1>
+}
+// CHECK-LABEL: func @failing_issue_59135
+//       CHECK:   %[[TRUE:.+]] = arith.constant true
+//       CHECK:   %[[OP:.+]] = test.cse_of_single_block_op
+//       CHECK:     test.region_yield %[[TRUE]]
+//       CHECK:   return %[[OP]], %[[OP]]