[Mlir-commits] [mlir] [llvm] [openmp] [OpenMP][OMPIRBuilder] Handle replace uses of ConstantExpr's inside of Target regions (PR #71891)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Nov 9 20:30:34 PST 2023
https://github.com/agozillon updated https://github.com/llvm/llvm-project/pull/71891
>From 3cdce16f440fae7dd3723783b5b91a54139f10cb Mon Sep 17 00:00:00 2001
From: Andrew Gozillon <Andrew.Gozillon at amd.com>
Date: Thu, 9 Nov 2023 22:01:09 -0600
Subject: [PATCH] [OpenMP][OMPIRBuilder] Handle replace uses of ConstantExpr's
inside of Target regions
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 neccessary.
This patch handles these by replacing 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 (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).
---
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 34 +++++++++++++--
...arget-constant-indexing-device-region.mlir | 41 +++++++++++++++++++
.../offloading/fortran/constant-arr-index.f90 | 27 ++++++++++++
3 files changed, 99 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
create mode 100644 openmp/libomptarget/test/offloading/fortran/constant-arr-index.f90
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index eb6c120b2a6dcbd..e28e2c66528eb05 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4817,11 +4817,39 @@ static Function *createOutlinedFunction(
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
More information about the Mlir-commits
mailing list