Skip to content

Commit 5e465ee

Browse files
committed
We've 2 ops:
1. nvvm.griddepcontrol.wait 1. nvvm.griddepcontrol.launch_dependents They are related to FDL and same concept. This PR unifies both ops into a single one.
1 parent 69d0bd5 commit 5e465ee

File tree

3 files changed

+47
-19
lines changed

3 files changed

+47
-19
lines changed

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2995,30 +2995,47 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
29952995
// NVVM Griddepcontrol Ops
29962996
//===----------------------------------------------------------------------===//
29972997

2998-
def NVVM_GriddepcontrolWaitOp : NVVM_IntrOp<"griddepcontrol.wait", [], 0> {
2999-
let assemblyFormat = "attr-dict";
2998+
def FDLWait : I32EnumAttrCase<"wait", 0>;
2999+
def FDLLaunchDependent : I32EnumAttrCase<"launch_dependents", 1>;
30003000

3001+
def FDLKind : I32EnumAttr<"FDLKind", "Fast dependenct launch kind",
3002+
[FDLWait, FDLLaunchDependent]> {
3003+
let genSpecializedAttr = 0;
3004+
let cppNamespace = "::mlir::NVVM";
3005+
}
3006+
3007+
def FDLKindAttr : EnumAttr<NVVM_Dialect, FDLKind, "fdl_kind">;
3008+
3009+
def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
30013010
let description = [{
3002-
Causes the executing thread to wait until all prerequisite grids in flight
3011+
If the $kind attribute is set to `wait`, the instruction, it causes the
3012+
executing thread to wait until all prerequisite grids in flight
30033013
have completed and all the memory operations from the prerequisite grids
30043014
are performed and made visible to the current grid.
30053015

3016+
When the $kind is launch_dependents, it signals that specific dependents
3017+
the runtime system designated to react to this instruction can be scheduled
3018+
as soon as all other CTAs in the grid issue the same instruction or have
3019+
completed.
30063020

30073021
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
30083022
}];
3009-
}
30103023

3011-
def NVVM_GriddepcontrolLaunchDependentsOp
3012-
: NVVM_IntrOp<"griddepcontrol.launch.dependents", [], 0> {
3013-
let assemblyFormat = "attr-dict";
3014-
3015-
let description = [{
3016-
Signals that specific dependents the runtime system designated to react to
3017-
this instruction can be scheduled as soon as all other CTAs in the grid
3018-
issue the same instruction or have completed.
3024+
let arguments = (ins FDLKindAttr:$kind);
30193025

3026+
let assemblyFormat = "$kind attr-dict";
30203027

3021-
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
3028+
string llvmBuilder = [{
3029+
llvm::Intrinsic::ID id;
3030+
switch ($kind) {
3031+
case NVVM::FDLKind::wait:
3032+
id = llvm::Intrinsic::nvvm_griddepcontrol_wait;
3033+
break;
3034+
case NVVM::FDLKind::launch_dependents:
3035+
id = llvm::Intrinsic::nvvm_griddepcontrol_launch_dependents;
3036+
break;
3037+
}
3038+
createIntrinsicCall(builder, id);
30223039
}];
30233040
}
30243041

mlir/test/Dialect/LLVMIR/nvvm.mlir

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -535,15 +535,15 @@ func.func @wgmma_wait_group_sync_aligned() {
535535
}
536536

537537
func.func @griddepcontrol_wait() {
538-
// CHECK: nvvm.griddepcontrol.wait
539-
nvvm.griddepcontrol.wait
538+
// CHECK: nvvm.griddepcontrol wait
539+
nvvm.griddepcontrol wait
540540
return
541541
}
542542

543543
func.func @griddepcontrol_launch_dependents()
544544
{
545-
// CHECK: nvvm.griddepcontrol.launch.dependents
546-
nvvm.griddepcontrol.launch.dependents
545+
// CHECK: nvvm.griddepcontrol launch_dependents
546+
nvvm.griddepcontrol launch_dependents
547547
return
548548
}
549549

mlir/test/Target/LLVMIR/nvvmir.mlir

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -766,15 +766,15 @@ llvm.func @nvvm_wgmma_wait_group_aligned() {
766766
// CHECK-LABEL: @nvvm_griddepcontrol_wait
767767
llvm.func @nvvm_griddepcontrol_wait() {
768768
// CHECK: call void @llvm.nvvm.griddepcontrol.wait()
769-
nvvm.griddepcontrol.wait
769+
nvvm.griddepcontrol wait
770770
llvm.return
771771
}
772772

773773
// -----
774774
// CHECK-LABEL: @nvvm_griddepcontrol_launch_dependents
775775
llvm.func @nvvm_griddepcontrol_launch_dependents() {
776776
// CHECK: call void @llvm.nvvm.griddepcontrol.launch.dependents()
777-
nvvm.griddepcontrol.launch.dependents
777+
nvvm.griddepcontrol launch_dependents
778778
llvm.return
779779
}
780780

@@ -918,3 +918,14 @@ llvm.func @nvvm_dot_accumulate_2way(%a: vector<2xi16>, %b: vector<4xi8>, %c: i32
918918
%7 = nvvm.dot.accumulate.2way %a <signed>, %b <signed>, %c {b_hi = true}: vector<2xi16>, vector<4xi8>
919919
llvm.return
920920
}
921+
922+
// -----
923+
924+
// CHECK-LABEL: @nvvm_pmevent
925+
llvm.func @nvvm_pmevent() {
926+
// CHECK: call void @llvm.nvvm.pm.event.mask(i16 15000)
927+
nvvm.pmevent mask = 15000
928+
// CHECK: call void @llvm.nvvm.pm.event.mask(i16 4)
929+
nvvm.pmevent mask = 4
930+
llvm.return
931+
}

0 commit comments

Comments
 (0)