[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

Author: Wolfram70

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 = [{

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.

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,

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.

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>}

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
+}