Skip to content

MetalCompilerPlugin support #313

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gin66 wants to merge 3 commits into
base: main
Choose a base branch
from
+66 −52
Open

MetalCompilerPlugin support #313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e4d2bc2
use MetalCompilerPlugin to create default.metallib
gin66 Dec 9, 2025
047706e
fix the warning about missing metal-cpp.patch
gin66 Dec 9, 2025
b499477
add info for default.metallib handling and ensure swift test passes w…
gin66 Dec 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions Package.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
// swift-tools-version: 5.12
// swift-tools-version: 6.2
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am concerned about this -- it locks out developers who use older Xcode. Why do we need this?

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to check, if this 6.2 is really necessary.

// The swift-tools-version declares the minimum version of Swift required to build this package.
// Copyright © 2024 Apple Inc.

import Foundation
import PackageDescription

let inXcode = ProcessInfo.processInfo.environment["XCODE_VERSION_ACTUAL"] != nil
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this won't work correctly: Xcode can't build the shaders with the correct arguments (in particular the NAX shaders). That won't be a problem once those move to JIT builds, but I think it isn't a good idea to have to different build paths where bugs might show up in one and not in the other (I realize we have swiftPM, Xcode and Cmake so this is already an issue, but the swiftpm and Xcode paths do use the same build process).

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For Xcode builds, the MetalShaderCompiler is not needed. And for spm builds it is required.
Not a major difference, because the underlying metal compiler is the same. Not sure, what you mean with inability to build with correct arguments.

Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The plugin allows you to specify arguments when building the shaders. xcodebuild does not -- the metal compile uses some default arguments, which have been fine up to the 0.30.0 mlx release but are no longer sufficient. We either need to move the shaders that require different arguments into JIT compiles (the current plan) or we need a solution like this that can vary them from file to file (for both swiftpm AND Xcode builds).


