[llvm] [NVPTX] Support 32-bit size operand for st.bulk (PR #217703)

[via llvm-commits]

https://github.com/robertvirany created https://github.com/llvm/llvm-project/pull/217703

PTX ISA 9.0 extends the `st.bulk` instruction to accept a 32-bit `size` operand.

Add `i32` variants of the generic and `.shared::cta` NVVM intrinsics and select them to `st.bulk` instructions using 32-bit registers. The new patterns require PTX ISA 9.0 and `sm_100`.

The existing intrinsics continue to represent the original 64-bit forms.

Add code-generation coverage for both the default and short-pointer ABIs.


>From 2ea11f9983877f96933859fe70cea1456ef59e44 Mon Sep 17 00:00:00 2001
From: Robert Virany <robertvirany at gmail.com>
Date: Thu, 20 Aug 2026 11:12:57 -0600
Subject: [PATCH] [NVPTX] Support 32-bit size operand for st.bulk

---
 llvm/include/llvm/IR/IntrinsicsNVVM.td   |  5 ++-
 llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 13 ++++++-
 llvm/test/CodeGen/NVPTX/st_bulk_i32.ll   | 47 ++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/st_bulk_i32.ll

diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index 6d6671e79dd24..325e50f3e17aa 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -3346,9 +3346,12 @@ let IntrProperties = [IntrArgMemOnly, IntrWriteMem, WriteOnly<ArgIndex<0>>,
                       NoCapture<ArgIndex<0>>, ImmArg<ArgIndex<2>>] in {
   def int_nvvm_st_bulk :
       DefaultAttrsIntrinsic<[], [llvm_ptr_ty, llvm_i64_ty, llvm_i64_ty]>;
-
+  def int_nvvm_st_bulk_i32 :
+      DefaultAttrsIntrinsic<[], [llvm_ptr_ty, llvm_i32_ty, llvm_i64_ty]>;
   def int_nvvm_st_bulk_shared_cta :
       DefaultAttrsIntrinsic<[], [llvm_shared_ptr_ty, llvm_i64_ty, llvm_i64_ty]>;
+  def int_nvvm_st_bulk_shared_cta_i32 :
+      DefaultAttrsIntrinsic<[], [llvm_shared_ptr_ty, llvm_i32_ty, llvm_i64_ty]>;
 }
 
 //
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index df5248a5bf7b5..3a2457f26f7eb 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -6181,11 +6181,22 @@ let Predicates = [SM100] in {
     BasicNVPTXInst<(outs), (ins ADDR:$dest_addr, B64:$size, i64imm:$value),
               "st.bulk",
               [(int_nvvm_st_bulk addr:$dest_addr, i64:$size, st_bulk_imm:$value)]>;
-
   def INT_NVVM_ST_BULK_SHARED_CTA:
     BasicNVPTXInst<(outs), (ins ADDR:$dest_addr, B64:$size, i64imm:$value),
               "st.bulk.shared::cta",
               [(int_nvvm_st_bulk_shared_cta addr:$dest_addr, i64:$size, st_bulk_imm:$value)]>;
+
+}
+
+let Predicates = [PTX90, SM100] in {
+  def INT_NVVM_ST_BULK_GENERIC_I32 :
+    BasicNVPTXInst<(outs), (ins ADDR:$dest_addr, B32:$size, i64imm:$value),
+              "st.bulk",
+              [(int_nvvm_st_bulk_i32 addr:$dest_addr, i32:$size, st_bulk_imm:$value)]>;
+  def INT_NVVM_ST_BULK_SHARED_CTA_I32:
+    BasicNVPTXInst<(outs), (ins ADDR:$dest_addr, B32:$size, i64imm:$value),
+              "st.bulk.shared::cta",
+              [(int_nvvm_st_bulk_shared_cta_i32 addr:$dest_addr, i32:$size, st_bulk_imm:$value)]>;
 }
 
 //
diff --git a/llvm/test/CodeGen/NVPTX/st_bulk_i32.ll b/llvm/test/CodeGen/NVPTX/st_bulk_i32.ll
new file mode 100644
index 0000000000000..62814dfa505d7
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/st_bulk_i32.ll
@@ -0,0 +1,47 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_100 -mattr=+ptx90 | FileCheck --check-prefixes=CHECK,CHECK-PTX64 %s
+; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_100 -mattr=+ptx90 -target-abi=shortptr | FileCheck --check-prefixes=CHECK,CHECK-PTX-SHARED32 %s
+; RUN: %if ptxas-sm_100 && ptxas-isa-9.0 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_100 -mattr=+ptx90 | %ptxas-verify -arch=sm_100 %}
+; RUN: %if ptxas-sm_100 && ptxas-isa-9.0 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_100 -mattr=+ptx90 -target-abi=shortptr | %ptxas-verify -arch=sm_100 %}
+
+declare void @llvm.nvvm.st.bulk.i32(ptr, i32, i64)
+define void @st_bulk_i32(ptr %dest_addr, i32 %size) {
+; CHECK-LABEL: st_bulk_i32(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<2>;
+; CHECK-NEXT:    .reg .b64 %rd<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param::func.b64 %rd1, [st_bulk_i32_param_0];
+; CHECK-NEXT:    ld.param::func.b32 %r1, [st_bulk_i32_param_1];
+; CHECK-NEXT:    st.bulk [%rd1], %r1, 0;
+; CHECK-NEXT:    ret;
+  call void @llvm.nvvm.st.bulk.i32(ptr %dest_addr, i32 %size, i64 0)
+  ret void
+}
+
+declare void @llvm.nvvm.st.bulk.shared.cta.i32(ptr addrspace(3), i32, i64)
+define void @st_bulk_shared_cta_i32(ptr addrspace(3) %dest_addr, i32 %size) {
+; CHECK-PTX64-LABEL: st_bulk_shared_cta_i32(
+; CHECK-PTX64:       {
+; CHECK-PTX64-NEXT:    .reg .b32 %r<2>;
+; CHECK-PTX64-NEXT:    .reg .b64 %rd<2>;
+; CHECK-PTX64-EMPTY:
+; CHECK-PTX64-NEXT:  // %bb.0:
+; CHECK-PTX64-NEXT:    ld.param::func.b64 %rd1, [st_bulk_shared_cta_i32_param_0];
+; CHECK-PTX64-NEXT:    ld.param::func.b32 %r1, [st_bulk_shared_cta_i32_param_1];
+; CHECK-PTX64-NEXT:    st.bulk.shared::cta [%rd1], %r1, 0;
+; CHECK-PTX64-NEXT:    ret;
+;
+; CHECK-PTX-SHARED32-LABEL: st_bulk_shared_cta_i32(
+; CHECK-PTX-SHARED32:       {
+; CHECK-PTX-SHARED32-NEXT:    .reg .b32 %r<3>;
+; CHECK-PTX-SHARED32-EMPTY:
+; CHECK-PTX-SHARED32-NEXT:  // %bb.0:
+; CHECK-PTX-SHARED32-NEXT:    ld.param::func.b32 %r1, [st_bulk_shared_cta_i32_param_0];
+; CHECK-PTX-SHARED32-NEXT:    ld.param::func.b32 %r2, [st_bulk_shared_cta_i32_param_1];
+; CHECK-PTX-SHARED32-NEXT:    st.bulk.shared::cta [%r1], %r2, 0;
+; CHECK-PTX-SHARED32-NEXT:    ret;
+   call void @llvm.nvvm.st.bulk.shared.cta.i32(ptr addrspace(3) %dest_addr, i32 %size, i64 0)
+   ret void
+}