[Mlir-commits] [mlir] [MLIR][NVVM] Add asynchronous store Op (PR #210931)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 21 03:19:14 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Srinivasa Ravi (Wolfram70)
<details>
<summary>Changes</summary>
This change adds the `store.async` op 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)
---
Full diff: https://github.com/llvm/llvm-project/pull/210931.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td (+14)
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+19)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+80)
- (added) mlir/test/Target/LLVMIR/nvvm/store_async.mlir (+58)
- (added) mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir (+63)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
index 42d196c5662d1..9d6a7f3bf1a6c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
@@ -69,4 +69,18 @@ def SaturationModeAttr : EnumAttr<NVVM_Dialect, SaturationMode, "sat_mode"> {
let assemblyFormat = "`<` $value `>`";
}
+def AsyncStoreScopeNone : I32EnumCase<"NONE", 0, "none">;
+def AsyncStoreScopeSys : I32EnumCase<"SYS", 1, "sys">;
+def AsyncStoreScopeGpu : I32EnumCase<"GPU", 2, "gpu">;
+
+def AsyncStoreScope :
+ I32Enum<"AsyncStoreScope", "NVVM Asynchronous Store Scope",
+ [AsyncStoreScopeNone, AsyncStoreScopeSys, AsyncStoreScopeGpu]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def AsyncStoreScopeAttr : EnumAttr<NVVM_Dialect, AsyncStoreScope,
+ "async_store_scope"> {
+ let assemblyFormat = "`<` $value `>`";
+}
#endif // NVVMIR_ENUMS
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 1f76218936e16..d905a5b212450 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5113,6 +5113,25 @@ def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// NVVM Asynchronous Store Op
+//===----------------------------------------------------------------------===//
+
+def NVVM_AsyncStoreOp: NVVM_VoidIntrinsicOp<"store.async", [NVVMRequiresSM<90>]> {
+ let summary = "Asynchronous Store Op";
+ let arguments = (ins AnyTypeOf<[LLVM_PointerGlobal,
+ LLVM_PointerSharedCluster]>:$addr,
+ AnyTypeOf<[I8, I16, I32, I64, I128]>:$value,
+ Optional<LLVM_PointerSharedCluster>:$mbarrier,
+ DefaultValuedAttr<AsyncStoreScopeAttr,
+ "AsyncStoreScope::NONE">:$scope,
+ DefaultValuedAttr<BoolAttr, "false">:$is_multimem,
+ DefaultValuedAttr<BoolAttr, "false">:$is_mmio);
+ let results = (outs );
+ let assemblyFormat = "$addr `,` $value (`,` `mbarrier` `=` $mbarrier^)? attr-dict `:` type($addr) `,` type($value) (`,` type($mbarrier)^)?";
+ let hasVerifier = 1;
+}
+
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 b29cba96a410f..603d28957176f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -668,6 +668,50 @@ LogicalResult BulkStoreOp::verify() {
return success();
}
+LogicalResult AsyncStoreOp::verify() {
+ unsigned addrSpace =
+ llvm::cast<LLVM::LLVMPointerType>(getAddr().getType()).getAddressSpace();
+ NVVM::AsyncStoreScope scope = getScope();
+ mlir::Type valueType = getValue().getType();
+ bool isMmio = getIsMmio();
+ bool isMultimem = getIsMultimem();
+
+ if (addrSpace == NVVMMemorySpace::Global) {
+ if (getMbarrier())
+ return emitOpError("mbarrier is not supported for global address space");
+
+ if (valueType.isInteger(128))
+ return emitOpError("only 8, 16, 32, and 64 bit values are supported for "
+ "global address space");
+
+ if (scope == AsyncStoreScope::NONE)
+ return emitOpError(
+ "scope must be set for async store to global address space");
+
+ if (isMmio && scope != AsyncStoreScope::SYS)
+ return emitOpError("mmio is only supported for SYS scope");
+
+ if (isMmio && isMultimem)
+ return emitOpError("multimem is not supported for mmio");
+ }
+
+ if (addrSpace == NVVMMemorySpace::SharedCluster) {
+ if (valueType.isInteger(8) || valueType.isInteger(16))
+ return emitOpError("only 32, 64, and 128 bit values are supported for "
+ "shared cluster address space");
+
+ if (isMultimem || isMmio)
+ return emitOpError("multimem and mmio are not supported for shared "
+ "cluster address space");
+
+ if (scope != AsyncStoreScope::NONE)
+ return emitOpError("scope is not supported for async store to shared "
+ "cluster address space");
+ }
+
+ return success();
+}
+
LogicalResult PMEventOp::verify() {
auto eventId = getEventId();
auto maskedEventId = getMaskedEventId();
@@ -3662,6 +3706,42 @@ DivFOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
{mt.lookupValue(thisOp.getLhs()), mt.lookupValue(thisOp.getRhs())}};
}
+mlir::NVVM::IDArgPair
+AsyncStoreOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::AsyncStoreOp>(op);
+ NVVM::NVVMMemorySpace addrSpace = static_cast<NVVM::NVVMMemorySpace>(
+ llvm::cast<LLVM::LLVMPointerType>(thisOp.getAddr().getType())
+ .getAddressSpace());
+ mlir::NVVM::AsyncStoreScope scope = thisOp.getScope();
+
+ llvm::Value *addr = mt.lookupValue(thisOp.getAddr());
+ llvm::Value *value = mt.lookupValue(thisOp.getValue());
+ llvm::Value *mbarrier =
+ thisOp.getMbarrier() ? mt.lookupValue(thisOp.getMbarrier()) : nullptr;
+ llvm::Value *isMultimem =
+ thisOp.getIsMultimem() ? builder.getInt1(true) : builder.getInt1(false);
+
+ switch (addrSpace) {
+ case NVVMMemorySpace::Global:
+ if (scope == AsyncStoreScope::SYS) {
+ if (thisOp.getIsMmio())
+ return {llvm::Intrinsic::nvvm_st_async_mmio_sys, {addr, value}};
+ else
+ return {llvm::Intrinsic::nvvm_st_async_sys, {addr, value, isMultimem}};
+ } else if (scope == AsyncStoreScope::GPU) {
+ return {llvm::Intrinsic::nvvm_st_async_gpu, {addr, value, isMultimem}};
+ }
+ llvm_unreachable("unsupported async store scope for global address space");
+ case NVVMMemorySpace::SharedCluster:
+ return {llvm::Intrinsic::nvvm_st_async, {addr, value, mbarrier}};
+ default:
+ llvm_unreachable("unsupported address space");
+ }
+
+ return {llvm::Intrinsic::not_intrinsic, {}};
+}
+
mlir::NVVM::IDArgPair
PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async.mlir
new file mode 100644
index 0000000000000..8e59f48c987ac
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async.mlir
@@ -0,0 +1,58 @@
+// 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 %addr, %value {scope = #nvvm.async_store_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 %addr, %value {scope = #nvvm.async_store_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 %addr, %value {scope = #nvvm.async_store_scope<sys>, is_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 %addr, %value {scope = #nvvm.async_store_scope<gpu>, is_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 %addr, %value {scope = #nvvm.async_store_scope<sys>, is_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 %addr, %v8 {scope = #nvvm.async_store_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 %addr, %v16 {scope = #nvvm.async_store_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 %addr, %v64 {scope = #nvvm.async_store_scope<gpu>} : !llvm.ptr<1>, i64
+ llvm.return
+}
+
+// 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 %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 %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 %addr, %v128, mbarrier = %mbar : !llvm.ptr<7>, i128, !llvm.ptr<7>
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
new file mode 100644
index 0000000000000..1fe4bbb85235f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
@@ -0,0 +1,63 @@
+// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
+
+llvm.func @st_async_global_mbarrier(%addr: !llvm.ptr<1>, %value: i32, %mbar: !llvm.ptr<7>) {
+ // expected-error @below {{mbarrier is not supported for global address space}}
+ nvvm.store.async %addr, %value, mbarrier = %mbar {scope = #nvvm.async_store_scope<sys>} : !llvm.ptr<1>, i32, !llvm.ptr<7>
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_global_i128(%addr: !llvm.ptr<1>, %value: i128) {
+ // expected-error @below {{only 8, 16, 32, and 64 bit values are supported for global address space}}
+ nvvm.store.async %addr, %value {scope = #nvvm.async_store_scope<sys>} : !llvm.ptr<1>, i128
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_global_no_scope(%addr: !llvm.ptr<1>, %value: i32) {
+ // expected-error @below {{scope must be set for async store to global address space}}
+ nvvm.store.async %addr, %value : !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 %addr, %value {scope = #nvvm.async_store_scope<gpu>, is_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 for mmio}}
+ nvvm.store.async %addr, %value {scope = #nvvm.async_store_scope<sys>, is_mmio = true, is_multimem = true} : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_shared_cluster_i8(%addr: !llvm.ptr<7>, %value: i8, %mbar: !llvm.ptr<7>) {
+ // expected-error @below {{only 32, 64, and 128 bit values are supported for shared cluster address space}}
+ nvvm.store.async %addr, %value, mbarrier = %mbar : !llvm.ptr<7>, i8, !llvm.ptr<7>
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_shared_cluster_multimem(%addr: !llvm.ptr<7>, %value: i32, %mbar: !llvm.ptr<7>) {
+ // expected-error @below {{multimem and mmio are not supported for shared cluster address space}}
+ nvvm.store.async %addr, %value, mbarrier = %mbar {is_multimem = true} : !llvm.ptr<7>, i32, !llvm.ptr<7>
+ llvm.return
+}
+
+// -----
+
+llvm.func @st_async_shared_cluster_scope(%addr: !llvm.ptr<7>, %value: i32, %mbar: !llvm.ptr<7>) {
+ // expected-error @below {{scope is not supported for async store to shared cluster address space}}
+ nvvm.store.async %addr, %value, mbarrier = %mbar {scope = #nvvm.async_store_scope<sys>} : !llvm.ptr<7>, i32, !llvm.ptr<7>
+ llvm.return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/210931
More information about the Mlir-commits
mailing list