{"id":50549297,"url":"https://github.com/eugenehp/amx-rs","last_synced_at":"2026-06-04T01:30:51.798Z","repository":{"id":346038763,"uuid":"1187495738","full_name":"eugenehp/amx-rs","owner":"eugenehp","description":"Rust reimplementation of Apple AMX (Matrix eXtensions)","archived":false,"fork":false,"pushed_at":"2026-03-22T19:19:39.000Z","size":295,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"master","last_synced_at":"2026-05-22T14:45:34.559Z","etag":null,"topics":["amx","apple","m1","m3","m4","rust","silicon"],"latest_commit_sha":null,"homepage":"https://crates.io/crates/amx-rs","language":"Rust","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/eugenehp.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-03-20T19:45:59.000Z","updated_at":"2026-03-22T19:19:42.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/eugenehp/amx-rs","commit_stats":null,"previous_names":["eugenehp/amx-rs"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/eugenehp/amx-rs","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eugenehp%2Famx-rs","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eugenehp%2Famx-rs/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eugenehp%2Famx-rs/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eugenehp%2Famx-rs/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/eugenehp","download_url":"https://codeload.github.com/eugenehp/amx-rs/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eugenehp%2Famx-rs/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":33886160,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-03T02:00:06.370Z","response_time":59,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["amx","apple","m1","m3","m4","rust","silicon"],"created_at":"2026-06-04T01:30:51.079Z","updated_at":"2026-06-04T01:30:51.789Z","avatar_url":"https://github.com/eugenehp.png","language":"Rust","funding_links":[],"categories":[],"sub_categories":[],"readme":"# AMX Rust Implementation\n\n[![Benchmark](https://github.com/eugenehp/amx-rs/actions/workflows/bench.yml/badge.svg)](https://github.com/eugenehp/amx-rs/actions/workflows/bench.yml)\n\n**Rust implementation of Apple AMX (Apple Matrix eXtensions) — achieving up to 90% of Apple Accelerate's sgemm performance through reverse-engineered undocumented instruction encodings and cache-optimized algorithms.**\n\n## Performance vs Apple Accelerate\n\nMeasured on Apple M4 Pro (Virtual), 10 cores:\n\n| Size | amx-rs (GFLOPS) | Accelerate (GFLOPS) | % of Accelerate |\n|------|-----------------|---------------------|-----------------|\n| 32×32 | 97 | 277 | 35% |\n| 64×64 | 107 | 776 | 14% |\n| 128×128 | 330 | 1350 | 24% |\n| 256×256 | 660 | 1650 | 40% |\n| 512×512 | 695 | 1680 | 41% |\n| 1024×1024 | **1500** | **1630** | **92%** |\n\n## Architecture\n\n### Multi-Backend Dispatch\n\n```\nmatmul(A, B)\n├── m=1 or n=1 → Scalar (auto-vectorized, optimal for rank-1)\n├── max_dim ≤ 32 → NEON 8×8 µ-kernel (zero AMX setup cost)\n├── medium (64-512) → Persistent AMX Thread Pool\n│   ├── Pre-pack A (NEON column-gather) and B (row-copy)\n│   ├── Distribute (i,j) tile pairs across spin-waiting workers\n│   └── 4Y AMX µ-kernel: 4 fma32 per Y-preload cycle\n└── large (1024+) → GEBP with L1/L2/L3 cache blocking\n    ├── Parallel B̃ packing (rayon par_iter)\n    ├── Dynamic MC for thread utilization\n    └── QoS pinning to P-cores\n```\n\n### AMX µ-Kernel: 4Y Register Preloading\n\nThe inner kernel preloads 4 Y registers (A columns) then issues 4 fma32 outer products per X load (B row), giving 4× better compute-to-load ratio:\n\n```c\n// Load 4 A columns into Y[0..3]\nldy Y[0] ← A[k+0];  ldy Y[1] ← A[k+1]\nldy Y[2] ← A[k+2];  ldy Y[3] ← A[k+3]\n\n// Each B row: 1 load + 1 fma32 (using Y-row select encoding)\nldx X[0] ← B[k+0];  fma32(0x00)   // Z += X[0] ⊗ Y[0]\nldx X[0] ← B[k+1];  fma32(0x40)   // Z += X[0] ⊗ Y[1]\nldx X[0] ← B[k+2];  fma32(0x80)   // Z += X[0] ⊗ Y[2]\nldx X[0] ← B[k+3];  fma32(0xC0)   // Z += X[0] ⊗ Y[3]\n```\n\n### GEBP Cache Blocking (Goto-style, tuned for Apple Silicon)\n\n```\nfor jc in 0..n step NC=1024:       ← L2 blocking (B̃ = 2 MB)\n  Pack B̃ panel (parallel across j-tiles)\n  for pc in 0..k step KC=512:      ← L1 blocking (64 KB working set)\n    for ic in 0..m step mc_par:    ← dynamic MC for thread balance\n      Pack à panel (per-thread, private L2)\n      GEBP macro-kernel:\n        for (ir, jr) tile pairs → 4Y AMX µ-kernel\n```\n\nCache sizing for Apple Silicon M1-M4:\n- **L1D**: 64 KB → KC × (MR+NR) × 4 = 64 KB ✓\n- **L2**: 4 MB shared → B̃ panel KC × NC × 4 = 2 MB ✓\n- **Dynamic MC**: auto-tuned so each thread gets ≥ 2 work items\n\n### Persistent AMX Thread Pool\n\nWorkers spin-wait on atomic generation counters (no mutex, no condvar). This eliminates the ~20-100µs thread spawn + AMX init overhead that kills parallelism at small matrix sizes:\n\n```\nMain thread                    Worker threads (persistent)\n─────────────                  ──────────────────────────\nWrite jobs to slots            Spin on: generation.load(Acquire)\nfence(SeqCst)                  │\ngeneration.fetch_add(1)  ───►  Wake: gen increased!\n│                              fence(SeqCst)\n│                              amx_set()\n│                              Execute tiles\n│                              amx_clr()\n│                              done_gen.store(gen, Release)\nWait: done_gen \u003e= gen    ◄───  │\n```\n\n---\n\n## Reverse Engineering: AMX Instruction Encodings\n\n### fma32 Y-Row Select (Discovered)\n\nThe AMX `fma32` instruction's Y-register selection was **undocumented**. We found the correct encoding through brute-force testing on M4 Pro hardware:\n\n```\nfma32 operand bits [8:6] = Y register row (0-7)\n\n  Y[0] = 0x000   Y[1] = 0x040   Y[2] = 0x080   Y[3] = 0x0C0\n  Y[4] = 0x100   Y[5] = 0x140   Y[6] = 0x180   Y[7] = 0x1C0\n```\n\nThis was not documented in any public source, including the [corsix/amx](https://github.com/corsix/amx) reverse-engineering project. Previous attempts using bits [19:10] or [21:20] produced incorrect results.\n\n### fma32 Operand Layout (Outer Product Mode)\n\n```\nBit 63:    0 = outer product mode (matrix), 1 = vector mode\nBits 8:6:  Y register row (0-7) — selects which Y row for the outer product\nBits 5:0:  Other flags (Z row offset, etc.)\n```\n\n### Load/Store Operand Layout\n\n```\nBits 55:0:   Memory address (pointer)\nBits 62:56:  Register row (0-63 for Z, 0-7 for X/Y)\n```\n\n### How Apple Accelerate Uses AMX (Reverse Engineered)\n\nWe disassembled `cblas_sgemm_singlecore` from Apple's Accelerate.framework:\n\n```\ncblas_sgemm\n└→ APPLE_NTHREADS (persistent thread pool dispatch)\n   └→ cblas_sgemm_singlecore (3516 insns, 287 AMX ops, ZERO loops)\n```\n\n**AMX instruction breakdown in Apple's kernel:**\n\n| Instruction | Count | Purpose |\n|-------------|------:|---------|\n| fma32 | 185 | Outer product compute |\n| extrx | 34 | Extract Z → X (register reuse) |\n| extry | 8 | Extract Z → Y (register reuse) |\n| vecfp | 30 | Vector add/sub (combine sub-results) |\n| ldy | 26 | Load A data |\n| ldx | 2 | Load B data |\n| stz | 2 | Store results |\n\n**Key insight:** Apple achieves **185 fma32 with only 28 loads** (ratio 6.6:1) by using `extrx`/`extry` to extract intermediate Z results back into X/Y registers for reuse. This is a **register-level recursive algorithm** — not traditional GEBP.\n\nApple's operand encoding (`0x38000000 | (pair \u003c\u003c 20)`) uses bits 29:27=111 which activates a different AMX mode than our outer-product mode. This encoding produces zeros on M4 Pro Virtual — it may require bare-metal execution or chip-specific features.\n\n**Apple's algorithm structure:**\n1. 48 initial fma32 with standard operand (compute sub-products)\n2. 4 extrx (extract Z rows → X registers)\n3. 32 fma32 with various operands (combine using extracted data)\n4. More extrx/extry + fma32 cycles (recursive combination)\n5. vecfp operations (vector add/subtract for final assembly)\n6. Pattern repeats for second half\n\nThis is consistent with a **Winograd-Strassen variant** operating at the register level.\n\n---\n\n## Workspace Structure\n\n### [`amx-sys`](crates/amx-sys) — Low-Level AMX Bindings\n\nC-compiled AMX instruction wrappers using the exact `AMX_OP_GPR` encoding from the reference implementation. Includes:\n- All 23 AMX instructions\n- Runtime AMX availability detection (fork + SIGILL probe)\n- Optimized µ-kernels: `amx_f32_tile_kernel`, `amx_f32_tile_kernel_4y`\n- NEON helper functions for packing and dot products\n- Strided sgemm kernel for zero-copy matmul\n\n### [`amx-rs`](crates/amx-rs) — High-Level API\n\nErgonomic matrix/vector operations with automatic backend dispatch:\n- `Matrix\u003cT\u003e`, `Vector\u003cT\u003e` generic types\n- `matmul()` — smart dispatch across all backends\n- `matmul_amx()` — single-threaded AMX with pre-packing\n- `matmul_pool()` — persistent thread pool dispatch\n- `matmul_gebp()` / `matmul_gebp_parallel()` — full GEBP\n- `matmul_recursive()` — cache-oblivious recursive (experimental)\n- `matmul_neon()` — NEON 8×8 for tiny matrices\n- `matmul_scalar()` — auto-vectorized fallback\n- `no_std` compatible (without `std` feature)\n\n## Quick Start\n\n```rust\nuse amx::Matrix;\n\nlet a = Matrix::from_data(vec![1.0f32; 1024*1024], 1024, 1024)?;\nlet b = Matrix::from_data(vec![1.0f32; 1024*1024], 1024, 1024)?;\n\n// Automatic dispatch: NEON → Pool → GEBP based on size\nlet c = a.matmul(\u0026b)?;\n```\n\n## Building \u0026 Testing\n\n```bash\n# Run all 33 tests\ncargo test --workspace --release\n\n# Benchmark vs Accelerate\ncargo bench -p amx-rs --bench scalar_vs_amx -- --nocapture\n\n# Control benchmark parameters\nBENCH_MAX_N=1024 BENCH_ITERS=20 cargo bench -p amx-rs --bench scalar_vs_amx -- --nocapture\n```\n\n## Optimization History\n\nEach optimization was validated through ablation studies:\n\n| Optimization | Impact | Status |\n|-------------|--------|--------|\n| AMX 16×16 tile kernel | Baseline 266 GF at 256 | ✅ Shipped |\n| Double-buffered X/Y loads | +5% µ-kernel throughput | ✅ Shipped |\n| NEON dispatch for N≤32 | 113 GF (beats Accelerate!) | ✅ Shipped |\n| Persistent AMX thread pool | 2× at N=128-512 | ✅ Shipped |\n| GEBP L1/L2/L3 cache blocking | 1400+ GF at N=1024 | ✅ Shipped |\n| Dynamic MC for thread balance | +40% at N=1024 | ✅ Shipped |\n| NC=1024 (B̃ fits in shared L2) | +60-90% GEBP single-thread | ✅ Shipped |\n| Parallel B̃ packing | Up to 97.8% of Accelerate | ✅ Shipped |\n| QoS pinning to P-cores | +3-5% parallel | ✅ Shipped |\n| NEON A-panel packing | +5-10% GEBP single-core | ✅ Shipped |\n| 4Y µ-kernel (multi-Y fma32) | 4 fma32 per Y-preload | ✅ Shipped |\n| Strided AMX kernel (no-copy) | Too slow (gather dominates) | ❌ Kept as reference |\n| Recursive cache-oblivious | Slower than GEBP (pack overhead per leaf) | ❌ Experimental |\n| KC=256 (smaller L1 block) | Worse (more packing) | ❌ Reverted |\n| NC=512 (smaller L2 block) | Worse (too much repacking) | ❌ Reverted |\n| Lower AMX-par threshold | Worse at N≤256 (thread overhead) | ❌ Reverted |\n\n## Coverage\n\n| Category | Count | Status |\n|----------|-------|--------|\n| AMX Instructions | 23/23 | ✅ |\n| Tests | 33 | ✅ |\n| Pass Rate | 100% | ✅ |\n| Matmul Backends | 7 | ✅ |\n| Data Types | f32 (primary), f64 (precision) | ✅ |\n\n## License\n\nMIT OR Apache-2.0\n\n## Acknowledgments\n\n- [corsix/amx](https://github.com/corsix/amx) — AMX instruction set reverse engineering\n- Apple Accelerate.framework — disassembly analysis for algorithm insights\n- Goto \u0026 Van de Geijn — GEBP algorithm design\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Feugenehp%2Famx-rs","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Feugenehp%2Famx-rs","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Feugenehp%2Famx-rs/lists"}