Adding a rustlang port of ktir-cpu

[fabianlim]

Rust Port — Landing Plan

TL;DR

PR #117 adds a complete Rust reimplementation of the KTIR CPU emulator under rust/,
parallel to the existing Python reference. 104 changed files, ~55k lines — broken
down as:

What	Files	Lines
Rust source (3 crates, 41 .rs files)	41	~31,800
Rust integration tests (35 .rs files)	35	~16,700
Cargo manifests	4 .toml	~100
rust/docs/port-map.md (design doc)	1	~5,000
New MLIR examples, CI workflow, bench scripts, Python fixes	~23	~1,400

The plan lands it in four commits: CI plumbing → ktir-core → full emulator + tests
→ equivalence workflow. Before merging, two decisions are needed:

1. Checkout strategy — adding rust/ to the repo will confuse Python contributors
   who clone and suddenly see unfamiliar Rust files (the concern is discoverability,
   not size). The fix is a .sparse-checkout file that hides rust/ by default, plus
   resolving the compile-time examples/ dependency so rust/ can be cloned alone.
   See the Checkout and packaging concern section.
2. Stale files — four files at the repo root (bench_*.py, bench_*.sh,
   mlx_fidelity_probe.py) contain hardcoded absolute paths and must be fixed or
   removed before Commit 3 lands.

---

This document covers the monorepo layout, the test harness, the commit sequence, and
the equivalence-testing strategy.

---

Monorepo Layout

The Rust port lives parallel to the Python reference. Logic-wise neither tree
depends on the other at build time — but see the checkout concern below.

ktir-cpu/
├── ktir_cpu/               # Python reference implementation (unchanged)
├── tests/                  # Python test suite
├── examples/               # shared MLIR examples (used by both Python and Rust tests)
├── rust/                   # Rust implementation
│   ├── Cargo.toml          # workspace root (members: ktir-core, ktir-emulator, ktir-optimizer)
│   └── crates/
│       ├── ktir-core/          # IR types, parser, affine, dtypes, codec, tile
│       │   └── src/            # inline #[cfg(test)] unit tests (57 total)
│       ├── ktir-emulator/      # interpreter, machine state, BLAS/Metal backends
│       │   ├── src/
│       │   └── tests/          # integration tests (port_*.rs, e2e_*.rs, golden tests)
│       └── ktir-optimizer/     # IR→IR passes: fusion, flash-attn, head-rewrite
│           └── src/
├── .github/workflows/
│   ├── ci.yml              # Python CI — must add paths-ignore for rust/**
│   ├── ci-rust.yml         # Rust CI — path-scoped to rust/**
│   └── equiv-test.yml      # Python↔Rust equivalence check (workflow_dispatch only)
├── docs/
│   └── gap_analysis.md
├── README.md               # top-level — add a "Rust port" section linking to rust/README.md

What's in rust/ — size breakdown

41 source files, 35 test files, 4 Cargo manifests. The heaviest files:

File	Lines	What it is
ktir-emulator/src/metal.rs	5,800	Apple Silicon GPU/AMX backend (Metal Shading Language codegen + runtime dispatch)
ktir-optimizer/src/flash_attn.rs	2,500	Flash-attention cap-tiling optimizer pass
ktir-core/src/parser.rs	2,030	Full MLIR text grammar parser (no MLIR dependency)
ktir-emulator/src/ops_memory.rs	1,924	ktdp.load/store data path (HBM/LX, distributed, indirect)
ktir-optimizer/src/fusion.rs	1,539	Function fusion pass (eliminates HBM round-trips between nodes)
ktir-emulator/src/dialects/linalg.rs	1,976	Linalg dialect handlers (matmul, reduce, generic, …)
ktir-emulator/src/dialects/arith.rs	1,885	Arith dialect handlers (all scalar + elementwise ops)
rust/docs/port-map.md	5,000	Design doc mapping every Python class/method to its Rust equivalent

Tests account for ~16,700 of the ~48,500 Rust lines — roughly a 2:1 source-to-test
ratio, which is healthy.

