[RISCV] Add a check in lowerSELECT after foldBinOpIntoSelectIfProfitable (#97391)

In certain case foldBinOpIntoSelectIfProfitable may return a constant
node, the node will be lowered in lowerSELECT and lead to crash.
This patch fix the bug by adding an extra check before lowerSELECT that
do lowerSELECT as before when foldBinOpIntoSelectIfProfitable returns a
select node, and return the node directly when
foldBinOpIntoSelectIfProfitable returns a constant node.

Fixes https://github.com/llvm/llvm-project/issues/97390

2 files changed, 19 insertions, 1 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 45368a0..c0a9ecf 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -7685,7 +7685,11 @@ SDValue RISCVTargetLowering::lowerSELECT(SDValue Op, SelectionDAG &DAG) const {
       if (SDValue NewSel = foldBinOpIntoSelectIfProfitable(*Op->use_begin(),
                                                            DAG, Subtarget)) {
         DAG.ReplaceAllUsesWith(BinOp, &NewSel);
-        return lowerSELECT(NewSel, DAG);
+        // Opcode check is necessary because foldBinOpIntoSelectIfProfitable
+        // may return a constant node and cause crash in lowerSELECT.
+        if (NewSel.getOpcode() == ISD::SELECT)
+          return lowerSELECT(NewSel, DAG);
+        return NewSel;
       }
     }
   }
diff --git a/llvm/test/CodeGen/RISCV/fold-binop-into-select-return-constant.ll b/llvm/test/CodeGen/RISCV/fold-binop-into-select-return-constant.ll
new file mode 100644
index 0000000..000da23
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/fold-binop-into-select-return-constant.ll
@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=riscv64 -mattr=+m < %s  | FileCheck  %s
+
+define i64 @fold_binop_into_select_return_constant(i1 %c) {
+; CHECK-LABEL: fold_binop_into_select_return_constant:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ret
+entry:
+  %select1 = select i1 %c, i32 4, i32 8
+  %select2 = sext i32 %select1 to i64
+  %div1 = sdiv i64 %select2, -5141143369814759789
+  ret i64 %div1
+}