[RISCV][SelectionDAG] Add a hook to sign extend i32 ConstantInt operands of phis on RV64.

[craig.topper]

Summary

Materializing constants on RISCV is simpler if the constant is sign
extended from i32. By default i32 constant operands of phis are
zero extended.

This patch adds a hook to allow RISCV to override this for i32. We
have an existing isSExtCheaperThanZExt, but it operates on EVT which
we don't have at these places in the code.

[luismarques]

Makes sense to me.

[jrtc27]

llvm/test/CodeGen/RISCV/aext-to-sext.ll

80

Also marking %c as signext (which is what the psABI gives you) would avoid the sext.w noise

[frasercrmck]

llvm/include/llvm/CodeGen/TargetLowering.h

2682

The comment about "this" constant is a little incongruous as no constant is provided. Even if RISC-V doesn't need it, I can imagine situations where targets may choose to sign-extend some specific constants and not others, so it may make sense to pass the Constant to this API.

[craig.topper]

Address review comments

[luismarques]

llvm/test/CodeGen/RISCV/aext-to-sext.ll

80

Also marking %c as signext (which is what the psABI gives you) would avoid the sext.w noise

In external conversations, I've commented that IMO we should be using signext more extensively in the RISC-V tests (e.g. default to adding it to every parameter unless it doesn't actually make sense). I didn't seem to have managed to persuade my audience but, until I see further arguments against it, I still stand by my position :)

[craig.topper]

Ping.

[efriedma]

LGTM

llvm/include/llvm/CodeGen/TargetLowering.h

 //===- llvm/CodeGen/TargetLowering.h - Target Lowering Info -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///
public:
   }
 
   /// Return true if sign-extension from FromTy to ToTy is cheaper than
   /// zero-extension.
   virtual bool isSExtCheaperThanZExt(EVT FromTy, EVT ToTy) const {
     return false;
   }
 
