[Mlir-commits] [mlir] aa4e54f - [mlir][transforms] CSE ops with multiple regions

[Matthias Springer]

Author: Matthias Springer
Date: 2023-01-27T12:14:45+01:00
New Revision: aa4e54f2f444c266310105fccbdbb07fc71894da

URL: https://github.com/llvm/llvm-project/commit/aa4e54f2f444c266310105fccbdbb07fc71894da
DIFF: https://github.com/llvm/llvm-project/commit/aa4e54f2f444c266310105fccbdbb07fc71894da.diff

LOG: [mlir][transforms] CSE ops with multiple regions

There were issues with the CSE equivalence analysis that have been fixed with D142558. This makes it possible to CSE ops with multiple regions.

Differential Revision: https://reviews.llvm.org/D142562

Added: 
    

Modified: 
    mlir/lib/Transforms/CSE.cpp
    mlir/test/Transforms/cse.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/CSE.cpp b/mlir/lib/Transforms/CSE.cpp
index 86debe7271fc..93e5c95cb199 100644
--- a/mlir/lib/Transforms/CSE.cpp
+++ b/mlir/lib/Transforms/CSE.cpp
@@ -199,11 +199,11 @@ LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
     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

diff  --git a/mlir/test/Transforms/cse.mlir b/mlir/test/Transforms/cse.mlir
index 7fdbf9551fe2..7086f5f462f5 100644
--- a/mlir/test/Transforms/cse.mlir
+++ b/mlir/test/Transforms/cse.mlir
@@ -468,3 +468,28 @@ func.func @failing_issue_59135(%arg0: tensor<2x2xi1>, %arg1: f32, %arg2 : tensor
 //       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]]