[llvm] [AMDGPU] Fix crash folding cos/sin table lookups on poison/undef vector lanes (PR #216972)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/216972

>From c9b0c9420a51079d963ae7eadd8c1dbf9eaf856b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 18 Aug 2026 11:49:53 +0200
Subject: [PATCH 1/2] [AMDGPU] Fix crash folding cos/sin table lookups on
 poison/undef vector lanes

TDOFold cast the lane directly to ConstantFP, which aborts under assertions when a lane is poison or undef instead of a table lookup failure
---
 llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp     |  8 +++++--
 .../AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp b/llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp
index 7291fdaca0d69..52bb6feb94229 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp
@@ -844,8 +844,12 @@ bool AMDGPULibCalls::TDOFold(CallInst *CI, const FuncInfo &FInfo) {
       SmallVector<APFloat, 4> Values;
       Values.reserve(vecSize);
       for (int eltNo = 0; eltNo < vecSize; ++eltNo) {
-        ConstantFP *eltval =
-            cast<ConstantFP>(CV->getAggregateElement((unsigned)eltNo));
+        // A lane may be undef or poison, in which case there is nothing to
+        // look up in the table.
+        ConstantFP *eltval = dyn_cast_or_null<ConstantFP>(
+            CV->getAggregateElement((unsigned)eltNo));
+        if (!eltval)
+          return false;
         auto MatchingRow = llvm::find_if(tr, [eltval](const TableEntry &entry) {
           return eltval->isExactlyValue(entry.input);
         });
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
index 936471bac4ca9..4f784d4a40046 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
@@ -61,6 +61,29 @@ entry:
   ret <2 x double> %c
 }
 
+; A lane that is not a ConstantFP has nothing to look up in the table.
+define <2 x float> @test_tdo_v2_f32_cos_poison_lane() {
+; CHECK-LABEL: define <2 x float> @test_tdo_v2_f32_cos_poison_lane() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[C:%.*]] = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float poison>)
+; CHECK-NEXT:    ret <2 x float> [[C]]
+;
+entry:
+  %c = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float poison>)
+  ret <2 x float> %c
+}
+
+define <2 x float> @test_tdo_v2_f32_cos_undef_lane() {
+; CHECK-LABEL: define <2 x float> @test_tdo_v2_f32_cos_undef_lane() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[C:%.*]] = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float undef>)
+; CHECK-NEXT:    ret <2 x float> [[C]]
+;
+entry:
+  %c = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float undef>)
+  ret <2 x float> %c
+}
+
 declare float        @_Z3cosf(float)
 declare <2 x float>  @_Z3cosDv2_f(<2 x float>)
 declare half         @_Z3cosDh(half)

>From d8cc0b343c18df7ee22b9787d76719d980a84205 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 18 Aug 2026 12:08:54 +0200
Subject: [PATCH 2/2] fix undef deprecator remark

---
 .../AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll      | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
index 4f784d4a40046..14ec399a18808 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll
@@ -61,7 +61,8 @@ entry:
   ret <2 x double> %c
 }
 
-; A lane that is not a ConstantFP has nothing to look up in the table.
+; A lane that is not a ConstantFP (e.g. poison or undef) has nothing to look
+; up in the table.
 define <2 x float> @test_tdo_v2_f32_cos_poison_lane() {
 ; CHECK-LABEL: define <2 x float> @test_tdo_v2_f32_cos_poison_lane() {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
@@ -73,17 +74,6 @@ entry:
   ret <2 x float> %c
 }
 
-define <2 x float> @test_tdo_v2_f32_cos_undef_lane() {
-; CHECK-LABEL: define <2 x float> @test_tdo_v2_f32_cos_undef_lane() {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[C:%.*]] = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float undef>)
-; CHECK-NEXT:    ret <2 x float> [[C]]
-;
-entry:
-  %c = call <2 x float> @_Z3cosDv2_f(<2 x float> <float 0.000000e+00, float undef>)
-  ret <2 x float> %c
-}
-
 declare float        @_Z3cosf(float)
 declare <2 x float>  @_Z3cosDv2_f(<2 x float>)
 declare half         @_Z3cosDh(half)