[OpenMP][OMPIRBuilder] Handle replace uses of ConstantExpr's inside of Target regions

[agozillon]

Currently there's an edge cases where constant indexing in target regions can lead to incorrect results as we do not correctly replace uses of mapped variables in generated target functions with the target arguments (and accessor instructions) that replace them. This patch seeks to fix that by extending the current logic in the OMPIRBuilder.

Things like GEP's can come in the form of Constants/ConstantExprs, Constants and ConstantExpr's do not have access to the knowledge of what they're contained in, so we must dig a little to find an instruction so we can tell if they're used inside of the function we're outlining so we can be sure they are replaceable and we are not accidentally replacing a usage somewhere else in the module that's still necessary.

This patch handles these by replacing the original constant expression with a new instruction equivalent; an instruction as it allows easy modification in the following loop, as we can now know the constant (instruction) is owned by our target function (as it holds this knowledge) and replaceUsesOfWith can now be invoked on it (cannot do this with constants it seems), a brand new one also allows us to be cautious as it is perhaps possible the old expression was used inside of the function but exists and is used externally (unlikely by the nature of a Constant, but still a positive side affect).

[llvmbot]

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

@llvm/pr-subscribers-openmp

Author: None (agozillon)

Changes

Currently there's an edge cases where constant indexing in target regions can lead to incorrect results as we do not correctly replace uses of mapped variables in generated target functions with the target arguments (and accessor instructions) that replace them. This patch seeks to fix that by extending the current logic in the OMPIRBuilder.

Things like GEP's can come in the form of Constants/ConstantExprs, Constants and ConstantExpr's do not have access to the knowledge of what they're contained in, so we must dig a little to find an instruction so we can tell if they're used inside of the function we're outlining so we can be sure they are replaceable and we are not accidentally replacing a usage somewhere else in the module that's still necessary.

This patch handles these by replacing the original constant expression with a new instruction equivalent; an instruction as it allows easy modification in the following loop, as we can now know the constant (instruction) is owned by our target function (as it holds this knowledge) and replaceUsesOfWith can now be invoked on it (cannot do this with constants it seems), a brand new one also allows us to be cautious as it is perhaps possible the old expression was used inside of the function but exists and is used externally (unlikely by the nature of a Constant, but still a positive side affect).

---

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

3 Files Affected:

• (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+32-4)
• (added) mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir (+41)
• (added) openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90 (+27)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index eb6c120b2a6dcbd..5b6960c70d51b39 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4813,15 +4813,43 @@ static Function *createOutlinedFunction(
     Value *Input = std::get<0>(InArg);
     Argument &Arg = std::get<1>(InArg);
     Value *InputCopy = nullptr;
-
+    
     Builder.restoreIP(
         ArgAccessorFuncCB(Arg, Input, InputCopy, AllocaIP, Builder.saveIP()));
 
-    // Collect all the instructions
+    // Things like GEP's can come in the form of Constants, constants and
+    // ConstantExpr's do not have access to the knowledge of what they're
+    // contained in, so we must dig a little to find an instruction so we can
+    // tell if they're used inside of the function we're outlining. We also
+    // replace the original constant expression with a new instruction
+    // equivelant; an instruction as it allows easy modification in the
+    // following loop, as we can now know the constant (instruction) is owned by
+    // our target function and replaceUsesOfWith can now be invoked on it
+    // (cannot do this with constants it seems), a brand new one also allows us
+    // to be cautious as it is perhaps possible the old expression was used
+    // inside of the function but exists and is used externally (unlikely by the
+    // nature of a Constant, but still)
+    auto ReplaceConstantUsedInFunction = [](Constant *Const, Function *Func) {
+      if (auto *ConstExpr = dyn_cast<ConstantExpr>(Const))
+        for (User *User : make_early_inc_range(ConstExpr->users()))
+          if (auto *Instr = dyn_cast<Instruction>(User))
+            if (Instr->getFunction() == Func)
+              Instr->replaceUsesOfWith(ConstExpr,
+                                       ConstExpr->getAsInstruction(Instr));
+    };
+
     for (User *User : make_early_inc_range(Input->users()))
-      if (auto Instr = dyn_cast<Instruction>(User))
-        if (Instr->getFunction() == Func)
+      if (auto *Const = dyn_cast<Constant>(User))
+          ReplaceConstantUsedInFunction(Const, Func);
+
+    // Collect all the instructions
+    for (User *User : make_early_inc_range(Input->users())) {
+      if (auto *Instr = dyn_cast<Instruction>(User)) {
+        if (Instr->getFunction() == Func) {
           Instr->replaceUsesOfWith(Input, InputCopy);
+        }
+      }
+    }
   }
 
   // Restore insert point.
