[Mlir-commits] [mlir] [OpenMP][OpenMPIRBuilder][NFC] Move copyInput to a passed in lambda function (PR #68124)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 3 09:18:45 PDT 2023
https://github.com/agozillon created https://github.com/llvm/llvm-project/pull/68124
This patch moves the existing copyInput function
into a lambda argument that can be defined
by a caller to the function.
This allows more flexibility in how the function
is defined, allowing Clang and MLIR to utilise
their own respective functions and types inside
of the lamba without affecting the OMPIRBuilder
itself.
The idea is to eventually replace/build on
the existing copyInput function that's used
and moved into OpenMPToLLVMIRTranslation.cpp
to a slightly more complex implementation
that uses MLIRs map information (primarily
ByRef and ByCapture information at the
moment).
For now this should be an NFC as far as
lowering behavior is concerned.
>From 4cf87ac5f73aa96d51f297c6483ee4cfad571276 Mon Sep 17 00:00:00 2001
From: Andrew Gozillon <Andrew.Gozillon at amd.com>
Date: Tue, 3 Oct 2023 10:48:03 -0500
Subject: [PATCH] [OpenMP][OpenMPIRBuilder][NFC] Move copyInput to a passed in
lambda function
This patch moves the existing copyInput function
into a lambda argument that can be defined
by a caller to the function.
This allows more flexibility in how the function
is defined, allowing Clang and Fortran to utilise
their own respective functions and types inside
of the lamba without affecting the OMPIRBuilder
itself.
The idea is to eventually replace/build on
the existing copyInput function that's used
and moved into OpenMPToLLVMIRTranslation.cpp
to a slightly more complex implementation
that uses Flang's map information (primarily
ByRef and ByCapture information at the
moment).
For now this should be an NFC as far as
lowering behavior is concerned.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 8 ++-
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 57 +++++++------------
.../Frontend/OpenMPIRBuilderTest.cpp | 48 ++++++++++++++--
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 40 ++++++++++++-
4 files changed, 109 insertions(+), 44 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 75da461cfd8d95e..e3f1cddb72fa03d 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -2166,6 +2166,9 @@ class OpenMPIRBuilder {
using TargetBodyGenCallbackTy = function_ref<InsertPointTy(
InsertPointTy AllocaIP, InsertPointTy CodeGenIP)>;
+ using TargetGenArgAccessorsCallbackTy = function_ref<Value *(
+ Argument &Arg, Value *Input, IRBuilderBase &Builder)>;
+
/// Generator for '#omp target'
///
/// \param Loc where the target data construct was encountered.
@@ -2177,6 +2180,8 @@ class OpenMPIRBuilder {
/// \param Inputs The input values to the region that will be passed.
/// as arguments to the outlined function.
/// \param BodyGenCB Callback that will generate the region code.
+ /// \param ArgAccessorFuncCB Callback that will generate accessors
+ /// instructions for passed in target arguments where neccessary
InsertPointTy createTarget(const LocationDescription &Loc,
OpenMPIRBuilder::InsertPointTy AllocaIP,
OpenMPIRBuilder::InsertPointTy CodeGenIP,
@@ -2184,7 +2189,8 @@ class OpenMPIRBuilder {
int32_t NumThreads,
SmallVectorImpl<Value *> &Inputs,
GenMapInfoCallbackTy GenMapInfoCB,
- TargetBodyGenCallbackTy BodyGenCB);
+ TargetBodyGenCallbackTy BodyGenCB,
+ TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB);
/// Returns __kmpc_for_static_init_* runtime function for the specified
/// size \a IVSize and sign \a IVSigned. Will create a distribute call
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 72e1af55fe63f60..bcb5d07bbf7edc8 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4539,25 +4539,11 @@ FunctionCallee OpenMPIRBuilder::createDispatchFiniFunction(unsigned IVSize,
return getOrCreateRuntimeFunction(M, Name);
}
-// Copy input from pointer or i64 to the expected argument type.
-static Value *copyInput(IRBuilderBase &Builder, unsigned AddrSpace,
- Value *Input, Argument &Arg) {
- auto Addr = Builder.CreateAlloca(Arg.getType()->isPointerTy()
- ? Arg.getType()
- : Type::getInt64Ty(Builder.getContext()),
- AddrSpace);
- auto AddrAscast =
- Builder.CreatePointerBitCastOrAddrSpaceCast(Addr, Input->getType());
- Builder.CreateStore(&Arg, AddrAscast);
- auto Copy = Builder.CreateLoad(Arg.getType(), AddrAscast);
-
- return Copy;
-}
-
-static Function *
-createOutlinedFunction(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
- StringRef FuncName, SmallVectorImpl<Value *> &Inputs,
- OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc) {
+static Function *createOutlinedFunction(
+ OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder, StringRef FuncName,
+ SmallVectorImpl<Value *> &Inputs,
+ OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc,
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB) {
SmallVector<Type *> ParameterTypes;
if (OMPBuilder.Config.isTargetDevice()) {
// All parameters to target devices are passed as pointers
@@ -4603,12 +4589,7 @@ createOutlinedFunction(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
Value *Input = std::get<0>(InArg);
Argument &Arg = std::get<1>(InArg);
- Value *InputCopy =
- OMPBuilder.Config.isTargetDevice()
- ? copyInput(Builder,
- OMPBuilder.M.getDataLayout().getAllocaAddrSpace(),
- Input, Arg)
- : &Arg;
+ Value *InputCopy = ArgAccessorFuncCB(Arg, Input, Builder);
// Collect all the instructions
for (User *User : make_early_inc_range(Input->users()))
@@ -4623,18 +4604,19 @@ createOutlinedFunction(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
return Func;
}
-static void
-emitTargetOutlinedFunction(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
- TargetRegionEntryInfo &EntryInfo,
- Function *&OutlinedFn, Constant *&OutlinedFnID,
- int32_t NumTeams, int32_t NumThreads,
- SmallVectorImpl<Value *> &Inputs,
- OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc) {
+static void emitTargetOutlinedFunction(
+ OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
+ TargetRegionEntryInfo &EntryInfo, Function *&OutlinedFn,
+ Constant *&OutlinedFnID, int32_t NumTeams, int32_t NumThreads,
+ SmallVectorImpl<Value *> &Inputs,
+ OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc,
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB) {
OpenMPIRBuilder::FunctionGenCallback &&GenerateOutlinedFunction =
- [&OMPBuilder, &Builder, &Inputs, &CBFunc](StringRef EntryFnName) {
+ [&OMPBuilder, &Builder, &Inputs, &CBFunc,
+ &ArgAccessorFuncCB](StringRef EntryFnName) {
return createOutlinedFunction(OMPBuilder, Builder, EntryFnName, Inputs,
- CBFunc);
+ CBFunc, ArgAccessorFuncCB);
};
OMPBuilder.emitTargetRegionFunction(EntryInfo, GenerateOutlinedFunction,
@@ -4698,7 +4680,9 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
const LocationDescription &Loc, InsertPointTy AllocaIP,
InsertPointTy CodeGenIP, TargetRegionEntryInfo &EntryInfo, int32_t NumTeams,
int32_t NumThreads, SmallVectorImpl<Value *> &Args,
- GenMapInfoCallbackTy GenMapInfoCB, TargetBodyGenCallbackTy CBFunc) {
+ GenMapInfoCallbackTy GenMapInfoCB,
+ OpenMPIRBuilder::TargetBodyGenCallbackTy CBFunc,
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB) {
if (!updateToLocation(Loc))
return InsertPointTy();
@@ -4707,7 +4691,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
Function *OutlinedFn;
Constant *OutlinedFnID;
emitTargetOutlinedFunction(*this, Builder, EntryInfo, OutlinedFn,
- OutlinedFnID, NumTeams, NumThreads, Args, CBFunc);
+ OutlinedFnID, NumTeams, NumThreads, Args, CBFunc,
+ ArgAccessorFuncCB);
if (!Config.isTargetDevice())
emitTargetCall(*this, Builder, AllocaIP, OutlinedFn, OutlinedFnID, NumTeams,
NumThreads, Args, GenMapInfoCB);
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index fd524f6067ee0ea..81733f1a2790287 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -5236,6 +5236,24 @@ TEST_F(OpenMPIRBuilderTest, TargetRegion) {
Inputs.push_back(BPtr);
Inputs.push_back(CPtr);
+ auto SimpleArgAccessorCB = [&OMPBuilder](llvm::Argument &Arg,
+ llvm::Value *Input,
+ IRBuilderBase &Builder) {
+ if (!OMPBuilder.Config.isTargetDevice())
+ return cast<llvm::Value>(&Arg);
+
+ llvm::Value *Addr = Builder.CreateAlloca(
+ Arg.getType()->isPointerTy() ? Arg.getType()
+ : Type::getInt64Ty(Builder.getContext()),
+ OMPBuilder.M.getDataLayout().getAllocaAddrSpace());
+ llvm::Value *AddrAscast =
+ Builder.CreatePointerBitCastOrAddrSpaceCast(Addr, Input->getType());
+ Builder.CreateStore(&Arg, AddrAscast);
+ llvm::Value *Copy = Builder.CreateLoad(Arg.getType(), AddrAscast);
+
+ return Copy;
+ };
+
llvm::OpenMPIRBuilder::MapInfosTy CombinedInfos;
auto GenMapInfoCB = [&](llvm::OpenMPIRBuilder::InsertPointTy codeGenIP)
-> llvm::OpenMPIRBuilder::MapInfosTy & {
@@ -5245,9 +5263,9 @@ TEST_F(OpenMPIRBuilderTest, TargetRegion) {
TargetRegionEntryInfo EntryInfo("func", 42, 4711, 17);
OpenMPIRBuilder::LocationDescription OmpLoc({Builder.saveIP(), DL});
- Builder.restoreIP(OMPBuilder.createTarget(OmpLoc, Builder.saveIP(),
- Builder.saveIP(), EntryInfo, -1, 0,
- Inputs, GenMapInfoCB, BodyGenCB));
+ Builder.restoreIP(OMPBuilder.createTarget(
+ OmpLoc, Builder.saveIP(), Builder.saveIP(), EntryInfo, -1, 0, Inputs,
+ GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB));
OMPBuilder.finalize();
Builder.CreateRetVoid();
@@ -5301,6 +5319,23 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
Constant::getNullValue(PointerType::get(Ctx, 0)),
Constant::getNullValue(PointerType::get(Ctx, 0))};
+ auto SimpleArgAccessorCB = [&](llvm::Argument &Arg, llvm::Value *Input,
+ IRBuilderBase &Builder) {
+ if (!OMPBuilder.Config.isTargetDevice())
+ return cast<llvm::Value>(&Arg);
+
+ llvm::Value *Addr = Builder.CreateAlloca(
+ Arg.getType()->isPointerTy() ? Arg.getType()
+ : Type::getInt64Ty(Builder.getContext()),
+ OMPBuilder.M.getDataLayout().getAllocaAddrSpace());
+ llvm::Value *AddrAscast =
+ Builder.CreatePointerBitCastOrAddrSpaceCast(Addr, Input->getType());
+ Builder.CreateStore(&Arg, AddrAscast);
+ llvm::Value *Copy = Builder.CreateLoad(Arg.getType(), AddrAscast);
+
+ return Copy;
+ };
+
llvm::OpenMPIRBuilder::MapInfosTy CombinedInfos;
auto GenMapInfoCB = [&](llvm::OpenMPIRBuilder::InsertPointTy codeGenIP)
-> llvm::OpenMPIRBuilder::MapInfosTy & {
@@ -5322,9 +5357,10 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
TargetRegionEntryInfo EntryInfo("parent", /*DeviceID=*/1, /*FileID=*/2,
/*Line=*/3, /*Count=*/0);
- Builder.restoreIP(OMPBuilder.createTarget(
- Loc, EntryIP, EntryIP, EntryInfo, /*NumTeams=*/-1,
- /*NumThreads=*/0, CapturedArgs, GenMapInfoCB, BodyGenCB));
+ Builder.restoreIP(
+ OMPBuilder.createTarget(Loc, EntryIP, EntryIP, EntryInfo, /*NumTeams=*/-1,
+ /*NumThreads=*/0, CapturedArgs, GenMapInfoCB,
+ BodyGenCB, SimpleArgAccessorCB));
Builder.CreateRetVoid();
OMPBuilder.finalize();
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 14bcbc3018f72bd..8326a20a15c4629 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -1993,6 +1993,27 @@ handleDeclareTargetMapVar(llvm::ArrayRef<Value> mapOperands,
}
}
+static llvm::Value *
+createDeviceArgumentAccessor(llvm::Argument &arg, llvm::Value *input,
+ llvm::IRBuilderBase &builder,
+ llvm::OpenMPIRBuilder &ompBuilder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ if (!ompBuilder.Config.isTargetDevice())
+ return cast<llvm::Value>(&arg);
+
+ llvm::Value *addr =
+ builder.CreateAlloca(arg.getType()->isPointerTy()
+ ? arg.getType()
+ : llvm::Type::getInt64Ty(builder.getContext()),
+ ompBuilder.M.getDataLayout().getAllocaAddrSpace());
+ llvm::Value *addrAscast =
+ builder.CreatePointerBitCastOrAddrSpaceCast(addr, input->getType());
+ builder.CreateStore(&arg, addrAscast);
+ llvm::Value *copy = builder.CreateLoad(arg.getType(), addrAscast);
+
+ return copy;
+}
+
static LogicalResult
convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation) {
@@ -2084,9 +2105,26 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
return combinedInfos;
};
+ auto argAccessorCB = [&moduleTranslation](llvm::Argument &arg,
+ llvm::Value *input,
+ llvm::IRBuilderBase &builder) {
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+ // We just return the unaltered argument for the host function
+ // for now, some alterations may be required in the future to
+ // keep host fallback functions working identically to the device
+ // version (e.g. pass ByCopy values should be treated as such on
+ // host and device, currently not always the case)
+ if (!ompBuilder->Config.isTargetDevice())
+ return cast<llvm::Value>(&arg);
+
+ return createDeviceArgumentAccessor(arg, input, builder, *ompBuilder,
+ moduleTranslation);
+ };
+
builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTarget(
ompLoc, allocaIP, builder.saveIP(), entryInfo, defaultValTeams,
- defaultValThreads, inputs, genMapInfoCB, bodyCB));
+ defaultValThreads, inputs, genMapInfoCB, bodyCB, argAccessorCB));
// Remap access operations to declare target reference pointers for the
// device, essentially generating extra loadop's as necessary
More information about the Mlir-commits
mailing list