Skip to content

Commit ad70f3e

Browse files
ergawyergawy
authored
[flang][OpenMP] Support target enter|update|exit .. nowait (#113305)
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 27158ed commit ad70f3e

File tree

6 files changed

+241
-110
lines changed

6 files changed

+241
-110
lines changed

‎clang/lib/CodeGen/CGOpenMPRuntime.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9672,8 +9672,8 @@ static void emitTargetCallKernelLaunch(
96729672
DynCGGroupMem, HasNoWait);
96739673

96749674
CGF.Builder.restoreIP(OMPRuntime->getOMPBuilder().emitKernelLaunch(
9675-
CGF.Builder, OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, Args,
9676-
DeviceID, RTLoc, AllocaIP));
9675+
CGF.Builder, OutlinedFnID, EmitTargetCallFallbackCB, Args, DeviceID,
9676+
RTLoc, AllocaIP));
96779677
};
96789678

96799679
if (RequiresOuterTask)

‎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 target 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 TargetTaskBodyCallbackTy =
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+
TargetTaskBodyCallbackTy 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: 82 additions & 43 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;
@@ -6403,16 +6403,44 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetData(
64036403
SrcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
64046404
}
64056405

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

64126413
if (IsStandAlone) {
64136414
assert(MapperFunc && "MapperFunc missing for standalone target data");
6414-
Builder.CreateCall(getOrCreateRuntimeFunctionPtr(*MapperFunc),
6415-
OffloadingArgs);
6415+
6416+
auto TaskBodyCB = [&](Value *, Value *, IRBuilderBase::InsertPoint) {
6417+
if (Info.HasNoWait) {
6418+
OffloadingArgs.append({llvm::Constant::getNullValue(Int32),
6419+
llvm::Constant::getNullValue(VoidPtr),
6420+
llvm::Constant::getNullValue(Int32),
6421+
llvm::Constant::getNullValue(VoidPtr)});
6422+
}
6423+
6424+
Builder.CreateCall(getOrCreateRuntimeFunctionPtr(*MapperFunc),
6425+
OffloadingArgs);
6426+
6427+
if (Info.HasNoWait) {
6428+
BasicBlock *OffloadContBlock =
6429+
BasicBlock::Create(Builder.getContext(), "omp_offload.cont");
6430+
Function *CurFn = Builder.GetInsertBlock()->getParent();
6431+
emitBlock(OffloadContBlock, CurFn, /*IsFinished=*/true);
6432+
Builder.restoreIP(Builder.saveIP());
6433+
}
6434+
};
6435+
6436+
bool RequiresOuterTargetTask = Info.HasNoWait;
6437+
6438+
if (!RequiresOuterTargetTask)
6439+
TaskBodyCB(/*DeviceID=*/nullptr, /*RTLoc=*/nullptr,
6440+
/*TargetTaskAllocaIP=*/{});
6441+
else
6442+
emitTargetTask(TaskBodyCB, DeviceID, SrcLocInfo, AllocaIP,
6443+
/*Dependencies=*/{}, Info.HasNoWait);
64166444
} else {
64176445
Function *BeginMapperFunc = getOrCreateRuntimeFunctionPtr(
64186446
omp::OMPRTL___tgt_target_data_begin_mapper);
@@ -6836,13 +6864,18 @@ static void emitTargetOutlinedFunction(
68366864
OMPBuilder.emitTargetRegionFunction(EntryInfo, GenerateOutlinedFunction,
68376865
IsOffloadEntry, OutlinedFn, OutlinedFnID);
68386866
}
6867+
68396868
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,
6869+
TargetTaskBodyCallbackTy TaskBodyCB, Value *DeviceID, Value *RTLoc,
6870+
OpenMPIRBuilder::InsertPointTy AllocaIP,
6871+
const SmallVector<llvm::OpenMPIRBuilder::DependData> &Dependencies,
68446872
bool HasNoWait) {
68456873

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

69516984
Builder.restoreIP(TargetTaskBodyIP);
69526985

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-
}
6986+
TaskBodyCB(DeviceID, RTLoc, TargetTaskAllocaIP);
69696987

69706988
OI.ExitBB = Builder.saveIP().getBlock();
69716989
OI.PostOutlineCB = [this, ToBeDeleted, Dependencies, HasNoWait,
@@ -7153,18 +7171,40 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
71537171
bool HasDependencies = Dependencies.size() > 0;
71547172
bool RequiresOuterTargetTask = HasNoWait || HasDependencies;
71557173

7174+
OpenMPIRBuilder::TargetKernelArgs KArgs;
7175+
7176+
auto TaskBodyCB = [&](Value *DeviceID, Value *RTLoc,
7177+
IRBuilderBase::InsertPoint TargetTaskAllocaIP) {
7178+
if (OutlinedFnID) {
7179+
// emitKernelLaunch makes the necessary runtime call to offload the
7180+
// kernel. We then outline all that code into a separate function
7181+
// ('kernel_launch_function' in the pseudo code above). This function is
7182+
// then called by the target task proxy function (see
7183+
// '@.omp_target_task_proxy_func' in the pseudo code above)
7184+
// "@.omp_target_task_proxy_func' is generated by
7185+
// emitTargetTaskProxyFunction.
7186+
Builder.restoreIP(OMPBuilder.emitKernelLaunch(
7187+
Builder, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID,
7188+
RTLoc, TargetTaskAllocaIP));
7189+
} else {
7190+
// When OutlinedFnID is set to nullptr, then it's not an offloading
7191+
// call. In this case, we execute the host implementation directly.
7192+
OMPBuilder.Builder.restoreIP(
7193+
EmitTargetCallFallbackCB(OMPBuilder.Builder.saveIP()));
7194+
}
7195+
};
7196+
71567197
// If we don't have an ID for the target region, it means an offload entry
71577198
// wasn't created. In this case we just run the host fallback directly.
71587199
if (!OutlinedFnID) {
71597200
if (RequiresOuterTargetTask) {
71607201
// Arguments that are intended to be directly forwarded to an
71617202
// emitKernelLaunch call are pased as nullptr, since OutlinedFnID=nullptr
71627203
// 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));
7204+
Builder.restoreIP(OMPBuilder.emitTargetTask(TaskBodyCB,
7205+
/*DeviceID=*/nullptr,
7206+
/*RTLoc=*/nullptr, AllocaIP,
7207+
Dependencies, HasNoWait));
71687208
} else {
71697209
Builder.restoreIP(EmitTargetCallFallbackCB(Builder.saveIP()));
71707210
}
@@ -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)