[Mlir-commits] [mlir] [mlir] [mem2reg] Fix Mem2Reg attempting to promote in graph regions (PR #104910)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Théo Degioanni (Moxinilian)

<details>
<summary>Changes</summary>

Mem2Reg assumes SSA dependencies but did not check for graph regions. This fixes it.

---
Full diff: https://github.com/llvm/llvm-project/pull/104910.diff


4 Files Affected:

- (modified) mlir/lib/Transforms/Mem2Reg.cpp (+13) 
- (modified) mlir/test/Transforms/mem2reg.mlir (+14-1) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+8) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+12) 


``````````diff
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 1f6998709ae02e..406787e2f591ca 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -13,6 +13,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
@@ -255,6 +256,18 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
   // delete itself). We thus need to start from the use of the slot pointer and
   // propagate further requests through the forward slice.
 
+  // Because this pass currently only supports analysing the parent region of
+  // the slot pointer, if a promotable memory op that needs promotion is within
+  // a graph region, the slot may only be used in a graph region and should
+  // therefore be ignored.
+  Region *slotPtrRegion = slot.ptr.getParentRegion();
+  auto slotPtrRegionOp =
+      llvm::dyn_cast<RegionKindInterface>(slotPtrRegion->getParentOp());
+  if (slotPtrRegionOp &&
+      slotPtrRegionOp.getRegionKind(slotPtrRegion->getRegionNumber()) ==
+          RegionKind::Graph)
+    return failure();
+
   // First insert that all immediate users of the slot pointer must no longer
   // use it.
   for (OpOperand &use : slot.ptr.getUses()) {
diff --git a/mlir/test/Transforms/mem2reg.mlir b/mlir/test/Transforms/mem2reg.mlir
index daeaa2da076341..db11e2d83249a7 100644
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg),test.isolated_graph_region(mem2reg))' --split-input-file | FileCheck %s
 
 // Verifies that allocators with mutliple slots are handled properly.
 
@@ -26,3 +26,16 @@ func.func @multi_slot_alloca_only_second() -> (i32, i32) {
   %4 = memref.load %2[] : memref<i32>
   return %3, %4 : i32, i32
 }
+
+// -----
+
+// Checks that slots are not promoted if used in a graph region.
+
+// CHECK-LABEL: test.isolated_graph_region
+test.isolated_graph_region {
+  // CHECK: %{{[[:alnum:]]+}} = test.multi_slot_alloca
+  %slot = test.multi_slot_alloca : () -> (memref<i32>)
+  memref.store %a, %slot[] : memref<i32>
+  %a = memref.load %slot[] : memref<i32>
+  "test.foo"() : () -> ()
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index fbaa102d3e33cc..69091fb893fad6 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -126,6 +126,14 @@ RegionKind GraphRegionOp::getRegionKind(unsigned index) {
   return RegionKind::Graph;
 }
 
+//===----------------------------------------------------------------------===//
+// IsolatedGraphRegionOp
+//===----------------------------------------------------------------------===//
+
+RegionKind IsolatedGraphRegionOp::getRegionKind(unsigned index) {
+  return RegionKind::Graph;
+}
+
 //===----------------------------------------------------------------------===//
 // AffineScopeOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 2b55bff3538d39..9e19966414d1d7 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2048,6 +2048,18 @@ def GraphRegionOp : TEST_Op<"graph_region",  [
   let assemblyFormat = "attr-dict-with-keyword $region";
 }
 
+def IsolatedGraphRegionOp : TEST_Op<"isolated_graph_region",  [
+    DeclareOpInterfaceMethods<RegionKindInterface>,
+    IsolatedFromAbove]> {
+  let summary =  "isolated from above operation with a graph region";
+  let description = [{
+    Test op that defines a graph region which is isolated from above.
+  }];
+
+  let regions = (region AnyRegion:$region);
+  let assemblyFormat = "attr-dict-with-keyword $region";
+}
+
 def AffineScopeOp : TEST_Op<"affine_scope", [AffineScope]> {
   let summary =  "affine scope operation";
   let description = [{

``````````

</details>


https://github.com/llvm/llvm-project/pull/104910