[Mlir-commits] [mlir] 40815be - [MLIR][NVVM] Add support for st.bulk Op (#131727)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Mar 25 20:01:09 PDT 2025
Author: Srinivasa Ravi
Date: 2025-03-26T08:30:58+05:30
New Revision: 40815be30a04c4646d74cf6a70bfa4596846f5a5
URL: https://github.com/llvm/llvm-project/commit/40815be30a04c4646d74cf6a70bfa4596846f5a5
DIFF: https://github.com/llvm/llvm-project/commit/40815be30a04c4646d74cf6a70bfa4596846f5a5.diff
LOG: [MLIR][NVVM] Add support for st.bulk Op (#131727)
This change adds the `st.bulk` NVVM Op for the `st.bulk` instruction
introduced in ptx8.6 for sm_100.
PTX Spec Reference:
https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-st-bulk
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
mlir/test/Dialect/LLVMIR/nvvm.mlir
mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
mlir/test/Target/LLVMIR/nvvmir.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7c2e042b55248..8a54804b220a1 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2630,6 +2630,36 @@ def NVVM_MatchSyncOp : NVVM_Op<"match.sync">,
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// NVVM Bulk Store Op
+//===----------------------------------------------------------------------===//
+
+def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
+ let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr, I64:$size, DefaultValuedAttr<I64Attr, "0">:$initVal);
+
+ let summary = "Bulk Store Op";
+ let description = [{
+ Initializes a region of shared memory at the address given by `addr`.
+ The `size` operand specifies the number of bytes to initialize and must be
+ a multiple of 8.
+ The `initVal` operand specifies the value to initialize the memory to. The
+ only supported value is 0.
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-st-bulk)
+ }];
+
+ string llvmBuilder = [{
+ auto intId = getStBulkIntrinsicId(
+ llvm::cast<LLVM::LLVMPointerType>(op.getAddr().getType()));
+ createIntrinsicCall(builder, intId,
+ {$addr, $size, builder.getInt64($initVal)});
+ }];
+
+ let assemblyFormat = "$addr `,` `size` `=` $size (`,` `init` `=` $initVal^)? attr-dict `:` type($addr)";
+
+ 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 ce93cb1ca4297..556114f4370b3 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -160,6 +160,12 @@ LogicalResult CvtFloatToTF32Op::verify() {
return success();
}
+LogicalResult BulkStoreOp::verify() {
+ if (getInitVal() != 0)
+ return emitOpError("only 0 is supported for initVal, got ") << getInitVal();
+ return success();
+}
+
// Given the element type of an operand and whether or not it is an accumulator,
// this function returns the PTX type (`NVVM::MMATypes`) that corresponds to the
// operand's element type.
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 8f0e89bc66096..9d14ff09ab434 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -150,6 +150,15 @@ static llvm::Intrinsic::ID getLdMatrixIntrinsicId(NVVM::MMALayout layout,
}
}
+/// Return the intrinsic ID associated with st.bulk for the given address type.
+static llvm::Intrinsic::ID
+getStBulkIntrinsicId(LLVM::LLVMPointerType addrType) {
+ bool isSharedMemory =
+ addrType.getAddressSpace() == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+ return isSharedMemory ? llvm::Intrinsic::nvvm_st_bulk_shared_cta
+ : llvm::Intrinsic::nvvm_st_bulk;
+}
+
static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
NVVM::ProxyKind toProxy,
NVVM::MemScopeKind scope,
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 903dbb1bb5c8e..18bf39424f0bf 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -563,6 +563,15 @@ func.func @match_sync(%val32: i32, %val64: i64, %thread_mask: i32) {
return
}
+// CHECK-LABEL: @st_bulk
+func.func @st_bulk(%addr_gen: !llvm.ptr, %addr_shared: !llvm.ptr<3>, %size: i64) {
+ // CHECK: nvvm.st.bulk %{{.*}}, size = %{{.*}} : !llvm.ptr
+ nvvm.st.bulk %addr_gen, size = %size, init = 0 : !llvm.ptr
+ // CHECK: nvvm.st.bulk %{{.*}}, size = %{{.*}} : !llvm.ptr<3>
+ nvvm.st.bulk %addr_shared, size = %size, init = 0 : !llvm.ptr<3>
+ return
+}
+
// -----
// Just check these don't emit errors.
diff --git a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
index 97fbbe2fe5fa3..f87f11daeef54 100644
--- a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
@@ -125,6 +125,14 @@ llvm.func @convert_float_to_tf32_no_rnd_mode(%src : f32) -> i32 {
// -----
+llvm.func @nvvm_st_bulk_initval_nonzero(%addr : !llvm.ptr, %size : i64) {
+ // expected-error @below {{only 0 is supported for initVal, got 1}}
+ nvvm.st.bulk %addr, size = %size, init = 1 : !llvm.ptr
+ llvm.return
+}
+
+// -----
+
llvm.func @nvvm_tcgen05_cp_128x256b_mc(%taddr : !llvm.ptr<6>, %smem_desc : i64) {
// expected-error @below {{Invalid multicast type for tcgen05.cp Op}}
nvvm.tcgen05.cp %taddr, %smem_desc {shape = #nvvm.tcgen05_cp_shape<shape_128x256b>, multicast = #nvvm.tcgen05_cp_multicast<warpx2_02_13>}
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index c113cd2fcf5f7..c3ec88db1d694 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -811,6 +811,7 @@ llvm.func @nvvm_redux_sync_f32(%value: f32, %offset: i32) {
llvm.return
}
+// -----
// CHECK-LABEL: @nvvm_match_sync
llvm.func @nvvm_match_sync(%mask: i32, %val32: i32, %val64: i64) {
// CHECK: call i32 @llvm.nvvm.match.any.sync.i32(i32 %{{.*}}, i32 %{{.*}})
@@ -823,3 +824,17 @@ llvm.func @nvvm_match_sync(%mask: i32, %val32: i32, %val64: i64) {
%3 = nvvm.match.sync all %mask, %val64 : i64 -> !llvm.struct<(i32, i1)>
llvm.return
}
+
+// -----
+// CHECK-LABEL: @nvvm_st_bulk
+llvm.func @nvvm_st_bulk(%addr_gen: !llvm.ptr, %addr_shared: !llvm.ptr<3>, %size: i64) {
+ // CHECK: call void @llvm.nvvm.st.bulk(ptr %{{.*}}, i64 %{{.*}}, i64 0)
+ nvvm.st.bulk %addr_gen, size = %size : !llvm.ptr
+ // CHECK: call void @llvm.nvvm.st.bulk.shared.cta(ptr addrspace(3) %{{.*}}, i64 %{{.*}}, i64 0)
+ nvvm.st.bulk %addr_shared, size = %size: !llvm.ptr<3>
+ // CHECK: call void @llvm.nvvm.st.bulk(ptr %{{.*}}, i64 %{{.*}}, i64 0)
+ nvvm.st.bulk %addr_gen, size = %size, init = 0 : !llvm.ptr
+ // CHECK: call void @llvm.nvvm.st.bulk.shared.cta(ptr addrspace(3) %{{.*}}, i64 %{{.*}}, i64 0)
+ nvvm.st.bulk %addr_shared, size = %size, init = 0: !llvm.ptr<3>
+ llvm.return
+}
More information about the Mlir-commits
mailing list