onnxruntime-ep-mlx 0.27.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. onnxruntime_ep_mlx-0.27.6/.gitignore +80 -0
  2. onnxruntime_ep_mlx-0.27.6/PKG-INFO +267 -0
  3. onnxruntime_ep_mlx-0.27.6/README.md +1 -0
  4. onnxruntime_ep_mlx-0.27.6/hatch_build.py +199 -0
  5. onnxruntime_ep_mlx-0.27.6/pyproject.toml +79 -0
  6. onnxruntime_ep_mlx-0.27.6/rust/Cargo.toml +56 -0
  7. onnxruntime_ep_mlx-0.27.6/rust/README.md +251 -0
  8. onnxruntime_ep_mlx-0.27.6/rust/build.rs +100 -0
  9. onnxruntime_ep_mlx-0.27.6/rust/src/compiled.rs +999 -0
  10. onnxruntime_ep_mlx-0.27.6/rust/src/engine.rs +2171 -0
  11. onnxruntime_ep_mlx-0.27.6/rust/src/ep.rs +2157 -0
  12. onnxruntime_ep_mlx-0.27.6/rust/src/factory.rs +268 -0
  13. onnxruntime_ep_mlx-0.27.6/rust/src/lib.rs +167 -0
  14. onnxruntime_ep_mlx-0.27.6/rust/src/logging.rs +101 -0
  15. onnxruntime_ep_mlx-0.27.6/rust/src/mlx.rs +319 -0
  16. onnxruntime_ep_mlx-0.27.6/rust/src/ops/attention.rs +2287 -0
  17. onnxruntime_ep_mlx-0.27.6/rust/src/ops/controlflow.rs +396 -0
  18. onnxruntime_ep_mlx-0.27.6/rust/src/ops/conv.rs +1604 -0
  19. onnxruntime_ep_mlx-0.27.6/rust/src/ops/elementwise.rs +1049 -0
  20. onnxruntime_ep_mlx-0.27.6/rust/src/ops/image.rs +139 -0
  21. onnxruntime_ep_mlx-0.27.6/rust/src/ops/math.rs +739 -0
  22. onnxruntime_ep_mlx-0.27.6/rust/src/ops/matmul.rs +168 -0
  23. onnxruntime_ep_mlx-0.27.6/rust/src/ops/misc.rs +1297 -0
  24. onnxruntime_ep_mlx-0.27.6/rust/src/ops/mod.rs +23 -0
  25. onnxruntime_ep_mlx-0.27.6/rust/src/ops/norm.rs +969 -0
  26. onnxruntime_ep_mlx-0.27.6/rust/src/ops/onnx_ml_linear.rs +878 -0
  27. onnxruntime_ep_mlx-0.27.6/rust/src/ops/onnx_ml_preprocess.rs +666 -0
  28. onnxruntime_ep_mlx-0.27.6/rust/src/ops/quant.rs +2293 -0
  29. onnxruntime_ep_mlx-0.27.6/rust/src/ops/random.rs +597 -0
  30. onnxruntime_ep_mlx-0.27.6/rust/src/ops/recurrent.rs +663 -0
  31. onnxruntime_ep_mlx-0.27.6/rust/src/ops/reduction.rs +968 -0
  32. onnxruntime_ep_mlx-0.27.6/rust/src/ops/shape.rs +2421 -0
  33. onnxruntime_ep_mlx-0.27.6/rust/src/ops/signal.rs +610 -0
  34. onnxruntime_ep_mlx-0.27.6/rust/src/ops/ssm.rs +835 -0
  35. onnxruntime_ep_mlx-0.27.6/rust/src/ops/stragglers.rs +155 -0
  36. onnxruntime_ep_mlx-0.27.6/rust/src/ops/vision.rs +1508 -0
  37. onnxruntime_ep_mlx-0.27.6/rust/src/registry.rs +1233 -0
  38. onnxruntime_ep_mlx-0.27.6/rust/src/sys.rs +26 -0
  39. onnxruntime_ep_mlx-0.27.6/rust/src/trace.rs +1786 -0
  40. onnxruntime_ep_mlx-0.27.6/rust/wrapper_mlx.h +1 -0
  41. onnxruntime_ep_mlx-0.27.6/rust/wrapper_ort.h +1 -0
  42. onnxruntime_ep_mlx-0.27.6/src/onnxruntime_ep_mlx/__init__.py +166 -0
  43. onnxruntime_ep_mlx-0.27.6/src/onnxruntime_ep_mlx/py.typed +0 -0
