[Mlir-commits] [mlir] 64171f0 - [mlir][affine] Implement PromotableRegionOpInterface for AffineForOp (#221123)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Sep 6 08:23:51 PDT 2026
Author: Bhavesh M
Date: 2026-09-06T20:53:46+05:30
New Revision: 64171f09367ae70475eebe9184353ca5948681fb
URL: https://github.com/llvm/llvm-project/commit/64171f09367ae70475eebe9184353ca5948681fb
DIFF: https://github.com/llvm/llvm-project/commit/64171f09367ae70475eebe9184353ca5948681fb.diff
LOG: [mlir][affine] Implement PromotableRegionOpInterface for AffineForOp (#221123)
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 to not be eliminated.
This change implements `PromotableRegionOpInterface` for `AffineForOp`,
allowing `mem2reg` to promote memory slots through affine loops.
Added:
mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
mlir/test/Dialect/Affine/mem2reg.mlir
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Dialect/Affine/IR/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index c5cd8c6c31375..5cc7acfe83c26 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -22,6 +22,8 @@
#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..7b74b7eb7bb6b
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
@@ -0,0 +1,56 @@
+//===----------------------------------------------------------------------===//
+//
+// 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> ®ionsToProcess) {
+ 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
+}
More information about the Mlir-commits
mailing list