diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index 81dfef4265131..ad6b0188390d8 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -4752,6 +4752,22 @@ FunctionCallee OpenMPIRBuilder::createDispatchFiniFunction(unsigned IVSize, return getOrCreateRuntimeFunction(M, Name); } +static void replaceConstatExprUsesInFuncWithInstr(ConstantExpr *ConstExpr, + Function *Func) { + for (User *User : make_early_inc_range(ConstExpr->users())) + if (auto *Instr = dyn_cast(User)) + if (Instr->getFunction() == Func) + Instr->replaceUsesOfWith(ConstExpr, ConstExpr->getAsInstruction(Instr)); +} + +static void replaceConstantValueUsesInFuncWithInstr(llvm::Value *Input, + Function *Func) { + for (User *User : make_early_inc_range(Input->users())) + if (auto *Const = dyn_cast(User)) + if (auto *ConstExpr = dyn_cast(Const)) + replaceConstatExprUsesInFuncWithInstr(ConstExpr, Func); +} + static Function *createOutlinedFunction( OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder, StringRef FuncName, SmallVectorImpl &Inputs, @@ -4822,9 +4838,23 @@ static Function *createOutlinedFunction( Builder.restoreIP( ArgAccessorFuncCB(Arg, Input, InputCopy, AllocaIP, Builder.saveIP())); + // 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). + replaceConstantValueUsesInFuncWithInstr(Input, Func); + // Collect all the instructions for (User *User : make_early_inc_range(Input->users())) - if (auto Instr = dyn_cast(User)) + if (auto *Instr = dyn_cast(User)) if (Instr->getFunction() == Func) Instr->replaceUsesOfWith(Input, InputCopy); } 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 0000000000000..5c0ac79271a69 --- /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 0000000000000..91dc30cf9604c --- /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