Compiler Bypasses and Native Code¶
Research and local checks: October 2, 2026, Apple M5, macOS 26.2, MLX 0.32.0, Apple Metal compiler 32023.864.
Two meTile experiments work below the normal MSL compilation path. One modifies AIR before Apple’s backend runs; the other rewrites selected AGX instructions afterward. Both probes show that modified code can execute, but neither has established a new speedup over MLX. Neither is enabled in ordinary kernel compilation.
What each route controls¶
Route |
What it controls |
What remains |
|---|---|---|
Tile IR / Metal IR to MSL |
Fusion, reuse, layouts, vector operations and matrix operations |
Apple selects instructions, allocates registers and schedules them |
Textual AIR to |
LLVM-level operations, intrinsics, metadata and arithmetic flags |
Apple’s backend still selects and schedules AGX instructions |
Rewrite a binary archive’s |
Actual encoded instructions in a verified region |
Must preserve the surrounding executable’s ABI and resource metadata |
A complete AGX backend |
Instruction selection, register allocation and scheduling |
Requires far more ISA and executable-format coverage than exists here |
Apple documents the source-to-AIR-to-device-code pipeline and the use of binary
archives as compilation caches. failOnBinaryArchiveMiss prevents a missing
archive entry from silently falling back to compilation; it does not make
archive editing a supported API. See Apple’s
binary archive presentation.
The existing source-order experiment was rerun on this machine:
Source change |
Compiled result |
|---|---|
Serial versus interleaved independent FMA chains |
Identical, 190 bytes |
Loads at use versus hoisted loads |
Identical, 218 bytes |
Addition chain versus balanced tree |
Identical, 282 bytes |
Apple’s compiler produced identical code for each of these source-order pairs.
Other source transformations may still change the result. The register and
throughput numbers in metile/target/agx.py describe specific M5 probes and
need new measurements for different kernels or hardware.
Direct AIR experiments¶
The recorded toolchain accepted this sequence:
xcrun metal -S -emit-llvm -c kernel.metal -o kernel.ll
xcrun metal-as kernel.ll -o kernel.air
xcrun metallib kernel.air -o kernel.metallib
xcrun metal -x ir -c kernel.ll -o kernel.air also worked locally.
To inspect binary AIR, metal -x ir -S -emit-llvm worked; a separate
metal-dis executable was not installed.
benchmarks/hardware/agx_air_probe.py reproduces the experiment. It compiles a
small MSL seed, changes one FMA multiplier in textual AIR, assembles both
versions, and dispatches both through meTile’s Metal runtime. Each version is
checked against an independent numerical expectation. If the AIR format
changes, the probe fails explicitly rather than making an unchecked string
replacement.
This demonstrates execution of modified IR, not a standalone AIR backend or a
faster kernel. On the setup above, both fma(input, 2, 1) and the rewritten
fma(input, 3, 1) matched all 257 quarter-step test inputs exactly.
An initial direct-AIR backend could lower a restricted subset of Metal IR: loads/stores, integer address arithmetic, floating-point operations, then SIMD reductions. It would need to preserve the versioned kernel metadata and address spaces from a toolchain-generated seed. Arbitrary system LLVM bitcode cannot be assumed to use the dialect and metadata Metal requires. Metal.jl’s compiler is a useful primary implementation to study for this boundary.
Useful experiments include precise contraction/reassociation permissions,
proven alignment/range information, and vector unpack/conversion forms that
MSL obscures. Each candidate needs a comparison with equivalent MSL, both in
final code and execution, before the backend expands. AIR instruction order
still does not specify the final schedule. Handwritten air.* declarations
are toolchain-dependent: a local MSL declaration using __asm("air.fma.f32")
was rejected, so that syntax is not a working inline-assembly escape hatch here.
Native instruction rewriting¶
metile/target/agx_isa.py already assembles a restricted compact F32 FMA
form and executes modified binary archives. metile/compiler/agx_schedule.py
operates on caller-supplied instruction offsets. The execution tests verify
that synthesised arithmetic really reaches the GPU.
The scheduler preserves original encodings, tracks register read/write dependencies including FMA addends, and treats unknown controls and disabled instructions as barriers. Archive rewriting extracts code from the archive it will patch, rather than recompiling and risking different output.
Recognizing eight plausible bytes does not establish an instruction boundary or operand mode. Even a confirmed boundary leaves the surrounding loads, waits, stores, branches, and matrix instructions undecoded. The low-nibble length heuristic already failed on other kernels in the repository. These passes therefore remain outside ordinary kernel compilation and MLX dispatch.
The next native experiment should target one demonstrated missed peephole in a real generated kernel. Keep instruction length, register allocation, memory operations and control flow unchanged. Preserve every field whose meaning is unknown, and reject a replacement unless both its operand semantics and its dependencies are understood. A speedup must be measured after validating the output; retiring an instruction does not establish a cycle saving.
For broader decoding, study the Apple G13 architecture reference and the Mesa Asahi driver. They provide a starting point for understanding encoding and compiler structure, not evidence that G13 instructions or Linux driver interfaces work unchanged on this G17/macOS target. A full backend also needs spills, scheduling controls, executable metadata and generation-specific validation.
Research directions¶
The following are research priorities, not new benchmark results:
Reuse weights across small batches in the general lowering. Sweep 2, 4, 8, 16 and 32 rows with identical BF16 or INT4 representations. Carry row tiling and accumulator limits through the compiler rather than adding a model-specific source template. Existing matched-representation benchmarks identify this as a productive regime; register growth is the tradeoff.
Fuse projection epilogues and adjacent operations. Use graph semantics to eliminate an intermediate write/read or launch. Validate the whole fused boundary and retain the native candidate. This can save work even when the individual arithmetic instructions are already good.
Compare direct device access against staging for tensor operations. Apple’s MPP programming guide recommends direct device access for GEMM and discusses when thread parallelism makes manual software pipelining unnecessary. Treat this as a candidate policy, measured by tile and dtype. meTile’s existing NAX lowering is the appropriate place to experiment.
Try AIR vector unpack and conversion candidates on INT4 compute-heavy shapes. First inspect whether the candidate changes final code; then measure the full projection, including tails, scales and biases. Preserve the representation and numerical contract. A faster isolated unpack loop may be irrelevant to a bandwidth-bound projection.
Use native rewriting only when the previous comparison exposes a specific backend limitation. Start with proven arithmetic regions; general load scheduling or register renaming needs much more decoding and liveness data.
MLX’s NAX implementation
already uses mpp::tensor_ops::matmul2d and SIMD-group execution. Emitting
that primitive alone does not distinguish a meTile kernel. The useful questions
concern the surrounding workload: where weights can be reused without changing
their representation, and where operations can be fused.
The results in Benchmarks motivate experiments with small-batch weight reuse and projection fusion. They do not establish a hardware-wide performance ceiling.
Reproduce and evaluate a candidate¶
python -m benchmarks.hardware.agx_source_order
python -m benchmarks.hardware.agx_air_probe
python -m pytest tests/target/test_agx_isa.py tests/compiler/test_agx_schedule.py \
tests/compiler/test_agx_schedule_safety.py tests/target/test_target.py -q
Run GPU experiments serially; the archive harness uses shared filenames in its working directory. Its subprocess duration includes compilation, loading and dispatch overhead and must not be presented as GPU kernel latency.
Before a new backend becomes a runtime candidate, compile both arms once, check outputs on random values and edge cases, and measure interleaved GPU dispatches with a byte-identical control. Report compilation time separately from steady-state execution. Reassociation needs an explicit numerical policy: collapsing FMA chains is not bit-exact and can change overflow or underflow as well as rounding. The native identity-removal pass also lacks a general signed-zero and NaN-payload guarantee; its current ordinary-value probes do not establish those edge cases.
Use the existing tuning tournament and native fallback for promotion. Persist any edited archive only with its GPU architecture, OS build, compiler identity, source/options, launch layout and rewrite version. Record correctness failures and inconclusive timings as rejected candidates. The current probes establish that these compilation routes work. Claiming an improvement over MLX requires the additional comparison described here.