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

[Théo Degioanni]

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

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

>From e696290874ff653d8f1d726352a24569fc4a1c24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9o=20Degioanni?=
 <theo.degioanni.llvm.deluge062 at simplelogin.fr>
Date: Tue, 20 Aug 2024 11:52:51 +0200
Subject: [PATCH] fix mem2reg visiting graph regions

---
 mlir/lib/Transforms/Mem2Reg.cpp           | 13 +++++++++++++
 mlir/test/Transforms/mem2reg.mlir         | 15 ++++++++++++++-
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp |  8 ++++++++
 mlir/test/lib/Dialect/Test/TestOps.td     | 12 ++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

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