+  /// Return true if this constant should be sign extended when promoting to
+  /// a larger type.
+  virtual bool signExtendConstant(const ConstantInt *C) const { return false; }
+
   /// Return true if sinking I's operands to the same basic block as I is
   /// profitable, e.g. because the operands can be folded into a target
   /// instruction during instruction selection. After calling the function
   /// \p Ops contains the Uses to sink ordered by dominance (dominating users
   /// come first).
   virtual bool shouldSinkOperands(Instruction *I,
                                   SmallVectorImpl<Use *> &Ops) const {
     return false;

llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp

 //===-- FunctionLoweringInfo.cpp ------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void FunctionLoweringInfo::ComputePHILiveOutRegInfo(const PHINode *PN) {
   Value *V = PN->getIncomingValue(0);
   if (isa<UndefValue>(V) || isa<ConstantExpr>(V)) {
     DestLOI.NumSignBits = 1;
     DestLOI.Known = KnownBits(BitWidth);
     return;
   }
 
   if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
-    APInt Val = CI->getValue().zextOrSelf(BitWidth);
+    APInt Val;
+    if (TLI->signExtendConstant(CI))
+      Val = CI->getValue().sextOrSelf(BitWidth);
+    else
+      Val = CI->getValue().zextOrSelf(BitWidth);
     DestLOI.NumSignBits = Val.getNumSignBits();
     DestLOI.Known = KnownBits::makeConstant(Val);
   } else {
     assert(ValueMap.count(V) && "V should have been placed in ValueMap when its"
                                 "CopyToReg node was created.");
     Register SrcReg = ValueMap[V];
     if (!Register::isVirtualRegister(SrcReg)) {
       DestLOI.IsValid = false;
for (unsigned i = 1, e = PN->getNumIncomingValues(); i != e; ++i) {
     Value *V = PN->getIncomingValue(i);
     if (isa<UndefValue>(V) || isa<ConstantExpr>(V)) {
       DestLOI.NumSignBits = 1;
       DestLOI.Known = KnownBits(BitWidth);
       return;
     }
 
     if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
-      APInt Val = CI->getValue().zextOrSelf(BitWidth);
+      APInt Val;
+      if (TLI->signExtendConstant(CI))
+        Val = CI->getValue().sextOrSelf(BitWidth);
+      else
+        Val = CI->getValue().zextOrSelf(BitWidth);
       DestLOI.NumSignBits = std::min(DestLOI.NumSignBits, Val.getNumSignBits());
       DestLOI.Known.Zero &= ~Val;
       DestLOI.Known.One &= Val;
       continue;
     }
 
     assert(ValueMap.count(V) && "V should have been placed in ValueMap when "
                                 "its CopyToReg node was created.");

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 //===- SelectionDAGBuilder.cpp - Selection-DAG building -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 /// Handle PHI nodes in successor blocks.  Emit code into the SelectionDAG to
 /// ensure constants are generated when needed.  Remember the virtual registers
 /// that need to be added to the Machine PHI nodes as input.  We cannot just
 /// directly add them, because expansion might result in multiple MBB's for one
 /// BB.  As such, the start of the BB might correspond to a different MBB than
 /// the end.
 void
 SelectionDAGBuilder::HandlePHINodesInSuccessorBlocks(const BasicBlock *LLVMBB) {
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   const Instruction *TI = LLVMBB->getTerminator();
 
   SmallPtrSet<MachineBasicBlock *, 4> SuccsHandled;
 
   // Check PHI nodes in successors that expect a value to be available from this
   // block.
   for (unsigned succ = 0, e = TI->getNumSuccessors(); succ != e; ++succ) {
     const BasicBlock *SuccBB = TI->getSuccessor(succ);
for (const PHINode &PN : SuccBB->phis()) {
 
       unsigned Reg;
       const Value *PHIOp = PN.getIncomingValueForBlock(LLVMBB);
 
       if (const Constant *C = dyn_cast<Constant>(PHIOp)) {
         unsigned &RegOut = ConstantsOut[C];
         if (RegOut == 0) {
           RegOut = FuncInfo.CreateRegs(C);
-          // We need to zero extend ConstantInt phi operands to match
+          // We need to zero/sign extend ConstantInt phi operands to match
           // assumptions in FunctionLoweringInfo::ComputePHILiveOutRegInfo.
-          ISD::NodeType ExtendType =
-              isa<ConstantInt>(PHIOp) ? ISD::ZERO_EXTEND : ISD::ANY_EXTEND;
+          ISD::NodeType ExtendType = ISD::ANY_EXTEND;
+          if (auto *CI = dyn_cast<ConstantInt>(C))
+            ExtendType = TLI.signExtendConstant(CI) ? ISD::SIGN_EXTEND
+                                                    : ISD::ZERO_EXTEND;
           CopyValueToVirtualRegister(C, RegOut, ExtendType);
         }
         Reg = RegOut;
       } else {
         DenseMap<const Value *, Register>::iterator I =
           FuncInfo.ValueMap.find(PHIOp);
         if (I != FuncInfo.ValueMap.end())
           Reg = I->second;
         else {
           assert(isa<AllocaInst>(PHIOp) &&
                  FuncInfo.StaticAllocaMap.count(cast<AllocaInst>(PHIOp)) &&
                  "Didn't codegen value into a register!??");
           Reg = FuncInfo.CreateRegs(PHIOp);
           CopyValueToVirtualRegister(PHIOp, Reg);
         }
       }
 
       // Remember that this register needs to added to the machine PHI node as
       // the input for this MBB.
       SmallVector<EVT, 4> ValueVTs;
-      const TargetLowering &TLI = DAG.getTargetLoweringInfo();
       ComputeValueVTs(TLI, DAG.getDataLayout(), PN.getType(), ValueVTs);
       for (unsigned vti = 0, vte = ValueVTs.size(); vti != vte; ++vti) {
         EVT VT = ValueVTs[vti];
         unsigned NumRegisters = TLI.getNumRegisters(*DAG.getContext(), VT);
         for (unsigned i = 0, e = NumRegisters; i != e; ++i)
           FuncInfo.PHINodesToUpdate.push_back(
               std::make_pair(&*MBBI++, Reg + i));
         Reg += NumRegisters;

llvm/lib/Target/RISCV/RISCVISelLowering.h

 //===-- RISCVISelLowering.h - RISCV DAG Lowering Interface ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
                              unsigned AS,
                              Instruction *I = nullptr) const override;
   bool isLegalICmpImmediate(int64_t Imm) const override;
   bool isLegalAddImmediate(int64_t Imm) const override;
   bool isTruncateFree(Type *SrcTy, Type *DstTy) const override;
   bool isTruncateFree(EVT SrcVT, EVT DstVT) const override;
   bool isZExtFree(SDValue Val, EVT VT2) const override;
   bool isSExtCheaperThanZExt(EVT SrcVT, EVT DstVT) const override;
+  bool signExtendConstant(const ConstantInt *CI) const override;
   bool isCheapToSpeculateCttz() const override;
   bool isCheapToSpeculateCtlz() const override;
   bool hasAndNotCompare(SDValue Y) const override;
   bool hasBitTest(SDValue X, SDValue Y) const override;
   bool shouldSinkOperands(Instruction *I,
                           SmallVectorImpl<Use *> &Ops) const override;
   bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

 //===-- RISCVISelLowering.cpp - RISCV DAG Lowering Implementation  --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
bool RISCVTargetLowering::isZExtFree(SDValue Val, EVT VT2) const {
 
   return TargetLowering::isZExtFree(Val, VT2);
 }
 
 bool RISCVTargetLowering::isSExtCheaperThanZExt(EVT SrcVT, EVT DstVT) const {
   return Subtarget.is64Bit() && SrcVT == MVT::i32 && DstVT == MVT::i64;
 }
 
+bool RISCVTargetLowering::signExtendConstant(const ConstantInt *CI) const {
+  return Subtarget.is64Bit() && CI->getType()->isIntegerTy(32);
+}
+
 bool RISCVTargetLowering::isCheapToSpeculateCttz() const {
   return Subtarget.hasStdExtZbb();
 }
 
 bool RISCVTargetLowering::isCheapToSpeculateCtlz() const {
   return Subtarget.hasStdExtZbb();
 }
 

llvm/test/CodeGen/RISCV/aext-to-sext.ll

bb:
   %c = select i1 %a, i32 8, i32 0
   %d = icmp eq i32 %b, %c
   br i1 %d, label %bb, label %bar
 
 bar:
   ret i32 %b
 }
 
-; The code that inserts AssertZExt based on predecessor information assumes
-; constants are zero extended for phi incoming values so an AssertZExt is
-; created in 'merge' allowing the zext to be removed.
-; SelectionDAG::getNode treats any_extend of i32 constants as sext for RISCV.
-; This code used to miscompile because the code that creates phi incoming values
-; in the predecessors created any_extend for the constants which then gets
-; treated as a sext by getNode. This made the removal of the zext incorrect.
-; SelectionDAGBuilder now creates a zero_extend instead of an any_extend to
-; match the assumption.
-; FIXME: RISCV would prefer constant inputs to phis to be sign extended.
-define i64 @miscompile(i32 %c) {
-; RV64I-LABEL: miscompile:
+; We prefer to sign extend i32 constants for phis. The default behavior in
+; SelectionDAGBuilder is zero extend we have a target hook to override it.
+define i64 @sext_phi_constants(i32 signext %c) {
+; RV64I-LABEL: sext_phi_constants:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sext.w a0, a0
-; RV64I-NEXT:    beqz a0, .LBB2_2
-; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    li a0, -1
+; RV64I-NEXT:    li a1, -1
+; RV64I-NEXT:    bnez a0, .LBB2_2
+; RV64I-NEXT:  # %bb.1: # %iffalse
+; RV64I-NEXT:    li a1, -2
+; RV64I-NEXT:  .LBB2_2: # %merge
+; RV64I-NEXT:    slli a0, a1, 32
 ; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    ret
-; RV64I-NEXT:  .LBB2_2: # %iffalse
-; RV64I-NEXT:    li a0, 1
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    addi a0, a0, -2
-; RV64I-NEXT:    ret
   %a = icmp ne i32 %c, 0
   br i1 %a, label %iftrue, label %iffalse
 
 iftrue:
   br label %merge
 
 iffalse:
   br label %merge
 
 merge:
   %b = phi i32 [-1, %iftrue], [-2, %iffalse]
   %d = zext i32 %b to i64
   ret i64 %d
 }