[llvm] [RISCV] Fix musttail with indirect arguments by forwarding incoming pointers (PR #185094)

[Xavier Roche via llvm-commits]

================
@@ -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)
----------------
xroche wrote:

Good catch, you were right to suggest the swapped-args test: it revealed a real bug. It appears that Outs[OutIdx].OrigArgIndex is the position in the call's argument list (per callee perspective), but the IncomingIndirectArgs map was keyed by the caller's formal parameter index. So when
  musttail reorders arguments (e.g., musttail call @f(%b, %a)), these diverge, causing the wrong pointers to be forwarded. The code is a bit more klunky, as we need to find the reordering by scanning CLI.CB->args()). There might be a simpler way ?


https://github.com/llvm/llvm-project/pull/185094