Skip to content

Commit 38c5ca3

Browse files
metascroymetascroy
authored
MLX delegate (part 1) (#17803)
Takes changes from #16718, but strips all ops (except addmm) and examples. Part2 will add back ops, and part 3 will add back examples.
1 parent 19f7ff2 commit 38c5ca3

56 files changed

Lines changed: 12676 additions & 1 deletion

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/mlx.yml‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
name: MLX
2+
3+
on:
4+
push:
5+
branches:
6+
- main
7+
- release/*
8+
pull_request:
9+
paths:
10+
- .github/workflows/mlx.yml
11+
- backends/mlx/**
12+
workflow_dispatch:
13+
14+
permissions: {}
15+
16+
jobs:
17+
test-mlx:
18+
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
19+
with:
20+
job-name: test-mlx
21+
runner: macos-14-xlarge
22+
python-version: "3.12"
23+
submodules: recursive
24+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
25+
timeout: 90
26+
script: |
27+
set -eux
28+
29+
echo "::group::Install ExecuTorch and configure build"
30+
${CONDA_RUN} python install_executorch.py > /dev/null
31+
# The sanitizers fail on github VM runner, but pass on real device
32+
# TODO: figure out why
33+
${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
34+
echo "::endgroup::"
35+
36+
${CONDA_RUN} pip list
37+
38+
echo "::group::Build test runners"
39+
${CONDA_RUN} cmake --build cmake-out --target op_test_runner -j$(( $(sysctl -n hw.ncpu) - 1 ))
40+
echo "::endgroup::"
41+
42+
echo "::group::Run op unit tests"
43+
${CONDA_RUN} python -m executorch.backends.mlx.test.run_all_tests -j4 --max-tasks-per-worker 10 --clean-after
44+
echo "::endgroup::"
45+
46+
echo "::group::Run Python unit tests"
47+
${CONDA_RUN} python -m pytest \
48+
backends/mlx/test/test_passes.py \
49+
backends/mlx/test/test_pattern_utils.py \
50+
backends/mlx/test/test_partitioner.py \
51+
-v
52+
echo "::endgroup::"
53+
54+
backend-tester:
55+
strategy:
56+
fail-fast: false
57+
matrix:
58+
suite: [models, operators]
59+
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
60+
with:
61+
job-name: test-mlx-backend-${{ matrix.suite }}
62+
runner: macos-14-xlarge
63+
python-version: "3.12"
64+
submodules: recursive
65+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
66+
timeout: 120
67+
script: |
68+
set -eux
69+
70+
echo "::group::Install ExecuTorch"
71+
${CONDA_RUN} python install_executorch.py > /dev/null
72+
echo "::endgroup::"
73+
74+
${CONDA_RUN} pip list
75+
76+
echo "::group::Run backend test suite (${{ matrix.suite }})"
77+
${CONDA_RUN} pytest -c /dev/null backends/test/suite/${{ matrix.suite }}/ -m flow_mlx -n auto 2>&1 | tee pytest_output.txt || true
78+
echo "::endgroup::"
79+
80+
# Parse pytest summary and check failure threshold
81+
if grep -E "^=+ .* =+$" pytest_output.txt | tail -1 | grep -q "failed"; then
82+
FAILED=$(grep -E "^=+ .* =+$" pytest_output.txt | tail -1 | grep -oE "[0-9]+ failed" | grep -oE "[0-9]+")
83+
else
84+
FAILED=0
85+
fi
86+
87+
if [ "${{ matrix.suite }}" = "operators" ]; then
88+
MAX_FAILURES=0
89+
else
90+
MAX_FAILURES=3
91+
fi
92+
93+
echo "Failed tests: $FAILED (max allowed: $MAX_FAILURES)"
94+
if [ "$FAILED" -gt "$MAX_FAILURES" ]; then
95+
echo "::error::Too many test failures: $FAILED > $MAX_FAILURES"
96+
exit 1
97+
fi

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,5 +74,7 @@ xcuserdata/
7474
*.dll
7575
*.pyd
7676

77+
7778
# Agents
7879
.claude/*.local.*
80+
extension/pybindings/mlx.metallib

‎.gitmodules‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,7 @@
6767
[submodule "third-party/json"]
6868
path = third-party/json
6969
url = https://github.com/nlohmann/json.git
70+
[submodule "backends/mlx/third-party/mlx"]
71+
path = backends/mlx/third-party/mlx
72+
url = https://github.com/ml-explore/mlx.git
73+
shallow = true

‎CMakeLists.txt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -682,6 +682,11 @@ if(EXECUTORCH_BUILD_MPS)
682682
list(APPEND _executorch_backends mpsdelegate)
683683
endif()
684684

685+
if(EXECUTORCH_BUILD_MLX)
686+
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/mlx)
687+
list(APPEND _executorch_backends mlxdelegate)
688+
endif()
689+
685690
if(EXECUTORCH_BUILD_NEURON)
686691
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/mediatek)
687692
list(APPEND _executorch_backends neuron_backend)
@@ -982,6 +987,10 @@ if(EXECUTORCH_BUILD_PYBIND)
982987
list(APPEND _dep_libs mpsdelegate)
983988
endif()
984989

990+
if(EXECUTORCH_BUILD_MLX)
991+
list(APPEND _dep_libs mlxdelegate)
992+
endif()
993+
985994
if(EXECUTORCH_BUILD_OPENVINO)
986995
list(APPEND _dep_libs openvino_backend)
987996
endif()
@@ -1082,6 +1091,12 @@ if(EXECUTORCH_BUILD_PYBIND)
10821091
install(TARGETS data_loader
10831092
LIBRARY DESTINATION executorch/extension/pybindings
10841093
)
1094+
1095+
# Copy MLX metallib next to _portable_lib.so for editable installs. MLX uses
1096+
# dladdr() to find the directory containing the library with MLX code, then
1097+
# looks for mlx.metallib in that directory. When MLX is statically linked into
1098+
# _portable_lib.so, we need the metallib colocated with it.
1099+
executorch_target_copy_mlx_metallib(portable_lib)
10851100
endif()
10861101

10871102
if(EXECUTORCH_BUILD_WASM)

‎CMakePresets.json‎

Lines changed: 84 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@
110110
"inherits": ["common"],
111111
"cacheVariables": {
112112
"EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/pybind.cmake",
113-
"CMAKE_OSX_DEPLOYMENT_TARGET": "12.0"
113+
"CMAKE_OSX_DEPLOYMENT_TARGET": "14.0"
114114
},
115115
"condition": {
116116
"type": "inList",
@@ -310,6 +310,43 @@
310310
"EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/arm_ethosu_linux.cmake",
311311
"CMAKE_TOOLCHAIN_FILE": "${sourceDir}/examples/arm/ethos-u-setup/aarch64-linux-musl-toolchain.cmake"
312312
}
313+
},
314+
{
315+
"name": "mlx",
316+
"displayName": "Build MLX delegate",
317+
"inherits": ["common"],
318+
"cacheVariables": {
319+
"EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/mlx.cmake",
320+
"EXECUTORCH_ENABLE_LOGGING": "ON",
321+
"CMAKE_OSX_DEPLOYMENT_TARGET": "14.0"
322+
},
323+
"condition": {
324+
"lhs": "${hostSystemName}",
325+
"type": "equals",
326+
"rhs": "Darwin"
327+
}
328+
},
329+
{
330+
"name": "mlx-release",
331+
"displayName": "MLX delegate release build",
332+
"inherits": ["mlx"],
333+
"cacheVariables": {
334+
"CMAKE_BUILD_TYPE": "Release",
335+
"CMAKE_INSTALL_PREFIX": "${sourceDir}/cmake-out",
336+
"ET_MLX_ENABLE_OP_LOGGING": "OFF",
337+
"ET_MIN_LOG_LEVEL": "Error"
338+
}
339+
},
340+
{
341+
"name": "mlx-debug",
342+
"displayName": "MLX delegate debug build with op logging",
343+
"inherits": ["mlx"],
344+
"cacheVariables": {
345+
"CMAKE_BUILD_TYPE": "Debug",
346+
"CMAKE_INSTALL_PREFIX": "${sourceDir}/cmake-out",
347+
"ET_MLX_ENABLE_OP_LOGGING": "ON",
348+
"ET_MIN_LOG_LEVEL": "Debug"
349+
}
313350
}
314351
],
315352
"buildPresets": [
@@ -387,6 +424,24 @@
387424
"install"
388425
],
389426
"jobs": 0
427+
},
428+
{
429+
"name": "mlx-release-install",
430+
"displayName": "Build and install MLX delegate release artifacts",
431+
"configurePreset": "mlx-release",
432+
"targets": [
433+
"install"
434+
],
435+
"jobs": 0
436+
},
437+
{
438+
"name": "mlx-debug-install",
439+
"displayName": "Build and install MLX delegate debug artifacts",
440+
"configurePreset": "mlx-debug",
441+
"targets": [
442+
"install"
443+
],
444+
"jobs": 0
390445
}
391446
],
392447
"workflowPresets": [
@@ -501,6 +556,34 @@
501556
"name": "llm-metal-stats-install"
502557
}
503558
]
559+
},
560+
{
561+
"name": "mlx-release",
562+
"displayName": "Configure, build and install ExecuTorch MLX delegate",
563+
"steps": [
564+
{
565+
"type": "configure",
566+
"name": "mlx-release"
567+
},
568+
{
569+
"type": "build",
570+
"name": "mlx-release-install"
571+
}
572+
]
573+
},
574+
{
575+
"name": "mlx-debug",
576+
"displayName": "Configure, build and install ExecuTorch MLX delegate with op logging (Debug)",
577+
"steps": [
578+
{
579+
"type": "configure",
580+
"name": "mlx-debug"
581+
},
582+
{
583+
"type": "build",
584+
"name": "mlx-debug-install"
585+
}
586+
]
504587
}
505588
]
506589
}

0 commit comments

Comments
 (0)