[AMDGPU] Fix crash folding cos/sin table lookups on poison/undef vector lanes

[aobolensk]

TDOFold cast the lane directly to ConstantFP, which aborts under assertions when a lane is poison or undef instead of a table lookup failure

[llvmorg-github-actions]

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

Changes

TDOFold cast the lane directly to ConstantFP, which aborts under assertions when a lane is poison or undef instead of a table lookup failure

---

Full diff: https://github.com/llvm/llvm-project/pull/216972.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp (+6-2)
• (modified) llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll (+23)

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)

[github-actions]

⚠️ undef deprecator found issues in your code. ⚠️

You can test this locally with the following command:

git diff -U0 --pickaxe-regex -S '([^a-zA-Z0-9#_-]undef([^a-zA-Z0-9_-]|$)|UndefValue::get)' 'HEAD~1' HEAD llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll

The following files introduce new uses of undef:

• llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-tdo-cos.ll

Undef is now deprecated and should only be used in the rare cases where no replacement is possible. For example, a load of uninitialized memory yields undef. You should use poison values for placeholders instead.

In tests, avoid using undef and having tests that trigger undefined behavior. If you need an operand with some unimportant value, you can add a new argument to the function and use that instead.

For example, this is considered a bad practice:

define void @fn() {
  ...
  br i1 undef, ...
}

Please use the following instead:

define void @fn(i1 %cond) {
  ...
  br i1 %cond, ...
}

Please refer to the Undefined Behavior Manual for more information.

[arsenm]

Leave the comment, ignore the dumb bot

[aobolensk]

ok, reverted

[aobolensk]

Ignoring undef related remark