Target Features Mismatch Finder

[hildebrandmw]

Compiling x86 binaries capable of running AVX-512 is currently a large game of inline and target-features whack-a-mole. As with all Rust code, enabling higher level architecture than the compilation target (e.g. AVX-512 when the binary is compiled for x86-64-v3) requires annotating functions with the correct target features.

When target_feature(enable = "...") is used to annotate a function foo, then all functions inlined into foo inherit those target features. Calling a non-inlined function from foo can cause those target feature to get lost.

Currently, we use architecture tokens to apply target features to closures and other types implementing the Target[1-3]? traits.

The problem arises when functions are not inlined into the immediate callsite. As an example, if you apply the following diff

diff --git a/diskann-wide/src/lib.rs b/diskann-wide/src/lib.rs
index 4f6ca5c9..da8e4762 100644
--- a/diskann-wide/src/lib.rs
+++ b/diskann-wide/src/lib.rs
@@ -236,6 +236,26 @@ pub(crate) mod helpers;
 #[cfg(test)]
 pub(crate) mod test_utils;

+alias!(f32s = <arch::x86_64::V4>::f32x16);
+
+pub fn uninlined_context(x: f32s, y: f32s) -> f32s {
+    let z = x + y;
+    let z = z * z;
+    z
+}
+
+pub fn inlined_context(x: f32s, y: f32s) -> f32s {
+    let arch = x.arch();
+    arch.run2(
+        |a: f32s, b: f32s| {
+            let z = a + b;
+            z * z
+        },
+        x,
+        y,
+    )
+}
+
 ///////////
 // Tests //
 //////////

And compile the resulting assembly, we get

        vmovaps zmm0, zmmword ptr [rsi]
        vaddps zmm0, zmm0, zmmword ptr [rdx]
        vmulps zmm0, zmm0, zmm0
        vmovaps zmmword ptr [rdi], zmm0
        vzeroupper
        ret

when unapplied_features is called from inlined_context but

        .cfi_startproc
        push rbp
        .cfi_def_cfa_offset 16
        .cfi_offset rbp, -16
        mov rbp, rsp
        .cfi_def_cfa_register rbp
        push rbx
        and rsp, -64
        sub rsp, 256
        .cfi_offset rbx, -24
        mov rbx, rdi
        vmovaps ymm0, ymmword ptr [rsi]
        vmovaps ymm1, ymmword ptr [rsi + 32]
        vmovaps ymm2, ymmword ptr [rdx]
        vmovaps ymm3, ymmword ptr [rdx + 32]
        vmovaps ymmword ptr [rsp + 32], ymm1
        vmovaps ymmword ptr [rsp], ymm0
        vmovaps ymmword ptr [rsp + 96], ymm3
        vmovaps ymmword ptr [rsp + 64], ymm2
        lea rdi, [rsp + 128]
        mov rsi, rsp
        lea rdx, [rsp + 64]
        vzeroupper
        call core::core_arch::x86::avx512f::_mm512_add_ps
        vmovaps ymm0, ymmword ptr [rsp + 128]
        vmovaps ymm1, ymmword ptr [rsp + 160]
        vmovaps ymmword ptr [rsp + 32], ymm1
        vmovaps ymmword ptr [rsp], ymm0
        vmovaps ymmword ptr [rsp + 96], ymm1
        vmovaps ymmword ptr [rsp + 64], ymm0
        mov rsi, rsp
        lea rdx, [rsp + 64]
        mov rdi, rbx
        vzeroupper
        call core::core_arch::x86::avx512f::_mm512_mul_ps
        mov rax, rbx
        lea rsp, [rbp - 8]
        pop rbx
        pop rbp
        .cfi_def_cfa rsp, 8
        ret

when it is not. What's happening in the latter case is that unapplied_features is not compiled with avx512f, and thus calls to core::core_arch::x86::avx512f::_mm512_add_ps and core::core_arch::x86::avx512f::_mm512_mul_ps (which require that feature) cannot be inlined.

This leads to a dramatic performance cliff when we accidentally fall out of the proper target features.

Fortunately, it should be pretty straight-forward to write tooling to detect this situation. All we really need to do (in theory) is disassemble compiled binaries and look for uninlined calls to the starch intrinsics. That can at least either put us in the ballpark of something that got missed, or give use higher confidence that we got everything.