[Mlir-commits] [mlir] 6402abf - Add [mlir][nvvm] `mbarrier.test.wait`

[Guray Ozen]

Author: Guray Ozen
Date: 2023-06-30T11:14:46+02:00
New Revision: 6402abf3222a9ccd4069a96d52e811f315f3c4f2

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

LOG: Add [mlir][nvvm] `mbarrier.test.wait`

This work adds `mbarrier.test.wait` and `mbarrier.test.wait.shared` Ops in NVVM dialect. Since they are already implemented in the LLVM kernel, it only calls createIntrinsicCall.

Reviewed By: qcolombet

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

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 f28eb1d414f16b..7669ff1957afe1 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -249,6 +249,23 @@ def NVVM_MBarrierArriveNocompleteSharedOp : NVVM_Op<"mbarrier.arrive.nocomplete.
   let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
 }
 
+def NVVM_MBarrierTestWaitOp : NVVM_Op<"mbarrier.test.wait">,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_any:$addr, LLVM_Type:$token)> {
+  string llvmBuilder = [{
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_test_wait, {$addr, $token});
+  }];
+  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
+}
+
+def NVVM_MBarrierTestWaitSharedOp : NVVM_Op<"mbarrier.test.wait.shared">,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_shared:$addr, LLVM_Type:$token)> {
+  string llvmBuilder = [{
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_test_wait_shared, {$addr, $token});
+  }];
+  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
+}
 //===----------------------------------------------------------------------===//
 // NVVM synchronization op definitions
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 0e92318c38375a..1e1cd660a48e58 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -394,3 +394,16 @@ llvm.func private @mbarrier_arrive_nocomplete_shared(%barrier: !llvm.ptr<3>) {
   %0 = nvvm.mbarrier.arrive.nocomplete.shared %barrier, %count : !llvm.ptr<3>, i32  -> i64
   llvm.return
 }
+
+llvm.func private @mbarrier_test_wait(%barrier: !llvm.ptr, %token : i64) -> i1 {  
+  // CHECK:   nvvm.mbarrier.test.wait %{{.*}}
+  %isComplete = nvvm.mbarrier.test.wait %barrier, %token : !llvm.ptr, i64 -> i1
+  llvm.return %isComplete : i1
+}
+
+llvm.func private @mbarrier_test_wait_shared(%barrier: !llvm.ptr<3>, %token : i64) {
+  %count = nvvm.read.ptx.sreg.ntid.x : i32
+  // CHECK:   nvvm.mbarrier.test.wait.shared %{{.*}}
+  %isComplete = nvvm.mbarrier.test.wait.shared %barrier, %token : !llvm.ptr<3>, i64 -> i1
+  llvm.return
+}