[llvm] [NVPTX] Add missing Range attr to tensormap.replace intrinsics. (PR #207099)

[Yunqing Yu via llvm-commits]

https://github.com/yunqingy180 updated https://github.com/llvm/llvm-project/pull/207099

>From 753eb4554b9ea395184042154505013051d94b95 Mon Sep 17 00:00:00 2001
From: Yunqing Yu <yunqingy at nvidia.com>
Date: Wed, 1 Jul 2026 16:29:27 -0700
Subject: [PATCH] [NVPTX] Add missing Range attr to tensormap.replace
 intrinsics.

---
 llvm/docs/NVPTXUsage.rst                      |  8 +++---
 llvm/include/llvm/IR/IntrinsicsNVVM.td        |  2 +-
 llvm/test/Verifier/NVPTX/tensormap-replace.ll | 26 +++++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Verifier/NVPTX/tensormap-replace.ll

diff --git a/llvm/docs/NVPTXUsage.rst b/llvm/docs/NVPTXUsage.rst
index c61883f007afc..d6ca9154a1be9 100644
--- a/llvm/docs/NVPTXUsage.rst
+++ b/llvm/docs/NVPTXUsage.rst
@@ -2653,7 +2653,7 @@ Overview:
 
 The '``@llvm.nvvm.tensormap.replace.global.stride.*``' intrinsics replace the 
 ``%ord``-th element of the ``global_stride`` field of the tensor-map object 
-with ``%new_value``.
+with ``%new_value``. ``%ord`` must be in the range [0, 4].
 
 '``llvm.nvvm.tensormap.replace.element.stride``'
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -2671,7 +2671,7 @@ Overview:
 
 The '``@llvm.nvvm.tensormap.replace.element.stride.*``' intrinsics replace the 
 ``%ord``-th element of the ``element_stride`` field of the tensor-map object 
-with ``%new_value``.
+with ``%new_value``. ``%ord`` must be in the range [0, 4].
 
 '``llvm.nvvm.tensormap.replace.global.dim``'
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -2689,7 +2689,7 @@ Overview:
 
 The '``@llvm.nvvm.tensormap.replace.global.dim.*``' intrinsics replace the 
 ``%ord``-th element of the ``global_dim`` field of the tensor-map object 
-with ``%new_value``.
+with ``%new_value``. ``%ord`` must be in the range [0, 4].
 
 '``llvm.nvvm.tensormap.replace.box.dim``'
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -2707,7 +2707,7 @@ Overview:
 
 The '``@llvm.nvvm.tensormap.replace.box.dim.*``' intrinsics replace the 
 ``%ord``-th element of the ``box_dim`` field of the tensor-map object with 
-``%new_value``.
+``%new_value``. ``%ord`` must be in the range [0, 4].
 
 '``llvm.nvvm.tensormap.replace.elemtype``'
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index fc7abc50dce01..9d3d2ab15f9a0 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -3428,7 +3428,7 @@ let IntrProperties = [IntrArgMemOnly, IntrWriteMem, NoCapture<ArgIndex<0>>] in {
 }
 
 let IntrProperties = [IntrArgMemOnly, ImmArg<ArgIndex<1>>, IntrWriteMem,
-                      NoCapture<ArgIndex<0>>] in {
+                      NoCapture<ArgIndex<0>>, Range<ArgIndex<1>, 0, 5>] in {
   def int_nvvm_tensormap_replace_global_stride :
     DefaultAttrsIntrinsic<[], [llvm_anyptr_ty, llvm_i32_ty, llvm_i64_ty], []>;
   foreach tmap_field = ["box_dim", "global_dim", "element_stride"] in {
diff --git a/llvm/test/Verifier/NVPTX/tensormap-replace.ll b/llvm/test/Verifier/NVPTX/tensormap-replace.ll
new file mode 100644
index 0000000000000..223dbadcf4711
--- /dev/null
+++ b/llvm/test/Verifier/NVPTX/tensormap-replace.ll
@@ -0,0 +1,26 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+; The %ord operand of the indexed-field tensormap.replace intrinsics must be in
+; the range [0, 5), per the PTX ISA (`tensormap.replace`: "The only valid values
+; for `ord` are [0..4]").
+
+declare void @llvm.nvvm.tensormap.replace.box.dim.p1(ptr addrspace(1) %addr, i32 immarg %ord, i32 %new_value)
+declare void @llvm.nvvm.tensormap.replace.global.dim.p1(ptr addrspace(1) %addr, i32 immarg %ord, i32 %new_value)
+declare void @llvm.nvvm.tensormap.replace.element.stride.p1(ptr addrspace(1) %addr, i32 immarg %ord, i32 %new_value)
+declare void @llvm.nvvm.tensormap.replace.global.stride.p1(ptr addrspace(1) %addr, i32 immarg %ord, i64 %new_value)
+
+define void @test_tensormap_replace_ord_out_of_range(ptr addrspace(1) %addr, i32 %v32, i64 %v64) {
+  ; CHECK: immarg value 5 out of range [0, 5)
+  call void @llvm.nvvm.tensormap.replace.box.dim.p1(ptr addrspace(1) %addr, i32 5, i32 %v32)
+
+  ; CHECK: immarg value 6 out of range [0, 5)
+  call void @llvm.nvvm.tensormap.replace.global.dim.p1(ptr addrspace(1) %addr, i32 6, i32 %v32)
+
+  ; CHECK: immarg value 5 out of range [0, 5)
+  call void @llvm.nvvm.tensormap.replace.element.stride.p1(ptr addrspace(1) %addr, i32 5, i32 %v32)
+
+  ; CHECK: immarg value 6 out of range [0, 5)
+  call void @llvm.nvvm.tensormap.replace.global.stride.p1(ptr addrspace(1) %addr, i32 6, i64 %v64)
+
+  ret void
+}