[MLIR][NVVM] Add CTA Cluster barrier intrinsics for sm_90

This work adds CTA Cluster barrier intrinsics for sm_90 in NVVM dialect. They are already supported in LLVM core, so this work uses the existing intrinsics.

Differential Revision: https://reviews.llvm.org/D158720

Author: grypp

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

@@ -483,6 +483,34 @@ def NVVM_Barrier0Op : NVVM_Op<"barrier0"> {
   let assemblyFormat = "attr-dict";
 }
 
+def NVVM_ClusterArriveOp : NVVM_Op<"cluster.arrive"> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive);
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
+def NVVM_ClusterArriveRelaxedOp : NVVM_Op<"cluster.arrive.relaxed"> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive_relaxed);
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
+def NVVM_ClusterWaitOp : NVVM_Op<"cluster.wait"> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_wait);
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
+def NVVM_FenceScClusterOp : NVVM_Op<"fence.sc.cluster"> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_fence_sc_cluster);
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
 def ShflKindBfly : I32EnumAttrCase<"bfly", 0>;
 def ShflKindUp   : I32EnumAttrCase<"up", 1>;
 def ShflKindDown : I32EnumAttrCase<"down", 2>;

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -43,6 +43,34 @@ func.func @llvm_nvvm_barrier0() {
   llvm.return
 }
 
+// CHECK-LABEL: @llvm_nvvm_cluster_arrive
+func.func @llvm_nvvm_cluster_arrive() {
+  // CHECK: nvvm.cluster.arrive
+  nvvm.cluster.arrive
+  llvm.return
+}
+
+// CHECK-LABEL: @llvm_nvvm_cluster_arrive_relaxed
+func.func @llvm_nvvm_cluster_arrive_relaxed() {
+  // CHECK: nvvm.cluster.arrive.relaxed
+  nvvm.cluster.arrive.relaxed
+  llvm.return
+}
+
+// CHECK-LABEL: @llvm_nvvm_cluster_wait
+func.func @llvm_nvvm_cluster_wait() {
+  // CHECK: nvvm.cluster.wait
+  nvvm.cluster.wait
+  llvm.return
+}
+
+// CHECK-LABEL: @llvm_nvvm_fence_sc_cluster
+func.func @llvm_nvvm_fence_sc_cluster() {
+  // CHECK: nvvm.fence.sc.cluster
+  nvvm.fence.sc.cluster
+  llvm.return
+}
+
 // CHECK-LABEL: @nvvm_shfl
 func.func @nvvm_shfl(
     %arg0 : i32, %arg1 : i32, %arg2 : i32,