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

[via llvm-commits]

https://github.com/Yunzezhu94 updated https://github.com/llvm/llvm-project/pull/97391

>From 418536e8fefc7e94dd446fc98cd1fc75e831ccde Mon Sep 17 00:00:00 2001
From: Yunze Zhu <yunzezhu at linux.alibaba.com>
Date: Tue, 21 May 2024 20:55:51 +0800
Subject: [PATCH] [RISCV] Add a check in lowerSELECT after
 foldBinOpIntoSelectIfProfitable

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.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp        |  4 +++-
 .../fold-binop-into-select-return-constant.ll      | 14 ++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/RISCV/fold-binop-into-select-return-constant.ll

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index ee45f730dc450..3b6d4dcaa357b 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -7664,7 +7664,9 @@ 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);
+        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 0000000000000..000da23b37043
--- /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
+}