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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Ivan R. Ivanov (ivanradanov)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Transforms/Mem2Reg.cpp (+1-1) 
- (modified) mlir/test/Transforms/mem2reg.mlir (+28) 


``````````diff
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
+}

``````````

</details>


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