[Mlir-commits] [llvm] [mlir] [OpenMP][OMPIRBuilder] Collect users of a value before replacing them in target outlined function (PR #139064)

[Kareem Ergawy]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/139064

>From 3db2c56779f76838def2b8a5475d77233a004cc5 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 8 May 2025 05:50:41 -0500
Subject: [PATCH] [OpenMP][OMPIRBuilder] Collect users of a value before
 replacing them in target outlined function

This PR fixes a crash that curently happens given the following input:
```fortran
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.   ....
n.   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.
n+1. 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.
n+2. 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.
n+3. 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.
---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     |  6 ++-
 ...p-target-call-with-repeated-parameter.mlir | 37 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Target/LLVMIR/omp-target-call-with-repeated-parameter.mlir

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..ebaecc8cf203b
--- /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 this test.
+    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: }