[MLIR][Mem2Reg] Ensure dominance of default value in regions

[ivanradanov]

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.

For reference, this is what used to be the output for the test. Note the use of %1.

  "func.func"() <{function_type = (f64) -> (), sym_name = "poison_insertion_point"}> ({
  ^bb0(%arg0: f64):
    "cf.br"(%1)[^bb1] : (f64) -> ()
  ^bb1(%0: f64):  // 2 preds: ^bb0, ^bb2
    %1 = "ub.poison"() <{value = #ub.poison}> : () -> f64
    %2 = "test.get"() : () -> i1
    "cf.cond_br"(%2)[^bb2, ^bb3] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()
  ^bb2:  // pred: ^bb1
    %3 = "test.get"() : () -> i1
    %4 = "scf.if"(%3) ({
      "scf.yield"(%arg0) : (f64) -> ()
    }, {
      "scf.yield"(%0) : (f64) -> ()
    }) : (i1) -> f64
    "test.use"(%4) : (f64) -> ()
    "cf.br"(%4)[^bb1] : (f64) -> ()
  ^bb3:  // pred: ^bb1
    "func.return"() : () -> ()
  }) : () -> ()

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ivan R. Ivanov (ivanradanov)

Changes

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

[llvmbot]

@llvm/pr-subscribers-mlir-core

Author: Ivan R. Ivanov (ivanradanov)

Changes

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

[vzakhari]

Thank you, Ivan! This looks good to me, but please wait for Theo's approval.

[Moxinilian]

This makes sense to me, thanks!