[Mlir-commits] [mlir] [MLIR][NVVM] Combine griddepcontrol Ops (PR #152525)
Guray Ozen
llvmlistbot at llvm.org
Thu Aug 14 03:49:03 PDT 2025
https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/152525
>From 10926b1b3a7b60ef33770becdba74d6a27fd8527 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Thu, 7 Aug 2025 14:59:54 +0000
Subject: [PATCH 1/5] 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.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 43 ++++++++++++++-------
mlir/test/Dialect/LLVMIR/nvvm.mlir | 8 ++--
mlir/test/Target/LLVMIR/nvvmir.mlir | 15 ++++++-
3 files changed, 47 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 30df3b739e5ca..136984caa7724 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2995,30 +2995,47 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
// NVVM Griddepcontrol Ops
//===----------------------------------------------------------------------===//
-def NVVM_GriddepcontrolWaitOp : NVVM_IntrOp<"griddepcontrol.wait", [], 0> {
- let assemblyFormat = "attr-dict";
+def FDLWait : I32EnumAttrCase<"wait", 0>;
+def FDLLaunchDependent : I32EnumAttrCase<"launch_dependents", 1>;
+def FDLKind : I32EnumAttr<"FDLKind", "Fast dependenct launch kind",
+ [FDLWait, FDLLaunchDependent]> {
+ let genSpecializedAttr = 0;
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def FDLKindAttr : EnumAttr<NVVM_Dialect, FDLKind, "fdl_kind">;
+
+def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
let description = [{
- Causes the executing thread to wait until all prerequisite grids in flight
+ If the $kind attribute is set to `wait`, the instruction, it causes the
+ executing thread to wait until all prerequisite grids in flight
have completed and all the memory operations from the prerequisite grids
are performed and made visible to the current grid.
+ When the $kind is launch_dependents, it signals that specific dependents
+ the runtime system designated to react to this instruction can be scheduled
+ as soon as all other CTAs in the grid issue the same instruction or have
+ completed.
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
}];
-}
-def NVVM_GriddepcontrolLaunchDependentsOp
- : NVVM_IntrOp<"griddepcontrol.launch.dependents", [], 0> {
- let assemblyFormat = "attr-dict";
-
- let description = [{
- Signals that specific dependents the runtime system designated to react to
- this instruction can be scheduled as soon as all other CTAs in the grid
- issue the same instruction or have completed.
+ let arguments = (ins FDLKindAttr:$kind);
+ let assemblyFormat = "$kind attr-dict";
- [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
+ string llvmBuilder = [{
+ llvm::Intrinsic::ID id;
+ switch ($kind) {
+ case NVVM::FDLKind::wait:
+ id = llvm::Intrinsic::nvvm_griddepcontrol_wait;
+ break;
+ case NVVM::FDLKind::launch_dependents:
+ id = llvm::Intrinsic::nvvm_griddepcontrol_launch_dependents;
+ break;
+ }
+ createIntrinsicCall(builder, id);
}];
}
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index c7fa41c98ac92..cd14be5473432 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -535,15 +535,15 @@ func.func @wgmma_wait_group_sync_aligned() {
}
func.func @griddepcontrol_wait() {
- // CHECK: nvvm.griddepcontrol.wait
- nvvm.griddepcontrol.wait
+ // CHECK: nvvm.griddepcontrol wait
+ nvvm.griddepcontrol wait
return
}
func.func @griddepcontrol_launch_dependents()
{
- // CHECK: nvvm.griddepcontrol.launch.dependents
- nvvm.griddepcontrol.launch.dependents
+ // CHECK: nvvm.griddepcontrol launch_dependents
+ nvvm.griddepcontrol launch_dependents
return
}
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 5c2cfa4683104..6af347d6dcea3 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -766,7 +766,7 @@ llvm.func @nvvm_wgmma_wait_group_aligned() {
// CHECK-LABEL: @nvvm_griddepcontrol_wait
llvm.func @nvvm_griddepcontrol_wait() {
// CHECK: call void @llvm.nvvm.griddepcontrol.wait()
- nvvm.griddepcontrol.wait
+ nvvm.griddepcontrol wait
llvm.return
}
@@ -774,7 +774,7 @@ llvm.func @nvvm_griddepcontrol_wait() {
// CHECK-LABEL: @nvvm_griddepcontrol_launch_dependents
llvm.func @nvvm_griddepcontrol_launch_dependents() {
// CHECK: call void @llvm.nvvm.griddepcontrol.launch.dependents()
- nvvm.griddepcontrol.launch.dependents
+ nvvm.griddepcontrol launch_dependents
llvm.return
}
@@ -918,3 +918,14 @@ llvm.func @nvvm_dot_accumulate_2way(%a: vector<2xi16>, %b: vector<4xi8>, %c: i32
%7 = nvvm.dot.accumulate.2way %a <signed>, %b <signed>, %c {b_hi = true}: vector<2xi16>, vector<4xi8>
llvm.return
}
+
+// -----
+
+// CHECK-LABEL: @nvvm_pmevent
+llvm.func @nvvm_pmevent() {
+ // CHECK: call void @llvm.nvvm.pm.event.mask(i16 15000)
+ nvvm.pmevent mask = 15000
+ // CHECK: call void @llvm.nvvm.pm.event.mask(i16 4)
+ nvvm.pmevent mask = 4
+ llvm.return
+}
\ No newline at end of file
>From 1f8ec3e42f81ec3657d980b2a99abae8df04844a Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 8 Aug 2025 07:20:24 +0000
Subject: [PATCH 2/5] fx
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 16 ++++++++--------
1 file changed, 8 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 136984caa7724..7b70da78fd296 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2995,20 +2995,20 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
// NVVM Griddepcontrol Ops
//===----------------------------------------------------------------------===//
-def FDLWait : I32EnumAttrCase<"wait", 0>;
-def FDLLaunchDependent : I32EnumAttrCase<"launch_dependents", 1>;
+def GridDepActionWait : I32EnumAttrCase<"wait", 0>;
+def GridDepActionLaunchDependent : I32EnumAttrCase<"launch_dependents", 1>;
-def FDLKind : I32EnumAttr<"FDLKind", "Fast dependenct launch kind",
- [FDLWait, FDLLaunchDependent]> {
+def GridDepActionKind : I32EnumAttr<"GridDepKind", "Action kind for grid dependency control",
+ [GridDepActionWait, GridDepActionLaunchDependent]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::NVVM";
}
-def FDLKindAttr : EnumAttr<NVVM_Dialect, FDLKind, "fdl_kind">;
+def GridDepActionAttr : EnumAttr<NVVM_Dialect, GridDepActionKind, "grid_dep_action">;
def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
let description = [{
- If the $kind attribute is set to `wait`, the instruction, it causes the
+ If the $kind attribute is set to `wait`, it causes the
executing thread to wait until all prerequisite grids in flight
have completed and all the memory operations from the prerequisite grids
are performed and made visible to the current grid.
@@ -3021,11 +3021,11 @@ def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
}];
- let arguments = (ins FDLKindAttr:$kind);
+ let arguments = (ins GridDepActionAttr:$kind);
let assemblyFormat = "$kind attr-dict";
- string llvmBuilder = [{
+ string llvmBuilder = [{
llvm::Intrinsic::ID id;
switch ($kind) {
case NVVM::FDLKind::wait:
>From c0123335b79f20bd1d1d20bc8b06fcdf554fafc5 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 8 Aug 2025 07:26:49 +0000
Subject: [PATCH 3/5] add nl
---
mlir/test/Target/LLVMIR/nvvmir.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 6af347d6dcea3..5108f8bd942a0 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -928,4 +928,4 @@ llvm.func @nvvm_pmevent() {
// CHECK: call void @llvm.nvvm.pm.event.mask(i16 4)
nvvm.pmevent mask = 4
llvm.return
-}
\ No newline at end of file
+}
>From 265cdc7a5482dcbf15186ce575123267cc6d689e Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 8 Aug 2025 08:39:47 +0000
Subject: [PATCH 4/5] fx
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7b70da78fd296..7b864a7be2c5e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -3028,10 +3028,10 @@ def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
string llvmBuilder = [{
llvm::Intrinsic::ID id;
switch ($kind) {
- case NVVM::FDLKind::wait:
+ case NVVM::GridDepActionKind::wait:
id = llvm::Intrinsic::nvvm_griddepcontrol_wait;
break;
- case NVVM::FDLKind::launch_dependents:
+ case NVVM::GridDepActionKind::launch_dependents:
id = llvm::Intrinsic::nvvm_griddepcontrol_launch_dependents;
break;
}
>From b7ff33c11292db5d0446c9f012b097c379785f44 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Thu, 14 Aug 2025 10:48:43 +0000
Subject: [PATCH 5/5] fx
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7b864a7be2c5e..ba32cde78930f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2998,7 +2998,7 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
def GridDepActionWait : I32EnumAttrCase<"wait", 0>;
def GridDepActionLaunchDependent : I32EnumAttrCase<"launch_dependents", 1>;
-def GridDepActionKind : I32EnumAttr<"GridDepKind", "Action kind for grid dependency control",
+def GridDepActionKind : I32EnumAttr<"GridDepActionKind", "Action kind for grid dependency control",
[GridDepActionWait, GridDepActionLaunchDependent]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::NVVM";
More information about the Mlir-commits
mailing list