Skip to content

Commit 9eb4056

Browse files
NimishMishraNimishMishra
authored
[mlir][llvm] Translation support for task detach (#116601)
This PR adds translation support for task detach. Essentially, if the `detach` clause is present on a task, emit a `__kmpc_task_allow_completion_event` on it, and store its return (of type `kmp_event_t*`) into the `event_handle`.
1 parent 25dcee4 commit 9eb4056

File tree

9 files changed

+88
-18
lines changed

9 files changed

+88
-18
lines changed

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1262,12 +1262,15 @@ class OpenMPIRBuilder {
12621262
/// cannot be resumed until execution of the structured
12631263
/// block that is associated with the generated task is
12641264
/// completed.
1265+
/// \param EventHandle If present, signifies the event handle as part of
1266+
/// the detach clause
12651267
/// \param Mergeable If the given task is `mergeable`
12661268
InsertPointOrErrorTy
12671269
createTask(const LocationDescription &Loc, InsertPointTy AllocaIP,
12681270
BodyGenCallbackTy BodyGenCB, bool Tied = true,
12691271
Value *Final = nullptr, Value *IfCondition = nullptr,
1270-
SmallVector<DependData> Dependencies = {}, bool Mergeable = false);
1272+
SmallVector<DependData> Dependencies = {}, bool Mergeable = false,
1273+
Value *EventHandle = nullptr);
12711274

12721275
/// Generator for the taskgroup construct
12731276
///

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1818,7 +1818,7 @@ static Value *emitTaskDependencies(
18181818
OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
18191819
const LocationDescription &Loc, InsertPointTy AllocaIP,
18201820
BodyGenCallbackTy BodyGenCB, bool Tied, Value *Final, Value *IfCondition,
1821-
SmallVector<DependData> Dependencies, bool Mergeable) {
1821+
SmallVector<DependData> Dependencies, bool Mergeable, Value *EventHandle) {
18221822

18231823
if (!updateToLocation(Loc))
18241824
return InsertPointTy();
@@ -1864,7 +1864,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
18641864
Builder, AllocaIP, ToBeDeleted, TaskAllocaIP, "global.tid", false));
18651865

18661866
OI.PostOutlineCB = [this, Ident, Tied, Final, IfCondition, Dependencies,
1867-
Mergeable, TaskAllocaBB,
1867+
Mergeable, EventHandle, TaskAllocaBB,
18681868
ToBeDeleted](Function &OutlinedFn) mutable {
18691869
// Replace the Stale CI by appropriate RTL function call.
18701870
assert(OutlinedFn.getNumUses() == 1 &&
@@ -1935,6 +1935,20 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
19351935
/*sizeof_task=*/TaskSize, /*sizeof_shared=*/SharedsSize,
19361936
/*task_func=*/&OutlinedFn});
19371937

