Skip to content

Commit acd617a

Browse files
author
ergawy
committed
[flang][OpenMP] Support target enter|update|exit .. nowait
Extends `nowait` support for other device directives. This PR refactors the task generation utils used for the `target` directive so that they are general enough to be reused for other device directives as well.
1 parent b90ea5c commit acd617a

5 files changed

Lines changed: 244 additions & 113 deletions

File tree

‎llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2264,6 +2264,9 @@ class OpenMPIRBuilder {
22642264

22652265
bool EmitDebug = false;
22662266

2267+
/// Whether the `target ... data` directive has a `nowait` clause.
2268+
bool HasNoWait = false;
2269+
22672270
explicit TargetDataInfo() {}
22682271
explicit TargetDataInfo(bool RequiresDevicePointerInfo,
22692272
bool SeparateBeginEndCalls)
@@ -2342,26 +2345,34 @@ class OpenMPIRBuilder {
23422345
/// Generate a target region entry call and host fallback call.
23432346
///
23442347
/// \param Loc The location at which the request originated and is fulfilled.
2345-
/// \param OutlinedFn The outlined kernel function.
23462348
/// \param OutlinedFnID The ooulined function ID.
23472349
/// \param EmitTargetCallFallbackCB Call back function to generate host
23482350
/// fallback code.
23492351
/// \param Args Data structure holding information about the kernel arguments.
23502352
/// \param DeviceID Identifier for the device via the 'device' clause.
23512353
/// \param RTLoc Source location identifier
23522354
/// \param AllocaIP The insertion point to be used for alloca instructions.
2353-
InsertPointTy emitKernelLaunch(
2354-
const LocationDescription &Loc, Function *OutlinedFn, Value *OutlinedFnID,
2355-
EmitFallbackCallbackTy EmitTargetCallFallbackCB, TargetKernelArgs &Args,
2356-
Value *DeviceID, Value *RTLoc, InsertPointTy AllocaIP);
2355+
InsertPointTy
2356+
emitKernelLaunch(const LocationDescription &Loc, Value *OutlinedFnID,
2357+
EmitFallbackCallbackTy EmitTargetCallFallbackCB,
2358+
TargetKernelArgs &Args, Value *DeviceID, Value *RTLoc,
2359+
InsertPointTy AllocaIP);
2360+
2361+
/// Callback type for generating the bodies of device directives that require
2362+
/// outer tasks (e.g. in case of having `nowait` or `depend` clauses).
2363+
///
2364+
/// \param DeviceID The ID of the device on which the target region will
2365+
/// execute.
2366+
/// \param RTLoc Source location identifier
2367+
/// \Param TargetTaskAllocaIP Insertion point for the alloca block of the
2368+
/// generated task.
2369+
using TaskBodyCallbackTy =
2370+
function_ref<void(Value *DeviceID, Value *RTLoc,
2371+
IRBuilderBase::InsertPoint TargetTaskAllocaIP)>;
23572372

23582373
/// Generate a target-task for the target construct
23592374
///
2360-
/// \param OutlinedFn The outlined device/target kernel function.
2361-
/// \param OutlinedFnID The ooulined function ID.
2362-
/// \param EmitTargetCallFallbackCB Call back function to generate host
2363-
/// fallback code.
2364-
/// \param Args Data structure holding information about the kernel arguments.
2375+
/// \param TaskBodyCB Callback to generate the actual body of the target task.
23652376
/// \param DeviceID Identifier for the device via the 'device' clause.
23662377
/// \param RTLoc Source location identifier
23672378
/// \param AllocaIP The insertion point to be used for alloca instructions.
@@ -2370,10 +2381,10 @@ class OpenMPIRBuilder {
23702381
/// \param HasNoWait True if the target construct had 'nowait' on it, false
23712382
/// otherwise
23722383
InsertPointTy emitTargetTask(
2373-
Function *OutlinedFn, Value *OutlinedFnID,
2374-
EmitFallbackCallbackTy EmitTargetCallFallbackCB, TargetKernelArgs &Args,
2375-
Value *DeviceID, Value *RTLoc, InsertPointTy AllocaIP,
2376-
SmallVector<OpenMPIRBuilder::DependData> &Dependencies, bool HasNoWait);
2384+
TaskBodyCallbackTy TaskBodyCB, Value *DeviceID, Value *RTLoc,
2385+
OpenMPIRBuilder::InsertPointTy AllocaIP,
2386+
const SmallVector<llvm::OpenMPIRBuilder::DependData> &Dependencies,
2387+
bool HasNoWait);
23772388

23782389
/// Emit the arguments to be passed to the runtime library based on the
23792390
/// arrays of base pointers, pointers, sizes, map types, and mappers. If

‎llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp‎

Lines changed: 87 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1080,8 +1080,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitTargetKernel(
10801080
}
10811081

10821082
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitKernelLaunch(
1083-
const LocationDescription &Loc, Function *OutlinedFn, Value *OutlinedFnID,
1084-
EmitFallbackCallbackTy emitTargetCallFallbackCB, TargetKernelArgs &Args,
1083+
const LocationDescription &Loc, Value *OutlinedFnID,
1084+
EmitFallbackCallbackTy EmitTargetCallFallbackCB, TargetKernelArgs &Args,
10851085
Value *DeviceID, Value *RTLoc, InsertPointTy AllocaIP) {
10861086

10871087
if (!updateToLocation(Loc))
@@ -1134,7 +1134,7 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitKernelLaunch(
11341134

11351135
auto CurFn = Builder.GetInsertBlock()->getParent();
11361136
emitBlock(OffloadFailedBlock, CurFn);
1137-
Builder.restoreIP(emitTargetCallFallbackCB(Builder.saveIP()));
1137+
Builder.restoreIP(EmitTargetCallFallbackCB(Builder.saveIP()));
11381138
emitBranch(OffloadContBlock);
11391139
emitBlock(OffloadContBlock, CurFn, /*IsFinished=*/true);
11401140
return Builder.saveIP();
@@ -1736,7 +1736,7 @@ void OpenMPIRBuilder::createTaskyield(const LocationDescription &Loc) {
17361736
// - All code is inserted in the entry block of the current function.
17371737
static Value *emitTaskDependencies(
17381738
OpenMPIRBuilder &OMPBuilder,
1739-
SmallVectorImpl<OpenMPIRBuilder::DependData> &Dependencies) {
1739+
const SmallVectorImpl<OpenMPIRBuilder::DependData> &Dependencies) {
17401740
// Early return if we have no dependencies to process
17411741
if (Dependencies.empty())
17421742
return nullptr;
@@ -6386,12 +6386,10 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetData(
63866386
// closing of the region.
63876387
auto BeginThenGen = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
63886388
MapInfo = &GenMapInfoCB(Builder.saveIP());
6389-
emitOffloadingArrays(AllocaIP, Builder.saveIP(), *MapInfo, Info,
6390-
/*IsNonContiguous=*/true, DeviceAddrCB,
6391-
CustomMapperCB);
6392-
63936389
TargetDataRTArgs RTArgs;
6394-
emitOffloadingArraysArgument(Builder, RTArgs, Info);
6390+
emitOffloadingArraysAndArgs(AllocaIP, Builder.saveIP(), Info, RTArgs,
6391+
*MapInfo, /*IsNonContiguous=*/true,
6392+
/*ForEndCall=*/false);
63956393

63966394
// Emit the number of elements in the offloading arrays.
63976395
Value *PointerNum = Builder.getInt32(Info.NumberOfPtrs);
@@ -6403,16 +6401,45 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetData(
64036401
SrcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
64046402
}
64056403

6406-
Value *OffloadingArgs[] = {SrcLocInfo, DeviceID,
6407-
PointerNum, RTArgs.BasePointersArray,
6408-
RTArgs.PointersArray, RTArgs.SizesArray,
6409-
RTArgs.MapTypesArray, RTArgs.MapNamesArray,
6410-
RTArgs.MappersArray};
6404+
SmallVector<llvm::Value *, 13> OffloadingArgs = {
6405+
SrcLocInfo, DeviceID,
6406+
PointerNum, RTArgs.BasePointersArray,
6407+
RTArgs.PointersArray, RTArgs.SizesArray,
6408+
RTArgs.MapTypesArray, RTArgs.MapNamesArray,
6409+
RTArgs.MappersArray};
64116410

64126411
if (IsStandAlone) {
64136412
assert(MapperFunc && "MapperFunc missing for standalone target data");
6414-
Builder.CreateCall(getOrCreateRuntimeFunctionPtr(*MapperFunc),
6415-
OffloadingArgs);
6413+
6414+
auto TaskBodyCB = [&](Value *, Value *, IRBuilderBase::InsertPoint) {
6415+
if (Info.HasNoWait) {
6416+
OffloadingArgs.push_back(llvm::Constant::getNullValue(Int32));
6417+
OffloadingArgs.push_back(llvm::Constant::getNullValue(VoidPtr));
6418+
OffloadingArgs.push_back(llvm::Constant::getNullValue(Int32));
6419+
OffloadingArgs.push_back(llvm::Constant::getNullValue(VoidPtr));
6420+
}
6421+
6422+
Builder.CreateCall(getOrCreateRuntimeFunctionPtr(*MapperFunc),
6423+
OffloadingArgs);
6424+
6425+
if (Info.HasNoWait) {
6426+
BasicBlock *OffloadContBlock =
6427+
BasicBlock::Create(Builder.getContext(), "omp_offload.cont");
6428+
auto *CurFn = Builder.GetInsertBlock()->getParent();
6429+
emitBranch(OffloadContBlock);
6430+
emitBlock(OffloadContBlock, CurFn, /*IsFinished=*/true);
6431+
Builder.restoreIP(Builder.saveIP());
6432+
}
6433+
};
6434+
6435+
bool RequiresOuterTargetTask = Info.HasNoWait;
6436+
6437+
if (!RequiresOuterTargetTask)
6438+
TaskBodyCB(/*DeviceID=*/nullptr, /*RTLoc=*/nullptr,
6439+
/*TargetTaskAllocaIP=*/{});
6440+
else
6441+
emitTargetTask(TaskBodyCB, DeviceID, SrcLocInfo, AllocaIP,
6442+
/*Dependencies=*/{}, Info.HasNoWait);
64166443
} else {
64176444
Function *BeginMapperFunc = getOrCreateRuntimeFunctionPtr(
64186445
omp::OMPRTL___tgt_target_data_begin_mapper);
@@ -6836,13 +6863,18 @@ static void emitTargetOutlinedFunction(
68366863
OMPBuilder.emitTargetRegionFunction(EntryInfo, GenerateOutlinedFunction,
68376864
IsOffloadEntry, OutlinedFn, OutlinedFnID);
68386865
}
6866+
68396867
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitTargetTask(
6840-
Function *OutlinedFn, Value *OutlinedFnID,
6841-
EmitFallbackCallbackTy EmitTargetCallFallbackCB, TargetKernelArgs &Args,
6842-
Value *DeviceID, Value *RTLoc, OpenMPIRBuilder::InsertPointTy AllocaIP,
6843-
SmallVector<llvm::OpenMPIRBuilder::DependData> &Dependencies,
6868+
TaskBodyCallbackTy TaskBodyCB, Value *DeviceID, Value *RTLoc,
6869+
OpenMPIRBuilder::InsertPointTy AllocaIP,
6870+
const SmallVector<llvm::OpenMPIRBuilder::DependData> &Dependencies,
68446871
bool HasNoWait) {
68456872

6873+
// The following explains the code-gen scenario for the `target` directive. A
6874+
// similar scneario is followed for other device-related directives (e.g.
6875+
// `target enter data`) but in similar fashion since we only need to emit task
6876+
// that encapsulates the proper runtime call.
6877+
//
68466878
// When we arrive at this function, the target region itself has been
68476879
// outlined into the function OutlinedFn.
68486880
// So at ths point, for
@@ -6950,22 +6982,7 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitTargetTask(
69506982

69516983
Builder.restoreIP(TargetTaskBodyIP);
69526984

6953-
if (OutlinedFnID) {
6954-
// emitKernelLaunch makes the necessary runtime call to offload the kernel.
6955-
// We then outline all that code into a separate function
6956-
// ('kernel_launch_function' in the pseudo code above). This function is
6957-
// then called by the target task proxy function (see
6958-
// '@.omp_target_task_proxy_func' in the pseudo code above)
6959-
// "@.omp_target_task_proxy_func' is generated by
6960-
// emitTargetTaskProxyFunction.
6961-
Builder.restoreIP(emitKernelLaunch(Builder, OutlinedFn, OutlinedFnID,
6962-
EmitTargetCallFallbackCB, Args, DeviceID,
6963-
RTLoc, TargetTaskAllocaIP));
6964-
} else {
6965-
// When OutlinedFnID is set to nullptr, then it's not an offloading call. In
6966-
// this case, we execute the host implementation directly.
6967-
Builder.restoreIP(EmitTargetCallFallbackCB(Builder.saveIP()));
6968-
}
6985+
TaskBodyCB(DeviceID, RTLoc, TargetTaskAllocaIP);
69696986

69706987
OI.ExitBB = Builder.saveIP().getBlock();
69716988
OI.PostOutlineCB = [this, ToBeDeleted, Dependencies, HasNoWait,
@@ -7153,21 +7170,44 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
71537170
bool HasDependencies = Dependencies.size() > 0;
71547171
bool RequiresOuterTargetTask = HasNoWait || HasDependencies;
71557172

7173+
OpenMPIRBuilder::TargetKernelArgs KArgs;
7174+
7175+
auto TaskBodyCB = [&](Value *DeviceID, Value *RTLoc,
7176+
IRBuilderBase::InsertPoint TargetTaskAllocaIP) {
7177+
if (OutlinedFnID) {
7178+
// emitKernelLaunch makes the necessary runtime call to offload the
7179+
// kernel. We then outline all that code into a separate function
7180+
// ('kernel_launch_function' in the pseudo code above). This function is
7181+
// then called by the target task proxy function (see
7182+
// '@.omp_target_task_proxy_func' in the pseudo code above)
7183+
// "@.omp_target_task_proxy_func' is generated by
7184+
// emitTargetTaskProxyFunction.
7185+
Builder.restoreIP(OMPBuilder.emitKernelLaunch(
7186+
Builder, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID,
7187+
RTLoc, TargetTaskAllocaIP));
7188+
} else {
7189+
// When OutlinedFnID is set to nullptr, then it's not an offloading
7190+
// call. In this case, we execute the host implementation directly.
7191+
OMPBuilder.Builder.restoreIP(
7192+
EmitTargetCallFallbackCB(OMPBuilder.Builder.saveIP()));
7193+
}
7194+
};
7195+
71567196
// If we don't have an ID for the target region, it means an offload entry
71577197
// wasn't created. In this case we just run the host fallback directly.
71587198
if (!OutlinedFnID) {
71597199
if (RequiresOuterTargetTask) {
71607200
// Arguments that are intended to be directly forwarded to an
71617201
// emitKernelLaunch call are pased as nullptr, since OutlinedFnID=nullptr
71627202
// results in that call not being done.
7163-
OpenMPIRBuilder::TargetKernelArgs KArgs;
7164-
Builder.restoreIP(OMPBuilder.emitTargetTask(
7165-
OutlinedFn, /*OutlinedFnID=*/nullptr, EmitTargetCallFallbackCB, KArgs,
7166-
/*DeviceID=*/nullptr, /*RTLoc=*/nullptr, AllocaIP, Dependencies,
7167-
HasNoWait));
7203+
Builder.restoreIP(OMPBuilder.emitTargetTask(TaskBodyCB,
7204+
/*DeviceID=*/nullptr,
7205+
/*RTLoc=*/nullptr, AllocaIP,
7206+
Dependencies, HasNoWait));
71687207
} else {
71697208
Builder.restoreIP(EmitTargetCallFallbackCB(Builder.saveIP()));
71707209
}
7210+
71717211
return;
71727212
}
71737213

@@ -7201,20 +7241,19 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
72017241
// TODO: Use correct DynCGGroupMem
72027242
Value *DynCGGroupMem = Builder.getInt32(0);
72037243

7204-
OpenMPIRBuilder::TargetKernelArgs KArgs(NumTargetItems, RTArgs, NumIterations,
7205-
NumTeamsC, NumThreadsC, DynCGGroupMem,
7206-
HasNoWait);
7244+
KArgs = OpenMPIRBuilder::TargetKernelArgs(
7245+
NumTargetItems, RTArgs, NumIterations, NumTeamsC, NumThreadsC,
7246+
DynCGGroupMem, HasNoWait);
72077247

72087248
// The presence of certain clauses on the target directive require the
72097249
// explicit generation of the target task.
72107250
if (RequiresOuterTargetTask) {
72117251
Builder.restoreIP(OMPBuilder.emitTargetTask(
7212-
OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID,
7213-
RTLoc, AllocaIP, Dependencies, HasNoWait));
7252+
TaskBodyCB, DeviceID, RTLoc, AllocaIP, Dependencies, HasNoWait));
72147253
} else {
72157254
Builder.restoreIP(OMPBuilder.emitKernelLaunch(
7216-
Builder, OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs,
7217-
DeviceID, RTLoc, AllocaIP));
7255+
Builder, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID, RTLoc,
7256+
AllocaIP));
72187257
}
72197258
}
72207259

‎mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2886,6 +2886,8 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
28862886
DataLayout DL = DataLayout(op->getParentOfType<ModuleOp>());
28872887

28882888
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
2889+
llvm::OpenMPIRBuilder::TargetDataInfo info(/*RequiresDevicePointerInfo=*/true,
2890+
/*SeparateBeginEndCalls=*/true);
28892891

28902892
LogicalResult result =
28912893
llvm::TypeSwitch<Operation *, LogicalResult>(op)
@@ -2905,9 +2907,9 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
29052907
return success();
29062908
})
29072909
.Case([&](omp::TargetEnterDataOp enterDataOp) {
2908-
if (enterDataOp.getNowait())
2910+
if (!enterDataOp.getDependVars().empty())
29092911
return (LogicalResult)(enterDataOp.emitError(
2910-
"`nowait` is not supported yet"));
2912+
"`depend` is not supported yet"));
29112913

29122914
if (auto ifVar = enterDataOp.getIfExpr())
29132915
ifCond = moduleTranslation.lookupValue(ifVar);
@@ -2917,14 +2919,18 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
29172919
dyn_cast<LLVM::ConstantOp>(devId.getDefiningOp()))
29182920
if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue()))
29192921
deviceID = intAttr.getInt();
2920-
RTLFn = llvm::omp::OMPRTL___tgt_target_data_begin_mapper;
2922+
RTLFn =
2923+
enterDataOp.getNowait()
2924+
? llvm::omp::OMPRTL___tgt_target_data_begin_nowait_mapper
2925+
: llvm::omp::OMPRTL___tgt_target_data_begin_mapper;
29212926
mapVars = enterDataOp.getMapVars();
2927+
info.HasNoWait = enterDataOp.getNowait();
29222928
return success();
29232929
})
29242930
.Case([&](omp::TargetExitDataOp exitDataOp) {
2925-
if (exitDataOp.getNowait())
2931+
if (!exitDataOp.getDependVars().empty())
29262932
return (LogicalResult)(exitDataOp.emitError(
2927-
"`nowait` is not supported yet"));
2933+
"`depend` is not supported yet"));
29282934

