Fixed a bug in -fsanitize-kcfi-arity

[scottconstable]

Compiling with fsanitize-kcfi-arity can crash the compiler if a function has more than 6 arguments, including floating-point arguments passed in XMM registers. This patch fixes the feature by only counter integer and stack arguments toward kCFI arity. For example, the compiler crashed when it attempted to generate kCFI arity information for this function: https://github.com/torvalds/linux/blob/16b70698aa3ae7888826d0c84567c72241cf6713/drivers/gpu/drm/amd/display/dc/dml/dcn32/display_mode_vba_util_32.h#L680

As noted in a comment, floating-point registers are not relevant to enforcing kCFI at this time.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Scott Constable (scottconstable)

Changes

Compiling with fsanitize-kcfi-arity can crash the compiler if a function has more than 6 arguments, including floating-point arguments passed in XMM registers. This patch fixes the feature by only counter integer and stack arguments toward kCFI arity. For example, the compiler crashed when it attempted to generate kCFI arity information for this function: https://github.com/torvalds/linux/blob/16b70698aa3ae7888826d0c84567c72241cf6713/drivers/gpu/drm/amd/display/dc/dml/dcn32/display_mode_vba_util_32.h#L680

As noted in a comment, floating-point registers are not relevant to enforcing kCFI at this time.

---

Full diff: https://github.com/llvm/llvm-project/pull/142867.diff

2 Files Affected:

• (modified) llvm/lib/Target/X86/X86AsmPrinter.cpp (+14-5)
• (modified) llvm/test/CodeGen/X86/kcfi-arity.ll (+27-2)

diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
index 24eda602effd1..7a8e3ef8ce3f7 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.cpp
+++ b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -198,14 +198,23 @@ 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;
+        Arity += !(Reg >= X86::XMM0 && Reg <= X86::XMM7);
+      }
+    }
     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}

[phoebewang]

Check

  if (X86::GR8RegClass.contains(Reg) ||
      X86::GR16RegClass.contains(Reg) ||
      X86::GR32RegClass.contains(Reg) ||
      X86::GR64RegClass.contains(Reg))

[scottconstable]

Thank you for the feedback, @phoebewang! I updated the PR.

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff HEAD~1 HEAD --extensions cpp -- llvm/lib/Target/X86/X86AsmPrinter.cpp

View the diff from clang-format here.

diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
index 6fe5862d5..c7238839c 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.cpp
+++ b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -212,8 +212,7 @@ void X86AsmPrinter::emitKCFITypeId(const MachineFunction &MF) {
       Arity = 0;
       for (const auto &LI : MF.getRegInfo().liveins()) {
         auto Reg = LI.first;
-        if (X86::GR8RegClass.contains(Reg) ||
-            X86::GR16RegClass.contains(Reg) ||
+        if (X86::GR8RegClass.contains(Reg) || X86::GR16RegClass.contains(Reg) ||
             X86::GR32RegClass.contains(Reg) ||
             X86::GR64RegClass.contains(Reg)) {
           ++Arity;

[phoebewang]

LGTM except the format issue.

[phoebewang]

Format the code.

[scottconstable]

Done, thanks @phoebewang !

[scottconstable]

@phoebewang Do we need more reviews before we merge?

[phoebewang]

@phoebewang Do we need more reviews before we merge?

Added @scottconstable for review and CCing @sirmc @Darksonn who commented prior patch. We can merge it if no objections in a few days.

[sirmc]

From my side this looks good to merge. I built kernel v6.15 successfully with this patch.