[RISCV] Do tail call when parameters are passed by stack

[StephenFan]

Summary

If parameters are passed by stack and their stack offsets are consistent
with stack offsets of caller's passed by stack parameters, tail call
optimization is feasible.

Fix https://github.com/llvm/llvm-project/issues/56908

[StephenFan]

Delete comments.

[StephenFan]

clang-format

[StephenFan]

Ping.

[luismarques]

I haven't noticed any specific issue so far but these things are fiddly, the patch needs more eye balls.
See if you can think of additional test coverage scenarios that might be relevant. Does this handle, say, when part of a value is passed on a register and part on the stack?

(OT: I just noticed that Phabricator is not syntax highlighting RISCVISelLowering.cpp. Is it because it's too big?)

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

11278

Is this for sure the only problematic flag? Having more test coverage of weird scenarios would help.

llvm/test/CodeGen/RISCV/tail-calls.ll

116

Nit: keep the comment but update it to explain that tail call optimization is possible if the stack layout is compatible

llvm/lib/Target/RISCV/RISCVISelLowering.h

private:
 
   SDValue lowerEH_DWARF_CFA(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue expandUnalignedRVVLoad(SDValue Op, SelectionDAG &DAG) const;
   SDValue expandUnalignedRVVStore(SDValue Op, SelectionDAG &DAG) const;
 
   bool isEligibleForTailCallOptimization(
       CCState &CCInfo, CallLoweringInfo &CLI, MachineFunction &MF,
-      const SmallVector<CCValAssign, 16> &ArgLocs) const;
+      const SmallVector<CCValAssign, 16> &ArgLocs, const SmallVectorImpl<SDValue> &OutVals,
+      const SmallVectorImpl<ISD::InputArg> &Ins) const;
 
   /// Generate error diagnostics if any register used by CC has been marked
   /// reserved.
   void validateCCReservedRegs(
       const SmallVectorImpl<std::pair<llvm::Register, llvm::SDValue>> &Regs,
       MachineFunction &MF) const;
 
   bool useRVVForFixedLengthVectorVT(MVT VT) const;

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

SDValue RISCVTargetLowering::LowerFormalArguments(
   if (!OutChains.empty()) {
     OutChains.push_back(Chain);
     Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, OutChains);
   }
 
   return Chain;
 }
 