29292935
if (auto ifVar = exitDataOp.getIfExpr())
29302936
ifCond = moduleTranslation.lookupValue(ifVar);
@@ -2935,14 +2941,17 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
29352941
if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue()))
29362942
deviceID = intAttr.getInt();
29372943

2938-
RTLFn = llvm::omp::OMPRTL___tgt_target_data_end_mapper;
2944+
RTLFn = exitDataOp.getNowait()
2945+
? llvm::omp::OMPRTL___tgt_target_data_end_nowait_mapper
2946+
: llvm::omp::OMPRTL___tgt_target_data_end_mapper;
29392947
mapVars = exitDataOp.getMapVars();
2948+
info.HasNoWait = exitDataOp.getNowait();
29402949
return success();
29412950
})
29422951
.Case([&](omp::TargetUpdateOp updateDataOp) {
2943-
if (updateDataOp.getNowait())
2952+
if (!updateDataOp.getDependVars().empty())
29442953
return (LogicalResult)(updateDataOp.emitError(
2945-
"`nowait` is not supported yet"));
2954+
"`depend` is not supported yet"));
29462955

29472956
if (auto ifVar = updateDataOp.getIfExpr())
29482957
ifCond = moduleTranslation.lookupValue(ifVar);
@@ -2953,8 +2962,12 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
29532962
if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue()))
29542963
deviceID = intAttr.getInt();
29552964

2956-
RTLFn = llvm::omp::OMPRTL___tgt_target_data_update_mapper;
2965+
RTLFn =
2966+
updateDataOp.getNowait()
2967+
? llvm::omp::OMPRTL___tgt_target_data_update_nowait_mapper
2968+
: llvm::omp::OMPRTL___tgt_target_data_update_mapper;
29572969
mapVars = updateDataOp.getMapVars();
2970+
info.HasNoWait = updateDataOp.getNowait();
29582971
return success();
29592972
})
29602973
.Default([&](Operation *op) {
@@ -3005,9 +3018,6 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
30053018
: basePointer);
30063019
};
30073020

3008-
llvm::OpenMPIRBuilder::TargetDataInfo info(/*RequiresDevicePointerInfo=*/true,
3009-
/*SeparateBeginEndCalls=*/true);
3010-
30113021
using BodyGenTy = llvm::OpenMPIRBuilder::BodyGenTy;
30123022
LogicalResult bodyGenStatus = success();
30133023
auto bodyGenCB = [&](InsertPointTy codeGenIP, BodyGenTy bodyGenType) {

0 commit comments

Comments
 (0)