[Mlir-commits] [mlir] [MLIR][NVVM] Add support for st.bulk Op (PR #131727)

[Srinivasa Ravi]

https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/131727

>From 147b862a546c437a1dfda451b50b33f69012c75b Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Wed, 26 Feb 2025 15:20:40 +0530
Subject: [PATCH] [MLIR][NVVM] Add support for st.bulk Op

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
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 19 +++++++++++++++++++
 mlir/test/Dialect/LLVMIR/nvvm.mlir          |  9 +++++++++
 mlir/test/Target/LLVMIR/nvvmir.mlir         | 10 ++++++++++
 3 files changed, 38 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index ff6696f6bec40..0ada5795cbc19 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2583,6 +2583,25 @@ def NVVM_MapaOp: NVVM_Op<"mapa",
   let assemblyFormat = "$a`,` $b attr-dict `:` type($a) `->` type($res)";
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM Bulk Store Op
+//===----------------------------------------------------------------------===//
+
+def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
+  let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr, I64:$size, I64Attr:$initval);
+
+  string llvmBuilder = [{
+    int addrSpace = llvm::cast<LLVMPointerType>(op.getAddr().getType()).getAddressSpace();
+    
+    bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+
+    auto intId = isSharedMemory ? llvm::Intrinsic::nvvm_st_bulk_shared_cta : llvm::Intrinsic::nvvm_st_bulk;
+    createIntrinsicCall(builder, intId, {$addr, $size, builder.getInt64($initval)});
+  }];
+  
+  let assemblyFormat = "$addr`,` $size`,` $initval attr-dict `:` type($addr)";
+}
+
 def NVVM_Exit : NVVM_Op<"exit"> {
   let summary = "Exit Op";
   let description = [{
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 85998d4e66254..d61c4d0a965b5 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -550,6 +550,15 @@ func.func @mapa(%a: !llvm.ptr, %a_shared: !llvm.ptr<3>, %b : 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 %{{.*}}, %{{.*}}, 0 : !llvm.ptr
+  nvvm.st.bulk %addr_gen, %size, 0 : !llvm.ptr
+  // CHECK:   nvvm.st.bulk %{{.*}}, %{{.*}}, 0 : !llvm.ptr<3>
+  nvvm.st.bulk %addr_shared, %size, 0 : !llvm.ptr<3>
+  return
+}
+
 // -----
 
 // Just check these don't emit errors.
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index f39aca95b918f..07bcdcd369909 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -810,3 +810,13 @@ llvm.func @nvvm_redux_sync_f32(%value: f32, %offset: i32) {
   %7 = nvvm.redux.sync fmax %value, %offset {abs = true, nan = true}: f32 -> f32
   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, 0 : !llvm.ptr
+  // CHECK: call void @llvm.nvvm.st.bulk.shared.cta(ptr addrspace(3) %{{.*}}, i64 %{{.*}}, i64 0)
+  nvvm.st.bulk %addr_shared, %size, 0: !llvm.ptr<3>
+  llvm.return
+}