[Mlir-commits] [mlir] 215950 (PR #221123)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Sep 3 19:56:31 PDT 2026 

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Bhavesh M (beamandala)

<details>
<summary>Changes</summary>

```mlir
module {
  func.func @<!-- -->affine_example(%arg0: memref<i32>) {
    %alloca = memref.alloca() : memref<i32>
    %c33_i32 = arith.constant 33 : i32
    affine.for %arg1 = 0 to 4 {
      memref.store %c33_i32, %alloca[] : memref<i32>
      %0 = memref.load %alloca[] : memref<i32>
      memref.store %0, %arg0[] : memref<i32>
    }
    return
  }
}
```

The `mem2reg` pass couldn't promote memory slots accessed within an `affine.for` because `AffineForOp` did not implement `PromotableRegionOpInterface`, leading to the stack allocation and its accesses in the example above to not be eliminated. 

This PR implements `PromotableRegionOpInterface` for `AffineForOp`, allowing `mem2reg` to promote memory slots through affine loops.

Fixes #<!-- -->215950 

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


5 Files Affected:

- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.h (+1) 
- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+3-1) 
- (modified) mlir/lib/Dialect/Affine/IR/CMakeLists.txt (+2) 
- (added) mlir/lib/Dialect/Affine/IR/MemorySlot.cpp (+57) 
- (added) mlir/test/Dialect/Affine/mem2reg.mlir (+116) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index c5cd8c6c31375..ffc9e4be4bb53 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -22,6 +22,7 @@
 #include "mlir/Interfaces/AlignmentAttrInterface.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 namespace mlir {
 namespace affine {
 
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 994df5aef0d9a..cf2848f55ea09 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -21,6 +21,7 @@ include "mlir/Interfaces/InferIntDivisibilityOpInterface.td"
 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Affine_Dialect : Dialect {
@@ -139,7 +140,8 @@ def AffineForOp : Affine_Op<"for",
       "getLoopUpperBounds", "getYieldedValuesMutable",
       "replaceWithAdditionalYields", "getStaticTripCount"]>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface,
-     ["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
+     ["getEntrySuccessorOperands", "getSuccessorInputs"]>,
+     DeclareOpInterfaceMethods<PromotableRegionOpInterface>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
diff --git a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
index 1caf2fa396797..e24cd4013bca3 100644
--- a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
@@ -4,6 +4,7 @@ add_mlir_dialect_library(MLIRAffineDialect
   AffineValueMap.cpp
   InferIntDivisibilityOpInterfaceImpl.cpp
   InferIntRangeInterfaceImpls.cpp
+  MemorySlot.cpp
   ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
@@ -22,6 +23,7 @@ add_mlir_dialect_library(MLIRAffineDialect
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
+  MLIRMemorySlotUtils
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRUBDialect
diff --git a/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp b/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
new file mode 100644
index 0000000000000..169414b44e2af
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Interfaces/Utils/MemorySlotUtils.h"
+
+using namespace mlir;
+using namespace mlir::affine;
+
+//===----------------------------------------------------------------------===//
+// AffineForOp
+//===----------------------------------------------------------------------===//
+
+bool AffineForOp::isRegionPromotable(const MemorySlot &slot, Region *region,
+                                     bool hasValueStores) {
+  return true;
+}
+
+void AffineForOp::setupPromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
+
+  Region &bodyRegion = getBodyRegion();
+  if (!hasValueStores) {
+    regionsToProcess.insert({&bodyRegion, reachingDef});
+    return;
+  }
+
+  getInitsMutable().append(reachingDef);
+  bodyRegion.addArgument(slot.elemType, slot.ptr.getLoc());
+  regionsToProcess.insert({&bodyRegion, bodyRegion.getArguments().back()});
+}
+
+Value AffineForOp::finalizePromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
+    OpBuilder &builder) {
+  if (!hasValueStores)
+    return reachingDef;
+
+  // Update the yield terminator to return the newly defined reaching
+  // definition.
+  memoryslot::updateTerminator(getBody(), reachingDef, reachingAtBlockEnd);
+
+  SmallVector<Type> resultTypes(getResultTypes());
+  resultTypes.push_back(slot.elemType);
+
+  IRRewriter rewriter(builder);
+  Operation *newOp =
+      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
+  return newOp->getResults().back();
+}
diff --git a/mlir/test/Dialect/Affine/mem2reg.mlir b/mlir/test/Dialect/Affine/mem2reg.mlir
new file mode 100644
index 0000000000000..af890cdb7a7a0
--- /dev/null
+++ b/mlir/test/Dialect/Affine/mem2reg.mlir
@@ -0,0 +1,116 @@
+// RUN: mlir-opt %s --mem2reg --split-input-file | FileCheck %s \
+// RUN:   -implicit-check-not "memref.alloca" \
+// RUN:   -implicit-check-not "memref.load" \
+// RUN:   -implicit-check-not "memref.store"
+
+/// Check promotion through a for loop with a load and store in the body.
+
+// CHECK-LABEL: func.func @for_load_and_store
+// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index)
+// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+// CHECK: %[[RES:.*]] = affine.for %{{.*}} = %[[LB]] to %[[UB]] iter_args(%[[ARG:.*]] = %[[C5]]) -> (i32) {
+// CHECK:   %[[NEW:.*]] = arith.addi %[[ARG]], %[[C1]] : i32
+// CHECK:   affine.yield %[[NEW]] : i32
+// CHECK: }
+// CHECK: return %[[RES]] : i32
+func.func @for_load_and_store(%lb: index, %ub: index) -> i32 {
+  %c5 = arith.constant 5 : i32
+  %c1 = arith.constant 1 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c5, %alloca[] : memref<i32>
+  affine.for %i = %lb to %ub {
+    %load = memref.load %alloca[] : memref<i32>
+    %new = arith.addi %load, %c1 : i32
+    memref.store %new, %alloca[] : memref<i32>
+  }
+  %load2 = memref.load %alloca[] : memref<i32>
+  return %load2 : i32
+}
+
+// -----
+
+/// Check promotion adds a second iter_arg when one already exists.
+
+// CHECK-LABEL: func.func @for_existing_iter_arg
+// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index, %[[INIT:.*]]: i32)
+// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+// CHECK: %[[RES:.*]]:2 = affine.for %{{.*}} = %[[LB]] to %[[UB]] iter_args(%[[MUL_ARG:.*]] = %[[INIT]], %[[SLOT_ARG:.*]] = %[[C5]]) -> (i32, i32) {
+// CHECK:   %[[MUL:.*]] = arith.muli %[[MUL_ARG]], %[[MUL_ARG]] : i32
+// CHECK:   %[[NEW:.*]] = arith.addi %[[SLOT_ARG]], %[[C1]] : i32
+// CHECK:   affine.yield %[[MUL]], %[[NEW]] : i32, i32
+// CHECK: }
+// CHECK: return %[[RES]]#1 : i32
+func.func @for_existing_iter_arg(%lb: index, %ub: index, %init: i32) -> i32 {
+  %c5 = arith.constant 5 : i32
+  %c1 = arith.constant 1 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c5, %alloca[] : memref<i32>
+  %mul_res = affine.for %i = %lb to %ub iter_args(%mul_arg = %init) -> i32 {
+    %mul = arith.muli %mul_arg, %mul_arg : i32
+    %load = memref.load %alloca[] : memref<i32>
+    %new = arith.addi %load, %c1 : i32
+    memref.store %new, %alloca[] : memref<i32>
+    affine.yield %mul : i32
+  }
+  %load2 = memref.load %alloca[] : memref<i32>
+  return %load2 : i32
+}
+
+// -----
+
+/// Check load-only promotion through a for loop generates no iter_arg.
+
+func.func private @use(i32)
+
+// CHECK-LABEL: func.func @for_load_only
+// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index)
+// CHECK: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK: affine.for %{{.*}} = %[[LB]] to %[[UB]] {
+// CHECK:   call @use(%[[C5]])
+// CHECK: }
+// CHECK: return %[[C5]] : i32
+func.func @for_load_only(%lb: index, %ub: index) -> i32 {
+  %c5 = arith.constant 5 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c5, %alloca[] : memref<i32>
+  affine.for %i = %lb to %ub {
+    %load = memref.load %alloca[] : memref<i32>
+    func.call @use(%load) : (i32) -> ()
+  }
+  %load2 = memref.load %alloca[] : memref<i32>
+  return %load2 : i32
+}
+
+// -----
+
+/// Check promotion through nested for loops with a load and store in the inner loop.
+
+// CHECK-LABEL: func.func @for_nested_load_and_store
+// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index)
+// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+// CHECK: %[[OUTER:.*]] = affine.for %{{.*}} = %[[LB]] to %[[UB]] iter_args(%[[OUTER_ARG:.*]] = %[[C5]]) -> (i32) {
+// CHECK:   %[[INNER:.*]] = affine.for %{{.*}} = 0 to 4 iter_args(%[[INNER_ARG:.*]] = %[[OUTER_ARG]]) -> (i32) {
+// CHECK:     %[[NEW:.*]] = arith.addi %[[INNER_ARG]], %[[C1]] : i32
+// CHECK:     affine.yield %[[NEW]] : i32
+// CHECK:   }
+// CHECK:   affine.yield %[[INNER]] : i32
+// CHECK: }
+// CHECK: return %[[OUTER]] : i32
+func.func @for_nested_load_and_store(%lb: index, %ub: index) -> i32 {
+  %c5 = arith.constant 5 : i32
+  %c1 = arith.constant 1 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c5, %alloca[] : memref<i32>
+  affine.for %i = %lb to %ub {
+    affine.for %j = 0 to 4 {
+      %load = memref.load %alloca[] : memref<i32>
+      %new = arith.addi %load, %c1 : i32
+      memref.store %new, %alloca[] : memref<i32>
+    }
+  }
+  %load2 = memref.load %alloca[] : memref<i32>
+  return %load2 : i32
+}

``````````

</details>


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

More information about the Mlir-commits mailing list