rlx_cpu/
asm_check.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! FileCheck-style disassembly regression tests (plan #10).
6//!
7//! Borrowed from MAX's `bazel/internal/mojo_filecheck_test.bzl` +
8//! `bazel/internal/lit.bzl` pattern: assertions on emitted IR/asm
9//! catch optimizer regressions that unit tests miss. The classic
10//! one is "the optimizer changed and now we lost a SIMD intrinsic"
11//! — a benchmark would notice the slowdown later, but FileCheck
12//! catches it at PR time.
13//!
14//! The Rust spelling shells out to `objdump` (or `llvm-objdump` if
15//! preferred) on the running test binary, locates a named function
16//! in the disassembly, and asserts each requested pattern appears.
17//!
18//! Tests that use this should mark themselves `#[ignore]` so
19//! they're opt-in via `cargo test -- --ignored asm` — `objdump`
20//! isn't always available in CI, and disassembling a debug binary
21//! is slow.
22//!
23//! Pattern matching is substring + line-anchored regex via the
24//! standard library only — no extern dep on `regex`.
25
26use std::path::PathBuf;
27use std::process::Command;
28
29/// What objdump-equivalent to use. macOS ships `objdump` as
30/// `Apple LLVM`; Linux usually has `llvm-objdump` available, with
31/// GNU `objdump` as a fallback.
32fn locate_objdump() -> Option<PathBuf> {
33    for candidate in ["llvm-objdump", "objdump"] {
34        let probe = Command::new(candidate).arg("--version").output();
35        if probe.ok().filter(|o| o.status.success()).is_some() {
36            return Some(PathBuf::from(candidate));
37        }
38    }
39    None
40}
41
42/// Path to the currently-running test binary. Cargo runs each
43/// integration / library test from a known location; the env var
44/// `CARGO_BIN_EXE_<name>` points binaries at themselves, but for
45/// library tests we walk up from `std::env::current_exe()`.
46fn current_test_binary() -> std::io::Result<PathBuf> {
47    std::env::current_exe()
48}
49
50/// Disassemble the test binary; return the full text dump.
51/// Errors out (with a skippable label) when the tool isn't
52/// available so callers can short-circuit gracefully.
53pub fn disassemble_self() -> Result<String, AsmCheckError> {
54    let tool = locate_objdump().ok_or(AsmCheckError::ToolMissing)?;
55    let bin = current_test_binary().map_err(|e| AsmCheckError::IoError(e.to_string()))?;
56    let out = Command::new(&tool)
57        .arg("-d")
58        .arg("--no-show-raw-insn")
59        .arg(&bin)
60        .output()
61        .map_err(|e| AsmCheckError::IoError(e.to_string()))?;
62    if !out.status.success() {
63        return Err(AsmCheckError::ToolFailed {
64            stderr: String::from_utf8_lossy(&out.stderr).into_owned(),
65        });
66    }
67    Ok(String::from_utf8_lossy(&out.stdout).into_owned())
68}
69
70/// Look up a function by demangled-name substring and return the
71/// slice of disassembly belonging to it (until the next
72/// function header).
73pub fn function_section<'a>(disasm: &'a str, name_substr: &str) -> Option<&'a str> {
74    let mut start = None;
75    for (i, line) in disasm.lines().enumerate() {
76        if line.contains(name_substr) && line.trim_end().ends_with(':') {
77            start = Some(i);
78            break;
79        }
80    }
81    let start = start?;
82    let lines: Vec<&str> = disasm.lines().collect();
83    let mut end = lines.len();
84    for j in (start + 1)..lines.len() {
85        if lines[j].ends_with(':')
86            && !lines[j].trim_start().starts_with('0')  // not a target label
87            && !lines[j].is_empty()
88        {
89            end = j;
90            break;
91        }
92    }
93    let from = lines[..start].iter().map(|s| s.len() + 1).sum::<usize>();
94    let to = lines[..end]
95        .iter()
96        .map(|s| s.len() + 1)
97        .sum::<usize>()
98        .min(disasm.len());
99    Some(&disasm[from..to])
100}
101
102/// Assert each substring in `expected` appears at least once in
103/// the function `name`'s disassembly. Returns Err if the
104/// disassembler isn't present so callers can `eprintln!` skip
105/// without failing.
106pub fn assert_function_contains(name_substr: &str, expected: &[&str]) -> Result<(), AsmCheckError> {
107    let disasm = disassemble_self()?;
108    let body = function_section(&disasm, name_substr).ok_or(AsmCheckError::FunctionNotFound {
109        name: name_substr.into(),
110    })?;
111    for pat in expected {
112        if !body.contains(pat) {
113            return Err(AsmCheckError::PatternMissing {
114                function: name_substr.into(),
115                pattern: (*pat).into(),
116            });
117        }
118    }
119    Ok(())
120}
121
122#[derive(Debug)]
123pub enum AsmCheckError {
124    ToolMissing,
125    ToolFailed { stderr: String },
126    IoError(String),
127    FunctionNotFound { name: String },
128    PatternMissing { function: String, pattern: String },
129}
130
131impl std::fmt::Display for AsmCheckError {
132    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133        match self {
134            Self::ToolMissing => write!(f, "neither llvm-objdump nor objdump found in PATH"),
135            Self::ToolFailed { stderr } => write!(f, "objdump failed: {stderr}"),
136            Self::IoError(s) => write!(f, "io error: {s}"),
137            Self::FunctionNotFound { name } => {
138                write!(f, "function matching `{name}` not found in disassembly")
139            }
140            Self::PatternMissing { function, pattern } => write!(
141                f,
142                "function `{function}` is missing expected pattern `{pattern}`"
143            ),
144        }
145    }
146}
147
148impl std::error::Error for AsmCheckError {}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    /// basic test: disassembly works at all. Marked `#[ignore]`
155    /// because objdump isn't always around in CI (in which case
156    /// we'd want to log-and-skip rather than fail).
157    #[test]
158    #[ignore = "needs objdump in PATH; run with --ignored for codegen disassembly checks"]
159    fn disassemble_self_succeeds() {
160        match disassemble_self() {
161            Ok(d) => assert!(d.len() > 1024, "disassembly suspiciously small"),
162            Err(AsmCheckError::ToolMissing) => {
163                eprintln!("[asm-check] skipping: objdump not in PATH");
164            }
165            Err(e) => panic!("disassembly failed: {e}"),
166        }
167    }
168
169    /// Real check: the cumsum kernel must contain the expected
170    /// f32 multiply / fused-multiply on aarch64. If the optimizer
171    /// regresses and inlines/loses these, the test catches it.
172    #[test]
173    #[ignore = "needs objdump in PATH (aarch64 host); run with --ignored for SIMD codegen check"]
174    fn cumsum_kernel_keeps_simd_on_aarch64() {
175        // On targets where we don't expect SIMD here, accept the
176        // miss as "no expected pattern" and move on.
177        if !cfg!(target_arch = "aarch64") {
178            eprintln!("[asm-check] skipping: not aarch64");
179            return;
180        }
181        // The cumsum direct-execution path adds + stores in a
182        // tight inner loop. We check for `fmadd` (FP add fused
183        // with multiply) or just `fadd` since cumsum doesn't
184        // multiply. Pick a stable substring: `fadd s` (single
185        // f32 register-form) shows up in the inner loop.
186        match assert_function_contains("Cumsum", &["fadd"]) {
187            Ok(()) => {}
188            Err(AsmCheckError::ToolMissing | AsmCheckError::FunctionNotFound { .. }) => {
189                eprintln!("[asm-check] skipping: tool or symbol missing");
190            }
191            Err(e) => panic!("Cumsum kernel asm check failed: {e}"),
192        }
193    }
194}
rlx_cpu/asm_check.rs

rlx_cpu/
asm_check.rs
