diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -19,6 +19,8 @@ def LLVM_i8Ptr_global : LLVM_IntPtrBase<8, 1>; def LLVM_i8Ptr_shared : LLVM_IntPtrBase<8, 3>; +def LLVM_i64ptr_global : LLVM_IntPtrBase<64, 1>; +def LLVM_i64ptr_shared : LLVM_IntPtrBase<64, 3>; //===----------------------------------------------------------------------===// // NVVM dialect definitions @@ -188,7 +190,7 @@ /// mbarrier.init instruction with shared pointer type def NVVM_MBarrierInitSharedOp : NVVM_Op<"mbarrier.init.shared">, - Arguments<(ins LLVM_i8Ptr_shared:$addr, I32:$count)> { + Arguments<(ins LLVM_i64Ptr_shared:$addr, I32:$count)> { string llvmBuilder = [{ createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_init_shared, {$addr, $count}); }]; @@ -204,13 +206,45 @@ } def NVVM_MBarrierInvalSharedOp : NVVM_Op<"mbarrier.inval.shared">, - Arguments<(ins LLVM_i8Ptr_shared:$addr)> { + Arguments<(ins LLVM_i64Ptr_shared:$addr)> { string llvmBuilder = [{ createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_inval_shared, {$addr}); }]; let assemblyFormat = "$addr attr-dict `:` type(operands)"; } +def NVVM_MBarrierArriveOp : NVVM_Op<"mbarrier.arrive">, + Arguments<(ins LLVM_AnyPointer:$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_AnyPointer:$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 --- a/mlir/test/Dialect/LLVMIR/nvvm.mlir +++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir @@ -339,6 +339,7 @@ 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,32 @@ nvvm.mbarrier.inval.shared %barrier : !llvm.ptr<3> llvm.return } + +llvm.func private @mbarrier_arrive(%barrier: !llvm.ptr) { + // CHECK: nvvm.mbarrier.init %{{.*}} : !llvm.ptr + nvvm.mbarrier.arrive %barrier : !llvm.ptr + llvm.return +} + + +llvm.func private @mbarrier_arrive_shared(%barrier: !llvm.ptr<3>) { + // CHECK: nvvm.mbarrier.init %{{.*}} : !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.init %{{.*}} : !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.init %{{.*}} : !llvm.ptr<3> + nvvm.mbarrier.arrive.nocomplete.shared %barrier, %count : !llvm.ptr<3>, i32 + llvm.return +}