[RISCV] Fix musttail with indirect arguments by forwarding incoming pointers

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -24462,6 +24462,8 @@ SDValue RISCVTargetLowering::LowerFormalArguments(
       InVals.push_back(DAG.getLoad(VA.getValVT(), DL, Chain, ArgValue,
                                    MachinePointerInfo()));
       unsigned ArgIndex = Ins[InsIdx].OrigArgIndex;
+      // Save the incoming indirect pointer for musttail forwarding.
+      RVFI->setIncomingIndirectArg(ArgIndex, ArgValue);
       unsigned ArgPartOffset = Ins[InsIdx].PartOffset;
       assert(VA.getValVT().isVector() || ArgPartOffset == 0);
       while (i + 1 != e && Ins[InsIdx + 1].OrigArgIndex == ArgIndex) {
@@ -24578,6 +24580,19 @@ bool RISCVTargetLowering::isEligibleForTailCallOptimization(
   if (CCInfo.getStackSize() > RVFI->getArgumentStackSize())
     return false;
 
+  // Do not tail call optimize if any argument needs to be passed indirectly.
+  // The caller allocates stack space and passes a pointer to the callee. On a
+  // tail call the caller's stack frame is deallocated before the callee
+  // executes, invalidating the pointer (use-after-free).
+  // musttail is excluded: callers forward incoming indirect pointers that
+  // point to the caller's caller's frame, which remains valid.
+  if (!CLI.CB || !CLI.CB->isMustTailCall()) {
+    for (const auto &VA : ArgLocs) {
+      if (VA.getLocInfo() == CCValAssign::Indirect)
+        return false;
+    }
+  }
+
   // Do not tail call opt if either caller or callee uses struct return
   // semantics.
   auto IsCallerStructRet = Caller.hasStructRetAttr();
@@ -24765,51 +24780,65 @@ SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
     // Promote the value if needed.
     // For now, only handle fully promoted and indirect arguments.
     if (VA.getLocInfo() == CCValAssign::Indirect) {
-      // Store the argument in a stack slot and pass its address.
-      Align StackAlign =
-          std::max(getPrefTypeAlign(Outs[OutIdx].ArgVT, DAG),
-                   getPrefTypeAlign(ArgValue.getValueType(), DAG));
-      TypeSize StoredSize = ArgValue.getValueType().getStoreSize();
-      // If the original argument was split (e.g. i128), we need
-      // to store the required parts of it here (and pass just one address).
-      // Vectors may be partly split to registers and partly to the stack, in
-      // which case the base address is partly offset and subsequent stores are
-      // relative to that.
-      unsigned ArgIndex = Outs[OutIdx].OrigArgIndex;
-      unsigned ArgPartOffset = Outs[OutIdx].PartOffset;
-      assert(VA.getValVT().isVector() || ArgPartOffset == 0);
-      // Calculate the total size to store. We don't have access to what we're
-      // actually storing other than performing the loop and collecting the
-      // info.
-      SmallVector<std::pair<SDValue, SDValue>> Parts;
-      while (i + 1 != e && Outs[OutIdx + 1].OrigArgIndex == ArgIndex) {
-        SDValue PartValue = OutVals[OutIdx + 1];
-        unsigned PartOffset = Outs[OutIdx + 1].PartOffset - ArgPartOffset;
-        SDValue Offset = DAG.getIntPtrConstant(PartOffset, DL);
-        EVT PartVT = PartValue.getValueType();
-        if (PartVT.isScalableVector())
-          Offset = DAG.getNode(ISD::VSCALE, DL, XLenVT, Offset);
-        StoredSize += PartVT.getStoreSize();
-        StackAlign = std::max(StackAlign, getPrefTypeAlign(PartVT, DAG));
-        Parts.push_back(std::make_pair(PartValue, Offset));
-        ++i;
-        ++OutIdx;
-      }
-      SDValue SpillSlot = DAG.CreateStackTemporary(StoredSize, StackAlign);
-      int FI = cast<FrameIndexSDNode>(SpillSlot)->getIndex();
-      MemOpChains.push_back(
-          DAG.getStore(Chain, DL, ArgValue, SpillSlot,
-                       MachinePointerInfo::getFixedStack(MF, FI)));
-      for (const auto &Part : Parts) {
-        SDValue PartValue = Part.first;
-        SDValue PartOffset = Part.second;
-        SDValue Address =
-            DAG.getNode(ISD::ADD, DL, PtrVT, SpillSlot, PartOffset);
+      // For musttail calls, forward the incoming indirect pointer instead
+      // of creating a new stack temporary. The incoming pointer points to
+      // the caller's caller's frame, which remains valid after a tail call.
+      if (IsTailCall && CLI.CB && CLI.CB->isMustTailCall()) {
+        unsigned ArgIndex = Outs[OutIdx].OrigArgIndex;
+        ArgValue = RVFI->getIncomingIndirectArg(ArgIndex);
+        // Skip any split parts of this argument (they are covered by the
+        // forwarded pointer).
+        while (i + 1 != e && Outs[OutIdx + 1].OrigArgIndex == ArgIndex) {
+          ++i;
+          ++OutIdx;
+        }
+      } else {
+        // Store the argument in a stack slot and pass its address.
+        Align StackAlign =
+            std::max(getPrefTypeAlign(Outs[OutIdx].ArgVT, DAG),
+                     getPrefTypeAlign(ArgValue.getValueType(), DAG));
+        TypeSize StoredSize = ArgValue.getValueType().getStoreSize();
+        // If the original argument was split (e.g. i128), we need
+        // to store the required parts of it here (and pass just one address).
+        // Vectors may be partly split to registers and partly to the stack, in
+        // which case the base address is partly offset and subsequent stores
+        // are relative to that.
+        unsigned ArgIndex = Outs[OutIdx].OrigArgIndex;
+        unsigned ArgPartOffset = Outs[OutIdx].PartOffset;
+        assert(VA.getValVT().isVector() || ArgPartOffset == 0);
+        // Calculate the total size to store. We don't have access to what
+        // we're actually storing other than performing the loop and collecting
+        // the info.
+        SmallVector<std::pair<SDValue, SDValue>> Parts;
+        while (i + 1 != e && Outs[OutIdx + 1].OrigArgIndex == ArgIndex) {
+          SDValue PartValue = OutVals[OutIdx + 1];
+          unsigned PartOffset = Outs[OutIdx + 1].PartOffset - ArgPartOffset;
+          SDValue Offset = DAG.getIntPtrConstant(PartOffset, DL);
+          EVT PartVT = PartValue.getValueType();
+          if (PartVT.isScalableVector())
+            Offset = DAG.getNode(ISD::VSCALE, DL, XLenVT, Offset);
+          StoredSize += PartVT.getStoreSize();
+          StackAlign = std::max(StackAlign, getPrefTypeAlign(PartVT, DAG));
+          Parts.push_back(std::make_pair(PartValue, Offset));
+          ++i;
+          ++OutIdx;
+        }
+        SDValue SpillSlot = DAG.CreateStackTemporary(StoredSize, StackAlign);
+        int FI = cast<FrameIndexSDNode>(SpillSlot)->getIndex();
         MemOpChains.push_back(
-            DAG.getStore(Chain, DL, PartValue, Address,
+            DAG.getStore(Chain, DL, ArgValue, SpillSlot,
                          MachinePointerInfo::getFixedStack(MF, FI)));
+        for (const auto &Part : Parts) {
+          SDValue PartValue = Part.first;
+          SDValue PartOffset = Part.second;
+          SDValue Address =
+              DAG.getNode(ISD::ADD, DL, PtrVT, SpillSlot, PartOffset);
+          MemOpChains.push_back(
+              DAG.getStore(Chain, DL, PartValue, Address,
+                           MachinePointerInfo::getFixedStack(MF, FI)));
+        }
+        ArgValue = SpillSlot;
       }
-      ArgValue = SpillSlot;
     } else {
       ArgValue = convertValVTToLocVT(DAG, ArgValue, VA, DL, Subtarget);
     }

llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.h

@@ -14,6 +14,7 @@
 #define LLVM_LIB_TARGET_RISCV_RISCVMACHINEFUNCTIONINFO_H
 
 #include "RISCVSubtarget.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -73,6 +74,10 @@ class RISCVMachineFunctionInfo : public MachineFunctionInfo {
   /// Incoming ByVal arguments
   SmallVector<SDValue, 8> IncomingByValArgs;
 
+  /// Incoming indirect argument pointers, keyed by OrigArgIndex.
+  /// Used for musttail forwarding of indirect args.
+  DenseMap<unsigned, SDValue> IncomingIndirectArgs;
+
   /// Is there any vector argument or return?
   bool IsVectorCall = false;
 
@@ -157,6 +162,15 @@ class RISCVMachineFunctionInfo : public MachineFunctionInfo {
   SDValue getIncomingByValArgs(unsigned Idx) { return IncomingByValArgs[Idx]; }
   unsigned getIncomingByValArgsSize() const { return IncomingByValArgs.size(); }
 
+  void setIncomingIndirectArg(unsigned ArgIndex, SDValue Val) {
+    IncomingIndirectArgs[ArgIndex] = Val;
+  }
+  SDValue getIncomingIndirectArg(unsigned ArgIndex) const {
+    auto It = IncomingIndirectArgs.find(ArgIndex);
+    assert(It != IncomingIndirectArgs.end() && "No incoming indirect arg");
+    return It->second;
+  }
+
   enum class PushPopKind { None = 0, StdExtZcmp, VendorXqccmp };
 
   PushPopKind getPushPopKind(const MachineFunction &MF) const;

llvm/test/CodeGen/RISCV/musttail-indirect-args.ll

@@ -0,0 +1,143 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=riscv32 %s -o - | FileCheck %s --check-prefix=RV32
+; RUN: llc -mtriple=riscv64 %s -o - | FileCheck %s --check-prefix=RV64
+
+; Test that musttail with indirect args (fp128 on RV32) forwards the incoming
+; pointer instead of creating a new stack temporary. Without this fix, the
+; pointer would dangle after the tail call deallocates the caller's frame.
+
+declare i32 @callee_musttail_indirect(fp128 %a)
+
+; fp128 is indirect on RV32 (too large for registers), direct on RV64.
+; On RV32, musttail must forward the incoming indirect pointer (a0) directly.
+define i32 @caller_musttail_indirect(fp128 %a) nounwind {
+; RV32-LABEL: caller_musttail_indirect:
+; RV32:       # %bb.0:
+; RV32-NEXT:    tail callee_musttail_indirect
+;
+; RV64-LABEL: caller_musttail_indirect:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_indirect
+  %call = musttail call i32 @callee_musttail_indirect(fp128 %a)
+  ret i32 %call
+}
+
+; Verify that non-musttail tail call with indirect args does NOT tail call
+; (this is the PR #184972 fix - indirect args are unsafe for regular tail calls).
+define void @caller_no_musttail_indirect() nounwind {
+; RV32-LABEL: caller_no_musttail_indirect:
+; RV32:       # %bb.0:
+; RV32-NEXT:    addi sp, sp, -32
+; RV32-NEXT:    sw ra, 28(sp) # 4-byte Folded Spill
+; RV32-NEXT:    lui a1, 262128
+; RV32-NEXT:    mv a0, sp
+; RV32-NEXT:    sw zero, 0(sp)
+; RV32-NEXT:    sw zero, 4(sp)
+; RV32-NEXT:    sw zero, 8(sp)
+; RV32-NEXT:    sw a1, 12(sp)
+; RV32-NEXT:    call callee_musttail_indirect
+; RV32-NEXT:    lw ra, 28(sp) # 4-byte Folded Reload
+; RV32-NEXT:    addi sp, sp, 32
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: caller_no_musttail_indirect:
+; RV64:       # %bb.0:
+; RV64-NEXT:    lui a1, 16383
+; RV64-NEXT:    slli a1, a1, 36
+; RV64-NEXT:    li a0, 0
+; RV64-NEXT:    tail callee_musttail_indirect
+  %call = tail call i32 @callee_musttail_indirect(fp128 0xL00000000000000003FFF000000000000)
+  ret void
+}
+
+; Verify that non-musttail tail call forwarding an indirect arg from the
+; caller's own parameters also does NOT tail call (the arg lives on the
+; caller's frame, which would be deallocated).
+define i32 @caller_no_musttail_forward_indirect(fp128 %a) nounwind {
+; RV32-LABEL: caller_no_musttail_forward_indirect:
+; RV32:       # %bb.0:
+; RV32-NEXT:    addi sp, sp, -32
+; RV32-NEXT:    sw ra, 28(sp) # 4-byte Folded Spill
+; RV32-NEXT:    lw a1, 0(a0)
+; RV32-NEXT:    lw a2, 4(a0)
+; RV32-NEXT:    lw a3, 8(a0)
+; RV32-NEXT:    lw a4, 12(a0)
+; RV32-NEXT:    mv a0, sp
+; RV32-NEXT:    sw a1, 0(sp)
+; RV32-NEXT:    sw a2, 4(sp)
+; RV32-NEXT:    sw a3, 8(sp)
+; RV32-NEXT:    sw a4, 12(sp)
+; RV32-NEXT:    call callee_musttail_indirect
+; RV32-NEXT:    lw ra, 28(sp) # 4-byte Folded Reload
+; RV32-NEXT:    addi sp, sp, 32
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: caller_no_musttail_forward_indirect:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_indirect
+  %call = tail call i32 @callee_musttail_indirect(fp128 %a)
+  ret i32 %call
+}
+
+; Test musttail with two indirect fp128 args on RV32. Both pointers must be
+; forwarded. Exercises the DenseMap with two distinct OrigArgIndex values.
+declare i32 @callee_musttail_two_indirect(fp128 %a, fp128 %b)
+
+define i32 @caller_musttail_two_indirect(fp128 %a, fp128 %b) nounwind {
+; RV32-LABEL: caller_musttail_two_indirect:
+; RV32:       # %bb.0:
+; RV32-NEXT:    tail callee_musttail_two_indirect
+;
+; RV64-LABEL: caller_musttail_two_indirect:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_two_indirect
+  %call = musttail call i32 @callee_musttail_two_indirect(fp128 %a, fp128 %b)
+  ret i32 %call
+}
+
+; Test musttail with mixed direct (i32 in register) + indirect (fp128) args.
+; Confirms OrigArgIndex lookup works when not all args are indirect.
+declare i32 @callee_musttail_mixed(i32 %x, fp128 %a)
+
+define i32 @caller_musttail_mixed(i32 %x, fp128 %a) nounwind {
+; RV32-LABEL: caller_musttail_mixed:
+; RV32:       # %bb.0:
+; RV32-NEXT:    tail callee_musttail_mixed
+;
+; RV64-LABEL: caller_musttail_mixed:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_mixed
+  %call = musttail call i32 @callee_musttail_mixed(i32 %x, fp128 %a)
+  ret i32 %call
+}
+
+; Test musttail with i128 on RV32 (indirect, split into 4 x i32 parts).
+declare i64 @callee_musttail_i128(i128 %a)
+
+define i64 @caller_musttail_i128(i128 %a) nounwind {
+; RV32-LABEL: caller_musttail_i128:
+; RV32:       # %bb.0:
+; RV32-NEXT:    tail callee_musttail_i128
+;
+; RV64-LABEL: caller_musttail_i128:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_i128
+  %call = musttail call i64 @callee_musttail_i128(i128 %a)
+  ret i64 %call
+}
+
+; Test musttail with i128 (indirect+split on RV32) plus a trailing i32 direct arg.
+; Exercises the split-skip logic followed by a normal register arg.
+declare i64 @callee_musttail_i128_and_i32(i128 %a, i32 %x)
+
+define i64 @caller_musttail_i128_and_i32(i128 %a, i32 %x) nounwind {
+; RV32-LABEL: caller_musttail_i128_and_i32:
+; RV32:       # %bb.0:
+; RV32-NEXT:    tail callee_musttail_i128_and_i32
+;
+; RV64-LABEL: caller_musttail_i128_and_i32:
+; RV64:       # %bb.0:
+; RV64-NEXT:    tail callee_musttail_i128_and_i32
+  %call = musttail call i64 @callee_musttail_i128_and_i32(i128 %a, i32 %x)
+  ret i64 %call
+}

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

@@ -247,20 +247,24 @@ declare i32 @callee_indirect_args(fp128 %a)
 define void @caller_indirect_args() nounwind {
 ; CHECK-LABEL: caller_indirect_args:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    addi sp, sp, -16
+; CHECK-NEXT:    addi sp, sp, -32
+; CHECK-NEXT:    sw ra, 28(sp) # 4-byte Folded Spill
 ; CHECK-NEXT:    lui a1, 262128
 ; CHECK-NEXT:    mv a0, sp
 ; CHECK-NEXT:    sw zero, 0(sp)
 ; CHECK-NEXT:    sw zero, 4(sp)
 ; CHECK-NEXT:    sw zero, 8(sp)
 ; CHECK-NEXT:    sw a1, 12(sp)
-; CHECK-NEXT:    addi sp, sp, 16
-; CHECK-NEXT:    tail callee_indirect_args
+; CHECK-NEXT:    call callee_indirect_args
+; CHECK-NEXT:    lw ra, 28(sp) # 4-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 32
+; CHECK-NEXT:    ret
 ;
 ; CHECK-LARGE-ZICFILP-LABEL: caller_indirect_args:
 ; CHECK-LARGE-ZICFILP:       # %bb.0: # %entry
 ; CHECK-LARGE-ZICFILP-NEXT:    lpad 0
-; CHECK-LARGE-ZICFILP-NEXT:    addi sp, sp, -16
+; CHECK-LARGE-ZICFILP-NEXT:    addi sp, sp, -32
+; CHECK-LARGE-ZICFILP-NEXT:    sw ra, 28(sp) # 4-byte Folded Spill
 ; CHECK-LARGE-ZICFILP-NEXT:    lui a1, 262128
 ; CHECK-LARGE-ZICFILP-NEXT:  .Lpcrel_hi9:
 ; CHECK-LARGE-ZICFILP-NEXT:    auipc a0, %pcrel_hi(.LCPI7_0)
@@ -270,8 +274,10 @@ define void @caller_indirect_args() nounwind {
 ; CHECK-LARGE-ZICFILP-NEXT:    sw zero, 4(sp)
 ; CHECK-LARGE-ZICFILP-NEXT:    sw zero, 8(sp)
 ; CHECK-LARGE-ZICFILP-NEXT:    sw a1, 12(sp)
-; CHECK-LARGE-ZICFILP-NEXT:    addi sp, sp, 16
-; CHECK-LARGE-ZICFILP-NEXT:    jr t2
+; CHECK-LARGE-ZICFILP-NEXT:    jalr t2
+; CHECK-LARGE-ZICFILP-NEXT:    lw ra, 28(sp) # 4-byte Folded Reload
+; CHECK-LARGE-ZICFILP-NEXT:    addi sp, sp, 32
+; CHECK-LARGE-ZICFILP-NEXT:    ret
 entry:
   %call = tail call i32 @callee_indirect_args(fp128 0xL00000000000000003FFF000000000000)
   ret void