[llvm] [llvm][OpenMPIRBuilder] Allow to not register offload entries in the entry manager (PR #78222)
Fabian Mora via llvm-commits
llvm-commits at lists.llvm.org
Mon Jan 15 18:24:14 PST 2024
https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/78222
This patch adds an optional field in the create target method to store the offload entry in a custom location and not register the entry in the entry manager.
This change is required to enable JIT compilation in MLIR for OpenMP target offload ops, as arrays of entries are handled differently for standalone MLIR compilation.
>From f2df0a65b01310fc7c8c60b92f035ad506ebe341 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Mon, 15 Jan 2024 15:01:42 +0000
Subject: [PATCH] [llvm][OpenMPIRBuilder] Allow to not register offload entries
in the entry manager This patch adds an optional field in the create target
method to store the offload entry in a custom location and not register the
entry in the entry manager. This change is required to enable JIT compilation
in MLIR for OpenMP target offload ops, as arrays of entries are handled
differently for standalone MLIR compilation.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 53 ++++++++++++-------
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 36 ++++++++-----
2 files changed, 57 insertions(+), 32 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 669104307fa0e2..a584552d15954d 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -298,9 +298,10 @@ class OffloadEntriesInfoManager {
void initializeTargetRegionEntryInfo(const TargetRegionEntryInfo &EntryInfo,
unsigned Order);
/// Register target region entry.
- void registerTargetRegionEntryInfo(TargetRegionEntryInfo EntryInfo,
- Constant *Addr, Constant *ID,
- OMPTargetRegionEntryKind Flags);
+ void registerTargetRegionEntryInfo(
+ TargetRegionEntryInfo EntryInfo, Constant *Addr, Constant *ID,
+ OMPTargetRegionEntryKind Flags,
+ OffloadEntryInfoTargetRegion *EntryInfoStorage = nullptr);
/// Return true if a target region entry with the provided information
/// exists.
bool hasTargetRegionEntryInfo(TargetRegionEntryInfo EntryInfo,
@@ -2139,10 +2140,14 @@ class OpenMPIRBuilder {
/// \param GenerateFunctionCallback The callback function to generate the code
/// \param OutlinedFunction Pointer to the outlined function
/// \param EntryFnIDName Name of the ID o be created
- void emitTargetRegionFunction(TargetRegionEntryInfo &EntryInfo,
- FunctionGenCallback &GenerateFunctionCallback,
- bool IsOffloadEntry, Function *&OutlinedFn,
- Constant *&OutlinedFnID);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry.
+ void emitTargetRegionFunction(
+ TargetRegionEntryInfo &EntryInfo,
+ FunctionGenCallback &GenerateFunctionCallback, bool IsOffloadEntry,
+ Function *&OutlinedFn, Constant *&OutlinedFnID,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// Registers the given function and sets up the attribtues of the function
/// Returns the FunctionID.
@@ -2152,10 +2157,13 @@ class OpenMPIRBuilder {
/// \param OutlinedFunction Pointer to the outlined function
/// \param EntryFnName Name of the outlined function
/// \param EntryFnIDName Name of the ID o be created
- Constant *registerTargetRegionFunction(TargetRegionEntryInfo &EntryInfo,
- Function *OutlinedFunction,
- StringRef EntryFnName,
- StringRef EntryFnIDName);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry.
+ Constant *registerTargetRegionFunction(
+ TargetRegionEntryInfo &EntryInfo, Function *OutlinedFunction,
+ StringRef EntryFnName, StringRef EntryFnIDName,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// Type of BodyGen to use for region codegen
///
@@ -2225,15 +2233,20 @@ class OpenMPIRBuilder {
/// \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,
- TargetRegionEntryInfo &EntryInfo, int32_t NumTeams,
- int32_t NumThreads,
- SmallVectorImpl<Value *> &Inputs,
- GenMapInfoCallbackTy GenMapInfoCB,
- TargetBodyGenCallbackTy BodyGenCB,
- TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry. If the pointer is null the entry gets registered in the
+ /// `OffloadEntriesInfoManager`, otherwise it gets stored in
+ /// `EntryInfoStorage` and doesn't get registered in
+ /// `OffloadEntriesInfoManager`.
+ InsertPointTy createTarget(
+ const LocationDescription &Loc, OpenMPIRBuilder::InsertPointTy AllocaIP,
+ OpenMPIRBuilder::InsertPointTy CodeGenIP,
+ TargetRegionEntryInfo &EntryInfo, int32_t NumTeams, int32_t NumThreads,
+ SmallVectorImpl<Value *> &Inputs, GenMapInfoCallbackTy GenMapInfoCB,
+ TargetBodyGenCallbackTy BodyGenCB,
+ TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// 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 f6cf358119fb71..c5498d6c67b455 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4741,7 +4741,8 @@ Constant *OpenMPIRBuilder::createTargetRegionEntryAddr(Function *OutlinedFn,
void OpenMPIRBuilder::emitTargetRegionFunction(
TargetRegionEntryInfo &EntryInfo,
FunctionGenCallback &GenerateFunctionCallback, bool IsOffloadEntry,
- Function *&OutlinedFn, Constant *&OutlinedFnID) {
+ Function *&OutlinedFn, Constant *&OutlinedFnID,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
SmallString<64> EntryFnName;
OffloadInfoManager.getTargetRegionEntryFnName(EntryFnName, EntryInfo);
@@ -4761,20 +4762,22 @@ void OpenMPIRBuilder::emitTargetRegionFunction(
? std::string(EntryFnName)
: createPlatformSpecificName({EntryFnName, "region_id"});
- OutlinedFnID = registerTargetRegionFunction(EntryInfo, OutlinedFn,
- EntryFnName, EntryFnIDName);
+ OutlinedFnID = registerTargetRegionFunction(
+ EntryInfo, OutlinedFn, EntryFnName, EntryFnIDName, EntryInfoStorage);
}
Constant *OpenMPIRBuilder::registerTargetRegionFunction(
TargetRegionEntryInfo &EntryInfo, Function *OutlinedFn,
- StringRef EntryFnName, StringRef EntryFnIDName) {
+ StringRef EntryFnName, StringRef EntryFnIDName,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
if (OutlinedFn)
setOutlinedTargetRegionFunctionAttributes(OutlinedFn);
auto OutlinedFnID = createOutlinedFunctionID(OutlinedFn, EntryFnIDName);
auto EntryAddr = createTargetRegionEntryAddr(OutlinedFn, EntryFnName);
OffloadInfoManager.registerTargetRegionEntryInfo(
EntryInfo, EntryAddr, OutlinedFnID,
- OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
+ OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion,
+ EntryInfoStorage);
return OutlinedFnID;
}
@@ -5094,7 +5097,8 @@ static void emitTargetOutlinedFunction(
TargetRegionEntryInfo &EntryInfo, Function *&OutlinedFn,
Constant *&OutlinedFnID, SmallVectorImpl<Value *> &Inputs,
OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc,
- OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB) {
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
OpenMPIRBuilder::FunctionGenCallback &&GenerateOutlinedFunction =
[&OMPBuilder, &Builder, &Inputs, &CBFunc,
@@ -5104,7 +5108,8 @@ static void emitTargetOutlinedFunction(
};
OMPBuilder.emitTargetRegionFunction(EntryInfo, GenerateOutlinedFunction, true,
- OutlinedFn, OutlinedFnID);
+ OutlinedFn, OutlinedFnID,
+ EntryInfoStorage);
}
static void emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
@@ -5165,7 +5170,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
int32_t NumThreads, SmallVectorImpl<Value *> &Args,
GenMapInfoCallbackTy GenMapInfoCB,
OpenMPIRBuilder::TargetBodyGenCallbackTy CBFunc,
- OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB) {
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
if (!updateToLocation(Loc))
return InsertPointTy();
@@ -5174,7 +5180,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
Function *OutlinedFn;
Constant *OutlinedFnID;
emitTargetOutlinedFunction(*this, Builder, EntryInfo, OutlinedFn,
- OutlinedFnID, Args, CBFunc, ArgAccessorFuncCB);
+ OutlinedFnID, Args, CBFunc, ArgAccessorFuncCB,
+ EntryInfoStorage);
if (!Config.isTargetDevice())
emitTargetCall(*this, Builder, AllocaIP, OutlinedFn, OutlinedFnID, NumTeams,
NumThreads, Args, GenMapInfoCB);
@@ -6929,7 +6936,8 @@ void OffloadEntriesInfoManager::initializeTargetRegionEntryInfo(
void OffloadEntriesInfoManager::registerTargetRegionEntryInfo(
TargetRegionEntryInfo EntryInfo, Constant *Addr, Constant *ID,
- OMPTargetRegionEntryKind Flags) {
+ OMPTargetRegionEntryKind Flags,
+ OffloadEntryInfoTargetRegion *EntryInfoStorage) {
assert(EntryInfo.Count == 0 && "expected default EntryInfo");
// Update the EntryInfo with the next available count for this location.
@@ -6953,8 +6961,12 @@ void OffloadEntriesInfoManager::registerTargetRegionEntryInfo(
assert(!hasTargetRegionEntryInfo(EntryInfo) &&
"Target region entry already registered!");
OffloadEntryInfoTargetRegion Entry(OffloadingEntriesNum, Addr, ID, Flags);
- OffloadEntriesTargetRegion[EntryInfo] = Entry;
- ++OffloadingEntriesNum;
+ if (EntryInfoStorage) {
+ *EntryInfoStorage = Entry;
+ } else {
+ OffloadEntriesTargetRegion[EntryInfo] = Entry;
+ ++OffloadingEntriesNum;
+ }
}
incrementTargetRegionEntryInfoCount(EntryInfo);
}
More information about the llvm-commits
mailing list