[RISCV] Add a check in lowerSELECT after foldBinOpIntoSelectIfProfitable

[Yunzezhu94]

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 #97390

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: None (Yunzezhu94)

Changes

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: #97390

---

Full diff: https://github.com/llvm/llvm-project/pull/97391.diff

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+4-1)
• (modified) llvm/test/CodeGen/RISCV/fold-binop-into-select.ll (+12)

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
+}

[dtcxzyw]

LGTM. Nice catch!

[dtcxzyw]

Missing newline.

[dtcxzyw]

Thank you!

[mgudim]

Can you please add a comment explaining why the check is needed - it doesn't seem obvious.

[mgudim]

I think using ternary operator is preferred in LLVM

[topperc]

The preferred thing would be to do

if (NewSel.getOpcode() == ISD::SELECT)
  return lowerSELECT(NewSel, DAG);
return NewSel;

[mgudim]

@Yunzezhu94 thanks for working on this, looks good. I left some minor comments.

[topperc]

The title and the description don't say what the extra check is.

[mgudim]

The title and the description don't say what the extra check is.

The description says what the check is.

[topperc]

The title and the description don't say what the extra check is.

The description says what the check is.

It could be a little better in my opinion. The first sentence says that foldBinOpIntoSelectIfProfitable returned a constant node. The second sentence says a new check is added. Reading the two sentences together, you might think that the new check is that the node isn't a constant. But that's not exactly what the new check is.

[Yunzezhu94]

Updated code and description as advised.

[topperc]

Updated code and description as advised.

I think you only updated the commit message. You need to update the PR description too. The "squash and merge" will use the PR description not the commit message.

[mgudim]

LGTM, thanks!

[topperc]

LGTM