[llvm] [AMDGPU] Handle CreateBinOp not returning BinaryOperator (PR #137791)

[via llvm-commits]

https://github.com/anjenner updated https://github.com/llvm/llvm-project/pull/137791

>From 4342e6aa5156fc2b7b7a3a881f99ebe142d4e797 Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Tue, 29 Apr 2025 06:52:50 -0400
Subject: [PATCH 1/4] [AMDGPU] Handle CreateBinOp not returning BinaryOperator

AMDGPUCodeGenPrepareImpl::visitBinaryOperator() calls Builder.CreateBinOp()
and casts the resulting Value as a BinaryOperator without checking, leading
to an assert failure in a case found by fuzzing. In this case, the operands
are constant and CreateBinOp does constant folding so returns a Constant
instead of a BinaryOperator.
---
 llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp       |  5 ++++-
 .../print-pipeline-passes.AFLCustomIRMutator.ll       | 11 +++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 6617373f89c8b..53bea99a2e98c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -1685,7 +1685,10 @@ bool AMDGPUCodeGenPrepareImpl::visitBinaryOperator(BinaryOperator &I) {
             // return the new value. Just insert a scalar copy and defer
             // expanding it.
             NewElt = Builder.CreateBinOp(Opc, NumEltN, DenEltN);
-            Div64ToExpand.push_back(cast<BinaryOperator>(NewElt));
+            // CreateBinOp does constant folding. If the operands are constant,
+            // it will return a Constant instead of a BinaryOperator.
+            if (auto *NewEltBO = dyn_cast<BinaryOperator>(NewElt))
+              Div64ToExpand.push_back(NewEltBO);
           }
         }
 
diff --git a/llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll b/llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll
new file mode 100644
index 0000000000000..583ef3a8bb7c7
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll
@@ -0,0 +1,11 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -O1 < %s | FileCheck -check-prefix=GCN %s
+
+define amdgpu_kernel void @kernel() {
+; GCN-LABEL: kernel:
+; GCN:       ; %bb.0: ; %entry
+; GCN-NEXT:    s_endpgm
+entry:
+  %B = srem <32 x i64> zeroinitializer, zeroinitializer
+  ret void
+}

>From ff33dbd0df3bff4141c9c360a16b3d4f12bc2195 Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Thu, 22 May 2025 11:16:16 -0400
Subject: [PATCH 2/4] Rename print-pipeline-passes.AFLCustomIRMutator.ll test
 to llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll .

---
 ...nt-pipeline-passes.AFLCustomIRMutator.ll => srem_zero_zero.ll} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename llvm/test/CodeGen/AMDGPU/{print-pipeline-passes.AFLCustomIRMutator.ll => srem_zero_zero.ll} (100%)

diff --git a/llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll b/llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll
similarity index 100%
rename from llvm/test/CodeGen/AMDGPU/print-pipeline-passes.AFLCustomIRMutator.ll
rename to llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll

>From 240da178df098c8f57947796f143aede2b7a011d Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Tue, 27 May 2025 08:35:46 -0400
Subject: [PATCH 3/4] Move testcase as per review feedback.

---
 .../test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll |  9 +++++++++
 llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll            | 11 -----------
 2 files changed, 9 insertions(+), 11 deletions(-)
 delete mode 100644 llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll

diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
index 8e16889c72e65..746c18b04aa2b 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
@@ -10094,3 +10094,12 @@ define i64 @udiv_i64_9divbits(i8 %size) {
   %div = udiv i64 %num, 10
   ret i64 %div
 }
+
+define void @srem_zero_zero() {
+; GCN-LABEL: kernel:
+; GCN:       ; %bb.0: ; %entry
+; GCN-NEXT:    s_endpgm
+entry:
+  %B = srem <32 x i64> zeroinitializer, zeroinitializer
+  ret void
+}
diff --git a/llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll b/llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll
deleted file mode 100644
index 583ef3a8bb7c7..0000000000000
--- a/llvm/test/CodeGen/AMDGPU/srem_zero_zero.ll
+++ /dev/null
@@ -1,11 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -O1 < %s | FileCheck -check-prefix=GCN %s
-
-define amdgpu_kernel void @kernel() {
-; GCN-LABEL: kernel:
-; GCN:       ; %bb.0: ; %entry
-; GCN-NEXT:    s_endpgm
-entry:
-  %B = srem <32 x i64> zeroinitializer, zeroinitializer
-  ret void
-}

>From 0b606d0ec7c01541f9fde6f0c6e3baa53dcb9629 Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Thu, 29 May 2025 10:14:49 -0400
Subject: [PATCH 4/4] Address review feedback.

---
 llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
index 746c18b04aa2b..b7097a9557b75 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
@@ -10095,11 +10095,11 @@ define i64 @udiv_i64_9divbits(i8 %size) {
   ret i64 %div
 }
 
-define void @srem_zero_zero() {
+define <2 x i64> @srem_zero_zero() {
 ; GCN-LABEL: kernel:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_endpgm
 entry:
-  %B = srem <32 x i64> zeroinitializer, zeroinitializer
-  ret void
+  %B = srem <2 x i64> zeroinitializer, zeroinitializer
+  ret <2 x i64> %B
 }