From 37145d9575a5ef2c2bcf506f36271813410d99aa Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen@gmail.com>
Date: Thu, 29 Jun 2023 12:09:31 +0200
Subject: [PATCH] [mlir] Add result to mbarrier.arrive

`mbarrier.arrive` returns token. This PR adds result to these ops.

Reviewed By: qcolombet

Differential Revision: https://reviews.llvm.org/D154059
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 20 ++++++++++++--------
 mlir/test/Dialect/LLVMIR/nvvm.mlir          |  8 ++++----
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index dda72f4..f28eb1d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -214,35 +214,39 @@ def NVVM_MBarrierInvalSharedOp : NVVM_Op<"mbarrier.inval.shared">,
 }
 
 def NVVM_MBarrierArriveOp : NVVM_Op<"mbarrier.arrive">,
+  Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_any:$addr)> {
   string llvmBuilder = [{
-      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive, {$addr});
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive, {$addr});
   }];
-  let assemblyFormat = "$addr attr-dict `:` type(operands)";
+  let assemblyFormat = "$addr attr-dict `:` type($addr) `->` type($res)";
 }
 
 def NVVM_MBarrierArriveSharedOp : NVVM_Op<"mbarrier.arrive.shared">,
+  Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_shared:$addr)> {
   string llvmBuilder = [{
-      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_shared, {$addr});
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_shared, {$addr});
   }];
-  let assemblyFormat = "$addr attr-dict `:` type(operands)";
+  let assemblyFormat = "$addr attr-dict `:` type($addr) `->` type($res)";
 }
 
 def NVVM_MBarrierArriveNocompleteOp : NVVM_Op<"mbarrier.arrive.nocomplete">,
+  Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_any:$addr, I32:$count)> {
   string llvmBuilder = [{
-      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete, {$addr, $count});
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete, {$addr, $count});
   }];
-  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands)";
+  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
 }
 
 def NVVM_MBarrierArriveNocompleteSharedOp : NVVM_Op<"mbarrier.arrive.nocomplete.shared">,
+  Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_shared:$addr, I32:$count)> {
   string llvmBuilder = [{
-      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete_shared, {$addr, $count});
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete_shared, {$addr, $count});
   }];
-  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands)";
+  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 30632ad..0e92318 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -371,26 +371,26 @@ llvm.func private @mbarrier_inval_shared(%barrier: !llvm.ptr<3>) {
 
 llvm.func private @mbarrier_arrive(%barrier: !llvm.ptr) {
   // CHECK:   nvvm.mbarrier.arrive %{{.*}} : !llvm.ptr
-  nvvm.mbarrier.arrive %barrier : !llvm.ptr
+  %0 = nvvm.mbarrier.arrive %barrier : !llvm.ptr  -> i64
   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>
+  %0 = nvvm.mbarrier.arrive.shared %barrier : !llvm.ptr<3> -> i64
   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
+  %0 = nvvm.mbarrier.arrive.nocomplete %barrier, %count : !llvm.ptr, i32 -> i64
   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
+  %0 = nvvm.mbarrier.arrive.nocomplete.shared %barrier, %count : !llvm.ptr<3>, i32  -> i64
   llvm.return
 }
-- 
2.7.4