Checkout and packaging concern

There are two separate but related questions here: what gets installed by the Python
package, and what needs to be present for each project to build.

Python packaging — already solved. pyproject.toml has:

[tool.setuptools.packages.find]
include = ["ktir_cpu*"]

rust/ is already excluded from the Python wheel. Anyone doing pip install or
uv sync never sees the Rust code. No action needed here.

Git checkout — examples/ is a compile-time Rust dependency. The Rust test and
source files embed MLIR programs from examples/ via include_str!, resolved by
rustc at compile time relative to the source file:

include_str!("../../../../examples/triton-ktir/vector_add_ktir.mlir")

cargo build / cargo test will fail if examples/ is not present at the repo
root. This is a hard build-time coupling that affects two scenarios:

• Git worktrees — unaffected. Each worktree is a full checkout of the repo tree,
  so examples/ is always present. This is the normal multi-branch workflow.
• Sparse checkouts — blocked for rust/-only clones unless we resolve it.

---

Sparse checkout proposal

The packaging exclusion hides rust/ from pip install, but not from git clone.
A Python contributor who clones the repo sees a rust/ directory in their working
tree with unfamiliar Rust files — even though none of it affects their work. The
concern is not storage (12 MB is nothing); it's confusion when something new
appears that wasn't there before.

A .sparse-checkout patterns file committed to the repo root solves this: Python
contributors clone and rust/ simply never appears in their working tree. Rust
contributors opt in with one extra command. The Python experience is unchanged;
the Rust subtree is invisible until explicitly requested.

How it works: git cannot auto-apply sparse-checkout on git clone — it is always
a manual post-clone step. But by shipping a .sparse-checkout file we make the
Python-focused view the documented default, and Rust contributors opt in explicitly.

The .sparse-checkout file (committed to the repo root) would contain:

