[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 7cc3fbc3810965712a6f14b0dc867f51af7648c6 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.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp         |  5 ++++-
 .../RISCV/fold-binop-into-select-return-constant.ll | 13 +++++++++++++
 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..5b3ad99ce3711 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -7664,7 +7664,10 @@ 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);
+        else
+          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..8ec2f88ae915f
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/fold-binop-into-select-return-constant.ll
@@ -0,0 +1,13 @@
+; 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
+}
\ No newline at end of file