[Mlir-commits] [openmp] [mlir] [llvm] [OpenMP][OMPIRBuilder] Handle replace uses of ConstantExpr's inside of Target regions (PR #71891)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Nov 10 07:53:57 PST 2023
https://github.com/agozillon updated https://github.com/llvm/llvm-project/pull/71891
>From 0f0189256a4f8ed488f974dd0eb146be74725890 Mon Sep 17 00:00:00 2001
From: Andrew Gozillon <Andrew.Gozillon at amd.com>
Date: Fri, 10 Nov 2023 09:28:42 -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 | 35 ++++++++++++++--
...arget-constant-indexing-device-region.mlir | 42 +++++++++++++++++++
.../offloading/fortran/constant-arr-index.f90 | 27 ++++++++++++
3 files changed, 101 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..706a08940529ee7 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4817,11 +4817,40 @@ 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
+ // 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 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 = [](ConstantExpr *ConstExpr,
+ Function *Func) {
+ 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))
+ if (auto *ConstExpr = dyn_cast<ConstantExpr>(Const))
+ ReplaceConstantUsedInFunction(ConstExpr, 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..5c0ac79271a6938
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-constant-indexing-device-region.mlir
@@ -0,0 +1,42 @@
+// 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 {{.*}} void @__omp_offloading_{{.*}}_{{.*}}__QQmain_{{.*}}(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