[OpenMP][OMPIRBuilder] Collect users of a value before replacing them in target outlined function

[ergawy]

This PR fixes a crash that curently happens given the following input:

subroutine caller()
  real :: x
  integer :: i

  !$omp target
    x = i
    call callee(x,x)
  !$omp end target
endsubroutine caller

subroutine callee(x1,x2)
  real :: x1, x2
endsubroutine callee

The crash happens because the following sequence of events is taken by the OMPIRBuilder:

1. ....
2. An outlined function for the target region is created. At first the outlined function still refers to the SSA values from the original function of the target region.
3. The builder then iterates over the users of SSA values used in the target region to replace them with the corresponding function arguments of outlined function.
4. If the same instruction references the SSA value more than once (say m), all uses of that SSA value are replaced in the instruction. Deleting all m uses of the value.
5. The next m-1 iterations will still iterate over the same instruction dropping the last m-1 actual users of the value.

Hence, we collect all users first before modifying them.

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-mlir-llvm

Author: Kareem Ergawy (ergawy)

Changes

This PR fixes a crash that curently happens given the following input:

subroutine caller()
  real :: x, y, z
  integer :: i

  !$omp target
    x = i
    call callee(x,x)
  !$omp end target
endsubroutine caller

subroutine callee(x1,x2)
  real :: x1, x2
endsubroutine callee

The crash happens because the following sequence of events is taken by the OMPIRBuilder:

1. ....
2. An outlined function for the target region is created. At first the
   outlined function still refers to the SSA values from the original
   function of the target region.
3. The builder then iterates over the users of SSA values used in the
   target region to replace them with the corresponding function arguments
   of outlined function.
4. If the same instruction references the SSA value more than once (say m),
   all uses of that SSA value are replaced in the instruction.
   Deleting all m uses of the value.
5. The next m-1 iterations will still iterate over the same
   instruction dropping the last m-1 actual users of the value.

Hence, we call all users first before modifying them.

---

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

2 Files Affected:

• (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+5-1)
• (added) mlir/test/Target/LLVMIR/omp-target-call-with-repeated-parameter.mlir (+37)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 7c1b64677a011..3dfde6f5e236d 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -7035,8 +7035,12 @@ static Expected<Function *> createOutlinedFunction(
     if (auto *Const = dyn_cast<Constant>(Input))
       convertUsersOfConstantsToInstructions(Const, Func, false);
 
+    // Collect users before iterating over them to avoid invalidating the
+    // iteration in case a user uses Input more than once (e.g. a call
+    // instruction).
+    SetVector<User *> Users(Input->users().begin(), Input->users().end());
     // Collect all the instructions
-    for (User *User : make_early_inc_range(Input->users()))
+    for (User *User : make_early_inc_range(Users))
       if (auto *Instr = dyn_cast<Instruction>(User))
         if (Instr->getFunction() == Func)
           Instr->replaceUsesOfWith(Input, InputCopy);
diff --git a/mlir/test/Target/LLVMIR/omp-target-call-with-repeated-parameter.mlir b/mlir/test/Target/LLVMIR/omp-target-call-with-repeated-parameter.mlir
new file mode 100644
index 0000000000000..91b4129f6ff9d
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omp-target-call-with-repeated-parameter.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @caller_()  {
+  %c1 = llvm.mlir.constant(1 : i64) : i64
+  %x_host = llvm.alloca %c1 x f32 {bindc_name = "x"} : (i64) -> !llvm.ptr
+  %i_host = llvm.alloca %c1 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+  %x_map = omp.map.info var_ptr(%x_host : !llvm.ptr, f32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "x"}
+  %i_map = omp.map.info var_ptr(%i_host : !llvm.ptr, i32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "i"}
+  omp.target map_entries(%x_map -> %x_arg, %i_map -> %i_arg : !llvm.ptr, !llvm.ptr) {
+    %1 = llvm.load %i_arg : !llvm.ptr -> i32
+    %2 = llvm.sitofp %1 : i32 to f32
+    llvm.store %2, %x_arg : f32, !llvm.ptr
+    // The call instruction uses %x_arg more than once. Hence modifying users
+    // while iterating them invalidates the iteration. Which is what is tested
+    // by here.
+    llvm.call @callee_(%x_arg, %x_arg) : (!llvm.ptr, !llvm.ptr) -> ()
+    omp.terminator
+  }
+  llvm.return
+}
+
+llvm.func @callee_(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+  llvm.return
+}
+
+
+// CHECK: define internal void @__omp_offloading_{{.*}}_caller__{{.*}}(ptr %[[X_PARAM:.*]], ptr %[[I_PARAM:.*]]) {
+
+// CHECK:   %[[I_VAL:.*]] = load i32, ptr %[[I_PARAM]], align 4
+// CHECK:   %[[I_VAL_FL:.*]] = sitofp i32 %[[I_VAL]] to float
+// CHECK:   store float %[[I_VAL_FL]], ptr %[[X_PARAM]], align 4
+// CHECK:   call void @callee_(ptr %[[X_PARAM]], ptr %[[X_PARAM]])
+// CHECK:   br label %[[REGION_CONT:.*]]
+
+// CHECK: [[REGION_CONT]]:
+// CHECK:   ret void
+// CHECK: }

[skatrak]

Thank you Kareem, if buildbots are clean this LGTM!

[agozillon]

LGTM, thanks for the fix!