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

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

---------

Co-authored-by: Christian Ulmann <[email protected]>

Author: Moxinilian

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 =
+      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()) {

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(any(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"() : () -> ()
+}

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
 //===----------------------------------------------------------------------===//

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 = [{