Ship the MLX delegate as a linkable library in the wheel - #22184
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22184
Note: Links to docs will display an error until the docs builds have been completed. ⏳ 102 Pending, 1 Unrelated FailureAs of commit 0411ced with merge base 04b3446 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
c3583de to
3e0cda3
Compare
3e0cda3 to
584e306
Compare
584e306 to
38437da
Compare
|
Updated. The blocker and all four majors are fixed; the minor was a wrong comment The blocker is real and I confirmed it against a real wheel. The release check Added the name. The check now passes on that wheel, and still rejects No ownership rows. Added two, modelled on the XNNPACK pair, one for the Marked not-required, like the optimized-kernels row, because the delegate only Metallib shipped nowhere with MLX on and shared off. Correct, and the I did not use a negated flag for this. Docs said the component cannot exist. Fixed all three surfaces: the paragraph No metallib path for a wheel consumer. The source config published The On R2, the split cache registry: agreed it is a latent trap for a source install |
38437da to
0dc1ea8
Compare
|
C4 is now fixed too, and thank you for catching that it was only half done. I had corrected the commit message but left the wrong reason sitting in the code That is the mechanism your measurement established: |
MLX is reachable from Python but not from C++. The macOS wheel carries no MLX
library at all: the delegate is statically fused into the Python extension, so a
C++ application that links the wheel cannot use it.
`find_package(executorch)` offers no MLX component, and running an MLX program
fails at load:
Backend MLXBackend is not registered.
Build the delegate as a shared library in the wheel, the same way the XNNPACK
delegate already is, and package it beside the runtime. A C++ consumer can then
link `executorch::backend_mlx`.
The static MLX archive is bundled into the shared library rather than left for
the consumer to supply, since it carries the delegate's kernels. The archive is
kept out of the target's interface deliberately, so an imported target that only
exists while building cannot reach the installed export set.
The Metal kernel archive moves from beside the Python extension to beside the
shared library. MLX finds that archive by resolving the address of its own code
with `dladdr()` and looking in the resulting directory, so it has to sit next to
whichever image carries MLX. Building the delegate shared moves that code out of
the Python extension entirely, so both the Python path and a C++ consumer now
resolve to the library directory and one copy serves both. The extension shrinks
from 5.7 MB to 0.9 MB as a result.
The copy placed beside the Python extension during a build is kept, because
without the shared option the delegate is static and is absorbed into the
extension, which is then the image that carries MLX.
Nothing changes for a build that is not producing a wheel: without the shared
option the delegate stays a static library linking the static runtime.
Test Plan:
Built the wheel on macOS arm64 with MLX enabled, installed it into a new virtual
environment holding only torch and the wheel, and exercised both paths. The
build-tree copy of the Metal archive was hidden for every run, because MLX falls
back to a compile-time path that still exists on a build machine and would
otherwise mask a failure.
`find_package(executorch)` now reports `executorch::backend_mlx` alongside the
other components, where the list previously ended at `backend_xnnpack`. A C++
program linking it loads an MLX program and produces the expected output, where
the same program previously failed with the error above. The same program on an
XNNPACK model still passes, so the harness is sound.
The Python path still exports, runs and matches eager exactly.
Which directory the Metal archive has to live in was measured rather than
reasoned about, by removing each copy in turn. Only the copy beside the shared
library works; the one beside the Python extension is not reachable from either
path once the delegate is shared. That is why a single copy ships.
Also add the delegate to the wheel's own release checks, which assert both that
every shipped library has an expected name and that each component has exactly
one owning library. The second one matters here because this library bundles a
third-party runtime, which is the case those checks exist to police.
The Metal kernel archive ships beside whichever image carries MLX code, which is the
delegate when it is built shared and the Python extension when it is not. The build
records which case applies, so exactly one copy ships either way.
Document the new component: the C++ page previously said a C++ application on
macOS cannot link MLX at all, which is what this change makes false, and its
component table and probe loop omitted it. The wheel's package config now also
publishes the archive's path, the way the source package already does, so an
application that relocates the delegate can bring the archive with it.
Ran the wheel's own two release checks against that install. Both pass, and the
new ownership rows are doing work rather than skipping: each reports a single
owner for the delegate and for the bundled runtime. Confirmed they can fail by
planting a second copy of the library, which makes them report two definers.
Confirmed a consumer is handed the Metal archive's path by configuring a project
against the installed package, and that the archive is selected for every build
that produces the delegate, shared or static.
Corrected after review, all three verified rather than reasoned about:
The archive was shipping twice in a shared build, because the entry covering the
static case was gated on MLX alone. The build now records which image carries MLX and
exactly one copy ships in every configuration.
The package config located the archive with find_file, which is the pattern the file's
own comment explains it avoids: it applies the consumer's find-root rules, so a
cross-compiling toolchain reroots an absolute path into its sysroot and reports a
complete package as missing. It also cached the result, which goes stale across two
installs in one build tree. Now a plain existence test and an uncached set, which is
what the rest of the file does.
Documented the new component and the published path in the config's own header block,
and gave the delegate the same runtime-file caveat the OpenVINO one already has.
0dc1ea8 to
49b34ef
Compare
|
@claude review this code |
|
Claude finished @metascroy's task in 2m 13s —— View job Code Review: Ship the MLX delegate as a linkable library in the wheelI read the full diff (6 files) against the XNNPACK precedent and the CMake helpers. This is a careful, well-reasoned change and the commit message / discussion already anticipate most of the sharp edges. Findings below, none blocking.
What's solid
Points worth a second look
Verification I could not performI can't build the macOS/Apple-Silicon wheel in this environment, so I'm trusting the author's reported runs ( Overall: LGTM. The one thing I'd add is a short comment at |
MLX is reachable from Python but not from C++. The macOS wheel carries no MLX
library at all: the delegate is statically fused into the Python extension, so a
C++ application that links the wheel cannot use it.
find_package(executorch)offers no MLX component, and running an MLX programfails at load:
Build the delegate as a shared library in the wheel, the same way the XNNPACK
delegate already is, and package it beside the runtime. A C++ consumer can then
link
executorch::backend_mlx.The static MLX archive is bundled into the shared library rather than left for
the consumer to supply, since it carries the delegate's kernels. The archive is
kept out of the target's interface deliberately, so an imported target that only
exists while building cannot reach the installed export set.
The Metal kernel archive moves from beside the Python extension to beside the
shared library. MLX finds that archive by resolving the address of its own code
with
dladdr()and looking in the resulting directory, so it has to sit next towhichever image carries MLX. Building the delegate shared moves that code out of
the Python extension entirely, so both the Python path and a C++ consumer now
resolve to the library directory and one copy serves both. The extension shrinks
from 5.7 MB to 0.9 MB as a result.
The copy placed beside the Python extension during a build is kept, because
without the shared option the delegate is static and is absorbed into the
extension, which is then the image that carries MLX.
Nothing changes for a build that is not producing a wheel: without the shared
option the delegate stays a static library linking the static runtime.
Test Plan:
Built the wheel on macOS arm64 with MLX enabled, installed it into a new virtual
environment holding only torch and the wheel, and exercised both paths. The
build-tree copy of the Metal archive was hidden for every run, because MLX falls
back to a compile-time path that still exists on a build machine and would
otherwise mask a failure.
find_package(executorch)now reportsexecutorch::backend_mlxalongside theother components, where the list previously ended at
backend_xnnpack. A C++program linking it loads an MLX program and produces the expected output, where
the same program previously failed with the error above. The same program on an
XNNPACK model still passes, so the harness is sound.
The Python path still exports, runs and matches eager exactly.
Which directory the Metal archive has to live in was measured rather than
reasoned about, by removing each copy in turn. Only the copy beside the shared
library works; the one beside the Python extension is not reachable from either
path once the delegate is shared. That is why a single copy ships.
Also add the delegate to the wheel's own release checks, which assert both that
every shipped library has an expected name and that each component has exactly
one owning library. The second one matters here because this library bundles a
third-party runtime, which is the case those checks exist to police.
The Metal kernel archive ships beside whichever image carries MLX code, which is the
delegate when it is built shared and the Python extension when it is not. The build
records which case applies, so exactly one copy ships either way.
Document the new component: the C++ page previously said a C++ application on
macOS cannot link MLX at all, which is what this change makes false, and its
component table and probe loop omitted it. The wheel's package config now also
publishes the archive's path, the way the source package already does, so an
application that relocates the delegate can bring the archive with it.
Ran the wheel's own two release checks against that install. Both pass, and the
new ownership rows are doing work rather than skipping: each reports a single
owner for the delegate and for the bundled runtime. Confirmed they can fail by
planting a second copy of the library, which makes them report two definers.
Confirmed a consumer is handed the Metal archive's path by configuring a project
against the installed package, and that the archive is selected for every build
that produces the delegate, shared or static.
Corrected after review, all three verified rather than reasoned about:
The archive was shipping twice in a shared build, because the entry covering the
static case was gated on MLX alone. The build now records which image carries MLX and
exactly one copy ships in every configuration.
The package config located the archive with find_file, which is the pattern the file's
own comment explains it avoids: it applies the consumer's find-root rules, so a
cross-compiling toolchain reroots an absolute path into its sysroot and reports a
complete package as missing. It also cached the result, which goes stale across two
installs in one build tree. Now a plain existence test and an uncached set, which is
what the rest of the file does.
Documented the new component and the published path in the config's own header block,
and gave the delegate the same runtime-file caveat the OpenVINO one already has.