[mlir][transforms] CSE ops with multiple regions

[springerm]

Summary

There were issues with the CSE equivalence analysis that have been fixed with D142558. This makes it possible to CSE ops with multiple regions (multiple block maybe even?).

Depends On: D142558

[rriddle]

mlir/lib/Transforms/CSE.cpp

205

Please don't use size, it's O(N) with the size of the region.

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

205

Oh right, use llvm::hasSingleElement(r.getBlocks()) instead!

mlir/lib/Transforms/CSE.cpp

LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
 
   // If the operation is already trivially dead just add it to the erase list.
   if (isOpTriviallyDead(op)) {
     opsToErase.push_back(op);
     ++numDCE;
     return success();
   }
 
-  // Don't simplify operations with nested blocks. We don't currently model
-  // equality comparisons correctly among other things. It is also unclear
-  // whether we would want to CSE such operations.
-  if (op->getNumRegions() != 0 &&
-      (op->getNumRegions() != 1 || !llvm::hasSingleElement(op->getRegion(0))))
+  // Don't simplify operations with regions that have multiple blocks.
+  // TODO: We need additional tests to verify that we handle such IR correctly.
+  if (!llvm::all_of(op->getRegions(), [](Region &r) {
+        return r.getBlocks().empty() || llvm::hasSingleElement(r.getBlocks());
+      }))
     return failure();
 
   // Some simple use case of operation with memory side-effect are dealt with
   // here. Operations with no side-effect are done after.
   if (!isMemoryEffectFree(op)) {
     auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
     // TODO: Only basic use case for operations with MemoryEffects::Read can be
     // eleminated now. More work needs to be done for more complicated patterns

mlir/test/Transforms/cse.mlir

func.func @failing_issue_59135(%arg0: tensor<2x2xi1>, %arg1: f32, %arg2 : tensor<2xi1>) -> (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]]
+
+func.func @cse_multiple_regions(%c: i1, %t: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
+  %r1 = scf.if %c -> (tensor<5xf32>) {
+    %0 = tensor.empty() : tensor<5xf32>
+    scf.yield %0 : tensor<5xf32>
+  } else {
+    scf.yield %t : tensor<5xf32>
+  }
+  %r2 = scf.if %c -> (tensor<5xf32>) {
+    %0 = tensor.empty() : tensor<5xf32>
+    scf.yield %0 : tensor<5xf32>
+  } else {
+    scf.yield %t : tensor<5xf32>
+  }
+  return %r1, %r2 : tensor<5xf32>, tensor<5xf32>
+}
+// CHECK-LABEL: func @cse_multiple_regions
+//       CHECK:   %[[if:.*]] = scf.if {{.*}} {
+//       CHECK:     tensor.empty
+//       CHECK:     scf.yield
+//       CHECK:   } else {
+//       CHECK:     scf.yield
+//       CHECK:   }
+//   CHECK-NOT:   scf.if
+//       CHECK:   return %[[if]], %[[if]]