@@ -0,0 +1,80 @@
1
+ # Xcode
2
+ #
3
+ # gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore
4
+
5
+ ## User settings
6
+ xcuserdata/
7
+
8
+ ## Obj-C/Swift specific
9
+ *.hmap
10
+
11
+ ## App packaging
12
+ *.ipa
13
+ *.dSYM.zip
14
+ *.dSYM
15
+
16
+ # CocoaPods
17
+ #
18
+ # We recommend against adding the Pods directory to your .gitignore. However
19
+ # you should judge for yourself, the pros and cons are mentioned at:
20
+ # https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control
21
+ #
22
+ # Pods/
23
+ #
24
+ # Add this line if you want to avoid checking in source code from the Xcode workspace
25
+ # *.xcworkspace
26
+
27
+ # Carthage
28
+ #
29
+ # Add this line if you want to avoid checking in source code from Carthage dependencies.
30
+ # Carthage/Checkouts
31
+
32
+ Carthage/Build/
33
+
34
+ # fastlane
35
+ #
36
+ # It is recommended to not store the screenshots in the git repo.
37
+ # Instead, use fastlane to re-generate the screenshots whenever they are needed.
38
+ # For more information about the recommended setup visit:
39
+ # https://docs.fastlane.tools/best-practices/source-control/#source-control
40
+
41
+ fastlane/report.xml
42
+ fastlane/Preview.html
43
+ fastlane/screenshots/**/*.png
44
+ fastlane/test_output
45
+
46
+ # --- ONNX Runtime MLX EP (Rust) ---
47
+ # Cargo build output (the plugin dylib is rebuilt from source: `cargo build`).
48
+ /rust/target/
49
+ build/
50
+ # Scratch reference material downloaded during development.
51
+ .ref/
52
+ build-mlx/
53
+ .vscode/
54
+
55
+ # Bundled runtime artifacts dropped into the wheel package by the build hook
56
+ # (python/hatch_build.py). Rebuilt on every wheel build; never committed.
57
+ python/src/onnxruntime_ep_mlx/*.dylib
58
+ python/src/onnxruntime_ep_mlx/*.metallib
59
+
60
+ # --- Python (op tests, tooling) ---
61
+ __pycache__/
62
+ *.py[cod]
63
+ *$py.class
64
+ .pytest_cache/
65
+ .mypy_cache/
66
+ .ruff_cache/
67
+ .hypothesis/
68
+ .coverage
69
+ .coverage.*
70
+ htmlcov/
71
+ *.egg-info/
72
+ .eggs/
73
+ dist/
74
+ # Virtual environments
75
+ .venv/
76
+ venv/
77
+ env/
78
+ .python-version
79
+ scratch/
80
+ .DS_Store
@@ -0,0 +1,267 @@
1
+ Metadata-Version: 2.5
2
+ Name: onnxruntime-ep-mlx
3
+ Version: 0.27.6
4
+ Summary: MLX-native ONNX Runtime execution provider (plugin EP) for Apple Silicon
5
+ Project-URL: Homepage, https://github.com/justinchuby/onnxruntime-mlx
6
+ Project-URL: Repository, https://github.com/justinchuby/onnxruntime-mlx
7
+ Author: onnxruntime-ep-mlx
8
+ License: MIT
9
+ Keywords: apple-silicon,coreml,execution-provider,metal,mlx,onnxruntime
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Operating System :: MacOS :: MacOS X
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Programming Language :: Rust
20
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
21
+ Requires-Python: >=3.10
22
+ Requires-Dist: onnxruntime>=1.22
23
+ Description-Content-Type: text/markdown
24
+
25
+ # onnxruntime-mlx
26
+
27
+ > **PyPI package: [`onnxruntime-ep-mlx`](https://pypi.org/project/onnxruntime-ep-mlx/)** — `pip install onnxruntime-ep-mlx`, `import onnxruntime_ep_mlx`. (Formerly published as `onnxruntime-mlx`, now renamed.)
28
+
29
+ An **MLX-native execution provider** for ONNX Runtime on Apple Silicon, built as an out-of-tree
30
+ **plugin EP** (ORT plugin-EP C ABI, ORT 1.27 / `ORT_API_VERSION 27`). It ships as a standalone
31
+ `libonnxruntime_mlx_ep.dylib` loaded by a stock prebuilt `libonnxruntime.dylib` via
32
+ `RegisterExecutionProviderLibrary` — **no ONNX Runtime fork required**.
33
+
34
+ The EP translates fused ONNX subgraphs into [MLX](https://github.com/ml-explore/mlx) graphs for
35
+ encoders, LLM prefill, and token-at-a-time decode.
36
+
37
+ ## How it works
38
+
39
+ `ONNX fused subgraph → MLX graph → mlx_compile → mlx_eval → ORT outputs`
40
+
41
+ The EP translates supported ONNX regions into MLX, compiles reusable closures, and leaves unsupported
42
+ ops on ORT CPU. It covers common encoder and decoder operators, including quantized matmul, GQA/MHA,
43
+ PagedAttention, RoPE, normalization, convolution, pooling, reductions, and shape operations. See
44
+ [`docs/OP_ARCHITECTURE.md`](docs/OP_ARCHITECTURE.md) for the coverage table.
45
+
46
+ Large, fully claimed regions are fastest. Dynamic shapes are compiled per shape key; autoregressive
47
+ decode uses a shapeless path so growing KV length does not retrace. Use
48
+ `ONNXRUNTIME_EP_MLX_VERBOSE=1` or `ONNXRUNTIME_EP_MLX_CLAIM_DEBUG=1` to diagnose fallback and
49
+ fragmentation.
50
+
51
+ ## Requirements
52
+
53
+ - macOS on Apple Silicon, ORT 1.27 prebuilt (`ORT_API_VERSION >= 27`)
54
+ - **`mlx-c` (and `mlx`) — a HARD build dependency**: `brew install mlx-c`
55
+ - A **Rust toolchain** (`rustup`) to build the EP from source
56
+
57
+ ## Versioning (ORT compatibility)
58
+
59
+ A plugin EP targets one ORT C-ABI version. Package versions use
60
+ `0.<ORT_API_VERSION>.<patch>`:
61
+
62
+ | onnxruntime-ep-mlx | ONNX Runtime | `ORT_API_VERSION` |
63
+ |---|---|---|
64
+ | `0.27.x` | 1.27.x | 27 |
65
+
66
+ For example, ORT 1.28 moves the EP to `0.28.x`.
67
+
68
+ ## Build
69
+
70
+ The EP is a Rust `cdylib` crate under [`rust/`](rust/). Point it at an ONNX Runtime C-API
71
+ include directory and `cargo build`:
72
+
73
+ ```sh
74
+ brew install mlx-c # HARD dependency (mlx-c + mlx)
75
+ cd rust
76
+ # Either point ORT_INCLUDE_DIR at the ORT headers directly, or set ORT_HOME to an
77
+ # ONNX Runtime release root (build.rs will look in $ORT_HOME/include):
78
+ export ORT_INCLUDE_DIR=/path/to/onnxruntime/include # or: export ORT_HOME=/path/to/onnxruntime-osx-arm64-1.27.0
79
+ cargo build --release
80
+ # => rust/target/release/libonnxruntime_mlx_ep.dylib (registers the EP as "MLXExecutionProvider")
81
+ ```
82
+
83
+ Set `MLX_PREFIX` and `MLXC_PREFIX` to link and bundle a custom MLX runtime
84
+ instead of the Homebrew installation. Each prefix must contain `include/` and
85
+ `lib/`; the wheel builder also bundles `mlx.metallib` and an optional
86
+ `libjaccl.dylib` from `MLX_PREFIX/lib`.
87
+
88
+ The crate binds the ORT plugin-EP C ABI and `mlx-c` directly via `bindgen`; it does **not** link
89
+ `libonnxruntime` (ORT is reached through the `OrtApi` function-pointer table passed to
90
+ `CreateEpFactories`).
91
+
92
+ ## Install & use
93
+
94
+ ### Python (recommended)
95
+
96
+ ```sh
97
+ pip install -U onnxruntime-ep-mlx # macOS/Apple-Silicon wheel; bundles the mlx runtime
98
+ ```
99
+
100
+ ```python
101
+ import onnxruntime as ort
102
+ import onnxruntime_ep_mlx
103
+
104
+ # Register the plugin EP once, then select it (with CPU fallback) like any provider.
105
+ onnxruntime_ep_mlx.register_execution_provider_library() # name: "MLXExecutionProvider"
106
+ sess = ort.InferenceSession(
107
+ "model.onnx",
108
+ providers=["MLXExecutionProvider", "CPUExecutionProvider"],
109
+ )
110
+ out = sess.run(None, feeds)
111
+ ```
112
+
113
+ `onnxruntime_ep_mlx` also exposes `library_path()`, `ep_name()`, `version()`, and
114
+ `append_to_session_options(so)`.
115
+
116
+ ### C / C++ (or any onnxruntime binding)
117
+
118
+ Point onnxruntime at the built dylib and select the provider by name:
119
+
120
+ ```c
121
+ // 1. Register the plugin library with the environment (once).
122
+ RegisterExecutionProviderLibrary(env, "MLXExecutionProvider",
123
+ "/abs/path/libonnxruntime_mlx_ep.dylib");
124
+ // 2. Append it to a session's options (falls back to CPU for unclaimed ops).
125
+ const char* ep = "MLXExecutionProvider";
126
+ SessionOptionsAppendExecutionProvider_V2(options, env, &ep, /*count*/ 1, ...);
127
+ ```
128
+
129
+ From Rust via **onnx-genai**: `ONNX_GENAI_EP=metal` +
130
+ `ONNX_GENAI_METAL_EP_LIB=/abs/path/libonnxruntime_mlx_ep.dylib`.
131
+
132
+ ## Large-decoder partition metadata
133
+
134
+ The EP automatically infers residual layer boundaries from graph topology for decoders of about 24
135
+ layers or larger. Exporters can make the boundaries explicit with the
136
+ ONNX custom metadata key `onnxruntime_ep_mlx.layer_boundary_outputs`. Its value is a JSON array of
137
+ residual output tensor names, one per transformer layer:
138
+
139
+ ```python
140
+ import json
141
+ import onnx
142
+
143
+ model = onnx.load("decoder/model.onnx", load_external_data=False)
144
+ entry = model.metadata_props.add()
145
+ entry.key = "onnxruntime_ep_mlx.layer_boundary_outputs"
146
+ entry.value = json.dumps(["layer.0.output", "layer.1.output", "layer.2.output"])
147
+ onnx.save(model, "decoder/model.onnx")
148
+ ```
149
+
150
+ Metadata takes precedence over inference and does not rely on node names. The EP chooses a dynamic
151
+ group size of 4-8 layers, targeting about seven partitions. Override it with
152
+ `ONNXRUNTIME_EP_MLX_LAYER_PARTITIONS=<layers>`; use `0` or `off` to disable partitioning.
153
+
154
+ ## Performance (M1 Max, warm)
155
+
156
+ Real end-to-end models, median of 10 runs, MLX EP vs the ORT **CPU** EP on the same machine — top-1
157
+ identical and max abs diff ≤ 6e-5 in every case:
158
+
159
+ | Model | Workload | CPU EP | MLX EP | Speedup |
160
+ |---|---|---:|---:|---:|
161
+ | Perch v2 | audio encoder (with DFT front-end) | 64.0 ms | 12.0 ms | **5.3×** |
162
+ | Perch v2 (no DFT) | audio encoder | 56.5 ms | 12.0 ms | **4.7×** |
163
+ | BirdNET | audio classifier (CNN) | 14.9 ms | 7.3 ms | **2.0×** |
164
+ | gemma-4-E2B | vision encoder (fp16 ViT) | 267 ms | 47 ms | **5.7×** |
165
+
166
+ Feed-forward encoders (audio / CNN / vision) are the EP's sweet spot: the whole graph fuses into a
167
+ single MLX closure that is traced + `mlx_compile`d once and replayed, so a static-shape model runs
168
+ end-to-end on the GPU with one dispatch (e.g. Perch: 725/725 nodes claimed, 1 fused subgraph).
169
+
170
+ Eligible BF16 INT4/INT8 prefill matmuls (block size 32/64/128, at least 32 rows) explicitly use stock
171
+ MLX's FP16 compute path by default. Set `ONNXRUNTIME_EP_MLX_BF16_QMM_FP16=0` to disable it.
172
+
173
+ The **Foundry Local** q4f16 decoders below run on the same M1 Max, warm, MLX EP vs the ORT CPU EP
174
+ (decode = 1 token with 128 past; prefill = 128-token step):
175
+
176
+ | Model | Arch | Prefill | Decode |
177
+ |---|---|---:|---:|
178
+ | Qwen2.5-0.5B | GQA, external rotary | **5.2×** | dispatch-bound (CPU-favored) |
179
+ | Phi-3.5-mini | Phi3, GQA | **5.29×** | 1.19× |
180
+ | Phi-4-mini | Phi4, long-context RoPE | **5.78×** | 1.10× |
181
+ | Mistral-7B-Instruct | GQA, growing KV | **11.89×** | **3.30×** |
182
+ | gemma-4-E2B | Gemma3n, 15-layer | 3.3× | **3.3×** |
183
+
184
+ Muse-Glimmer-30B INT4, with the optimized bundled MLX runtime and automatic 8-layer groups, reaches
185
+ **138.67 prefill tok/s** at 512 tokens and **14.79 decode tok/s** over 200 generated tokens. The
186
+ same-quantization llama.cpp baseline reaches **137.84 / 13.50 tok/s**.
187
+
188
+ Decode is weight-bandwidth-bound: small models can favor CPU, while larger q4 decoders benefit from
189
+ MLX. Unclaimed ops fall back to ORT CPU.
190
+
191
+ ## Profiling & tracing (Perfetto)
192
+
193
+ The EP ships a built-in tracer (compiled in by default, **near-zero cost when off**). Recording is
194
+ gated entirely by environment variables — set one, run your model, and inspect the result.
195
+
196
+ **Get a Perfetto/Chrome trace.** Point `ONNXRUNTIME_EP_MLX_TRACE` at an output path; the JSON trace is
197
+ written when the inference session is torn down:
198
+
199
+ ```bash
200
+ ONNXRUNTIME_EP_MLX_TRACE=/tmp/mlx_trace.json python your_script.py
201
+ # then open https://ui.perfetto.dev (or chrome://tracing) and load /tmp/mlx_trace.json
202
+ ```
203
+
204
+ The timeline shows one span per fused subgraph (`mlx.subgraph`), a nested span around the synchronous
205
+ `mlx_eval` (`mlx.eval` — its CPU wall time is the GPU-inclusive time of the whole fused subgraph),
206
+ per-op build spans with shapes/dtype/bytes, and counter tracks for GPU memory / utilisation. Ops that
207
+ fell back to a slower *composed* path (despite a fused kernel existing) are coloured distinctly with a
208
+ `reason=…`, and a top-10 slowest-ops summary is emitted at teardown.
209
+
210
+ **Lighter options** (no JSON file):
211
+
212
+ | Env var | Effect |
213
+ |---|---|
214
+ | `ONNXRUNTIME_EP_MLX_VERBOSE=1` | Print the end-of-run session summary (claim rate, compute-path breakdown, time attribution) to stderr. |
215
+ | `ONNXRUNTIME_EP_MLX_CLAIM_DEBUG=1` | Print each unclaimed node + the actionable reason (why the graph fragmented). |
216
+ | `ONNXRUNTIME_EP_MLX_SIGNPOST=1` | Emit `os_signpost` intervals so an Instruments *Metal System Trace* correlates. |
217
+ | `ONNXRUNTIME_EP_MLX_NO_STABLE_CROSS_CACHE=1` | Disable per-generation MLX reuse of immutable MHA cross-attention K/V inputs for performance A/B. |
218
+
219
+ **Per-kernel GPU detail (Xcode).** MLX hides its Metal command buffers inside one fused `mlx_eval`, so
220
+ the JSON trace times the fused eval as a whole. To see *inside* it, capture a boundary eval to a
221
+ `.gputrace` bundle (full per-kernel timing / occupancy / bandwidth) and open it in Xcode:
222
+
223
+ ```bash
224
+ MTL_CAPTURE_ENABLED=1 \
225
+ ONNXRUNTIME_EP_MLX_GPU_CAPTURE=/tmp/mlx.gputrace \
226
+ ONNXRUNTIME_EP_MLX_GPU_CAPTURE_EVAL=5 \
227
+ python your_script.py
228
+ ```
229
+
230
+ `MTL_CAPTURE_ENABLED=1` must be set before process start. `…_GPU_CAPTURE_EVAL` picks which eval to
231
+ capture (0-based, default 0); for decode, eval 0 is prefill/warmup, so pick a steady-state token.
232
+
233
+ ## Concurrency
234
+
235
+ MLX evaluation is thread-affine. Use one `InferenceSession` per thread; do not call `Run()` on one
236
+ shared session from multiple threads.
237
+
238
+ ## Numerical accuracy
239
+
240
+ Outputs are tolerance-matched against ORT CPU but are not bit-identical. Long greedy generations can
241
+ diverge after near-tied logits because MLX and CPU use different floating-point reduction orders.
242
+
243
+ ## Layout
244
+
245
+ ```
246
+ docs/ design docs (DESIGN, OP_ARCHITECTURE, COMPILED_CAPTURE, MLX_EVALUATION)
247
+ rust/ the Rust EP: plugin-EP C-ABI vtables (factory/ep) + the modular ONNX->MLX
248
+ translator (engine, registry, ops/*.rs) over a mlx-c RAII layer (mlx.rs)
249
+ python/ pure-Python pip package (onnxruntime-ep-mlx): a locator that bundles + registers
250
+ the cargo-built dylib (hatchling build hook, hatch_build.py)
251
+ tests/ MLX op-correctness (tests/ops, pytest) + ONNX-standard conformance (tests/conformance)
252
+ .github/ CI (cargo build + op tests) and PyPI trusted-publishing workflows
253
+ ```
254
+
255
+ ## Testing
256
+
257
+ Build the EP (above), then run the pytest op-correctness suite (MLX vs ORT CPU reference):
258
+
259
+ ```sh
260
+ export ONNXRUNTIME_MLX_EP_LIB=$PWD/rust/target/release/libonnxruntime_mlx_ep.dylib
261
+ export DYLD_LIBRARY_PATH=<ort-prebuilt/lib>
262
+ python -m pytest tests/ops -q
263
+ ```
264
+
265
+ - `tests/ops` — each translated decoder op via MLX vs. ORT CPU reference (tolerance-gated, pytest)
266
+ - `tests/conformance` — opt-in fuzz-conformance of the MLX EP against the ONNX standard
267
+ (`cbourjau/onnx-tests`); see [`tests/conformance/README.md`](tests/conformance/README.md)
@@ -0,0 +1 @@
1
+ ../README.md
@@ -0,0 +1,199 @@
1
+ """Hatchling build hook for onnxruntime-ep-mlx.
2
+
3
+ Builds the Rust execution provider (`cargo build --release` in ../rust), then
4
+ bundles the resulting `libonnxruntime_mlx_ep.dylib` together with its mlx-c/mlx
5
+ runtime dependencies into the wheel package, relinked so they load from
6
+ ``@loader_path`` (a self-contained wheel). Finally it forces a platform wheel
7
+ tag: the package ships no CPython-ABI extension, so a single
8
+ ``py3-none-macosx_*_arm64`` wheel installs on 3.12, 3.13 and the free-threaded
9
+ builds alike.
10
+
11
+ The onnxruntime dependency is intentionally NOT bundled — it is resolved at
12
+ runtime from the host ``onnxruntime`` package (two-level namespace), matching the
13
+ EP's design.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import os
19
+ import shutil
20
+ import subprocess
21
+ import sys
22
+ import sysconfig
23
+ from pathlib import Path
24
+
25
+ from hatchling.builders.hooks.plugin.interface import BuildHookInterface
26
+
27
+ PLUGIN_DYLIB = "libonnxruntime_mlx_ep.dylib"
28
+
29
+
30
+ def _brew_prefix(pkg: str) -> Path:
31
+ out = subprocess.run(
32
+ ["brew", "--prefix", pkg],
33
+ check=True,
34
+ capture_output=True,
35
+ text=True,
36
+ )
37
+ return Path(out.stdout.strip())
38
+
39
+
40
+ def _dependency_prefix(env_name: str, brew_pkg: str) -> Path:
41
+ value = os.environ.get(env_name)
42
+ return Path(value) if value else _brew_prefix(brew_pkg)
43
+
44
+
45
+ def _linked_dependency(binary: Path, basename: str) -> str | None:
46
+ out = subprocess.run(
47
+ ["otool", "-L", str(binary)],
48
+ check=True,
49
+ capture_output=True,
50
+ text=True,
51
+ ).stdout
52
+ for line in out.splitlines()[1:]:
53
+ dependency = line.strip().split(" ", 1)[0]
54
+ if Path(dependency).name == basename:
55
+ return dependency
56
+ return None
57
+
58
+
59
+ def _run(cmd: list[str], **kw) -> None:
60
+ print("[onnxruntime-ep-mlx build] $", " ".join(cmd), flush=True)
61
+ subprocess.run(cmd, check=True, **kw)
62
+
63
+
64
+ def _resolve_ort_include() -> str:
65
+ """Mirror rust/build.rs: ORT_INCLUDE_DIR, else $ORT_HOME/include."""
66
+ inc = os.environ.get("ORT_INCLUDE_DIR")
67
+ if not inc:
68
+ home = os.environ.get("ORT_HOME")
69
+ if home:
70
+ inc = str(Path(home) / "include")
71
+ if inc and (Path(inc) / "onnxruntime_c_api.h").is_file():
72
+ return inc
73
+ raise RuntimeError(
74
+ "Could not locate the ONNX Runtime headers. Set ORT_INCLUDE_DIR to the "
75
+ "ORT C-API include dir, or ORT_HOME to an ONNX Runtime release root "
76
+ "(expects $ORT_HOME/include/onnxruntime_c_api.h)."
77
+ )
78
+
79
+
80
+ class CustomBuildHook(BuildHookInterface):
81
+ PLUGIN_NAME = "custom"
82
+
83
+ def initialize(self, version: str, build_data: dict) -> None:
84
+ if sys.platform != "darwin":
85
+ raise RuntimeError("onnxruntime-ep-mlx only builds on macOS (Apple Silicon).")
86
+
87
+ project_root = Path(self.root) # the python/ dir
88
+ # In a repository checkout the crate is ../rust. In an sdist it is
89
+ # included as ./rust so installers can build the wheel from source.
90
+ rust_dir = project_root / "rust"
91
+ if not (rust_dir / "Cargo.toml").is_file():
92
+ rust_dir = project_root.parent / "rust"
93
+ if not (rust_dir / "Cargo.toml").is_file():
94
+ raise RuntimeError(
95
+ "Rust crate not found. Expected ./rust in an sdist or ../rust "
96
+ "in a repository checkout."
97
+ )
98
+
99
+ pkg_dir = project_root / "src" / "onnxruntime_ep_mlx"
100
+
101
+ # 1) Build the Rust EP dylib.
102
+ env = dict(os.environ)
103
+ env["ORT_INCLUDE_DIR"] = _resolve_ort_include()
104
+ _run(["cargo", "build", "--release"], cwd=str(rust_dir), env=env)
105
+
106
+ built = rust_dir / "target" / "release" / PLUGIN_DYLIB
107
+ if not built.is_file():
108
+ raise RuntimeError(f"cargo did not produce {built}")
109
+
110
+ # 2) Copy the plugin into the package.
111
+ dest_plugin = pkg_dir / PLUGIN_DYLIB
112
+ shutil.copy2(built, dest_plugin)
113
+ os.chmod(dest_plugin, 0o755)
114
+
115
+ # 3) Bundle + relink the mlx runtime next to the plugin (self-contained).
116
+ self._bundle_mlx(pkg_dir, dest_plugin)
117
+
118
+ # 4) This is a platform wheel with no Python-ABI extension: force a
119
+ # py3-none-macosx_*_arm64 tag so ONE wheel serves every interpreter
120
+ # (CPython 3.10+, free-threaded 3.13t/3.14t, ...).
121
+ plat = sysconfig.get_platform().replace("-", "_").replace(".", "_")
122
+ # Honour MACOSX_DEPLOYMENT_TARGET for the platform floor: the bundled
123
+ # dylibs (and mlx) target it, so the tag should advertise it rather than
124
+ # whatever floor the running interpreter was built against.
125
+ dep_target = os.environ.get("MACOSX_DEPLOYMENT_TARGET")
126
+ if dep_target and plat.startswith("macosx_"):
127
+ arch = plat.rsplit("_", 1)[-1]
128
+ major, _, minor = dep_target.partition(".")
129
+ plat = f"macosx_{major}_{minor or '0'}_{arch}"
130
+ build_data["pure_python"] = False
131
+ build_data["infer_tag"] = False
132
+ build_data["tag"] = f"py3-none-{plat}"
133
+
134
+ # -- mlx bundling ---------------------------------------------------------
135
+ def _bundle_mlx(self, pkg_dir: Path, plugin: Path) -> None:
136
+ mlxc_pfx = _dependency_prefix("MLXC_PREFIX", "mlx-c")
137
+ mlx_pfx = _dependency_prefix("MLX_PREFIX", "mlx")
138
+ mlxc_src = mlxc_pfx / "lib" / "libmlxc.dylib"
139
+ mlx_src = mlx_pfx / "lib" / "libmlx.dylib"
140
+ metallib_src = mlx_pfx / "lib" / "mlx.metallib"
141
+ for f in (mlxc_src, mlx_src, metallib_src):
142
+ if not f.is_file():
143
+ raise RuntimeError(f"Required mlx artifact missing: {f} (brew install mlx-c)")
144
+
145
+ mlxc_dst = pkg_dir / "libmlxc.dylib"
146
+ mlx_dst = pkg_dir / "libmlx.dylib"
147
+ for src, dst in ((mlxc_src, mlxc_dst), (mlx_src, mlx_dst), (metallib_src, pkg_dir / "mlx.metallib")):
148
+ shutil.copy2(src, dst)
149
+ os.chmod(dst, 0o644)
150
+ os.chmod(mlxc_dst, 0o755)
151
+ os.chmod(mlx_dst, 0o755)
152
+
153
+ jaccl_src = mlx_pfx / "lib" / "libjaccl.dylib"
154
+ jaccl_dst = pkg_dir / "libjaccl.dylib"
155
+ if jaccl_src.is_file():
156
+ shutil.copy2(jaccl_src, jaccl_dst)
157
+ os.chmod(jaccl_dst, 0o755)
158
+
159
+ def name_tool(*args: str) -> None:
160
+ _run(["install_name_tool", *args])
161
+
162
+ def resign(f: Path) -> None:
163
+ subprocess.run(["codesign", "--force", "--sign", "-", str(f)], check=False)
164
+
165
+ # Bundled mlx install ids -> @loader_path.
166
+ name_tool("-id", "@loader_path/libmlxc.dylib", str(mlxc_dst))
167
+ name_tool("-id", "@loader_path/libmlx.dylib", str(mlx_dst))
168
+ # Relink MLX dependencies by basename so both Homebrew absolute paths and
169
+ # custom-build @rpath install names work.
170
+ mlxc_mlx = _linked_dependency(mlxc_dst, "libmlx.dylib")
171
+ if mlxc_mlx:
172
+ name_tool("-change", mlxc_mlx, "@loader_path/libmlx.dylib", str(mlxc_dst))
173
+
174
+ # Plugin's mlx deps -> colocated copies.
175
+ plugin_mlxc = _linked_dependency(plugin, "libmlxc.dylib")
176
+ plugin_mlx = _linked_dependency(plugin, "libmlx.dylib")
177
+ if plugin_mlxc:
178
+ name_tool("-change", plugin_mlxc, "@loader_path/libmlxc.dylib", str(plugin))
179
+ if plugin_mlx:
180
+ name_tool("-change", plugin_mlx, "@loader_path/libmlx.dylib", str(plugin))
181
+
182
+ if jaccl_src.is_file():
183
+ name_tool("-id", "@loader_path/libjaccl.dylib", str(jaccl_dst))
184
+ mlx_jaccl = _linked_dependency(mlx_dst, "libjaccl.dylib")
185
+ if mlx_jaccl:
186
+ name_tool("-change", mlx_jaccl, "@loader_path/libjaccl.dylib", str(mlx_dst))
187
+
188
+ # The Rust EP does NOT link libonnxruntime — it reaches ORT purely through
189
+ # the OrtApi function-pointer table handed to CreateEpFactories (see
190
+ # rust/build.rs). So there is no onnxruntime dependency to relink here;
191
+ # ORT dlopen()s the plugin by the absolute path library_path() returns.
192
+
193
+ # Re-sign everything we mutated (install_name_tool voids the ad-hoc sig;
194
+ # dyld SIGKILLs unsigned/invalid arm64 images).
195
+ bundled = [mlxc_dst, mlx_dst, plugin]
196
+ if jaccl_src.is_file():
197
+ bundled.append(jaccl_dst)
198
+ for f in bundled:
199
+ resign(f)
@@ -0,0 +1,79 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.24", "delocate>=0.11"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "onnxruntime-ep-mlx"
7
+ version = "0.27.6"
8
+ description = "MLX-native ONNX Runtime execution provider (plugin EP) for Apple Silicon"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "onnxruntime-ep-mlx" }]
13
+ keywords = ["onnxruntime", "mlx", "execution-provider", "apple-silicon", "metal", "coreml"]
14
+ dependencies = ["onnxruntime>=1.22"]
15
+ classifiers = [
16
+ "Development Status :: 3 - Alpha",
17
+ "Intended Audience :: Developers",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Operating System :: MacOS :: MacOS X",
20
+ "Programming Language :: Rust",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Programming Language :: Python :: 3.13",
26
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
27
+ ]
28
+
29
+ [project.urls]
30
+ Homepage = "https://github.com/justinchuby/onnxruntime-mlx"
31
+ Repository = "https://github.com/justinchuby/onnxruntime-mlx"
32
+
33
+ # --- Hatchling wheel build ----------------------------------------------------
34
+ # The execution provider is a Rust `cdylib` (../rust). It is NOT a Python
35
+ # extension module, so there is no compiled Python surface to build: the Python
36
+ # layer (onnxruntime_ep_mlx/__init__.py) is a *pure-Python locator* that only finds
37
+ # and registers the bundled dylib. The custom build hook (hatch_build.py):
38
+ # 1. runs `cargo build --release` in ../rust (honouring ORT_INCLUDE_DIR / ORT_HOME),
39
+ # 2. bundles the cargo-built libonnxruntime_mlx_ep.dylib + the mlx-c/mlx runtime
40
+ # (via delocate) into the package, relinked to @loader_path,
41
+ # 3. forces a platform wheel tag (py3-none-macosx_*_arm64) — one wheel that
42
+ # installs on CPython 3.10-3.13 AND the free-threaded (3.13t/3.14t) builds,
43
+ # because it ships no CPython-ABI-bound extension. abi3audit has nothing to
44
+ # audit (no extension modules) so `abi3audit --strict` is trivially clean.
45
+ #
46
+ # Why pure-Python instead of maturin+pyo3: `_core` was only ever a locator
47
+ # (ep_name/version/library_path). A compiled pyo3 abi3 extension cannot also be
48
+ # free-threaded (abi3 wheels are unsupported under Py_GIL_DISABLED), so "abi3 AND
49
+ # abi3t" would force *separate* per-version cp313t/cp314t wheels. Shipping zero
50
+ # native Python code sidesteps that entirely: one platform wheel serves every
51
+ # interpreter (GIL and free-threaded alike) and is abi3audit-clean by construction.
52
+
53
+ [tool.hatch.build.targets.wheel]
54
+ packages = ["src/onnxruntime_ep_mlx"]
55
+ # The dylibs / metallib are placed into the package dir by the build hook.
56
+ artifacts = [
57
+ "src/onnxruntime_ep_mlx/*.dylib",
58
+ "src/onnxruntime_ep_mlx/*.metallib",
59
+ ]
60
+
61
+ [tool.hatch.build.targets.wheel.hooks.custom]
62
+ path = "hatch_build.py"
63
+
64
+ [tool.hatch.build.targets.sdist]
65
+ # Carry the Rust crate too, so `pip wheel` can build the macOS wheel directly
66
+ # from the PyPI sdist rather than requiring a full repository checkout.
67
+ include = [
68
+ "src/onnxruntime_ep_mlx",
69
+ "hatch_build.py",
70
+ "README.md",
71
+ ]
72
+
73
+ [tool.hatch.build.targets.sdist.force-include]
74
+ "../rust/Cargo.toml" = "rust/Cargo.toml"
75
+ "../rust/build.rs" = "rust/build.rs"
76
+ "../rust/README.md" = "rust/README.md"
77
+ "../rust/src" = "rust/src"
78
+ "../rust/wrapper_mlx.h" = "rust/wrapper_mlx.h"
79
+ "../rust/wrapper_ort.h" = "rust/wrapper_ort.h"
@@ -0,0 +1,56 @@
1
+ [package]
2
+ name = "onnxruntime-ep-mlx"
3
+ version = "0.27.6"
4
+ edition = "2024"
5
+ description = "MLX-native ONNX Runtime execution provider (plugin EP) for Apple Silicon — binds mlx-c directly, no mlx-rs."
6
+ license = "MIT"
7
+ repository = "https://github.com/justinchuby/onnxruntime-mlx"
8
+ readme = "README.md"
9
+ keywords = ["onnx", "onnxruntime", "mlx", "apple-silicon", "metal"]
10
+ categories = ["science", "api-bindings", "hardware-support"]
11
+ # Keep repo-only tooling out of the published crate: the *.py files are leak-check /
12
+ # tracing scripts run by CI from the repo (not part of the Rust crate).
13
+ exclude = ["*.py"]
14
+
15
+ [lib]
16
+ # Pin the lib (and therefore dylib) name so it stays `libonnxruntime_mlx_ep.dylib`
17
+ # regardless of the crate name — the Python package, tests, CI and onnx-genai all
18
+ # load the plugin by that exact filename.
19
+ name = "onnxruntime_mlx_ep"
20
+ crate-type = ["cdylib"]
21
+
22
+ [dependencies]
23
+ # Structured logging facade — zero-cost when the log level is disabled at
24
+ # compile time (via `max_level_*` features) or when no logger is installed.
25
+ # This is the FACADE only; the in-crate `logging` module provides the minimal
26
+ # subscriber (a 20-line stderr logger gated by env vars). Using `log` rather
27
+ # than `tracing` because this is a cdylib plugin loaded into someone else's
28
+ # process: the `log` facade is lighter, has no global subscriber conflict risk,
29
+ # and the plugin's statically-linked copy of `log` is fully independent of
30
+ # whatever the host process uses.
31
+ log = { version = "0.4", features = ["std"] }
32
+ half = "2"
33
+ serde_json = "1"
34
+ image = { version = "0.25", default-features = false, features = ["bmp", "jpeg", "png", "pnm", "tiff", "webp"] }
35
+ jpeg2k = { version = "0.10", default-features = true }
36
+
37
+ # Pure-Rust Chrome/Perfetto tracer, from crates.io. `default-features = false`
38
+ # drops the `perfetto` (prost) feature — we export Chrome Trace JSON only.
39
+ #
40
+ # Pinned to a release with the absolute UNIX-microsecond clock. Every copy of
41
+ # the crate -- the host's and this dylib's, which has its own statics and so
42
+ # could never have shared a process-global epoch -- reports the same instant as
43
+ # the same number, which is what lets a trace from this plugin be laid over a
44
+ # host trace with no offset to negotiate.
45
+ onnx-runtime-tracer = { version = "0.1.0-dev.5", default-features = false }
46
+
47
+ [build-dependencies]
48
+ bindgen = "0.72"
49
+
50
+ [profile.release]
51
+ # Unwind (not abort) so panics can be caught at every C-ABI boundary (see
52
+ # `guard_ffi_status` in lib.rs and the trace_thunk catch_unwind): a plugin loaded into a
53
+ # host process must degrade to an error, never take the host down. Every extern-"C" entry
54
+ # that runs real logic is wrapped, and the MLX trace callback catches locally, so a panic
55
+ # never unwinds across the FFI boundary (which would be UB).
56
+ panic = "unwind"