[Mlir-commits] [mlir] 040805d - [mlir] Add bar.warp.sync to NVVM

[Guray Ozen]

Author: Guray Ozen
Date: 2022-10-06T12:27:53+02:00
New Revision: 040805dc4748f5b21261acdf562a424dd2f565c8

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

LOG: [mlir] Add bar.warp.sync to NVVM

It adds the missing `bar.warp.sync` to the nvvm dialect. It is a barrier to synchronize for threads in a warp.

Reviewed By: ftynse

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

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 861ae0cc62622..9a94cf745b7db 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -168,6 +168,15 @@ def NVVM_VoteBallotOp :
   let hasCustomAssemblyFormat = 1;
 }
 
+def NVVM_SyncWarpOp : 
+  NVVM_Op<"bar.warp.sync">,
+  Arguments<(ins LLVM_Type:$mask)> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_bar_warp_sync, {$mask});
+  }];
+  let assemblyFormat = "$mask attr-dict `:` type($mask)";
+}
+
 
 def NVVM_CpAsyncOp : NVVM_Op<"cp.async.shared.global">,
   Arguments<(ins LLVM_i8Ptr_shared:$dst,

diff  --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 283127382df1b..e43cc7806e5cc 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -78,6 +78,13 @@ func.func @nvvm_vote(%arg0 : i32, %arg1 : i1) -> i32 {
   llvm.return %0 : i32
 }
 
+// CHECK-LABEL: @llvm_nvvm_bar_warp_sync
+func.func @llvm_nvvm_bar_warp_sync(%mask : i32) {
+  // CHECK: nvvm.bar.warp.sync %{{.*}}
+  nvvm.bar.warp.sync %mask : i32
+  llvm.return
+}  
+
 // CHECK-LABEL: @nvvm_mma_m8n8k4_row_col_f32_f32
 func.func @nvvm_mma_m8n8k4_row_col_f32_f32(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
                %b0 : vector<2xf16>, %b1 : vector<2xf16>,