[SelectionDAG] Use correct boolean representation in FoldConstantArithmetic

[bjope]

Summary

The description of SETCC says

/// SetCC operator - This evaluates to a true value iff the condition is
/// true.  If the result value type is not i1 then the high bits conform
/// to getBooleanContents.

Without this patch, we sign extended the i1 to the used larger type
regardless of getBooleanContents. This resulted in miscompiles, as
shown in the attached testcase that ended up returning -1 instead of
1 when using -mattr=+v.

Fixes https://github.com/llvm/llvm-project/issues/55168

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5587

"ExtendToBooleanContent" should be sufficient

[bjope]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5649

Not sure if it perhaps would be OK (or even more correct) to always extend according to getBooleanContents here if ScalarResult:s type is i1?

(when preparing this patch it felt less intrusive to just consider the case when SVT has been adjusted down for a SETCC above)

[bjope]

Shorten a variable name according to review feedback.

[craig.topper]

I want to look more closely at this. Aren't the extended bits are going into a build_vector or splat_vector that is supposed to ignore the extra bits because they are past the size of the result element type?

[craig.topper]

In D124618#3480230, @craig.topper wrote:

I want to look more closely at this. Aren't the extended bits are going into a build_vector or splat_vector that is supposed to ignore the extra bits because they are past the size of the result element type?

Oh, I see. I missed that we produce an i1 setcc result. So it's not extending past the element result type.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5654

Can we pick the extend opcode outside the loop?

[bjope]

Calculate the ExtendCode outside the loop (also no longer needing the extra bool
to determine how we extend inside the loop).

(Thanks Craig for this idea!)

[craig.topper]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
 
   // All operands must be vector types with the same number of elements as
   // the result type and must be either UNDEF or a build/splat vector
   // or UNDEF scalars.
   if (!llvm::all_of(Ops, IsBuildVectorSplatVectorOrUndef) ||
       !llvm::all_of(Ops, IsScalarOrSameVectorSize))
     return SDValue();
 
-  // If we are comparing vectors, then the result needs to be a i1 boolean
-  // that is then sign-extended back to the legal result type.
+  // If we are comparing vectors, then the result needs to be a i1 boolean that
+  // is then extended back to the legal result type depending on how booleans
+  // are represented.
   EVT SVT = (Opcode == ISD::SETCC ? MVT::i1 : VT.getScalarType());
+  ISD::NodeType ExtendCode =
+      (Opcode == ISD::SETCC && SVT != VT.getScalarType())
+          ? TargetLowering::getExtendForContent(TLI->getBooleanContents(VT))
+          : ISD::SIGN_EXTEND;
 
   // Find legal integer scalar type for constant promotion and
   // ensure that its scalar size is at least as large as source.
   EVT LegalSVT = VT.getScalarType();
   if (NewNodesMustHaveLegalTypes && LegalSVT.isInteger()) {
     LegalSVT = TLI->getTypeToTransformTo(*getContext(), LegalSVT);
     if (LegalSVT.bitsLT(VT.getScalarType()))
       return SDValue();
for (SDValue Op : Ops) {
       ScalarOps.push_back(ScalarOp);
     }
 
     // Constant fold the scalar operands.
     SDValue ScalarResult = getNode(Opcode, DL, SVT, ScalarOps);
 
     // Legalize the (integer) scalar constant if necessary.
     if (LegalSVT != SVT)
-      ScalarResult = getNode(ISD::SIGN_EXTEND, DL, LegalSVT, ScalarResult);
+      ScalarResult = getNode(ExtendCode, DL, LegalSVT, ScalarResult);
 
     // Scalar folding only succeeded if the result is a constant or UNDEF.
     if (!ScalarResult.isUndef() && ScalarResult.getOpcode() != ISD::Constant &&
         ScalarResult.getOpcode() != ISD::ConstantFP)
       return SDValue();
     ScalarResults.push_back(ScalarResult);
   }
 
   SDValue V = NumElts.isScalable() ? getSplatVector(VT, DL, ScalarResults[0])
                                    : getBuildVector(VT, DL, ScalarResults);
   NewSDValueDbgMsg(V, "New node fold constant vector: ", this);
   return V;
 }
 

llvm/test/CodeGen/RISCV/fold-vector-cmp.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -start-after codegenprepare -mtriple=riscv64 -mattr=-v -o - %s | FileCheck %s
+; RUN: llc -start-after codegenprepare -mtriple=riscv64 -mattr=+v -o - %s | FileCheck %s
+
+; Reproducer for https://github.com/llvm/llvm-project/issues/55168.
+; We should always return 1 (and not -1).
+define i32 @test(i32 %call.i) {
+; CHECK-LABEL: test:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 1
+; CHECK-NEXT:    ret
+  %t2 = insertelement <2 x i32> <i32 poison, i32 -2147483648>, i32 %call.i, i64 0
+  %t3 = icmp slt <2 x i32> %t2, <i32 -2147483646, i32 -2147483646>
+  %t4 = zext <2 x i1> %t3 to <2 x i32>
+  %t6 = extractelement <2 x i32> %t4, i64 1
+  ret i32 %t6
+}