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

[via llvm-commits]

https://github.com/Yunzezhu94 created https://github.com/llvm/llvm-project/pull/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.

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

>From ebdc2aa48dc151431c081e13721b500fafb42a95 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 ++++-
 llvm/test/CodeGen/RISCV/fold-binop-into-select.ll | 12 ++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

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.ll b/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll
index 1512db87b9311..8b9bafbbf4476 100644
--- a/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll
+++ b/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll
@@ -58,3 +58,15 @@ entry:
   %res = sub i64 2, %select_
   ret i64 %res
 }
+
+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
+}