diff --git a/llvm/include/llvm/IR/ReplaceConstant.h b/llvm/include/llvm/IR/ReplaceConstant.h index 56027d7fba6f8..e1af352f295ae 100644 --- a/llvm/include/llvm/IR/ReplaceConstant.h +++ b/llvm/include/llvm/IR/ReplaceConstant.h @@ -18,10 +18,21 @@ namespace llvm { template class ArrayRef; class Constant; +class Function; /// Replace constant expressions users of the given constants with /// instructions. Return whether anything was changed. -bool convertUsersOfConstantsToInstructions(ArrayRef Consts); +/// +/// Passing RestrictToFunc will restrict the constant replacement +/// to the passed in functions scope, as opposed to the replacements +/// occurring at module scope. +/// +/// RemoveDeadConstants by default will remove all dead constants as +/// the final step of the function after replacement, when passed +/// false it will skip this final step. +bool convertUsersOfConstantsToInstructions(ArrayRef Consts, + Function *RestrictToFunc = nullptr, + bool RemoveDeadConstants = true); } // end namespace llvm diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index 92213e19c9d9d..ccec66fcb7bac 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -40,6 +40,7 @@ #include "llvm/IR/MDBuilder.h" #include "llvm/IR/Metadata.h" #include "llvm/IR/PassManager.h" +#include "llvm/IR/ReplaceConstant.h" #include "llvm/IR/Value.h" #include "llvm/MC/TargetRegistry.h" #include "llvm/Support/CommandLine.h" @@ -5092,27 +5093,6 @@ 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) { - Instruction *ConstInst = ConstExpr->getAsInstruction(); - ConstInst->insertBefore(*Instr->getParent(), Instr->getIterator()); - Instr->replaceUsesOfWith(ConstExpr, ConstInst); - } - } - } -} - -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, @@ -5191,17 +5171,23 @@ static Function *createOutlinedFunction( // 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 + // 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); + // 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). + // NOTE: We cannot remove dead constants that have been rewritten to + // instructions at this stage, we run the risk of breaking later lowering + // by doing so as we could still be in the process of lowering the module + // from MLIR to LLVM-IR and the MLIR lowering may still require the original + // constants we have created rewritten versions of. + if (auto *Const = dyn_cast(Input)) + convertUsersOfConstantsToInstructions(Const, Func, false); // Collect all the instructions for (User *User : make_early_inc_range(Input->users())) diff --git a/llvm/lib/IR/ReplaceConstant.cpp b/llvm/lib/IR/ReplaceConstant.cpp index 9b07bd8040492..67b6fe6fda3b9 100644 --- a/llvm/lib/IR/ReplaceConstant.cpp +++ b/llvm/lib/IR/ReplaceConstant.cpp @@ -49,7 +49,9 @@ static SmallVector expandUser(BasicBlock::iterator InsertPt, return NewInsts; } -bool convertUsersOfConstantsToInstructions(ArrayRef Consts) { +bool convertUsersOfConstantsToInstructions(ArrayRef Consts, + Function *RestrictToFunc, + bool RemoveDeadConstants) { // Find all expandable direct users of Consts. SmallVector Stack; for (Constant *C : Consts) @@ -74,7 +76,8 @@ bool convertUsersOfConstantsToInstructions(ArrayRef Consts) { for (Constant *C : ExpandableUsers) for (User *U : C->users()) if (auto *I = dyn_cast(U)) - InstructionWorklist.insert(I); + if (!RestrictToFunc || I->getFunction() == RestrictToFunc) + InstructionWorklist.insert(I); // Replace those expandable operands with instructions bool Changed = false; @@ -102,8 +105,9 @@ bool convertUsersOfConstantsToInstructions(ArrayRef Consts) { } } - for (Constant *C : Consts) - C->removeDeadConstantUsers(); + if (RemoveDeadConstants) + for (Constant *C : Consts) + C->removeDeadConstantUsers(); return Changed; } diff --git a/offload/test/offloading/fortran/dtype-array-constant-index-map.f90 b/offload/test/offloading/fortran/dtype-array-constant-index-map.f90 new file mode 100644 index 0000000000000..580eef42e51a8 --- /dev/null +++ b/offload/test/offloading/fortran/dtype-array-constant-index-map.f90 @@ -0,0 +1,51 @@ +! Offloading test which maps a specific element of a +! derived type to the device and then accesses the +! element alongside an individual element of an array +! that the derived type contains. In particular, this +! test helps to check that we can replace the constants +! within the kernel with instructions and then replace +! these instructions with the kernel parameters. +! REQUIRES: flang +! 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 +module test_0 + type dtype + integer elements(20) + integer value + end type dtype + + type (dtype) array_dtype(5) + contains + + subroutine assign() + implicit none +!$omp target map(tofrom: array_dtype(5)) + array_dtype(5)%elements(5) = 500 +!$omp end target + end subroutine + + subroutine add() + implicit none + +!$omp target map(tofrom: array_dtype(5)) + array_dtype(5)%elements(5) = array_dtype(5)%elements(5) + 500 +!$omp end target + end subroutine +end module test_0 + +program main + use test_0 + + call assign() + call add() + + print *, array_dtype(5)%elements(5) +end program + +! CHECK: 1000 \ No newline at end of file