Skip to content

feat(cpu_jit): add CPUID-aware vectorized JIT and cache topology tiling engine - #2467

Open
presentyourlove wants to merge 1 commit into
exo-explore:mainfrom
presentyourlove:feature/cpuid-vectorized-jit-cache-tiling
Open

presentyourlove wants to merge 1 commit into
exo-explore:mainfrom
presentyourlove:feature/cpuid-vectorized-jit-cache-tiling

Conversation

@presentyourlove

Copy link
Copy Markdown

Description

This PR implements Frontier Dimension 5.1: CPUID-Aware Vectorized JIT & Cache Topology Tiling Engine, resolving compatibility crashes and performance bottlenecks on legacy edge CPUs lacking advanced vector instruction sets.

Background & Motivation

Many edge devices and older PCs (e.g. Intel Core 2 Duo, 1st–4th Gen Core, legacy AMD FX/Phenom, or older ARM SBCs) lack modern vector extensions such as AVX-512, AVX2, or even AVX. Running pre-compiled matrix math binaries directly triggers Illegal Instruction (SIGILL) crashes. Falling back to non-vectorized scalar operations results in severe performance collapse. Furthermore, small L1/L2 caches on legacy hardware frequently suffer from cache thrashing during matrix multiplication.

This PR introduces an architectural CPU JIT subsystem that dynamically detects hardware vector capabilities, gracefully degrades to specialized 128-bit unrolled SSE4 or safe scalar kernels, and calculates optimal 3-level loop tiling dimensions based on hardware L1/L2 cache capacity.

Key Changes

  1. CPU Capability & Cache Geometry Types (src/exo/shared/types/cpu_jit.py):

    • CPUCapabilityInfo: Captures active instruction set levels (AVX512, AVX2, AVX, SSE4_2, SSE4_1, ARM_NEON, SCALAR), hardware flags, vendor, and core counts.
    • CPUCacheGeometry: Models L1D, L2, L3 cache capacities and cache line sizes with validation.
    • CacheTilingPlan & JITExecutionReport: Captures dynamic M/N/K tile blocks, working set footprints, estimated cache hit rates, and kernel GFLOPS.
  2. Cross-Platform Hardware Probers (src/exo/worker/engines/cpu_jit/cpu_prober.py):

    • CPUIDInstructionProber: Cross-platform probing (Linux /proc/cpuinfo, macOS sysctl, Windows processor identifiers) with support for simulated feature overrides and infallible fallback to SCALAR.
    • CacheTopologyProber: Queries OS sysfs cache topologies with empirical fallback models.
  3. Micro-Architecture Vectorized Dispatcher (src/exo/worker/engines/cpu_jit/vector_jit.py):

    • VectorizedJITDispatcher: Dynamically binds optimal compute kernels:
      • 256-bit AVX2 8-wide unrolled tiles.
      • 128-bit SSE4 4-wide unrolled tiles (for legacy Intel/AMD CPUs lacking AVX2).
      • ARM NEON vector tiles.
      • Pure scalar safe kernel ensuring 100% crash-free execution on legacy CPUs.
    • unpack_ternary_weights: Real-time on-CPU unpacking for 1.58-bit ternary {-1, 0, +1} quantized weights.
  4. Cache Tiling & Execution Engine (src/exo/worker/engines/cpu_jit/cache_tiling_engine.py):

    • CacheTilingPlanner: Solves optimal tile dimensions ($Tile_M, Tile_N, Tile_K$) such that $B_{tile} \le 60% \text{ L1D}$ and $A_{tile} + C_{tile} \le 75% \text{ L2}$, preventing cache thrashing.
    • CacheTiledGEMMEngine: Executes 3-level cache-partitioned matrix multiplication with micro-kernel dispatch and telemetry reporting.
    • CoreAffinityManager: Safe, non-throwing CPU thread and core affinity pinning.
  5. Unit Tests (src/exo/shared/tests/test_cpu_jit.py):

    • 8 comprehensive test cases verifying CPUID hardware probing, simulated legacy degradation, cache tiling calculation, numerical equivalence across scalar/SSE4/AVX2 kernels, ternary weight unpacking, end-to-end GEMM computation, and core affinity management.

Verification

  • Unit Tests:
    .venv/Scripts/pytest.exe src/exo/shared/tests/test_cpu_jit.py -v
    # 8 passed in 0.12s
  • Code Linting:
    .venv/Scripts/ruff.exe check src/exo/shared/types/cpu_jit.py src/exo/worker/engines/cpu_jit src/exo/shared/tests/test_cpu_jit.py
    # All checks passed!

Related Issues

  • Resolves: Frontier Dimension 5.1 in 人人都能跑llm.md (Legacy CPU Microarchitecture Vectorized JIT & Cache Tiling)

…ng engine

- Implement CPUIDInstructionProber with cross-platform SIMD vector flag detection and graceful degradation to prevent Illegal Instruction crashes on legacy hardware

- Implement CacheTopologyProber and CacheTilingPlanner for L1/L2 cache capacity matching

- Implement VectorizedJITDispatcher supporting AVX512, AVX2, SSE4.1/4.2 unrolled, NEON, and scalar safe kernels with ternary weight unpacking

- Implement CacheTiledGEMMEngine with 3-level loop tiling and CoreAffinityManager

- Add comprehensive unit test suite in test_cpu_jit.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant