[MLIR][NVVM] Add support for griddepcontrol Ops (#124603)

Adds `griddepcontrol.wait` and `griddepcontrol.launch.dependents`
MLIR Ops to generate griddepcontrol instructions.

`griddepcontrol` - Allows dependent and prerequisite grids as defined by
the runtime to control execution in the following ways:

- `griddepcontrol.wait` - 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.

- `griddepcontrol.launch.dependents` - 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.

PTX Spec Reference:

https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol

Author: Wolfram70

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

@@ -2512,6 +2512,33 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM Griddepcontrol Ops
+//===----------------------------------------------------------------------===//
+
+def NVVM_GriddepcontrolWaitOp : NVVM_IntrOp<"griddepcontrol.wait", [], 0> {
+  let assemblyFormat = "attr-dict";
+
+  let description = [{
+    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.
+    [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.
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
+  }];
+}
+
 def NVVM_Exit : NVVM_Op<"exit"> {
   let summary = "Exit Op";
   let description = [{

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -509,6 +509,19 @@ func.func @wgmma_wait_group_sync_aligned() {
   return
 }
 
+func.func @griddepcontrol_wait() {
+  // CHECK: nvvm.griddepcontrol.wait
+  nvvm.griddepcontrol.wait
+  return
+}
+
+func.func @griddepcontrol_launch_dependents()
+{
+  // CHECK: nvvm.griddepcontrol.launch.dependents
+  nvvm.griddepcontrol.launch.dependents
+  return
+}
+
 // -----
 
 // Just check these don't emit errors.

mlir/test/Target/LLVMIR/nvvmir.mlir

@@ -757,3 +757,19 @@ llvm.func @nvvm_wgmma_wait_group_aligned() {
   nvvm.wgmma.wait.group.sync.aligned 20
   llvm.return
 }
+
+// -----
+// CHECK-LABEL: @nvvm_griddepcontrol_wait
+llvm.func @nvvm_griddepcontrol_wait() {
+  // CHECK: call void @llvm.nvvm.griddepcontrol.wait()
+  nvvm.griddepcontrol.wait
+  llvm.return
+}
+
+// -----
+// CHECK-LABEL: @nvvm_griddepcontrol_launch_dependents
+llvm.func @nvvm_griddepcontrol_launch_dependents() {
+  // CHECK: call void @llvm.nvvm.griddepcontrol.launch.dependents()
+  nvvm.griddepcontrol.launch.dependents
+  llvm.return
+}