[llvm] Fixed a bug in `-fsanitize-kcfi-arity` (PR #142867)

[Scott Constable via llvm-commits]

https://github.com/scottconstable updated https://github.com/llvm/llvm-project/pull/142867

>From d5adc1c0142cc22278e650ee0b1dcd3dc2b71962 Mon Sep 17 00:00:00 2001
From: Scott D Constable <scott.d.constable at intel.com>
Date: Mon, 9 Jun 2025 09:10:52 -0700
Subject: [PATCH] Fixed a bug in `-fsanitize-kcfi-arity`

---
 llvm/lib/Target/X86/X86AsmPrinter.cpp | 23 ++++++++++++++++-----
 llvm/test/CodeGen/X86/kcfi-arity.ll   | 29 +++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
index 24eda602effd1..c7238839c26b5 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.cpp
+++ b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -198,14 +198,27 @@ void X86AsmPrinter::emitKCFITypeId(const MachineFunction &MF) {
     // Determine the function's arity (i.e., the number of arguments) at the ABI
     // level by counting the number of parameters that are passed
     // as registers, such as pointers and 64-bit (or smaller) integers. The
-    // Linux x86-64 ABI allows up to 6 parameters to be passed in GPRs.
+    // Linux x86-64 ABI allows up to 6 integer parameters to be passed in GPRs.
     // Additional parameters or parameters larger than 64 bits may be passed on
-    // the stack, in which case the arity is denoted as 7.
+    // the stack, in which case the arity is denoted as 7. Floating-point
+    // arguments passed in XMM0-XMM7 are not counted toward arity because
+    // floating-point values are not relevant to enforcing kCFI at this time.
     const unsigned ArityToRegMap[8] = {X86::EAX, X86::ECX, X86::EDX, X86::EBX,
                                        X86::ESP, X86::EBP, X86::ESI, X86::EDI};
-    int Arity = MF.getInfo<X86MachineFunctionInfo>()->getArgumentStackSize() > 0
-                    ? 7
-                    : MF.getRegInfo().liveins().size();
+    int Arity;
+    if (MF.getInfo<X86MachineFunctionInfo>()->getArgumentStackSize() > 0) {
+      Arity = 7;
+    } else {
+      Arity = 0;
+      for (const auto &LI : MF.getRegInfo().liveins()) {
+        auto Reg = LI.first;
+        if (X86::GR8RegClass.contains(Reg) || X86::GR16RegClass.contains(Reg) ||
+            X86::GR32RegClass.contains(Reg) ||
+            X86::GR64RegClass.contains(Reg)) {
+          ++Arity;
+        }
+      }
+    }
     DestReg = ArityToRegMap[Arity];
   }
 
diff --git a/llvm/test/CodeGen/X86/kcfi-arity.ll b/llvm/test/CodeGen/X86/kcfi-arity.ll
index 68d90adaf2a17..009fa7d2dc0a4 100644
--- a/llvm/test/CodeGen/X86/kcfi-arity.ll
+++ b/llvm/test/CodeGen/X86/kcfi-arity.ll
@@ -192,9 +192,33 @@ entry:
   ret void
 }
 
+;; Ensure that floating-point values are not counted toward the arity
+; ASM-LABEL: __cfi_f12:
+; ASM: movl $2253188362, %ebp
+define dso_local void @f12(i32 noundef %v1, i32 noundef %v2, float noundef %v3, double noundef %v4, float noundef %v5, i32 noundef %v6, i32 noundef %v7, i32 noundef %v8) #0 !kcfi_type !7 {
+entry:
+  %v1.addr = alloca i32, align 4
+  %v2.addr = alloca i32, align 4
+  %v3.addr = alloca float, align 4
+  %v4.addr = alloca double, align 4
+  %v5.addr = alloca float, align 4
+  %v6.addr = alloca i32, align 4
+  %v7.addr = alloca i32, align 4
+  %v8.addr = alloca i32, align 4
+  store i32 %v1, ptr %v1.addr, align 4
+  store i32 %v2, ptr %v2.addr, align 4
+  store float %v3, ptr %v3.addr, align 4
+  store double %v4, ptr %v4.addr, align 4
+  store float %v5, ptr %v5.addr, align 4
+  store i32 %v6, ptr %v6.addr, align 4
+  store i32 %v7, ptr %v7.addr, align 4
+  store i32 %v8, ptr %v8.addr, align 4
+  ret void
+}
+
 attributes #0 = { "target-features"="+retpoline-indirect-branches,+retpoline-indirect-calls" }
 
-!llvm.module.flags = !{!0, !7}
+!llvm.module.flags = !{!0, !8}
 !0 = !{i32 4, !"kcfi", i32 1}
 !1 = !{i32 12345678}
 !2 = !{i32 4196274163}
@@ -202,4 +226,5 @@ attributes #0 = { "target-features"="+retpoline-indirect-branches,+retpoline-ind
 !4 = !{i32 199571451}
 !5 = !{i32 1046421190}
 !6 = !{i32 1342488295}
-!7 = !{i32 4, !"kcfi-arity", i32 1}
+!7 = !{i32 2253188362}
+!8 = !{i32 4, !"kcfi-arity", i32 1}