From d0233ccbf8807f86edef3aa1fb90b3786a9ce017 Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen@gmail.com>
Date: Fri, 16 Jun 2023 13:52:29 +0200
Subject: [PATCH] [mlir][nvvm] Introduce `mbarrier.arrive`

It introduces `mbarrier.arrive` that are in ptx78.

Reviewed By: qcolombet

Differential Revision: https://reviews.llvm.org/D153021
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 32 +++++++++++++++++++++++++++++
 mlir/test/Dialect/LLVMIR/nvvm.mlir          | 27 ++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 44aa4aa..dda72f4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -213,6 +213,38 @@ def NVVM_MBarrierInvalSharedOp : NVVM_Op<"mbarrier.inval.shared">,
   let assemblyFormat = "$addr attr-dict `:` type(operands)";
 }
 
+def NVVM_MBarrierArriveOp : NVVM_Op<"mbarrier.arrive">,
+  Arguments<(ins LLVM_i64ptr_any:$addr)> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive, {$addr});
+  }];
+  let assemblyFormat = "$addr attr-dict `:` type(operands)";
+}
+
+def NVVM_MBarrierArriveSharedOp : NVVM_Op<"mbarrier.arrive.shared">,
+  Arguments<(ins LLVM_i64ptr_shared:$addr)> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_shared, {$addr});
+  }];
+  let assemblyFormat = "$addr attr-dict `:` type(operands)";
+}
+
+def NVVM_MBarrierArriveNocompleteOp : NVVM_Op<"mbarrier.arrive.nocomplete">,
+  Arguments<(ins LLVM_i64ptr_any:$addr, I32:$count)> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete, {$addr, $count});
+  }];
+  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands)";
+}
+
+def NVVM_MBarrierArriveNocompleteSharedOp : NVVM_Op<"mbarrier.arrive.nocomplete.shared">,
+  Arguments<(ins LLVM_i64ptr_shared:$addr, I32:$count)> {
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete_shared, {$addr, $count});
+  }];
+  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands)";
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM synchronization op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index a32b33d..30632ad 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -339,6 +339,7 @@ llvm.func @redux_sync(%value : i32, %offset : i32) -> i32 {
 func.func private @expected_llvm_func() attributes { nvvm.kernel }
 
 // -----
+
 llvm.func private @mbarrier_init_generic(%barrier: !llvm.ptr) {
   %count = nvvm.read.ptx.sreg.ntid.x : i32
   // CHECK:   nvvm.mbarrier.init %{{.*}}, %{{.*}} : !llvm.ptr, i32
@@ -367,3 +368,29 @@ llvm.func private @mbarrier_inval_shared(%barrier: !llvm.ptr<3>) {
   nvvm.mbarrier.inval.shared %barrier : !llvm.ptr<3>
   llvm.return
 }
+
+llvm.func private @mbarrier_arrive(%barrier: !llvm.ptr) {
+  // CHECK:   nvvm.mbarrier.arrive %{{.*}} : !llvm.ptr
+  nvvm.mbarrier.arrive %barrier : !llvm.ptr
+  llvm.return
+}
+
+llvm.func private @mbarrier_arrive_shared(%barrier: !llvm.ptr<3>) {
+  // CHECK:   nvvm.mbarrier.arrive.shared %{{.*}} : !llvm.ptr<3>
+  nvvm.mbarrier.arrive.shared %barrier : !llvm.ptr<3>
+  llvm.return
+}
+
+llvm.func private @mbarrier_arrive_nocomplete(%barrier: !llvm.ptr) {
+  %count = nvvm.read.ptx.sreg.ntid.x : i32
+  // CHECK:   nvvm.mbarrier.arrive.nocomplete %{{.*}} : !llvm.ptr
+  nvvm.mbarrier.arrive.nocomplete %barrier, %count : !llvm.ptr, i32
+  llvm.return
+}
+
+llvm.func private @mbarrier_arrive_nocomplete_shared(%barrier: !llvm.ptr<3>) {
+  %count = nvvm.read.ptx.sreg.ntid.x : i32
+  // CHECK:   nvvm.mbarrier.arrive.nocomplete.shared %{{.*}} : !llvm.ptr<3>
+  nvvm.mbarrier.arrive.nocomplete.shared %barrier, %count : !llvm.ptr<3>, i32
+  llvm.return
+}
-- 
2.7.4