1938+
// Emit detach clause initialization.
1939+
// evt = (typeof(evt))__kmpc_task_allow_completion_event(loc, tid,
1940+
// task_descriptor);
1941+
if (EventHandle) {
1942+
Function *TaskDetachFn = getOrCreateRuntimeFunctionPtr(
1943+
OMPRTL___kmpc_task_allow_completion_event);
1944+
llvm::Value *EventVal =
1945+
Builder.CreateCall(TaskDetachFn, {Ident, ThreadID, TaskData});
1946+
llvm::Value *EventHandleAddr =
1947+
Builder.CreatePointerBitCastOrAddrSpaceCast(EventHandle,
1948+
Builder.getPtrTy(0));
1949+
EventVal = Builder.CreatePtrToInt(EventVal, Builder.getInt64Ty());
1950+
Builder.CreateStore(EventVal, EventHandleAddr);
1951+
}
19381952
// Copy the arguments for outlined function
19391953
if (HasShareds) {
19401954
Value *Shareds = StaleCI->getArgOperand(1);

‎mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -909,6 +909,34 @@ class OpenMP_ParallelizationLevelClauseSkip<
909909

910910
def OpenMP_ParallelizationLevelClause : OpenMP_ParallelizationLevelClauseSkip<>;
911911

912+
//===----------------------------------------------------------------------===//
913+
// OpenMPV5.2: [12.5.2] `detach` clause
914+
//===----------------------------------------------------------------------===//
915+
916+
class OpenMP_DetachClauseSkip<
917+
bit traits = false, bit arguments = false, bit assemblyFormat = false,
918+
bit description = false, bit extraClassDeclaration = false>
919+
: OpenMP_Clause<traits, arguments, assemblyFormat, description,
920+
extraClassDeclaration> {
921+
922+
let traits = [BlockArgOpenMPOpInterface];
923+
924+
let arguments = (ins Optional<OpenMP_PointerLikeType>:$event_handle);
925+
926+
let optAssemblyFormat = [{
927+
`detach` `(` $event_handle `:` type($event_handle) `)`
928+
}];
929+
930+
let description = [{
931+
The detach clause specifies that the task generated by the construct on which it appears is a
932+
detachable task. A new allow-completion event is created and connected to the completion of the
933+
associated task region. The original event-handle is updated to represent that allow-completion
934+
event before the task data environment is created.
935+
}];
936+
}
937+
938+
def OpenMP_DetachClause : OpenMP_DetachClauseSkip<>;
939+
912940
//===----------------------------------------------------------------------===//
913941
// V5.2: [12.4] `priority` clause
914942
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -657,15 +657,18 @@ def DistributeOp : OpenMP_Op<"distribute", traits = [
657657
// 2.10.1 task Construct
658658
//===----------------------------------------------------------------------===//
659659

660-
def TaskOp : OpenMP_Op<"task", traits = [
661-
AttrSizedOperandSegments, AutomaticAllocationScope,
662-
OutlineableOpenMPOpInterface
663-
], clauses = [
664-
// TODO: Complete clause list (affinity, detach).
665-
OpenMP_AllocateClause, OpenMP_DependClause, OpenMP_FinalClause,
666-
OpenMP_IfClause, OpenMP_InReductionClause, OpenMP_MergeableClause,
667-
OpenMP_PriorityClause, OpenMP_PrivateClause, OpenMP_UntiedClause
668-
], singleRegion = true> {
660+
def TaskOp
661+
: OpenMP_Op<"task",
662+
traits = [AttrSizedOperandSegments, AutomaticAllocationScope,
663+
OutlineableOpenMPOpInterface],
664+
clauses = [
665+
// TODO: Complete clause list (affinity, detach).
666+
OpenMP_AllocateClause, OpenMP_DependClause,
667+
OpenMP_FinalClause, OpenMP_IfClause,
668+
OpenMP_InReductionClause, OpenMP_MergeableClause,
669+
OpenMP_PriorityClause, OpenMP_PrivateClause,
670+
OpenMP_UntiedClause, OpenMP_DetachClause],
671+
singleRegion = true> {
669672
let summary = "task construct";
670673
let description = [{
671674
The task construct defines an explicit task.

‎mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2303,7 +2303,7 @@ void TaskOp::build(OpBuilder &builder, OperationState &state,
23032303
makeArrayAttr(ctx, clauses.inReductionSyms), clauses.mergeable,
23042304
clauses.priority, /*private_vars=*/clauses.privateVars,
23052305
/*private_syms=*/makeArrayAttr(ctx, clauses.privateSyms),
2306-
clauses.untied);
2306+
clauses.untied, clauses.eventHandle);
23072307
}
23082308

23092309
LogicalResult TaskOp::verify() {

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1703,7 +1703,8 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
17031703
ompLoc, allocaIP, bodyCB, !taskOp.getUntied(),
17041704
moduleTranslation.lookupValue(taskOp.getFinal()),
17051705
moduleTranslation.lookupValue(taskOp.getIfExpr()), dds,
1706-
taskOp.getMergeable());
1706+
taskOp.getMergeable(),
1707+
moduleTranslation.lookupValue(taskOp.getEventHandle()));
17071708

17081709
if (failed(handleError(afterIP, *taskOp)))
17091710
return failure();

‎mlir/test/Dialect/OpenMP/invalid.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1634,7 +1634,7 @@ func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
16341634
// -----
16351635

16361636
func.func @omp_task_depend(%data_var: memref<i32>) {
1637-
// expected-error @below {{op expected as many depend values as depend variables}}
1637+
// expected-error @below {{'omp.task' op operand count (1) does not match with the total size (0) specified in attribute 'operandSegmentSizes'}}
16381638
"omp.task"(%data_var) ({
16391639
"omp.terminator"() : () -> ()
16401640
}) {depend_kinds = [], operandSegmentSizes = array<i32: 0, 0, 1, 0, 0, 0, 0, 0>} : (memref<i32>) -> ()

‎mlir/test/Dialect/OpenMP/ops.mlir

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1975,8 +1975,8 @@ func.func @omp_single_copyprivate(%data_var: memref<i32>) {
19751975
}
19761976

19771977
// CHECK-LABEL: @omp_task
1978-
// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref<i32>)
1979-
func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref<i32>) {
1978+
// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref<i32>, %[[event_handle:.*]]: !llvm.ptr)
1979+
func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref<i32>, %event_handle : !llvm.ptr) {
19801980

19811981
// Checking simple task
19821982
// CHECK: omp.task {
@@ -2054,7 +2054,11 @@ func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memr
20542054
// CHECK: omp.terminator
20552055
omp.terminator
20562056
}
2057-
2057+
// Checking detach clause
2058+
// CHECK: omp.task detach(%[[event_handle]] : !llvm.ptr)
2059+
omp.task detach(%event_handle : !llvm.ptr){
2060+
omp.terminator
2061+
}
20582062
// Checking multiple clauses
20592063
// CHECK: omp.task allocate(%[[data_var]] : memref<i32> -> %[[data_var]] : memref<i32>)
20602064
omp.task allocate(%data_var : memref<i32> -> %data_var : memref<i32>)

‎mlir/test/Target/LLVMIR/openmp-llvm.mlir

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2505,6 +2505,23 @@ llvm.mlir.global internal @_QFsubEx() : i32
25052505

25062506
// -----
25072507

2508+
// CHECK-LABEL: define void @omp_task_detach
2509+
// CHECK-SAME: (ptr %[[event_handle:.*]])
2510+
llvm.func @omp_task_detach(%event_handle : !llvm.ptr){
2511+
// CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
2512+
// CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
2513+
// CHECK: %[[return_val:.*]] = call ptr @__kmpc_task_allow_completion_event(ptr {{.*}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
2514+
// CHECK: %[[conv:.*]] = ptrtoint ptr %[[return_val]] to i64
2515+
// CHECK: store i64 %[[conv]], ptr %[[event_handle]], align 4
2516+
// CHECK: call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
2517+
omp.task detach(%event_handle : !llvm.ptr){
2518+
omp.terminator
2519+
}
2520+
llvm.return
2521+
}
2522+
2523+
// -----
2524+
25082525
// CHECK-LABEL: define void @omp_task
25092526
// CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
25102527
llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr) {

0 commit comments

Comments
 (0)