[flang-commits] [flang] [flang] Restrict mem2reg promotion through fir.declare to single-block case (PR #182933)

Mon Feb 23 12:47:00 PST 2026

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/182933

The PromotableOpInterface on fir.declare allows mem2reg to promote allocas accessed through declare ops. However, MLIR's mem2reg computes defining blocks and live-in sets only from direct users of the slot pointer. Stores through fir.declare are users of the declare result, not the alloca, so they are not registered as defining blocks. This causes missing phi nodes at join points (loop headers, merge blocks), which silently drops conditional updates to promoted variables.
This was observed in CUDA Fortran kernels where a loop variable updated conditionally (e.g., mywatch = max(1, mywatch-32)) became constant after promotion, producing incorrect results at runtime.
The fix restricts promotion through fir.declare to cases where all users of the declare are in the same block. In single-block cases no phi nodes are needed, so the MLIR limitation does not apply. Cross-block cases are left unpromoted until the MLIR mem2reg infrastructure is extended to track defining blocks through PromotableOpInterface results.

With the current behavior, this would be the result. 
```
func.func @loop_conditional_update(%arg0: i32, %cdt: i1) -> i32 {
  %c1 = arith.constant 1 : i32
  %alloca = fir.alloca i32 {bindc_name = "mywatch", uniq_name = "_QFkernelEmywatch"}
  %declare = fir.declare %alloca {uniq_name = "_QFkernelEmywatch"} : (!fir.ref<i32>) -> !fir.ref<i32>
  fir.store %arg0 to %declare : !fir.ref<i32>
  llvm.br ^loop
^loop:
  %val = fir.load %declare : !fir.ref<i32>
  llvm.cond_br %cdt, ^update, ^exit
^update:
  %new = arith.subi %val, %c1 : i32
  fir.store %new to %declare : !fir.ref<i32>
  llvm.br ^loop
^exit:
  %result = fir.load %declare : !fir.ref<i32>
  return %result : i32
}
```

```
  func.func @loop_conditional_update(%arg0: i32, %arg1: i1) -> i32 {
    %c1_i32 = arith.constant 1 : i32
    fir.declare_value %arg0 {uniq_name = "_QFkernelEmywatch"} : i32
    llvm.br ^bb1
  ^bb1:  // 2 preds: ^bb0, ^bb2
    llvm.cond_br %arg1, ^bb2, ^bb3
  ^bb2:  // pred: ^bb1
    %0 = arith.subi %arg0, %c1_i32 : i32 // Doesn't use current value. 
    fir.declare_value %0 {uniq_name = "_QFkernelEmywatch"} : i32
    llvm.br ^bb1
  ^bb3:  // pred: ^bb1
    return %arg0 : i32 // always return $arg0
  }
```

A better fix should probably be done in mem2reg to support these cases better. I'll look into that later this week. 

>From 0694d6f438b1cf15d7893047d9a80e659a6b6462 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Mon, 23 Feb 2026 12:43:23 -0800
Subject: [PATCH] [flang] Restrict mem2reg promotion through fir.declare to
 single-block cases

---
 flang/lib/Optimizer/Dialect/FIROps.cpp | 12 +++++-
 flang/test/Fir/mem2reg.mlir            | 56 ++++++++++++++++++++------
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 03deba420ec4c..17aa042bda60b 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -5306,9 +5306,17 @@ bool fir::DeclareOp::canUsesBeRemoved(
     const mlir::DataLayout &dataLayout) {
   if (!isLegalTypeForValueDeclare(fir::unwrapRefType(getType())))
     return false;
-  // Forward uses to the users of the fir.declare.
-  for (mlir::OpOperand &use : getResult().getUses())
+  // MLIR's mem2reg computes defining blocks only from direct users of
+  // the slot pointer. Stores through fir.declare are not direct users,
+  // so they are not registered as defining blocks. This causes missing
+  // phi nodes at join points (e.g., loop headers). Restrict promotion
+  // to the single-block case where no phi nodes are needed.
+  mlir::Block *declBlock = getOperation()->getBlock();
+  for (mlir::OpOperand &use : getResult().getUses()) {
+    if (use.getOwner()->getBlock() != declBlock)
+      return false;
     newBlockingUses.push_back(&use);
+  }
   return true;
 }
 
diff --git a/flang/test/Fir/mem2reg.mlir b/flang/test/Fir/mem2reg.mlir
index 0329dfc6f26d2..5109c2ac3001b 100644
--- a/flang/test/Fir/mem2reg.mlir
+++ b/flang/test/Fir/mem2reg.mlir
@@ -158,18 +158,17 @@ func.func @box_not_mem2reg(%arg0: !fir.ref<!fir.box<f32>> {fir.bindc_name = "i"}
 
 // -----
 
-// CHECK-LABEL:   func.func @block_argument_value(
-// CHECK-SAME:      %[[ARG0:.*]]: i32,
-// CHECK-SAME:      %[[ARG1:.*]]: i1) -> i32 {
-// CHECK:           %[[CONSTANT_0:.*]] = arith.constant 42 : i32
-// CHECK:           fir.declare_value %[[CONSTANT_0]] {uniq_name = "_QFfooEjlocal"} : i32
-// CHECK:           llvm.cond_br %[[ARG1]], ^bb1, ^bb2
-// CHECK:         ^bb1:
-// CHECK:           fir.declare_value %[[ARG0]] {uniq_name = "_QFfooEjlocal"} : i32
-// CHECK:           llvm.br ^bb2
-// CHECK:         ^bb2:
-// CHECK:           return %[[CONSTANT_0]] : i32
-// CHECK:         }
+// Conditional store in a different block through fir.declare is not promoted
+// because MLIR mem2reg would not place the needed phi nodes correctly.
+
+// CHECK-LABEL: func.func @block_argument_value(
+// CHECK-SAME: %[[ARG0:.*]]: i32,
+// CHECK-SAME: %[[ARG1:.*]]: i1) -> i32 {
+// CHECK: fir.alloca i32
+// CHECK: fir.declare
+// CHECK: fir.store
+// CHECK: fir.store
+// CHECK: fir.load
 func.func @block_argument_value(%arg0: i32, %cdt: i1) -> i32 {
   %c42_i32 = arith.constant 42 : i32
   %3 = fir.alloca i32 {bindc_name = "jlocal", uniq_name = "_QFfooEjlocal"}
@@ -183,3 +182,36 @@ func.func @block_argument_value(%arg0: i32, %cdt: i1) -> i32 {
   %6 = fir.load %4 : !fir.ref<i32>
   return %6 : i32
 }
+
+// -----
+
+// Conditional store inside a loop through fir.declare must not be promoted.
+// MLIR's mem2reg does not register stores through declares as defining blocks,
+// so phi nodes at the loop header would be missing, losing the update.
+
+// CHECK-LABEL: func.func @loop_conditional_update(
+// CHECK-SAME: %[[ARG0:.*]]: i32,
+// CHECK-SAME: %[[ARG1:.*]]: i1) -> i32 {
+// CHECK: fir.alloca i32
+// CHECK: fir.declare
+// CHECK: fir.store
+// CHECK: fir.load
+// CHECK: fir.store
+// CHECK: fir.load
+func.func @loop_conditional_update(%arg0: i32, %cdt: i1) -> i32 {
+  %c1 = arith.constant 1 : i32
+  %alloca = fir.alloca i32 {bindc_name = "mywatch", uniq_name = "_QFkernelEmywatch"}
+  %declare = fir.declare %alloca {uniq_name = "_QFkernelEmywatch"} : (!fir.ref<i32>) -> !fir.ref<i32>
+  fir.store %arg0 to %declare : !fir.ref<i32>
+  llvm.br ^loop
+^loop:
+  %val = fir.load %declare : !fir.ref<i32>
+  llvm.cond_br %cdt, ^update, ^exit
+^update:
+  %new = arith.subi %val, %c1 : i32
+  fir.store %new to %declare : !fir.ref<i32>
+  llvm.br ^loop
+^exit:
+  %result = fir.load %declare : !fir.ref<i32>
+  return %result : i32
+}