diff --git a/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir b/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
new file mode 100644
index 000000000000000..ef7ff696b35567e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
@@ -0,0 +1,41 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {omp.is_target_device = true} {
+  llvm.func @_QQmain() attributes {bindc_name = "main"} {
+    %0 = llvm.mlir.addressof @_QFEsp : !llvm.ptr
+    %1 = llvm.mlir.constant(10 : index) : i64
+    %2 = llvm.mlir.constant(1 : index) : i64
+    %3 = llvm.mlir.constant(0 : index) : i64
+    %4 = llvm.mlir.constant(9 : index) : i64
+    %5 = omp.bounds lower_bound(%3 : i64) upper_bound(%4 : i64) extent(%1 : i64) stride(%2 : i64) start_idx(%2 : i64)
+    %6 = omp.map_info var_ptr(%0 : !llvm.ptr, !llvm.array<10 x i32>) map_clauses(tofrom) capture(ByRef) bounds(%5) -> !llvm.ptr {name = "sp"}
+    omp.target map_entries(%6 -> %arg0 : !llvm.ptr) {
+    ^bb0(%arg0: !llvm.ptr):
+      %7 = llvm.mlir.constant(20 : i32) : i32
+      %8 = llvm.mlir.constant(0 : i64) : i64
+      %9 = llvm.getelementptr %arg0[0, %8] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i32>
+      llvm.store %7, %9 : i32, !llvm.ptr
+      %10 = llvm.mlir.constant(10 : i32) : i32
+      %11 = llvm.mlir.constant(4 : i64) : i64
+      %12 = llvm.getelementptr %arg0[0, %11] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i32>
+      llvm.store %10, %12 : i32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+  llvm.mlir.global internal @_QFEsp(dense<0> : tensor<10xi32>) {addr_space = 0 : i32} : !llvm.array<10 x i32>
+  llvm.mlir.global external constant @_QQEnvironmentDefaults() {addr_space = 0 : i32} : !llvm.ptr {
+    %0 = llvm.mlir.zero : !llvm.ptr
+    llvm.return %0 : !llvm.ptr
+  }
+}
+
+// CHECK: define weak_odr protected void @__omp_offloading_fd00_494504a__QQmain_l12(ptr %{{.*}}, ptr %[[ARG1:.*]]) {
+
+// CHECK: %[[ARG1_ALLOCA:.*]] = alloca ptr, align 8
+// CHECK: store ptr %[[ARG1]], ptr %[[ARG1_ALLOCA]], align 8
+// CHECK: %[[LOAD_ARG1_ALLOCA:.*]] = load ptr, ptr %[[ARG1_ALLOCA]], align 8
+// CHECK: store i32 20, ptr %[[LOAD_ARG1_ALLOCA]], align 4
+// CHECK: %[[GEP_ARG1_ALLOCA:.*]] = getelementptr inbounds [10 x i32], ptr %[[LOAD_ARG1_ALLOCA]], i32 0, i64 4
+// CHECK: store i32 10, ptr %[[GEP_ARG1_ALLOCA]], align 4
+
diff --git a/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90 b/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90
new file mode 100644
index 000000000000000..91dc30cf9604c4b
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90
@@ -0,0 +1,27 @@
+! Basic offloading test with a target region
+! that checks constant indexing on device 
+! correctly works (regression test for prior
+! bug).
+! REQUIRES: flang, amdgcn-amd-amdhsa
+! UNSUPPORTED: nvptx64-nvidia-cuda
+! UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+! UNSUPPORTED: aarch64-unknown-linux-gnu
+! UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+! UNSUPPORTED: x86_64-pc-linux-gnu
+! UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+program main
+    INTEGER :: sp(10) = (/0,0,0,0,0,0,0,0,0,0/)
+
+  !$omp target map(tofrom:sp)
+     sp(1) = 20
+     sp(5) = 10
+  !$omp end target
+
+   ! print *, sp(1)
+   ! print *, sp(5)
+end program
+
+! CHECK: 20
+! CHECK: 10

[llvmbot]

@llvm/pr-subscribers-mlir-openmp

Author: None (agozillon)

Changes

Currently there's an edge cases where constant indexing in target regions can lead to incorrect results as we do not correctly replace uses of mapped variables in generated target functions with the target arguments (and accessor instructions) that replace them. This patch seeks to fix that by extending the current logic in the OMPIRBuilder.

Things like GEP's can come in the form of Constants/ConstantExprs, Constants and ConstantExpr's do not have access to the knowledge of what they're contained in, so we must dig a little to find an instruction so we can tell if they're used inside of the function we're outlining so we can be sure they are replaceable and we are not accidentally replacing a usage somewhere else in the module that's still necessary.

This patch handles these by replacing the original constant expression with a new instruction equivalent; an instruction as it allows easy modification in the following loop, as we can now know the constant (instruction) is owned by our target function (as it holds this knowledge) and replaceUsesOfWith can now be invoked on it (cannot do this with constants it seems), a brand new one also allows us to be cautious as it is perhaps possible the old expression was used inside of the function but exists and is used externally (unlikely by the nature of a Constant, but still a positive side affect).

---

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

3 Files Affected:

• (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+32-4)
• (added) mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir (+41)
• (added) openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90 (+27)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index eb6c120b2a6dcbd..5b6960c70d51b39 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4813,15 +4813,43 @@ static Function *createOutlinedFunction(
     Value *Input = std::get<0>(InArg);
     Argument &Arg = std::get<1>(InArg);
     Value *InputCopy = nullptr;
-
+    
     Builder.restoreIP(
         ArgAccessorFuncCB(Arg, Input, InputCopy, AllocaIP, Builder.saveIP()));
 
-    // Collect all the instructions
+    // Things like GEP's can come in the form of Constants, constants and
+    // ConstantExpr's do not have access to the knowledge of what they're
+    // contained in, so we must dig a little to find an instruction so we can
+    // tell if they're used inside of the function we're outlining. We also
+    // replace the original constant expression with a new instruction
+    // equivelant; an instruction as it allows easy modification in the
+    // following loop, as we can now know the constant (instruction) is owned by
+    // our target function and replaceUsesOfWith can now be invoked on it
+    // (cannot do this with constants it seems), a brand new one also allows us
+    // to be cautious as it is perhaps possible the old expression was used
+    // inside of the function but exists and is used externally (unlikely by the
+    // nature of a Constant, but still)
+    auto ReplaceConstantUsedInFunction = [](Constant *Const, Function *Func) {
+      if (auto *ConstExpr = dyn_cast<ConstantExpr>(Const))
+        for (User *User : make_early_inc_range(ConstExpr->users()))
+          if (auto *Instr = dyn_cast<Instruction>(User))
+            if (Instr->getFunction() == Func)
+              Instr->replaceUsesOfWith(ConstExpr,
+                                       ConstExpr->getAsInstruction(Instr));
+    };
+
     for (User *User : make_early_inc_range(Input->users()))
-      if (auto Instr = dyn_cast<Instruction>(User))
-        if (Instr->getFunction() == Func)
+      if (auto *Const = dyn_cast<Constant>(User))
+          ReplaceConstantUsedInFunction(Const, Func);
+
+    // Collect all the instructions
+    for (User *User : make_early_inc_range(Input->users())) {
+      if (auto *Instr = dyn_cast<Instruction>(User)) {
+        if (Instr->getFunction() == Func) {
           Instr->replaceUsesOfWith(Input, InputCopy);
+        }
+      }
+    }
   }
 
   // Restore insert point.
diff --git a/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir b/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
new file mode 100644
index 000000000000000..ef7ff696b35567e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
@@ -0,0 +1,41 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {omp.is_target_device = true} {
+  llvm.func @_QQmain() attributes {bindc_name = "main"} {
+    %0 = llvm.mlir.addressof @_QFEsp : !llvm.ptr
+    %1 = llvm.mlir.constant(10 : index) : i64
+    %2 = llvm.mlir.constant(1 : index) : i64
+    %3 = llvm.mlir.constant(0 : index) : i64
+    %4 = llvm.mlir.constant(9 : index) : i64
+    %5 = omp.bounds lower_bound(%3 : i64) upper_bound(%4 : i64) extent(%1 : i64) stride(%2 : i64) start_idx(%2 : i64)
+    %6 = omp.map_info var_ptr(%0 : !llvm.ptr, !llvm.array<10 x i32>) map_clauses(tofrom) capture(ByRef) bounds(%5) -> !llvm.ptr {name = "sp"}
+    omp.target map_entries(%6 -> %arg0 : !llvm.ptr) {
+    ^bb0(%arg0: !llvm.ptr):
+      %7 = llvm.mlir.constant(20 : i32) : i32
+      %8 = llvm.mlir.constant(0 : i64) : i64
+      %9 = llvm.getelementptr %arg0[0, %8] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i32>
+      llvm.store %7, %9 : i32, !llvm.ptr
+      %10 = llvm.mlir.constant(10 : i32) : i32
+      %11 = llvm.mlir.constant(4 : i64) : i64
+      %12 = llvm.getelementptr %arg0[0, %11] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i32>
+      llvm.store %10, %12 : i32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+  llvm.mlir.global internal @_QFEsp(dense<0> : tensor<10xi32>) {addr_space = 0 : i32} : !llvm.array<10 x i32>
+  llvm.mlir.global external constant @_QQEnvironmentDefaults() {addr_space = 0 : i32} : !llvm.ptr {
+    %0 = llvm.mlir.zero : !llvm.ptr
+    llvm.return %0 : !llvm.ptr
+  }
+}
+
+// CHECK: define weak_odr protected void @__omp_offloading_fd00_494504a__QQmain_l12(ptr %{{.*}}, ptr %[[ARG1:.*]]) {
+
+// CHECK: %[[ARG1_ALLOCA:.*]] = alloca ptr, align 8
+// CHECK: store ptr %[[ARG1]], ptr %[[ARG1_ALLOCA]], align 8
+// CHECK: %[[LOAD_ARG1_ALLOCA:.*]] = load ptr, ptr %[[ARG1_ALLOCA]], align 8
+// CHECK: store i32 20, ptr %[[LOAD_ARG1_ALLOCA]], align 4
+// CHECK: %[[GEP_ARG1_ALLOCA:.*]] = getelementptr inbounds [10 x i32], ptr %[[LOAD_ARG1_ALLOCA]], i32 0, i64 4
+// CHECK: store i32 10, ptr %[[GEP_ARG1_ALLOCA]], align 4
+
diff --git a/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90 b/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90
new file mode 100644
index 000000000000000..91dc30cf9604c4b
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90
@@ -0,0 +1,27 @@
+! Basic offloading test with a target region
+! that checks constant indexing on device 
+! correctly works (regression test for prior
+! bug).
+! REQUIRES: flang, amdgcn-amd-amdhsa
+! UNSUPPORTED: nvptx64-nvidia-cuda
+! UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+! UNSUPPORTED: aarch64-unknown-linux-gnu
+! UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+! UNSUPPORTED: x86_64-pc-linux-gnu
+! UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+program main
+    INTEGER :: sp(10) = (/0,0,0,0,0,0,0,0,0,0/)
+
+  !$omp target map(tofrom:sp)
+     sp(1) = 20
+     sp(5) = 10
+  !$omp end target
+
+   ! print *, sp(1)
+   ! print *, sp(5)
+end program
+
+! CHECK: 20
+! CHECK: 10

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[skatrak]

Thanks Andrew! LGTM, just a couple nits. I'd wait for a second reviewer before merging.

[agozillon]

Thank you for the review Sergio! I've updated the PR with your feedback. I'll await one more reviewer as you suggest.

[agozillon]

@jsjodin your comments are hopefully addressed now!

[jsjodin]

LGTM