[Mlir-commits] [mlir] 79679be - [MLIR][NVVM] Add asynchronous store Ops (#210931)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 17 02:11:35 PDT 2026
Author: Srinivasa Ravi
Date: 2026-08-17T14:41:30+05:30
New Revision: 79679be4d194bf28f25e93d2782384cd7f2606b5
URL: https://github.com/llvm/llvm-project/commit/79679be4d194bf28f25e93d2782384cd7f2606b5
DIFF: https://github.com/llvm/llvm-project/commit/79679be4d194bf28f25e93d2782384cd7f2606b5.diff
LOG: [MLIR][NVVM] Add asynchronous store Ops (#210931)
This change adds the `store.async.global` and `store.async.shared`
ops to the NVVM dialect to perform asynchronous stores to global
or shared-cluster address spaces.
PTX Spec References:
1.
[`st.async`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-st-async)
2.
[`multimem.st.async`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-multimem-st-async)
Added:
mlir/test/Target/LLVMIR/nvvm/store_async_global.mlir
mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
mlir/test/Target/LLVMIR/nvvm/store_async_shared.mlir
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 46c38bcb5475d..6241606122f43 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5124,6 +5124,57 @@ def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// NVVM Asynchronous Store Ops
+//===----------------------------------------------------------------------===//
+
+def NVVM_AsyncStoreGlobalOp: NVVM_VoidIntrinsicOp<"store.async.global",
+ [NVVMRequiresSM<100>]> {
+ let description = [{
+ Performs an asynchronous store to global memory to the address given by
+ `addr`.
+ The `value` operand specifies the value to store.
+ The `scope` operand specifies the scope of the store and must be one of the
+ following:
+ - `sys`: Synchronization with all threads in the system.
+ - `gpu`: Synchronization with all threads in the same GPU.
+ The `multimem` operand specifies whether the store is performed on a
+ multimem address.
+ The `mmio` operand specifies whether this is an MMIO operation.
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-st-async)
+ }];
+ let summary = "Asynchronous Store Op to global memory";
+ let arguments = (ins LLVM_PointerGlobal:$addr,
+ AnyTypeOf<[I8, I16, I32, I64]>:$value,
+ MemScopeKindAttr:$scope,
+ DefaultValuedAttr<BoolAttr, "false">:$multimem,
+ DefaultValuedAttr<BoolAttr, "false">:$mmio);
+ let results = (outs );
+ let assemblyFormat = "$addr `,` $value attr-dict `:` type($addr) `,` type($value)";
+ let hasVerifier = 1;
+}
+
+def NVVM_AsyncStoreSharedOp:
+ LLVM_IntrOpBase<NVVM_Dialect, "store.async.shared", "nvvm_st_async",
+ /*overloadedResults=*/[], /*overloadedOperands=*/[1],
+ [NVVMRequiresSM<90>], /*numResults=*/0> {
+ let description = [{
+ Performs an asynchronous store to shared cluster memory to the address
+ given by `addr`.
+ The `value` operand specifies the value to store.
+ The `mbarrier` operand specifies the mbarrier object which signals the
+ completion of the store.
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-st-async)
+ }];
+ let summary = "Asynchronous Store Op to shared cluster memory";
+ let arguments = (ins LLVM_PointerSharedCluster:$addr,
+ AnyTypeOf<[I32, I64, I128]>:$value,
+ LLVM_PointerSharedCluster:$mbarrier);
+ let assemblyFormat = "$addr `,` $value `,` `mbarrier` `=` $mbarrier attr-dict `:` type($addr) `,` type($value) `,` type($mbarrier)";
+}
+
def NVVM_Exit : NVVM_Op<"exit"> {
let summary = "Exit Op";
let description = [{
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 31fed0b25990d..3cb03e297d03c 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -671,6 +671,23 @@ LogicalResult BulkStoreOp::verify() {
return success();
}
+LogicalResult AsyncStoreGlobalOp::verify() {
+ NVVM::MemScopeKind scope = getScope();
+ bool isMmio = getMmio();
+ bool isMultimem = getMultimem();
+
+ if (scope != MemScopeKind::SYS && scope != MemScopeKind::GPU)
+ return emitOpError("scope must be either SYS or GPU");
+
+ if (isMmio && scope != MemScopeKind::SYS)
+ return emitOpError("mmio is only supported for SYS scope");
+
+ if (isMmio && isMultimem)
+ return emitOpError("multimem is not supported with mmio");
+
+ return success();
+}
+
LogicalResult PMEventOp::verify() {
auto eventId = getEventId();
auto maskedEventId = getMaskedEventId();
@@ -3665,6 +3682,29 @@ DivFOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
{mt.lookupValue(thisOp.getLhs()), mt.lookupValue(thisOp.getRhs())}};
}
+mlir::NVVM::IDArgPair AsyncStoreGlobalOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ using IDArgPair = mlir::NVVM::IDArgPair;
+ auto thisOp = cast<NVVM::AsyncStoreGlobalOp>(op);
+ mlir::NVVM::MemScopeKind scope = thisOp.getScope();
+ bool isMmio = thisOp.getMmio();
+
+ llvm::Value *addr = mt.lookupValue(thisOp.getAddr());
+ llvm::Value *value = mt.lookupValue(thisOp.getValue());
+ llvm::Value *isMultimem = builder.getInt1(thisOp.getMultimem());
+
+ if (scope == MemScopeKind::SYS) {
+ return isMmio ? IDArgPair(llvm::Intrinsic::nvvm_st_async_mmio_sys,
+ {addr, value})
+ : IDArgPair(llvm::Intrinsic::nvvm_st_async_sys,
+ {addr, value, isMultimem});
+ } else if (scope == MemScopeKind::GPU) {
+ return IDArgPair(llvm::Intrinsic::nvvm_st_async_gpu,
+ {addr, value, isMultimem});
+ }
+ llvm_unreachable("unsupported scope for AsyncStoreGlobalOp");
+}
+
mlir::NVVM::IDArgPair
PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async_global.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_global.mlir
new file mode 100644
index 0000000000000..11cbf424167f1
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async_global.mlir
@@ -0,0 +1,42 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @st_async_global_sys
+llvm.func @st_async_global_sys(%addr: !llvm.ptr<1>, %value: i32) {
+ // CHECK: call void @llvm.nvvm.st.async.sys.i32(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, /* isMultimem= */ i1 false)
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<sys>} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// CHECK-LABEL: define void @st_async_global_gpu
+llvm.func @st_async_global_gpu(%addr: !llvm.ptr<1>, %value: i32) {
+ // CHECK: call void @llvm.nvvm.st.async.gpu.i32(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, /* isMultimem= */ i1 false)
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<gpu>} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// CHECK-LABEL: define void @st_async_global_multimem
+llvm.func @st_async_global_multimem(%addr: !llvm.ptr<1>, %value: i32) {
+ // CHECK: call void @llvm.nvvm.st.async.sys.i32(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, /* isMultimem= */ i1 true)
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<sys>, multimem = true} : !llvm.ptr<1>, i32
+ // CHECK: call void @llvm.nvvm.st.async.gpu.i32(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, /* isMultimem= */ i1 true)
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<gpu>, multimem = true} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// CHECK-LABEL: define void @st_async_global_mmio
+llvm.func @st_async_global_mmio(%addr: !llvm.ptr<1>, %value: i32) {
+ // CHECK: call void @llvm.nvvm.st.async.mmio.sys.i32(ptr addrspace(1) %{{.*}}, i32 %{{.*}})
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<sys>, mmio = true} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// CHECK-LABEL: define void @st_async_global_types
+llvm.func @st_async_global_types(%addr: !llvm.ptr<1>, %v8: i8, %v16: i16, %v64: i64) {
+ // CHECK: call void @llvm.nvvm.st.async.gpu.i8(ptr addrspace(1) %{{.*}}, i8 %{{.*}}, /* isMultimem= */ i1 false)
+ nvvm.store.async.global %addr, %v8 {scope = #nvvm.mem_scope<gpu>} : !llvm.ptr<1>, i8
+ // CHECK: call void @llvm.nvvm.st.async.gpu.i16(ptr addrspace(1) %{{.*}}, i16 %{{.*}}, /* isMultimem= */ i1 false)
+ nvvm.store.async.global %addr, %v16 {scope = #nvvm.mem_scope<gpu>} : !llvm.ptr<1>, i16
+ // CHECK: call void @llvm.nvvm.st.async.gpu.i64(ptr addrspace(1) %{{.*}}, i64 %{{.*}}, /* isMultimem= */ i1 false)
+ nvvm.store.async.global %addr, %v64 {scope = #nvvm.mem_scope<gpu>} : !llvm.ptr<1>, i64
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
new file mode 100644
index 0000000000000..18fb78ff4341c
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
+
+llvm.func @st_async_global_invalid_scope(%addr: !llvm.ptr<1>, %value: i32) {
+ // expected-error @below {{scope must be either SYS or GPU}}
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<cta>} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_global_mmio_non_sys(%addr: !llvm.ptr<1>, %value: i32) {
+ // expected-error @below {{mmio is only supported for SYS scope}}
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<gpu>, mmio = true} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_global_mmio_multimem(%addr: !llvm.ptr<1>, %value: i32) {
+ // expected-error @below {{multimem is not supported with mmio}}
+ nvvm.store.async.global %addr, %value {scope = #nvvm.mem_scope<sys>, mmio = true, multimem = true} : !llvm.ptr<1>, i32
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async_shared.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_shared.mlir
new file mode 100644
index 0000000000000..a7aa611c58173
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async_shared.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @st_async_shared_cluster
+llvm.func @st_async_shared_cluster(%addr: !llvm.ptr<7>, %value: i32, %mbar: !llvm.ptr<7>) {
+ // CHECK: call void @llvm.nvvm.st.async.i32(ptr addrspace(7) %{{.*}}, i32 %{{.*}}, ptr addrspace(7) %{{.*}})
+ nvvm.store.async.shared %addr, %value, mbarrier = %mbar : !llvm.ptr<7>, i32, !llvm.ptr<7>
+ llvm.return
+}
+
+// CHECK-LABEL: define void @st_async_shared_cluster_types
+llvm.func @st_async_shared_cluster_types(%addr: !llvm.ptr<7>, %v64: i64, %v128: i128, %mbar: !llvm.ptr<7>) {
+ // CHECK: call void @llvm.nvvm.st.async.i64(ptr addrspace(7) %{{.*}}, i64 %{{.*}}, ptr addrspace(7) %{{.*}})
+ nvvm.store.async.shared %addr, %v64, mbarrier = %mbar : !llvm.ptr<7>, i64, !llvm.ptr<7>
+ // CHECK: call void @llvm.nvvm.st.async.i128(ptr addrspace(7) %{{.*}}, i128 %{{.*}}, ptr addrspace(7) %{{.*}})
+ nvvm.store.async.shared %addr, %v128, mbarrier = %mbar : !llvm.ptr<7>, i128, !llvm.ptr<7>
+ llvm.return
+}
More information about the Mlir-commits
mailing list