# Default: Python reference implementation only.
# To add the Rust port: git sparse-checkout add rust examples
/*
!rust/

The leading /* includes all root-level files (pyproject.toml, uv.lock,
README.md, etc.). !rust/ excludes the Rust subtree. examples/ and all other
directories are included by default, which is correct — Python tests use them.

Per-use-case clone recipe:

# Full clone — unchanged, always works, recommended if unsure:
git clone https://github.com/torch-spyre/ktir-cpu

# Python-focused sparse clone (default after .sparse-checkout is committed):
git clone --filter=blob:none --sparse https://github.com/torch-spyre/ktir-cpu
cd ktir-cpu
git sparse-checkout reapply    # reads .sparse-checkout; rust/ excluded by default

# Rust contributor — add rust/ to an existing sparse clone:
git sparse-checkout add rust   # + examples if Option B below (no copy)

# Rust contributor — fresh sparse clone:
git clone --filter=blob:none --sparse https://github.com/torch-spyre/ktir-cpu
cd ktir-cpu
git sparse-checkout set rust examples   # today; just 'rust' after Option A below

--filter=blob:none (partial clone) skips fetching file blobs until they are checked
out — the lightest possible clone. It composes with sparse-checkout: sparse-checkout
controls which paths appear in the working tree; partial clone controls when blobs are
fetched.

Git worktrees and sparse-checkout: the sparse-checkout config is per-worktree, so
a full-clone worktree and a sparse worktree can coexist from the same repo object.

---

Resolving the examples/ Rust build dependency — two options

Option	What changes	Trade-off
A — Copy examples into rust/	Copy the 13 referenced .mlir files into rust/examples/ and update include_str! paths. examples/ at repo root remains canonical for Python.	rust/ becomes fully self-contained; sparse rust/-only clones work. 13 small files exist in two places; a CI diff check prevents drift.
B — Always include examples/ in the Rust sparse set	No code change. Document that Rust sparse clones need examples/.	Zero code change; current contributors unaffected. rust/-only sparse clones fail with an opaque rustc error rather than a clear message.

Recommendation: Option A. The 13 files are small and stable; a one-line CI diff
catches drift. The payoff is a self-contained rust/ that can be sparse-cloned,
published to crates.io, or vendored without pulling in the Python tree.

Open item: resolve with the PR author before Commit 3 lands.

---

Test Harness

ktir-core

No separate tests/ directory. All unit tests are inline #[cfg(test)] blocks:

Source file	Test count
affine.rs	19
codec.rs	15
parser_ast.rs	14
parser.rs	11
memref.rs	5
dtypes.rs	4
Total	57

Run with: cargo test -p ktir-core

ktir-emulator integration tests (tests/*.rs)

All 36 files import ktir_emulator (which re-exports ktir-core). Categorized:

Parity tests (port_*.rs) are direct Rust ports of the Python test suite — each
file corresponds 1:1 to a tests/test_*.py. They re-express the same assertions in
Rust (same inputs, same expected outputs or behaviours), translated to Rust idioms
where needed (e.g. ctx.get_value(name) instead of a dict lookup). They do not
run both interpreters: they simply verify the Rust implementation is correct by
reproducing what the Python tests already assert.

E2E correctness tests (e2e_*.rs) are Rust-only (no Python counterpart). They
drive a complete MLIR program through the full Rust pipeline — parse → HBM marshal →
multi-core grid execution → result readback — and assert the numerical output against
a reference computed inline (e.g. a naive GEMM for e2e_matmul, a Python-generated
golden for e2e_smollm2). e2e_matmul specifically is the first end-to-end check
that the matmul output values are correct (the Python parity tests only check the
latency report, not the numeric result).

Category	Files	What they test	Run in CI?
Parity / port	port_affine, port_ast, port_dtypes, port_tile, port_parse, port_parser_errors, port_parser_utils, port_ops, port_dialects_exec, port_distributed_view, port_indirect_access, port_interpreter, port_ktir_cpu, port_ktir_simple, port_latency, port_latency_modeling, port_lx_scoping, port_grid_scheduler, port_spec_gaps, port_examples	Rust ports of Python tests/test_*.py — same assertions, Rust idioms	Yes
E2E correctness	e2e_matmul, e2e_layernorm, e2e_smollm2, end_to_end	Full pipeline (parse→execute→readback) with numeric output assertions	Yes
Real-model e2e	e2e_real_forward	Full forward of smollm2-135m and llama-3.2-1b vs vendored transformers goldens; fetches HF weights at test time	Yes (background binary in CI)
Optimizer correctness	fuse_run_e2e, fuse_run_smollm2, flash_attn_golden, head_rewrite_golden, head_rewrite_e2e	Fusion passes and flash-attn against numeric goldens	Yes
Dispatch / parse	dispatch_coverage, parser_exec	Every registered op dispatches; parser roundtrips	Yes
Benches	bench_amx_vs_metal, bench_py_vs_rust, flash_attn_timing	Timing only — all marked #[ignore]	No — excluded by #[ignore]

Benches are never run in CI. cargo test skips #[ignore] tests by default; the CI
does not pass --ignored. To run benches manually:

cargo test --release --test bench_py_vs_rust -- --ignored --nocapture

Python↔Rust numeric equivalence

There is no automated CI gate that runs the Python interpreter (KTIRInterpreter)
and the Rust interpreter (execute_function) on the same input and diffs their
outputs. The equivalence guarantee today is indirect:

• Parity tests (port_*.rs) reproduce the same assertions as the Python test suite.
• e2e_real_forward validates the Rust production path against a shared transformers
  golden (the same golden bench_e2e_hermetic.py uses for the Python interpreter).

A direct cross-language output-comparison test is planned as a workflow_dispatch
workflow (Commit 4 below).

---

Commit Sequence

Commit 1 — CI plumbing only (no Rust source)

This repo is a monorepo: ktir_cpu/ (Python) and rust/ (Rust) are parallel
implementations with no build-time dependency on each other. Each CI workflow must be
scoped so it only fires when its own subtree changes:

• ci.yml (Python) must not run on Rust-only commits
• ci-rust.yml (Rust) must not run on Python-only commits

What lands:

• .github/workflows/ci-rust.yml — fixed (see issues below), with paths: scoped
  to rust/** and .github/workflows/ci-rust.yml
• .github/workflows/ci.yml — add paths-ignore: ['rust/**'] to both push and
  pull_request triggers

Required fixes to ci-rust.yml before landing:

Issue	Fix
actions/checkout@v6 (does not exist — build-breaking)	→ actions/checkout@v4
macos-26 (pre-release runner tag)	→ macos-latest or macos-15
No rust-version in Cargo.toml	Add rust-version = "<tested msrv>" to rust/Cargo.toml

Background e2e trick (keep as-is): The real-model e2e test fetches ~2.5 GB of HF
weights and runs a 1B-parameter forward — up to ~44s in a debug build. Running it
sequentially after all unit tests would dominate CI wall-clock. Instead, ci-rust.yml
builds all test binaries first (cargo test --no-run --message-format=json), extracts
the e2e_real_forward binary path from the build JSON, launches that binary in the
background, then runs all other tests in the foreground with --skip real_forward.
Both processes complete independently; the step fails if either exits non-zero.

Result: Green CI with no Rust source in the repo yet. Path matching is validated
before any Rust code lands.

---

Commit 2 — ktir-core + workspace stub

What lands:

• rust/Cargo.toml — workspace with single member ["crates/ktir-core"]
• rust/crates/ktir-core/Cargo.toml + src/ — all 9 source files

Dependencies: None outside std. Truly dependency-free.

Tests that run: 57 inline unit tests via cargo test -p ktir-core.

CI job: cargo fmt --check → cargo clippy --release → cargo test --release.
The emulator and optimizer crates do not exist yet; the workspace is minimal.

---

Commit 3 — ktir-optimizer + ktir-emulator + all correctness tests

The main Rust code commit. Completes the workspace.

What lands:

• rust/Cargo.toml — updated to all three workspace members
• rust/crates/ktir-optimizer/ — fusion, flash_attn, head_rewrite passes
• rust/crates/ktir-emulator/src/ — interpreter, dialects, machine state, blas,
  metal, comm, program
• rust/crates/ktir-emulator/tests/ — all 36 test files
• rust/crates/ktir-emulator/tests/fixtures/*.tar.gz — four vendored model fixtures
• rust/crates/ktir-emulator/build.rs
• rust/README.md

Tests that run in CI (cargo test --release --workspace):

All tests except #[ignore]d ones — covers all parity, e2e, optimizer golden, and
the real-model e2e (as background binary). Timing benches are excluded by #[ignore].

On the vendored fixture archives:

The four tests/fixtures/*.tar.gz files (total ~10 MB) must be committed — they
cannot be generated in CI. Each archive contains two kinds of content:

• MLIR programs (manifest.json + node*.mlir) — emitted by scratchy (the
  internal cudaforge KTIR emitter via SCRATCHY_KTIR_DUMP). This tool has a
  ~/.cache/cudaforge dependency and is not available in CI.
• Goldens + runtime inputs (t<id>.f16.gz, golden.f16.gz) — generated by
  gen_golden.py, which requires torch + transformers + safetensors. This is
  a ~3 GB install and a real HF model forward; not suitable for a CI step that would
  run on every push.

The archives are small (largest is 7.8 MB, smollm2 decode is 362 KB) and content-
stable: they change only when the KTIR program or the reference model changes. The
vendoring approach is intentional and correct. Re-archiving instructions are in
tests/fixtures/README.md.

Files to clean up before landing:

These files exist in the PR but should not land on main as-is:

File	Issue	Action
bench_py_vs_rust.py (repo root)	Hardcoded absolute path /Users/moosevan/git/ktir-cpu	Fix to repo-relative or move to rust/
bench_e2e_py_vs_rust.py (repo root)	Same hardcoded path	Fix to repo-relative or move to rust/
bench_e2e_sweep.sh (repo root)	Same hardcoded path	Fix or move
mlx_fidelity_probe.py (repo root)	Scratch file, not part of stated scope	Remove or move to rust/

---

Commit 4 — Python↔Rust equivalence workflow (manual dispatch)

What lands:

• .github/workflows/equiv-test.yml — workflow_dispatch trigger only

Design:

• Trigger: manual (workflow_dispatch) — never runs automatically. No path matching
  needed. Used for ad-hoc cross-language output checks.
• Environment: same runner needs both uv (Python) and cargo (Rust).

What the equivalence test does:

1. Selects a small fixed set of MLIR programs from examples/
   (vector_add, matmul, layernorm — already shared by both bench_py_vs_rust.py
   and bench_py_vs_rust.rs)
2. Runs each through KTIRInterpreter (Python) and execute_function (Rust)
3. Asserts max-abs difference is within f16 tolerance

Implementation (Option A — recommended):
A Python script tests/equiv/test_py_vs_rust.py that:

• Runs the Python interpreter on each program and captures outputs
• Invokes cargo test --test bench_py_vs_rust -- --ignored --nocapture and parses
  the Rust outputs from stdout
• Diffs both and reports max-abs per kernel

This pattern already exists in rust/crates/ktir-emulator/tests/fixtures/bench_e2e_hermetic.py
for the real-model case — generalize it for simpler kernels.

---

Wiki + README

GitHub Wiki page: Rust Port — Build & Development

Suggested contents:

• Build instructions (cargo, feature flags, macOS vs Linux BLAS providers)
• How to run specific test categories
  • Parity tests: cargo test -p ktir-emulator --test port_affine etc.
  • Real-model e2e: cargo test --test e2e_real_forward -- --test-threads=1
  • Timing benches (manual): cargo test --release --ignored --nocapture
• How to regenerate fixtures: gen_golden.py runs a real transformers forward
  (public HF repos, no token) and writes the golden logits + runtime input tensors
  (t<id>.f16.gz, golden.f16.gz) into each fixture dir. Requires torch,
  transformers, safetensors (run via uv run --no-project --with ...). Only
  needed when the KTIR program changes or the reference model output needs refreshing.
• Port map summary (condensed from rust/port-map.md) — which port_*.rs maps to
  which Python test_*.py
• Known gaps (rust/TODOs.md summary) and deferred items

Top-level README.md: Add after the existing content:

## Rust port

A high-performance Rust reimplementation lives in [`rust/`](rust/). It targets
the same RFC 0682 spec as the Python reference and validates against the same
golden outputs. See [`rust/README.md`](rust/README.md) for build instructions,
or the [Rust Port wiki page](../../wiki/Rust-Port) for the full development guide.

rust/port-map.md (a 5000-line design document) stays in the repo as a reference
artifact for the port author but is not linked prominently. It will be condensed
into the wiki over time.

---

Summary

Step	Contents	CI trigger	Tests
Commit 1	ci-rust.yml (fixed) + ci.yml path-ignore	rust/** or CI file change	Format check only (no source yet)
Commit 2	ktir-core + workspace stub	rust/**	57 inline unit tests
Commit 3	ktir-optimizer + ktir-emulator + all tests + fixtures	rust/**	All parity + e2e + golden + real-model; benches excluded by #[ignore]
Commit 4	Equivalence workflow	workflow_dispatch (manual)	Python + Rust outputs compared on shared MLIR examples
—	Wiki page + README.md link	n/a	n/a

[starpit]

Status of the landing-plan items, re-checked against the current PR HEAD. Note the plan reads as if written against an earlier snapshot of the branch, so several items are already resolved.

Item	Status	Note
rust/docs/port-map.md design doc	✅ Done	Already in the PR (~5k lines, Python→Rust class/method map).
rust-version (MSRV) in Cargo.toml	✅ Done	Verified MSRV = 1.94 on all three crates. It's not the 1.85 the edition implies: edition-2024 → 1.85, the idna/icu deps → 1.86, let-chains → 1.88, and stdarch_neon_f16 (the f16 SIMD path in codec.rs) → 1.94. Proven by building 1.85 → 1.94, not guessed.
Python↔Rust equivalence test (Commit 4)	✅ Done — stronger than proposed	A direct differential harness: the same seeded inputs run through Python KTIRInterpreter and the Rust port, diffed head-to-head (not both-vs-a-hardcoded-answer-key, which is all the port_*.rs parity tests do). Seeded + fuzzed, gated per-push in rust-conformance.yml. CPU/AMX path: 100% across all 19 shared programs — 17 bit-exact (max-abs 0) + 2 matched-failure (both sides raise the same category). Metal fast path (NAX/simdgroup GEMM + fused map) is covered too: the same programs forced through the resident/GPU executor and diffed within a principled bf16/f16 band, each with a mandatory offload proof (a silent CPU fallback FAILS) — 9/9 Metal-eligible programs conform, zero divergence. It also caught real bugs the parity tests missed (f16-subnormal truncf; non-identity indirect subscripts; a Metal GEMM weight-reader; a fusion reduce-accumulator). See rust/PERFORMANCE.md.
Stale root files (bench_*.py, bench_*.sh, mlx_fidelity_probe.py)	✅ Already removed	All absent from the PR — the "cleanup" list reflects a pre-cleanup snapshot.
actions/checkout@v6 → v4	❌ Skip	checkout@v6 exists and CI already runs green on it; this would be a downgrade.
macos-26 → macos-latest/macos-15	❌ Skip — intentional pin	Not a stray pre-release tag — it's required. The Metal backend compiles Metal 4 (MTLLanguageVersion::Version4_0) and AOT-builds the NAX kernels with -mmacosx-version-min=26.2, both of which need macOS 26+. macos-latest currently resolves to macOS 15, and macos-15 obviously lacks Metal 4 — either would break the build. (macOS 26 exists as the macos-26 label but isn't GA/-latest yet, which is why we pin it explicitly.) Revisit macos-latest once GitHub promotes macOS 26 to -latest.
.sparse-checkout to hide rust/	❌ Skip	Git does not auto-apply sparse-checkout on git clone (as the plan itself notes), so committing the file does not achieve the stated goal — a normal clone still shows rust/. It adds IDE/worktree/tooling footguns for what is a non-problem (a rust/ folder in a monorepo). A one-line README pointer covers discoverability.
Copy examples/*.mlir into rust/ (self-containment)	❌ Skip	Only matters to enable sparse rust/-only clones (skipped above) or a crates.io publish (not planned). In a monorepo the include_str! resolves in every real scenario — full clone, git worktree, CI, and cd rust && cargo test — so copying just adds duplication + drift risk for no benefit.
4-commit split (CI / core / emulator / equiv)	❌ Skip	Landed as one signed commit. The split is cosmetic — "Commit 3" is still ~45k lines, so it doesn't aid review, and a CI-only "Commit 1" adds no value.
Wiki page + top-level README.md "Rust port" section	⏳ Optional / post-merge	rust/README.md + rust/docs/port-map.md already cover build/dev. A top-level README pointer and a condensed wiki page are easy follow-ups; we kept the PR scoped to rust/.

Net: the only "required fixes" worth doing were rust-version (done) — the CI version "fixes" would regress a green pipeline, and the sparse-checkout / examples-copying / 4-commit items don't pay for themselves. The headline addition the plan didn't anticipate: a directly-proven, per-push 100% Python↔Rust conformance gate, not just parity tests.

[fabianlim]

@raghukiran1224 @mudhakar would like your weigh-in on this. My proposal is have the rust port live independently of the ktir-cpu (since its a re-implmeentation) and we will CI them separately. My main concerns are

• testing: we will trigger rust CI only if there are changes to the rust code, pushes to python code will not run any rust-CI. Equivalence tests will be triggered also only on changes to rust
• I would like to have the rust checkout optional. If this were to be a real mono-repo, then I would be able to check out each folder separately, but its too much effort to move it that way, so my suggestion is to setup git sparse checkout. the rust port seems to have significant number of files, and in vscode this can slow down file searches, intellisense, etc. So I didnt want to burden users on carrying a rust folder with a lot of files.