[Mlir-commits] [mlir] [MLIR][NVVM] Add asynchronous store Op (PR #210931)
Srinivasa Ravi
llvmlistbot at llvm.org
Mon Aug 10 00:15:56 PDT 2026
https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/210931
>From 8e7054aebd54568d560c20efa9cd955f8806818b Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Fri, 17 Jul 2026 06:48:11 +0000
Subject: [PATCH 1/5] [MLIR][NVVM] Add asynchronous store Op
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
---
mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td | 14 ++++
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 19 +++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 78 +++++++++++++++++++
mlir/test/Target/LLVMIR/nvvm/store_async.mlir | 58 ++++++++++++++
.../LLVMIR/nvvm/store_async_invalid.mlir | 63 +++++++++++++++
5 files changed, 232 insertions(+)
create mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
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 75bc7527aba52..c108243ec797a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5122,6 +5122,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 67c4be9670c91..6a9ee48a31778 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -669,6 +669,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();
@@ -3663,6 +3707,40 @@ 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}};
+ }
+ 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
+}
>From 7352e6f9eb39310f8b1d59d53f11abe0057d31e2 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Tue, 21 Jul 2026 10:15:20 +0000
Subject: [PATCH 2/5] add missing llvm_unreachable
---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 6a9ee48a31778..3963e467d382f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3733,12 +3733,14 @@ AsyncStoreOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
} 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, {}};
}
+
+ return {llvm::Intrinsic::not_intrinsic, {}};
}
mlir::NVVM::IDArgPair
>From 77124d5361d037e4ae02304cf86765fbe5a0fe80 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Tue, 4 Aug 2026 11:23:15 +0000
Subject: [PATCH 3/5] address comments
---
mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td | 14 ---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 55 ++++++++--
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 103 +++++++-----------
mlir/test/Target/LLVMIR/nvvm/store_async.mlir | 58 ----------
.../LLVMIR/nvvm/store_async_global.mlir | 42 +++++++
.../nvvm/store_async_global_invalid.mlir | 23 ++++
.../LLVMIR/nvvm/store_async_invalid.mlir | 63 -----------
.../LLVMIR/nvvm/store_async_shared.mlir | 17 +++
8 files changed, 162 insertions(+), 213 deletions(-)
delete mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async_global.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
delete mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/store_async_shared.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
index 9d6a7f3bf1a6c..42d196c5662d1 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMEnums.td
@@ -69,18 +69,4 @@ 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 c108243ec797a..c477a28d4b7f2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5123,24 +5123,55 @@ def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
}
//===----------------------------------------------------------------------===//
-// NVVM Asynchronous Store Op
+// NVVM Asynchronous Store Ops
//===----------------------------------------------------------------------===//
-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);
+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 (`,` `mbarrier` `=` $mbarrier^)? attr-dict `:` type($addr) `,` type($value) (`,` type($mbarrier)^)?";
+ let assemblyFormat = "$addr `,` $value attr-dict `:` type($addr) `,` type($value)";
let hasVerifier = 1;
}
+def NVVM_AsyncStoreSharedOp: NVVM_VoidIntrinsicOp<"store.async.shared",
+ [NVVMRequiresSM<90>]> {
+ 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 results = (outs );
+ 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 3963e467d382f..2c112daa0e7bb 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -669,46 +669,19 @@ 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");
+LogicalResult AsyncStoreGlobalOp::verify() {
+ NVVM::MemScopeKind scope = getScope();
+ bool isMmio = getMmio();
+ bool isMultimem = getMultimem();
- 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 (scope != MemScopeKind::SYS && scope != MemScopeKind::GPU)
+ return emitOpError("scope must be either SYS or GPU");
- if (isMmio && isMultimem)
- return emitOpError("multimem is not supported for mmio");
- }
+ if (isMmio && scope != MemScopeKind::SYS)
+ return emitOpError("mmio is only supported for SYS scope");
- 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");
- }
+ if (isMmio && isMultimem)
+ return emitOpError("multimem is not supported for mmio");
return success();
}
@@ -3707,40 +3680,38 @@ 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();
+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 *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");
+ 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 NVVM::IDArgPair(llvm::Intrinsic::nvvm_st_async_gpu,
+ {addr, value, isMultimem});
}
+ llvm_unreachable("unsupported async store scope for global address space");
+}
+
+mlir::NVVM::IDArgPair AsyncStoreSharedOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::AsyncStoreSharedOp>(op);
+
+ llvm::Value *addr = mt.lookupValue(thisOp.getAddr());
+ llvm::Value *value = mt.lookupValue(thisOp.getValue());
+ llvm::Value *mbarrier = mt.lookupValue(thisOp.getMbarrier());
- return {llvm::Intrinsic::not_intrinsic, {}};
+ return {llvm::Intrinsic::nvvm_st_async, {addr, value, mbarrier}};
}
mlir::NVVM::IDArgPair
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async.mlir
deleted file mode 100644
index 8e59f48c987ac..0000000000000
--- a/mlir/test/Target/LLVMIR/nvvm/store_async.mlir
+++ /dev/null
@@ -1,58 +0,0 @@
-// 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_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..ebe83cc55f888
--- /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 for 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_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
deleted file mode 100644
index 1fe4bbb85235f..0000000000000
--- a/mlir/test/Target/LLVMIR/nvvm/store_async_invalid.mlir
+++ /dev/null
@@ -1,63 +0,0 @@
-// 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
-}
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
+}
>From 1a073abfcb9c225a93c9f708c73fa3bc66f5d365 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Wed, 5 Aug 2026 13:07:03 +0000
Subject: [PATCH 4/5] address comments
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 7 ++++---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 15 ++-------------
.../LLVMIR/nvvm/store_async_global_invalid.mlir | 2 +-
3 files changed, 7 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index c477a28d4b7f2..8c5712b0828ad 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5153,8 +5153,10 @@ def NVVM_AsyncStoreGlobalOp: NVVM_VoidIntrinsicOp<"store.async.global",
let hasVerifier = 1;
}
-def NVVM_AsyncStoreSharedOp: NVVM_VoidIntrinsicOp<"store.async.shared",
- [NVVMRequiresSM<90>]> {
+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`.
@@ -5168,7 +5170,6 @@ def NVVM_AsyncStoreSharedOp: NVVM_VoidIntrinsicOp<"store.async.shared",
let arguments = (ins LLVM_PointerSharedCluster:$addr,
AnyTypeOf<[I32, I64, I128]>:$value,
LLVM_PointerSharedCluster:$mbarrier);
- let results = (outs );
let assemblyFormat = "$addr `,` $value `,` `mbarrier` `=` $mbarrier attr-dict `:` type($addr) `,` type($value) `,` type($mbarrier)";
}
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 2c112daa0e7bb..4917a55637877 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -681,7 +681,7 @@ LogicalResult AsyncStoreGlobalOp::verify() {
return emitOpError("mmio is only supported for SYS scope");
if (isMmio && isMultimem)
- return emitOpError("multimem is not supported for mmio");
+ return emitOpError("multimem is not supported with mmio");
return success();
}
@@ -3700,18 +3700,7 @@ mlir::NVVM::IDArgPair AsyncStoreGlobalOp::getIntrinsicIDAndArgs(
return NVVM::IDArgPair(llvm::Intrinsic::nvvm_st_async_gpu,
{addr, value, isMultimem});
}
- llvm_unreachable("unsupported async store scope for global address space");
-}
-
-mlir::NVVM::IDArgPair AsyncStoreSharedOp::getIntrinsicIDAndArgs(
- Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
- auto thisOp = cast<NVVM::AsyncStoreSharedOp>(op);
-
- llvm::Value *addr = mt.lookupValue(thisOp.getAddr());
- llvm::Value *value = mt.lookupValue(thisOp.getValue());
- llvm::Value *mbarrier = mt.lookupValue(thisOp.getMbarrier());
-
- return {llvm::Intrinsic::nvvm_st_async, {addr, value, mbarrier}};
+ llvm_unreachable("unsupported scope for AsyncStoreGlobalOp");
}
mlir::NVVM::IDArgPair
diff --git a/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
index ebe83cc55f888..18fb78ff4341c 100644
--- a/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/store_async_global_invalid.mlir
@@ -17,7 +17,7 @@ llvm.func @st_async_global_mmio_non_sys(%addr: !llvm.ptr<1>, %value: i32) {
// -----
llvm.func @st_async_global_mmio_multimem(%addr: !llvm.ptr<1>, %value: i32) {
- // expected-error @below {{multimem is not supported for mmio}}
+ // 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
}
>From e2232a975a25529875712b061a098e02399ebc16 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Mon, 10 Aug 2026 07:15:19 +0000
Subject: [PATCH 5/5] address comment
---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 4917a55637877..908628a2ec0a4 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3697,8 +3697,8 @@ mlir::NVVM::IDArgPair AsyncStoreGlobalOp::getIntrinsicIDAndArgs(
: IDArgPair(llvm::Intrinsic::nvvm_st_async_sys,
{addr, value, isMultimem});
} else if (scope == MemScopeKind::GPU) {
- return NVVM::IDArgPair(llvm::Intrinsic::nvvm_st_async_gpu,
- {addr, value, isMultimem});
+ return IDArgPair(llvm::Intrinsic::nvvm_st_async_gpu,
+ {addr, value, isMultimem});
}
llvm_unreachable("unsupported scope for AsyncStoreGlobalOp");
}
More information about the Mlir-commits
mailing list