[Mlir-commits] [mlir] [MLIR][NVVM]: Add cp.async.mbarrier.arrive Op (PR #74241)

[Guray Ozen]

================
@@ -592,6 +592,42 @@ def NVVM_CpAsyncWaitGroupOp : NVVM_Op<"cp.async.wait.group">,
   let assemblyFormat = "$n attr-dict";
 }
 
+def NVVM_CpAsyncMBarrierArriveOp : NVVM_Op<"cp.async.mbarrier.arrive"> {
+  let summary = "NVVM Dialect Op for cp.async.mbarrier.arrive";
+  let description = [{
+    The `cp.async.mbarrier.arrive` Op makes the mbarrier object track
+    all prior cp.async operations initiated by the executing thread.
+    The `addr` operand specifies the address of the mbarrier object
+    in generic or shared memory space. The `noinc` attr impacts how
+    how the mbarrier's state is updated.
+    [For more information, refer PTX ISA]
+    (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive)
+  }];
+  let assemblyFormat = "$addr attr-dict `:` type(operands)";
+
+  let arguments = (ins
+    AnyTypeOf<[LLVM_AnyPointer, LLVM_PointerShared]>:$addr,
+    DefaultValuedAttr<I1Attr, "0">:$noinc);
+
+  string llvmBuilder = [{
+    unsigned addressSpace =
+      llvm::cast<LLVM::LLVMPointerType>(op.getAddr().getType()).getAddressSpace();
+    bool isShared = addressSpace == NVVM::kSharedMemorySpace;
+
+    llvm::Intrinsic::ID intId;
+    if ($noinc) {
+      intId = isShared ?
+        llvm::Intrinsic::nvvm_cp_async_mbarrier_arrive_noinc_shared :
+        llvm::Intrinsic::nvvm_cp_async_mbarrier_arrive_noinc;
+    } else {
+      intId = isShared ?
+        llvm::Intrinsic::nvvm_cp_async_mbarrier_arrive_shared :
+        llvm::Intrinsic::nvvm_cp_async_mbarrier_arrive;
+    }
+    createIntrinsicCall(builder, intId, {$addr});
+  }];
+}
+
----------------
grypp wrote:

> What is the nicest way to emit a NVVM operation for producers? How to avoid extra if/else or switch for them?

Chatted with @durga4github about this today. We will plan a design in the coming weeks. We can all meet up and do this together if you have time.

Let's keep NVVM as is for time being. We can do refactoring if it's necessary. The only user of the H100 NVVM dialect is now the NVGPU dialect. Even if we want refactoring, I expect it will go very smoothly.

https://github.com/llvm/llvm-project/pull/74241