[Mlir-commits] [mlir] [MLIR][Mem2Reg] Ensure dominance of default value in regions (PR #193708)

[Ivan R. Ivanov]

https://github.com/ivanradanov created https://github.com/llvm/llvm-project/pull/193708

When we promote an allocation, and a default value for a load from an uninitialized slot is required, this value used to get inserted in the same block as the allocation. However, in some cases, the default value needs to be available in the predecessor blocks so that they can pass it to the block of the allocation as an argument. For example, this is the case for loops containing an allocation where the promoted value will become and IV.

Make sure the default value is always available to all blocks by creating it in the entry block of the region.

>From b6af18ae6b4cec9b520cf4584a098dcad138ac1b Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Thu, 23 Apr 2026 02:56:30 -0700
Subject: [PATCH] [MLIR][Mem2Reg] Ensure dominance of default value in regions

When we promote an allocation, and a default value for a load from an
uninitialized slot is required, this value used to get inserted in the
same block as the allocation. However, in some cases, the default value
needs to be available in the predecessor blocks so that they can pass it
to the block of the allocation as an argument. For example, this is the
case for loops containing an allocation where the promoted value will
become and IV.

Make sure the default value is always available to all blocks by
creating it in the entry block of the region.
---
 mlir/lib/Transforms/Mem2Reg.cpp   |  2 +-
 mlir/test/Transforms/mem2reg.mlir | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index deaa226f3e809..40d08d869a9e2 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -338,7 +338,7 @@ Value MemorySlotPromoter::getOrCreateDefaultValue() {
     return defaultValue;
 
   OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPointToStart(slot.ptr.getParentBlock());
+  builder.setInsertionPointToStart(&slot.ptr.getParentRegion()->front());
   return defaultValue = allocator.getDefaultValue(slot, builder);
 }
 
diff --git a/mlir/test/Transforms/mem2reg.mlir b/mlir/test/Transforms/mem2reg.mlir
index 484d3e46881d8..94b721cf28dcf 100644
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@@ -153,3 +153,31 @@ func.func @cyclic_unused_merge_points(%cond: i1) -> i32 {
 ^exit:
   return %c0 : i32
 }
+
+// -----
+
+// This function contains a loop bb1 -> bb2 -> bb1, and to promote the alloca, we
+// need to insert a poison value that gets used initially instead of the stored
+// value in case %cond2 is false. This poison value becomes a bb arg in bb1 and
+// bb2, so the entry block must also pass the poison value to bb1. Make sure the
+// poison value is generated in the entry block where it dominates all uses.
+
+// CHECK-LABEL: func.func @poison_insertion_point
+// CHECK-NEXT: ub.poison
+func.func @poison_insertion_point(%val: f64) {
+  cf.br ^bb1
+^bb1:
+  %alloca = memref.alloca() : memref<f64>
+  %cond1 = "test.get"() : () -> i1
+  cf.cond_br %cond1, ^bb2, ^bb3
+^bb2:
+  %cond2 = "test.get"() : () -> i1
+  scf.if %cond2 {
+    memref.store %val, %alloca[] : memref<f64>
+  }
+  %reload = memref.load %alloca[] : memref<f64>
+  "test.use"(%reload) : (f64) -> ()
+  cf.br ^bb1
+^bb3:
+  return
+}