[Mlir-commits] [mlir] [mlir][affine] Implement PromotableRegionOpInterface for AffineIfOp (PR #225388)

[Paul Berg]

https://github.com/Pangoraw created https://github.com/llvm/llvm-project/pull/225388

Follow-up to #221123 to use the structured mem2reg pass on affine if operations.

>From 00e4812ed72dd2d0e43886abbf336a08b4720048 Mon Sep 17 00:00:00 2001
From: Paul Berg <9824244+Pangoraw at users.noreply.github.com>
Date: Tue, 22 Sep 2026 07:48:50 -0500
Subject: [PATCH] [mlir][affine] Implement PromotableRegionOpInterface for
 AffineIfOp

---
 .../mlir/Dialect/Affine/IR/AffineOps.td       |  3 +-
 mlir/lib/Dialect/Affine/IR/MemorySlot.cpp     | 45 +++++++++++++++++++
 mlir/test/Dialect/Affine/mem2reg.mlir         | 38 ++++++++++++++++
 3 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index cf2848f55ea09..100eb9b5ce2b3 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -355,7 +355,8 @@ def AffineIfOp : Affine_Op<"if",
                            [ImplicitAffineTerminator, RecursivelySpeculatable,
                             RecursiveMemoryEffects, NoRegionArguments,
                             DeclareOpInterfaceMethods<RegionBranchOpInterface,
-                                ["getSuccessorInputs"]>
+                                ["getSuccessorInputs"]>,
+                            DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
                            ]> {
   let summary = "if-then-else operation";
   let description = [{
diff --git a/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp b/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
index 7b74b7eb7bb6b..b98cb0d16f92d 100644
--- a/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
+++ b/mlir/lib/Dialect/Affine/IR/MemorySlot.cpp
@@ -54,3 +54,48 @@ Value AffineForOp::finalizePromotion(
       memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
   return newOp->getResults().back();
 }
+
+//===----------------------------------------------------------------------===//
+// AffineIfOp
+//===----------------------------------------------------------------------===//
+
+bool AffineIfOp::isRegionPromotable(const MemorySlot &slot, Region *region,
+                                    bool hasValueStores) {
+  return true;
+}
+
+void AffineIfOp::setupPromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
+  regionsToProcess.insert({&getThenRegion(), reachingDef});
+  regionsToProcess.insert({&getElseRegion(), reachingDef});
+}
+
+Value AffineIfOp::finalizePromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
+    OpBuilder &builder) {
+  if (!hasValueStores)
+    return reachingDef;
+
+  IRRewriter rewriter(builder);
+
+  memoryslot::updateTerminator(&getThenRegion().back(), reachingDef,
+                               reachingAtBlockEnd);
+
+  if (getElseRegion().hasOneBlock()) {
+    memoryslot::updateTerminator(&getElseRegion().back(), reachingDef,
+                                 reachingAtBlockEnd);
+  } else {
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.createBlock(&getElseRegion());
+    AffineYieldOp::create(rewriter, getOperation()->getLoc(), reachingDef);
+  }
+
+  SmallVector<Type> resultTypes(getResultTypes());
+  resultTypes.push_back(slot.elemType);
+
+  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
index af890cdb7a7a0..5920fab45d3b6 100644
--- a/mlir/test/Dialect/Affine/mem2reg.mlir
+++ b/mlir/test/Dialect/Affine/mem2reg.mlir
@@ -114,3 +114,41 @@ func.func @for_nested_load_and_store(%lb: index, %ub: index) -> i32 {
   %load2 = memref.load %alloca[] : memref<i32>
   return %load2 : i32
 }
+
+// -----
+
+#iseven = affine_set<(d0) : (d0 mod 2 == 0)>
+// CHECK-LABEL: func.func @affine_for_affine_if
+// CHECK-SAME: (%[[X:.*]]: f32)
+// CHECK-DAG: %[[CST:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK: %[[RES:.*]] = affine.for %[[IV:.*]] = 0 to 128 iter_args(%[[ARG:.*]] = %[[X]]) -> (f32) {
+// CHECK:   %[[IFRES:.*]] = affine.if #{{.*}}(%[[IV]]) -> f32 {
+// CHECK:     %[[NEW:.*]] = arith.mulf %[[ARG]], %[[CST]] : f32
+// CHECK:     affine.yield %[[NEW]] : f32
+// CHECK:   } else {
+// CHECK:     affine.yield %[[ARG]] : f32
+// CHECK:   }
+// CHECK:   arith.mulf %[[IFRES]], %[[CST]] : f32
+// CHECK:   affine.yield %[[IFRES]] : f32
+// CHECK: }
+// CHECK: return %[[RES]] : f32
+func.func @affine_for_affine_if(%x: f32) -> (f32) {
+  %cst = arith.constant 2.0 : f32
+  %alloc = memref.alloca() : memref<f32>
+  memref.store %x, %alloc[] : memref<f32>
+
+  affine.for %iv = 0 to 128 {
+
+    affine.if #iseven(%iv) {
+      %val = memref.load %alloc[] : memref<f32>
+      %new_val = arith.mulf %val, %cst : f32
+      memref.store %new_val, %alloc[] : memref<f32>
+    }
+
+    %cur_val = memref.load %alloc[] : memref<f32>
+    %prod_next = arith.mulf %cur_val, %cst : f32
+  }
+
+  %prod = memref.load %alloc[] : memref<f32>
+  return %prod : f32
+}