+/// MatchingStackOffset - Return true if the given stack call argument is
+/// already available in the same position (relatively) of the caller's
+/// incoming argument stack.
+static
+bool MatchingStackOffset(SDValue Arg, unsigned Offset, ISD::ArgFlagsTy Flags,
+                         MachineFrameInfo &MFI, const MachineRegisterInfo *MRI,
+                         const TargetInstrInfo *TII) {
+  unsigned Bytes = Arg.getValueSizeInBits() / 8;
+  int FI = std::numeric_limits<int>::max();
+  if (LoadSDNode *Ld = dyn_cast<LoadSDNode>(Arg)) {
+    if (Flags.isByVal())
+      // ByVal argument is passed in as a pointer but it's now being
+      // dereferenced. e.g.
+      // define @foo(%struct.X* %A) {
+      //   tail call @bar(%struct.X* byval %A)
+      // }
+      return false;
+    SDValue Ptr = Ld->getBasePtr();
+    FrameIndexSDNode *FINode = dyn_cast<FrameIndexSDNode>(Ptr);
+    if (!FINode)
+      return false;
+    FI = FINode->getIndex();
+  } else
+    return false;
+
+  assert(FI != std::numeric_limits<int>::max());
+  if (!MFI.isFixedObjectIndex(FI))
+    return false;
+  return Offset == MFI.getObjectOffset(FI) && Bytes == MFI.getObjectSize(FI);
+}
+
 /// isEligibleForTailCallOptimization - Check whether the call is eligible
 /// for tail call optimization.
 /// Note: This is modelled after ARM's IsEligibleForTailCallOptimization.
 bool RISCVTargetLowering::isEligibleForTailCallOptimization(
     CCState &CCInfo, CallLoweringInfo &CLI, MachineFunction &MF,
-    const SmallVector<CCValAssign, 16> &ArgLocs) const {
+    const SmallVector<CCValAssign, 16> &ArgLocs, const SmallVectorImpl<SDValue> &OutVals,
+    const SmallVectorImpl<ISD::InputArg> &Ins) const {
 
   auto &Callee = CLI.Callee;
   auto CalleeCC = CLI.CallConv;
   auto &Outs = CLI.Outs;
   auto &Caller = MF.getFunction();
   auto CallerCC = Caller.getCallingConv();
 
   // Exception-handling functions need a special set of instructions to
   // indicate a return to the hardware. Tail-calling another function would
   // probably break this.
   // TODO: The "interrupt" attribute isn't currently defined by RISC-V. This
   // should be expanded as new function attributes are introduced.
   if (Caller.hasFnAttribute("interrupt"))
     return false;
 
-  // Do not tail call opt if the stack is used to pass parameters.
-  if (CCInfo.getNextStackOffset() != 0)
+  if (CCInfo.getNextStackOffset() != 0) {
+    MachineFrameInfo &MFI = MF.getFrameInfo();
+    const MachineRegisterInfo *MRI = &MF.getRegInfo();
+    const TargetInstrInfo *TII = Subtarget.getInstrInfo();
+    for (unsigned I = 0, RealArgIdx = 0, E = ArgLocs.size();
+         I != E; ++I, ++RealArgIdx) {
+      const CCValAssign &VA = ArgLocs[I];
+      SDValue Arg = OutVals[RealArgIdx];
+      ISD::ArgFlagsTy Flags = Outs[RealArgIdx].Flags;
+      if (VA.getLocInfo() == CCValAssign::Indirect)
         return false;
+      else if (!VA.isRegLoc()) {
+        if (!MatchingStackOffset(Arg, VA.getLocMemOffset(), Flags, MFI, MRI, TII))
+          return false;
+      }
+    }
+
+  }
 
   // Do not tail call opt if any parameters need to be passed indirectly.
   // Since long doubles (fp128) and i128 are larger than 2*XLEN, they are
   // passed indirectly. So the address of the value will be passed in a
   // register, or if not available, then the address is put on the stack. In
   // order to pass indirectly, space on the stack often needs to be allocated
   // in order to store the value. In this case the CCInfo.getNextStackOffset()
   // != 0 check is not enough and we need to check if any CCValAssign ArgsLocs
if (CallConv == CallingConv::GHC)
     ArgCCInfo.AnalyzeCallOperands(Outs, CC_RISCV_GHC);
   else
     analyzeOutputArgs(MF, ArgCCInfo, Outs, /*IsRet=*/false, &CLI,
                       CallConv == CallingConv::Fast ? CC_RISCV_FastCC
                                                     : CC_RISCV);
 
   // Check if it's really possible to do a tail call.
   if (IsTailCall)
-    IsTailCall = isEligibleForTailCallOptimization(ArgCCInfo, CLI, MF, ArgLocs);
+    IsTailCall = isEligibleForTailCallOptimization(ArgCCInfo, CLI, MF, ArgLocs, OutVals, Ins);
 
   if (IsTailCall)
     ++NumTailCalls;
   else if (CLI.CB && CLI.CB->isMustTailCall())
     report_fatal_error("failed to perform tail call elimination on a call "
                        "site marked musttail");
 
   // Get a count of how many bytes are to be pushed on the stack.
for (unsigned i = 0, j = 0, e = ArgLocs.size(); i != e; ++i) {
 
     // Use local copy if it is a byval arg.
     if (Flags.isByVal())
       ArgValue = ByValArgs[j++];
 
     if (VA.isRegLoc()) {
       // Queue up the argument copies and emit them at the end.
       RegsToPass.push_back(std::make_pair(VA.getLocReg(), ArgValue));
-    } else {
+    } else if (!IsTailCall) {
       assert(VA.isMemLoc() && "Argument not register or memory");
-      assert(!IsTailCall && "Tail call not allowed if stack is used "
-                            "for passing parameters");
 
       // Work out the address of the stack slot.
       if (!StackPtr.getNode())
         StackPtr = DAG.getCopyFromReg(Chain, DL, RISCV::X2, PtrVT);
       SDValue Address =
           DAG.getNode(ISD::ADD, DL, PtrVT, StackPtr,
                       DAG.getIntPtrConstant(VA.getLocMemOffset(), DL));
 

llvm/test/CodeGen/RISCV/tail-calls.ll

 ; CHECK-NEXT:    lw ra, 12(sp) # 4-byte Folded Reload
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
 entry:
   %call = tail call i32 (i32, ...) @callee_varargs(i32 %a, i32 %b, i32 %b, i32 %a, i32 %a, i32 %b, i32 %b, i32 %a, i32 %a)
   ret void
 }
 
-; Do not tail call optimize if stack is used to pass parameters.
 declare i32 @callee_args(i32 %a, i32 %b, i32 %c, i32 %dd, i32 %e, i32 %ff, i32 %g, i32 %h, i32 %i, i32 %j, i32 %k, i32 %l, i32 %m, i32 %n)
 define i32 @caller_args(i32 %a, i32 %b, i32 %c, i32 %dd, i32 %e, i32 %ff, i32 %g, i32 %h, i32 %i, i32 %j, i32 %k, i32 %l, i32 %m, i32 %n) nounwind {
 ; CHECK-LABEL: caller_args:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    addi sp, sp, -32
-; CHECK-NEXT:    sw ra, 28(sp) # 4-byte Folded Spill
-; CHECK-NEXT:    lw t0, 32(sp)
-; CHECK-NEXT:    lw t1, 36(sp)
-; CHECK-NEXT:    lw t2, 40(sp)
-; CHECK-NEXT:    lw t3, 44(sp)
-; CHECK-NEXT:    lw t4, 48(sp)
-; CHECK-NEXT:    lw t5, 52(sp)
-; CHECK-NEXT:    sw t5, 20(sp)
-; CHECK-NEXT:    sw t4, 16(sp)
-; CHECK-NEXT:    sw t3, 12(sp)
-; CHECK-NEXT:    sw t2, 8(sp)
-; CHECK-NEXT:    sw t1, 4(sp)
-; CHECK-NEXT:    sw t0, 0(sp)
-; CHECK-NEXT:    call callee_args@plt
-; CHECK-NEXT:    lw ra, 28(sp) # 4-byte Folded Reload
-; CHECK-NEXT:    addi sp, sp, 32
-; CHECK-NEXT:    ret
+; CHECK-NEXT:    tail callee_args@plt
 entry:
   %r = tail call i32 @callee_args(i32 %a, i32 %b, i32 %c, i32 %dd, i32 %e, i32 %ff, i32 %g, i32 %h, i32 %i, i32 %j, i32 %k, i32 %l, i32 %m, i32 %n)
   ret i32 %r
 }
 
 ; Do not tail call optimize if parameters need to be passed indirectly.
 declare i32 @callee_indirect_args(fp128 %a)
 define void @caller_indirect_args() nounwind {