let package = Package(
name: "mlx-swift",

Expand All @@ -26,7 +29,8 @@ let package = Package(
],
dependencies: [
// for Complex type
.package(url: "https://github.com/apple/swift-numerics", from: "1.0.0")
.package(url: "https://github.com/apple/swift-numerics", from: "1.0.0"),
.package(url: "https://github.com/schwa/MetalCompilerPlugin", branch: "main"),
],
targets: [
.target(
Expand Down Expand Up @@ -155,7 +159,14 @@ let package = Package(
.linkedFramework("Foundation"),
.linkedFramework("Metal"),
.linkedFramework("Accelerate"),
]
],

plugins:
// Optional: Use plugin for custom Metal compilation
// needed for swift build. Xcode does it automatically
inXcode ? [] : [
.plugin(name: "MetalCompilerPlugin", package: "MetalCompilerPlugin")
],
),
.testTarget(
name: "CmlxTests",
Expand Down
53 changes: 4 additions & 49 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,6 @@ MLX is an array framework for machine learning on Apple silicon. MLX Swift
expands MLX to the Swift language, making research and experimentation easier
on Apple silicon.

## Language Models

LLM and VLM implementations are available in [mlx-swift-lm](https://github.com/ml-explore/mlx-swift-lm).
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are these removed from the README?

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

was not intentional


## Examples

MLX Swift has [many
Expand Down Expand Up @@ -42,7 +38,7 @@ from MLX Python.

## Installation

The ``MLX`` Swift package can be built and run from Xcode or SwiftPM. A CMake installation is also provided, featuring a native Linux build option.
The ``MLX`` Swift package can be built and run from Xcode or SwiftPM. A CMake install is also provided.

More details are in the [documentation](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/install).

Expand Down Expand Up @@ -72,6 +68,8 @@ dependencies: [.product(name: "MLX", package: "mlx-swift"),
> [!Note]
> SwiftPM (command line) cannot build the Metal shaders so the ultimate build has to be done
> via Xcode.
>
>Update: Using [Metal Compiler Plugin](https://github.com/schwa/MetalCompilerPlugin), the library will be compiled and stored as default.metallib inside the Clmx-bundle. It is not in the Resources of the main bundle! With [patch](https://github.com/ml-explore/mlx/pull/2885) mlx is able to load default.metallib from any bundle root. What is left to the application/bundle to ensure, the bundle is loaded. Check the `setUp()` function in ArrayAtTests.swift as an example for manual loading.

### xcodebuild

Expand Down Expand Up @@ -99,57 +97,14 @@ brew install ninja
With CMake:

```shell
mkdir -p build
mkdir build
cd build
cmake .. -G Ninja
ninja
./example1
./tutorial
```

<details>
<summary>Expand Native Linux Build Instructions</summary>

#### (1) Install Dependencies

RHEL/Fedora:
```shell
sudo dnf install -y blas-devel lapack-devel openblas-devel clang llvm cmake make ninja
# Then install Swift by following the instructions at https://swift.org
```

Ubuntu/Debian:
```shell
sudo apt update;
sudo apt install -y libblas-dev liblapack-dev libopenblas-dev clang llvm cmake make ninja;
# Then install Swift by following the instructions at https://swift.org
```

Refer to [swift.org](https://www.swift.org/install/linux/) for installation options and instructions specific to your Linux distribution.


#### (2) Build + Run Examples

On Linux, the examples use the CPU backend by default.

Note: GPU+CUDA support is a work in progress for `mlx-swift` on Linux, but is available in the Python-based MLX.

Note: SwiftPM builds are not currently supported for native Linux targets.

```shell
mkdir -p build
cd build
cmake -DMLX_BUILD_METAL=OFF .. -G Ninja
ninja
./example1
./tutorial
```


</details>

</br>
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this removed?

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same. I think, I have copied over the Readme instead of redoing the modification. Should be an easy fix.


## Contributing

Check out the [contribution guidelines](CONTRIBUTING.md) for more information
Expand Down
23 changes: 23 additions & 0 deletions Source/Cmlx/.metal-compiler-plugin.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
{
"xcrun": true,
"find-inputs": false,
"include-dependencies": false,
"dependency-path-suffix": "include",
"include-paths": ["mlx-generated"],
"inputs": [
"Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal",
"Source/Cmlx/mlx-generated/metal/arg_reduce.metal",
"Source/Cmlx/mlx-generated/metal/conv.metal",
"Source/Cmlx/mlx-generated/metal/rms_norm.metal",
"Source/Cmlx/mlx-generated/metal/random.metal",
"Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal",
"Source/Cmlx/mlx-generated/metal/gemv.metal",
"Source/Cmlx/mlx-generated/metal/layer_norm.metal",
"Source/Cmlx/mlx-generated/metal/rope.metal"],
"output": "default.metallib",
"flags": ["-gline-tables-only", "-frecord-sources", "-Wno-c++20-extensions", "-std=metal4.0"],
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure we can use metal4 for all the shaders, we just want it for the NAX shaders. If we do this then it won't run on pre-macOS 26. I think we need to vary this per file (which is why we can't currently support NAX in main).

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One way out would be, that swiftpm shader compilation is only supported from e.g. OS 26 onwards.
Otherwise it gets unnecessary complicated.

"plugin-logging": true,
"verbose-logging": false,
"metal-enable-logging": false,
"logging-prefix": "[Metal]"
}
25 changes: 25 additions & 0 deletions Tests/MLXTests/ArrayAtTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,33 @@
import Foundation
import MLX
import XCTest
import Cmlx

class ArrayAtTests: XCTestCase {
override class func setUp() {
super.setUp()

// Cmlx is not an automatically loaded bundle.
// Let's load Cmlx manually, so that the mlx-lib can find the default.metallib.
// This requires patch #2885 in ml-explore/mlx be applied.
Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a workaround for the fix from the mlx repo?

Copy link
Copy Markdown
Author

@gin66 gin66 Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No. Both are needed. The issue is, that the plugin can only write to its plugin dir and the module dir, which depends on it. There is no way to get this metallib out of the module in the main module. The mlx-repo already supports searching all bundles‘ resource URL, which is a major step. Now the issue is, that I doubt, that the plugin can write a valid resource. At least I was not successful. So the remedy is, that not only the bundles’ resourceURL are searched, but the bundles‘ root directories, too. This is the purpose of the patch in mlx.

Now the tricky thing is, that on swiftpm level modules are not autloaded. The patch in the test performs the manual loading of the bundle and then the mlx-code is (magically) able to find the bundle in the list.

There is a lot of automatisms in the Xcode build process.

Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, I think this will probably be a blocker until we can figure out how to get the resource production working -- applications can't be required to do this same work also.

Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought that producing resources was the only thing a plugin could do. We used to have a plugin that would prepare the include paths in the shaders:

it didn't compile them. I know it can't produce a .o that is linked in to the resulting binary.


// Find the test bundle (.xctest)
let testBundle = Bundle(for: self)

// Go to its parent directory: …/debug
let buildDir = testBundle.bundleURL.deletingLastPathComponent()

// Append the actual SwiftPM bundle name for Cmlx
let cmlxURL = buildDir.appendingPathComponent("mlx-swift_Cmlx.bundle")

if let cmlxBundle = Bundle(url: cmlxURL) {
print("Loaded Cmlx bundle at:", cmlxURL.path)
_ = cmlxBundle.load() // registers it globally
} else {
print("Failed to load Cmlx bundle at:", cmlxURL.path)
}

}

func testArrayAt() {
// from example at https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.array.at.html#mlx.core.array.at
Expand Down