[Mlir-commits] [mlir] de4d742 - [MLIR][NVVM] Add CTA Cluster barrier intrinsics for sm_90

[Guray Ozen]

Author: Guray Ozen
Date: 2023-09-01T16:39:52+02:00
New Revision: de4d742ae1cd21dc4141cdbd47f23a2bc37e6d42

URL: https://github.com/llvm/llvm-project/commit/de4d742ae1cd21dc4141cdbd47f23a2bc37e6d42
DIFF: https://github.com/llvm/llvm-project/commit/de4d742ae1cd21dc4141cdbd47f23a2bc37e6d42.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
    mlir/test/Dialect/LLVMIR/nvvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 55aa820314e89b..1eff22226e6669 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/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>;

diff  --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 6dce8eafc29d54..39516b5090d07b 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/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,