[llvm] [RISCV] Use t3 for static chain register when branch CFI is enabled (PR #142344)

[Jesse Huang via llvm-commits]

https://github.com/jaidTw created https://github.com/llvm/llvm-project/pull/142344

Use t3 for static chain register when branch CFI is enabled to align with gcc.[1]

[1] https://github.com/gcc-mirror/gcc/blob/master/gcc/config/riscv/riscv.h#L417

>From fbdda7cd46eb7455272bcc3921b9910014633737 Mon Sep 17 00:00:00 2001
From: Jesse Huang <jesse.huang at sifive.com>
Date: Mon, 2 Jun 2025 01:08:32 -0700
Subject: [PATCH] [RISCV] Use t3 for static chain register when branch CFI is
 enabled

---
 llvm/lib/Target/RISCV/RISCVCallingConv.cpp | 11 +++++++++--
 llvm/test/CodeGen/RISCV/nest-register.ll   | 23 ++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVCallingConv.cpp b/llvm/lib/Target/RISCV/RISCVCallingConv.cpp
index d3bfbb0943766..e0d1fb2facc87 100644
--- a/llvm/lib/Target/RISCV/RISCVCallingConv.cpp
+++ b/llvm/lib/Target/RISCV/RISCVCallingConv.cpp
@@ -13,6 +13,7 @@
 #include "RISCVCallingConv.h"
 #include "RISCVSubtarget.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Module.h"
 #include "llvm/MC/MCRegister.h"
 
 using namespace llvm;
@@ -333,9 +334,15 @@ bool llvm::CC_RISCV(unsigned ValNo, MVT ValVT, MVT LocVT,
   MVT XLenVT = Subtarget.getXLenVT();
 
   // Static chain parameter must not be passed in normal argument registers,
-  // so we assign t2 for it as done in GCC's __builtin_call_with_static_chain
+  // so we assign t2/t3 for it as done in GCC's __builtin_call_with_static_chain
+  bool HasCFBranch =
+      Subtarget.hasStdExtZicfilp() &&
+      MF.getFunction().getParent()->getModuleFlag("cf-protection-branch");
+  // Normal: t2, Branch control flow protection: t3
+  const auto StaticChainReg = HasCFBranch ? RISCV::X28 : RISCV::X7;
+
   if (ArgFlags.isNest()) {
-    if (MCRegister Reg = State.AllocateReg(RISCV::X7)) {
+    if (MCRegister Reg = State.AllocateReg(StaticChainReg)) {
       State.addLoc(CCValAssign::getReg(ValNo, ValVT, Reg, LocVT, LocInfo));
       return false;
     }
diff --git a/llvm/test/CodeGen/RISCV/nest-register.ll b/llvm/test/CodeGen/RISCV/nest-register.ll
index e222beee45783..9f8e4e1a2d8d3 100644
--- a/llvm/test/CodeGen/RISCV/nest-register.ll
+++ b/llvm/test/CodeGen/RISCV/nest-register.ll
@@ -3,6 +3,8 @@
 ; RUN:   | FileCheck -check-prefix=RV32I %s
 ; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck -check-prefix=RV64I %s
+; RUN: llc -mtriple=riscv64 -mattr=+experimental-zicfilp -verify-machineinstrs < %s \
+; RUN:   | FileCheck -check-prefix=RV64I-ZICFILP %s
 
 ; Tests that the 'nest' parameter attribute causes the relevant parameter to be
 ; passed in the right register.
@@ -17,6 +19,12 @@ define ptr @nest_receiver(ptr nest %arg) nounwind {
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    mv a0, t2
 ; RV64I-NEXT:    ret
+;
+; RV64I-ZICFILP-LABEL: nest_receiver:
+; RV64I-ZICFILP:       # %bb.0:
+; RV64I-ZICFILP-NEXT:    lpad 0
+; RV64I-ZICFILP-NEXT:    mv a0, t3
+; RV64I-ZICFILP-NEXT:    ret
   ret ptr %arg
 }
 
@@ -40,6 +48,21 @@ define ptr @nest_caller(ptr %arg) nounwind {
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
+;
+; RV64I-ZICFILP-LABEL: nest_caller:
+; RV64I-ZICFILP:       # %bb.0:
+; RV64I-ZICFILP-NEXT:    lpad 0
+; RV64I-ZICFILP-NEXT:    addi sp, sp, -16
+; RV64I-ZICFILP-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; RV64I-ZICFILP-NEXT:    mv t3, a0
+; RV64I-ZICFILP-NEXT:    call nest_receiver
+; RV64I-ZICFILP-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; RV64I-ZICFILP-NEXT:    addi sp, sp, 16
+; RV64I-ZICFILP-NEXT:    ret
   %result = call ptr @nest_receiver(ptr nest %arg)
   ret ptr %result
 }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 8, !"cf-protection-branch", i32 1}