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

[guraypp]

Summary

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.

[guraypp]

fix typo

[qcolombet]

BTW, LGTM :D.

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

 
 def NVVM_Barrier0Op : NVVM_Op<"barrier0"> {
   string llvmBuilder = [{
       createIntrinsicCall(builder, llvm::Intrinsic::nvvm_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>;
 def ShflKindIdx  : I32EnumAttrCase<"idx", 3>;
 
 /// Enum attribute of the different shuffle kinds.
 def ShflKind : I32EnumAttr<"ShflKind", "NVVM shuffle kind",
   [ShflKindBfly, ShflKindUp, ShflKindDown, ShflKindIdx]> {

mlir/test/Dialect/LLVMIR/nvvm.mlir

 
 // CHECK-LABEL: @llvm_nvvm_barrier0
 func.func @llvm_nvvm_barrier0() {
   // CHECK: nvvm.barrier0
   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,
     %arg3 : i32, %arg4 : f32) -> i32 {
   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32 -> i32
   %0 = nvvm.shfl.sync bfly %arg0, %arg3, %arg1, %arg2 : i32 -> i32
   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 -> f32
   %1 = nvvm.shfl.sync bfly %arg0, %arg4, %arg1, %arg2 : f32 -> f32