[llvm] [AMDGPU] Add verification for amdgcn.init.exec.from.input (PR #128172)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i created https://github.com/llvm/llvm-project/pull/128172

Check that the input register is an inreg argument to the parent function. (See the comment in `IntrinsicsAMDGPU.td`.)

This LLVM defect was identified via the AMD Fuzzing project.

>From dda8be4f0c7229d670e7e5e3058b28d7eebfc59a Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 21 Feb 2025 07:05:04 -0600
Subject: [PATCH] [AMDGPU] Add verification for amdgcn.init.exec.from.input

Check that the input register is an inreg argument to the parent
function. (See the comment in `IntrinsicsAMDGPU.td`.)

This LLVM defect was identified via the AMD Fuzzing project.
---
 llvm/lib/IR/Verifier.cpp                      | 15 +++++++++++
 .../intrinsic-amdgpu-init-exec-from-input.ll  | 27 +++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 llvm/test/Verifier/AMDGPU/intrinsic-amdgpu-init-exec-from-input.ll

diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 8432779c107de..55923fd56a109 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -6369,6 +6369,21 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
           "VGPR arguments must not have the `inreg` attribute", &Call);
     break;
   }
+  case Intrinsic::amdgcn_init_exec_from_input: {
+    const Value *InputVal = Call.getOperand(0);
+    bool InRegArgFound = false;
+    for (const Argument &Arg : Call.getCaller()->args()) {
+      if (Arg.hasInRegAttr() && &Arg == InputVal) {
+        InRegArgFound = true;
+        break;
+      }
+    }
+    Check(InRegArgFound,
+          "only inreg arguments to the parent function are valid as inputs to "
+          "this intrinsic",
+          &Call);
+    break;
+  }
   case Intrinsic::amdgcn_set_inactive_chain_arg: {
     auto CallerCC = Call.getCaller()->getCallingConv();
     switch (CallerCC) {
diff --git a/llvm/test/Verifier/AMDGPU/intrinsic-amdgpu-init-exec-from-input.ll b/llvm/test/Verifier/AMDGPU/intrinsic-amdgpu-init-exec-from-input.ll
new file mode 100644
index 0000000000000..5e4a6638dc170
--- /dev/null
+++ b/llvm/test/Verifier/AMDGPU/intrinsic-amdgpu-init-exec-from-input.ll
@@ -0,0 +1,27 @@
+; RUN: sed -e '/^; MARK$/,$d' %s | not llc -mtriple=amdgcn -mcpu=gfx942 2>&1 | FileCheck %s
+; RUN: sed -e '1,/^; MARK$/d' %s | llc -mtriple=amdgcn -mcpu=gfx942 -filetype=null
+
+; Function Attrs: convergent nocallback nofree nounwind willreturn
+declare void @llvm.amdgcn.init.exec.from.input(i32, i32 immarg) #0
+attributes #0 = { convergent nocallback nofree nounwind willreturn }
+
+; CHECK: only inreg arguments to the parent function are valid as inputs to this intrinsic
+; CHECK-NEXT: call void @llvm.amdgcn.init.exec.from.input(i32 0, i32 0)
+define amdgpu_ps void @init_exec_from_input_fail_immarg(i32 inreg %a, i32 %b) {
+  call void @llvm.amdgcn.init.exec.from.input(i32 0, i32 0)
+  ret void
+}
+
+; CHECK: only inreg arguments to the parent function are valid as inputs to this intrinsic
+; CHECK-NEXT: call void @llvm.amdgcn.init.exec.from.input(i32 %b, i32 0)
+define amdgpu_ps void @init_exec_from_input_fail_not_inreg(i32 inreg %a, i32 %b) {
+  call void @llvm.amdgcn.init.exec.from.input(i32 %b, i32 0)
+  ret void
+}
+
+; MARK
+
+define amdgpu_ps void @init_exec_from_input_success(i32 inreg %a, i32 %b) {
+  call void @llvm.amdgcn.init.exec.from.input(i32 %a, i32 0)
+  ret void
+}