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

[llvmlistbot at llvm.org]

Author: Théo Degioanni
Date: 2024-08-23T15:15:10+02:00
New Revision: b084111c8e26f96975f505c37d42e992066776f8

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

LOG: [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 <christianulmann at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Transforms/Mem2Reg.cpp
    mlir/test/Transforms/mem2reg.mlir
    mlir/test/lib/Dialect/Test/TestOpDefs.cpp
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 1f6998709ae02e..4bc2bbd12b0945 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 =
+      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..89472ac0ca2842 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(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"() : () -> ()
+}

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