From 869ae115e4a109ec82bc07bfc09dfec53347a3a7 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Thu, 15 Feb 2024 09:45:38 -0800
Subject: [PATCH] [RISCV] Use APInt in useInversedSetcc to prevent crashes when
 mask is larger than UINT64_MAX.

There are no checks that the type is legal so we need to handle
any type.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp |  4 +-
 llvm/test/CodeGen/RISCV/condops.ll          | 51 +++++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 8235b536c4e00a..4f809ccf7b3019 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -15000,8 +15000,8 @@ static SDValue useInversedSetcc(SDNode *N, SelectionDAG &DAG,
     ISD::CondCode CC = cast<CondCodeSDNode>(Cond.getOperand(2))->get();
     if (CC == ISD::SETEQ && LHS.getOpcode() == ISD::AND &&
         isa<ConstantSDNode>(LHS.getOperand(1)) && isNullConstant(RHS)) {
-      uint64_t MaskVal = LHS.getConstantOperandVal(1);
-      if (isPowerOf2_64(MaskVal) && !isInt<12>(MaskVal))
+      const APInt &MaskVal = LHS.getConstantOperandAPInt(1);
+      if (MaskVal.isPowerOf2() && !MaskVal.isSignedIntN(12))
         return DAG.getSelect(DL, VT,
                              DAG.getSetCC(DL, CondVT, LHS, RHS, ISD::SETNE),
                              False, True);
diff --git a/llvm/test/CodeGen/RISCV/condops.ll b/llvm/test/CodeGen/RISCV/condops.ll
index 8e53782b5dcd78..101cb5aeeb0940 100644
--- a/llvm/test/CodeGen/RISCV/condops.ll
+++ b/llvm/test/CodeGen/RISCV/condops.ll
@@ -3719,3 +3719,54 @@ entry:
   %cond = select i1 %tobool.not, i64 0, i64 %x
   ret i64 %cond
 }
+
+; Test that we don't crash on types larger than 64 bits.
+define i64 @single_bit3(i80 %x, i64 %y) {
+; RV32I-LABEL: single_bit3:
+; RV32I:       # %bb.0: # %entry
+; RV32I-NEXT:    lw a0, 8(a0)
+; RV32I-NEXT:    slli a0, a0, 31
+; RV32I-NEXT:    srai a3, a0, 31
+; RV32I-NEXT:    and a0, a3, a1
+; RV32I-NEXT:    and a1, a3, a2
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: single_bit3:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    slli a1, a1, 63
+; RV64I-NEXT:    srai a0, a1, 63
+; RV64I-NEXT:    and a0, a0, a2
+; RV64I-NEXT:    ret
+;
+; RV64XVENTANACONDOPS-LABEL: single_bit3:
+; RV64XVENTANACONDOPS:       # %bb.0: # %entry
+; RV64XVENTANACONDOPS-NEXT:    andi a1, a1, 1
+; RV64XVENTANACONDOPS-NEXT:    vt.maskc a0, a2, a1
+; RV64XVENTANACONDOPS-NEXT:    ret
+;
+; RV64XTHEADCONDMOV-LABEL: single_bit3:
+; RV64XTHEADCONDMOV:       # %bb.0: # %entry
+; RV64XTHEADCONDMOV-NEXT:    slli a1, a1, 63
+; RV64XTHEADCONDMOV-NEXT:    srai a0, a1, 63
+; RV64XTHEADCONDMOV-NEXT:    and a0, a0, a2
+; RV64XTHEADCONDMOV-NEXT:    ret
+;
+; RV32ZICOND-LABEL: single_bit3:
+; RV32ZICOND:       # %bb.0: # %entry
+; RV32ZICOND-NEXT:    lw a0, 8(a0)
+; RV32ZICOND-NEXT:    andi a3, a0, 1
+; RV32ZICOND-NEXT:    czero.eqz a0, a1, a3
+; RV32ZICOND-NEXT:    czero.eqz a1, a2, a3
+; RV32ZICOND-NEXT:    ret
+;
+; RV64ZICOND-LABEL: single_bit3:
+; RV64ZICOND:       # %bb.0: # %entry
+; RV64ZICOND-NEXT:    andi a1, a1, 1
+; RV64ZICOND-NEXT:    czero.eqz a0, a2, a1
+; RV64ZICOND-NEXT:    ret
+entry:
+  %and = and i80 %x, 18446744073709551616 ; 1 << 64
+  %tobool.not = icmp eq i80 %and, 0
+  %cond = select i1 %tobool.not, i64 0, i64 %y
+  ret i64 %cond
+}