[NVPTX] Add intrinsics for wgmma.fence PTX instructions

Adds NVVM intrinsics and NVPTX codegen for:

- wgmma.fence (https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-fence)
- wgmma.commit_group (https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-commit-group)
- wgmma.wait_group (https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-wait-group)

Author: Wolfram70

llvm/include/llvm/IR/IntrinsicsNVVM.td

@@ -4805,6 +4805,21 @@ def int_nvvm_redux_sync_or : ClangBuiltin<"__nvvm_redux_sync_or">,
   Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty],
             [IntrConvergent, IntrInaccessibleMemOnly, IntrNoCallback]>;
 
+//
+// WGMMA instructions
+//
+// wgmma.fence.sync.aligned;
+def int_nvvm_wgmma_fence_sync_aligned 
+  : Intrinsic<[], [], [IntrConvergent], "llvm.nvvm.wgmma.fence.sync.aligned">;
+
+// wgmma.commit_group.sync.aligned;
+def int_nvvm_wgmma_commit_group_sync_aligned
+  : Intrinsic<[], [], [IntrConvergent], "llvm.nvvm.wgmma.commit_group.sync.aligned">;
+
+// wgmma.wait_group.sync.aligned N;
+def int_nvvm_wgmma_wait_group_sync_aligned
+  : Intrinsic<[], [llvm_i32_ty], [IntrConvergent, ImmArg<ArgIndex<0>>], "llvm.nvvm.wgmma.wait_group.sync.aligned">;
+
 //
 // WMMA instructions
 //

llvm/lib/Target/NVPTX/NVPTXIntrinsics.td

@@ -7484,4 +7484,16 @@ defm INT_SET_MAXNREG_DEC : SET_MAXNREG<"dec", int_nvvm_setmaxnreg_dec_sync_align
 
 } // isConvergent
 
+//
+// WGMMA instructions
+//
+def INT_NVVM_WGMMA_FENCE_SYNC_ALIGNED : NVPTXInst<(outs), (ins), "wgmma.fence.sync.aligned;",
+                             [(int_nvvm_wgmma_fence_sync_aligned)]>, Requires<[hasSM90a, hasPTX<80>]>;
+
+def INT_NVVM_WGMMA_COMMIT_GROUP_SYNC_ALIGNED : NVPTXInst<(outs), (ins), "wgmma.commit_group.sync.aligned;",
+                             [(int_nvvm_wgmma_commit_group_sync_aligned)]>, Requires<[hasSM90a, hasPTX<80>]>;
+
+def INT_NVVM_WGMMA_WAIT_GROUP_SYNC_ALIGNED : NVPTXInst<(outs), (ins i32imm:$n), "wgmma.wait_group.sync.aligned \t$n;",
+                             [(int_nvvm_wgmma_wait_group_sync_aligned timm:$n)]>, Requires<[hasSM90a, hasPTX<80>]>;
+
 def INT_EXIT : NVPTXInst<(outs), (ins), "exit;", [(int_nvvm_exit)]>;

llvm/test/CodeGen/NVPTX/wgmma-commit-group.ll

@@ -0,0 +1,15 @@
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | FileCheck %s
+; RUN: %if ptxas-12.0 %{ llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | %ptxas-verify -arch=sm_90a %}
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare void @llvm.nvvm.wgmma.commit_group.sync.aligned()
+
+define void @test_wgmma_commit_group_sync_aligned() {
+  ; CHECK-LABEL:  test_wgmma_commit_group_sync_aligned(
+  ; CHECK:        // %bb.0:
+  ; CHECK-NEXT:     wgmma.commit_group.sync.aligned;
+  ; CHECK-NEXT:     ret;
+  call void @llvm.nvvm.wgmma.commit_group.sync.aligned()
+  ret void
+}

llvm/test/CodeGen/NVPTX/wgmma-fence.ll

@@ -0,0 +1,15 @@
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | FileCheck %s
+; RUN: %if ptxas-12.0 %{ llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | %ptxas-verify -arch=sm_90a %}
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare void @llvm.nvvm.wgmma.fence.sync.aligned()
+
+define void @test_wgmma_fence_sync_aligned() {
+  ; CHECK-LABEL:  test_wgmma_fence_sync_aligned(
+  ; CHECK:        // %bb.0:
+  ; CHECK-NEXT:     wgmma.fence.sync.aligned;
+  ; CHECK-NEXT:     ret;
+  call void @llvm.nvvm.wgmma.fence.sync.aligned()
+  ret void
+}

llvm/test/CodeGen/NVPTX/wgmma-wait-group.ll

@@ -0,0 +1,15 @@
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | FileCheck %s
+; RUN: %if ptxas-12.0 %{ llc < %s -march=nvptx64 -mcpu=sm_90a -mattr=+ptx80 | %ptxas-verify -arch=sm_90a %}
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare void @llvm.nvvm.wgmma.wait_group.sync.aligned(i32)
+
+define void @test_wgmma_wait_group_sync_aligned() {
+  ; CHECK-LABEL:  test_wgmma_wait_group_sync_aligned(
+  ; CHECK:        // %bb.0:
+  ; CHECK-NEXT:     wgmma.wait_group.sync.aligned   10;
+  ; CHECK-NEXT:     ret;
+  call void @llvm.nvvm.wgmma.wait_group.sync.aligned(i32 10)
+  ret void
+}