[llvm] [IR] Use immarg for preallocated intrinsics (NFC) (PR #155835)

[Nikita Popov via llvm-commits]

https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/155835

>From 708c01845a8474943ebf464eba97780043dd9101 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 28 Aug 2025 15:28:59 +0200
Subject: [PATCH 1/2] [IR] Use immarg for preallocated intrinsics (NFC)

Mark the attributes as immarg to indicate that they require a
constant integer. This was previously enforced with a manual
verifier check.
---
 llvm/include/llvm/IR/Intrinsics.td         | 8 ++++++--
 llvm/lib/IR/Verifier.cpp                   | 8 ++------
 llvm/test/Verifier/preallocated-invalid.ll | 2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index dba44e7c3c506..5ca7d24c8aaa3 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -977,8 +977,12 @@ def int_instrprof_mcdc_tvbitmap_update : Intrinsic<[],
                                         [llvm_ptr_ty, llvm_i64_ty,
                                          llvm_i32_ty, llvm_ptr_ty]>;
 
-def int_call_preallocated_setup : DefaultAttrsIntrinsic<[llvm_token_ty], [llvm_i32_ty]>;
-def int_call_preallocated_arg : DefaultAttrsIntrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_i32_ty]>;
+def int_call_preallocated_setup
+    : DefaultAttrsIntrinsic<[llvm_token_ty], [llvm_i32_ty],
+                            [ImmArg<ArgIndex<0>>]>;
+def int_call_preallocated_arg
+    : DefaultAttrsIntrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_i32_ty],
+                            [ImmArg<ArgIndex<1>>]>;
 def int_call_preallocated_teardown : DefaultAttrsIntrinsic<[], [llvm_token_ty]>;
 
 // This intrinsic is intentionally undocumented and users shouldn't call it;
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index b2e76cc7a8a90..27e78ecfa9288 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -5869,9 +5869,7 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
     break;
   }
   case Intrinsic::call_preallocated_setup: {
-    auto *NumArgs = dyn_cast<ConstantInt>(Call.getArgOperand(0));
-    Check(NumArgs != nullptr,
-          "llvm.call.preallocated.setup argument must be a constant");
+    auto *NumArgs = cast<ConstantInt>(Call.getArgOperand(0));
     bool FoundCall = false;
     for (User *U : Call.users()) {
       auto *UseCall = dyn_cast<CallBase>(U);
@@ -5879,9 +5877,7 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
             "Uses of llvm.call.preallocated.setup must be calls");
       Intrinsic::ID IID = UseCall->getIntrinsicID();
       if (IID == Intrinsic::call_preallocated_arg) {
-        auto *AllocArgIndex = dyn_cast<ConstantInt>(UseCall->getArgOperand(1));
-        Check(AllocArgIndex != nullptr,
-              "llvm.call.preallocated.alloc arg index must be a constant");
+        auto *AllocArgIndex = cast<ConstantInt>(UseCall->getArgOperand(1));
         auto AllocArgIndexInt = AllocArgIndex->getValue();
         Check(AllocArgIndexInt.sge(0) &&
                   AllocArgIndexInt.slt(NumArgs->getValue()),
diff --git a/llvm/test/Verifier/preallocated-invalid.ll b/llvm/test/Verifier/preallocated-invalid.ll
index 38ed1067c497d..921fa69dcb23b 100644
--- a/llvm/test/Verifier/preallocated-invalid.ll
+++ b/llvm/test/Verifier/preallocated-invalid.ll
@@ -65,7 +65,7 @@ define void @preallocated_one_call() {
     ret void
 }
 
-; CHECK: must be a constant
+; CHECK: immarg operand has non-immediate parameter
 define void @preallocated_setup_constant() {
     %ac = call i32 @blackbox()
     %cs = call token @llvm.call.preallocated.setup(i32 %ac)

>From cd339d8239dca3549da3307b2411b09967540117 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Fri, 26 Sep 2025 12:50:14 +0200
Subject: [PATCH 2/2] Add test for call.preallocated.arg and restore check

---
 llvm/lib/IR/Verifier.cpp                   | 4 +++-
 llvm/test/Verifier/preallocated-invalid.ll | 8 ++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 27e78ecfa9288..8c03d6f809d50 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -5877,7 +5877,9 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
             "Uses of llvm.call.preallocated.setup must be calls");
       Intrinsic::ID IID = UseCall->getIntrinsicID();
       if (IID == Intrinsic::call_preallocated_arg) {
-        auto *AllocArgIndex = cast<ConstantInt>(UseCall->getArgOperand(1));
+        auto *AllocArgIndex = dyn_cast<ConstantInt>(UseCall->getArgOperand(1));
+        Check(AllocArgIndex != nullptr,
+              "llvm.call.preallocated.alloc arg index must be a constant");
         auto AllocArgIndexInt = AllocArgIndex->getValue();
         Check(AllocArgIndexInt.sge(0) &&
                   AllocArgIndexInt.slt(NumArgs->getValue()),
diff --git a/llvm/test/Verifier/preallocated-invalid.ll b/llvm/test/Verifier/preallocated-invalid.ll
index 921fa69dcb23b..2c5aff231e1bd 100644
--- a/llvm/test/Verifier/preallocated-invalid.ll
+++ b/llvm/test/Verifier/preallocated-invalid.ll
@@ -72,6 +72,14 @@ define void @preallocated_setup_constant() {
     ret void
 }
 
+; CHECK: llvm.call.preallocated.alloc arg index must be a constant
+define void @preallocated_arg_constant() {
+    %ac = call i32 @blackbox()
+    %cs = call token @llvm.call.preallocated.setup(i32 3)
+    call token @llvm.call.preallocated.arg(token %cs, i32 %ac)
+    ret void
+}
+
 ; CHECK: must be between 0 and corresponding
 define void @preallocated_setup_arg_index_in_bounds() {
     %cs = call token @llvm.call.preallocated.setup(i32 2)