Skip to content

Support building MLX-Swift on Linux through SwiftPM (CPU only)#304

Merged
davidkoski merged 1 commit intoml-explore:mainfrom
Joannis:update/support-linux-native-build-cmake-cpu-backend
Jan 29, 2026
Merged

Support building MLX-Swift on Linux through SwiftPM (CPU only)#304
davidkoski merged 1 commit intoml-explore:mainfrom
Joannis:update/support-linux-native-build-cmake-cpu-backend

Conversation

@Joannis
Copy link
Contributor

@Joannis Joannis commented Dec 1, 2025

Based on #293

Proposed changes

  • Linux compilation works.
  • Doesn't support cross-compilation yet

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@Joannis Joannis marked this pull request as ready for review December 1, 2025 15:39
@incertum
Copy link
Contributor

incertum commented Dec 1, 2025

Amazing @Joannis! Let me check in with @davidkoski and ask when we might be ready to bump mlx-c in order to unblock current PRs.

@Joannis Joannis force-pushed the update/support-linux-native-build-cmake-cpu-backend branch 2 times, most recently from 878fc10 to e8dd62b Compare December 1, 2025 18:20
@incertum
Copy link
Contributor

incertum commented Dec 2, 2025

@Joannis you may rebase this PR now since the other PR got merged yesterday. Happy to help with the review once it's rebased.

@Joannis Joannis force-pushed the update/support-linux-native-build-cmake-cpu-backend branch from e8dd62b to 76494af Compare December 18, 2025 20:38
@Joannis
Copy link
Contributor Author

Joannis commented Dec 18, 2025

Rebased @incertum

@incertum
Copy link
Contributor

incertum commented Dec 18, 2025

Rebased @incertum

wohoo, could you run pre-commit run --all and re-push?

This commit should not exist f85906c after a clean rebase (since you worked on top of the previous PR). You could cherry delete it and or perhaps retry the rebase? Thanks a bunch!

}

open func callAsFunction(_ x: MLXArray) -> MLXArray {
#if canImport(MLXFast)
Copy link
Contributor

Choose a reason for hiding this comment

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

@davidkoski do you have more context around MLXFast issues? Intuitively it sounds like something you would like to have. It seems to just compile fine over CMake on Linux.

Copy link
Collaborator

Choose a reason for hiding this comment

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

I am not sure what this is for. MLXFast used to be a separate package, but was folded into MLX, see Sources/MLX/MLXFast.swift.

MLXFast remains as a stub to avoid breaking things and this function is:

@available(*, deprecated, message: "layerNorm is now available in the main MLX module")
@_disfavoredOverload
public func layerNorm(
    _ x: MLXArray, weight: MLXArray? = nil, bias: MLXArray? = nil, eps: Float,
    stream: StreamOrDevice = .default
) -> MLXArray {
    return MLXFast.layerNorm(x, weight: weight, bias: bias, eps: eps, stream: stream)
}

What issue is this working around?

Copy link
Collaborator

Choose a reason for hiding this comment

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

In this case MLXFast is an enum to give a namespace to look like the old module MLXFast. So we have:

  • MLXFast.layerNorm (the module MLXFast if you were to import that)
  • MLX.MLXFast.layerNorm aka MLXFast.layerNorm (the new home, you should not import MLXFast and it is not here)
  • MLX.layerNorm (just forwards to MLX.MLXFast.layerNorm for convenience)

OK, so the code here is referencing the second one as there is no import of MLXFast, thus we are getting the MLXFast scoped to the MLX module.

Dockerfile Outdated
@@ -0,0 +1,24 @@
FROM swift:6.2.1-jammy AS base
Copy link
Contributor

Choose a reason for hiding this comment

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

Could we move the Dockerfile somewhere under .github for CI use or for documentation purposes?

How much work would it be to add the swift build to https://github.com/ml-explore/mlx-swift/blob/main/.github/scripts/setup%2Bbuild-linux-container-cmake.sh

or better yet a new script / CI job (my preference)?

Package.swift Outdated
dependencies: [
"MLX",
"MLXRandom",
.target(name: "MLXFast", condition: .when(platforms: [.macOS, .iOS, .tvOS, .visionOS, .watchOS]))
Copy link
Collaborator

Choose a reason for hiding this comment

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

I think these references to MLXFast are:

  1. stale
  2. maybe causing the problems in Normalization.swift. I propose the following:
diff --git a/Package.swift b/Package.swift
index d88d9e4..0d1ed11 100644
--- a/Package.swift
+++ b/Package.swift
@@ -188,7 +188,7 @@ let package = Package(
         ),
         .target(
             name: "MLXNN",
-            dependencies: ["MLX", "MLXRandom", "MLXFast"],
+            dependencies: ["MLX"],
             swiftSettings: [
                 .enableExperimentalFeature("StrictConcurrency")
             ]
@@ -218,7 +218,7 @@ let package = Package(
         .testTarget(
             name: "MLXTests",
             dependencies: [
-                "MLX", "MLXRandom", "MLXNN", "MLXOptimizers", "MLXFFT", "MLXLinalg", "MLXFast",
+                "MLX", "MLXNN", "MLXOptimizers",
             ]
         ),
 

@davidkoski
Copy link
Collaborator

FYI #319 will have more Package.swift changes -- if we can get this merged first you won't have to figure out how to merge those in.

Outstanding issues. These have to be fixed before the merge:

  • swift-format
  • figure out this MLXFast issue
  • move Dockerfile under .github

Ideally this would be fixed but we can do a followup PR:

  • integrate swift build into GitHub actions

@Joannis
Copy link
Contributor Author

Joannis commented Dec 23, 2025

@davidkoski shall I use that as a base branch then?

@davidkoski
Copy link
Collaborator

@davidkoski shall I use that as a base branch then?

I think it is fine to get this prepared on its own. The mlx-c part isn't tagged yet, so it isn't ready to merge.

If we can get this merged then I will adopt it in that branch -- I think that will be easier.

@davidkoski
Copy link
Collaborator

@Joannis mlx/mlx-c have both been updated. Do you want to rebase and we can get this merged? I will maintain the Package.swift with these changes going forward.

@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

Hey @davidkoski , I missed your message. Will update it next week, as I'm currently prepping for FOSDEM

@davidkoski
Copy link
Collaborator

Hey @davidkoski , I missed your message. Will update it next week, as I'm currently prepping for FOSDEM

Awesome, no rush, take your time -- when we get this integrated I can cut a new release.

@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

@davidkoski I don't know what your policy is regarding this, but it would be good to have SwiftPM on Linux in CI

@Joannis Joannis requested a review from incertum January 28, 2026 18:27
@Joannis Joannis force-pushed the update/support-linux-native-build-cmake-cpu-backend branch from 7590346 to 94a0d2a Compare January 28, 2026 18:30
@Joannis Joannis requested a review from davidkoski January 28, 2026 18:30
@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

My Dockerfile to test:

FROM swift:6.2.1-jammy AS base

RUN apt-get update && apt-get install -y \
    libblas-dev \
    liblapack-dev \
    liblapacke-dev \
    libopenblas-dev \
    gfortran \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

FROM base AS builder

COPY . .
RUN swift build -c release --product Example1 --static-swift-stdlib -Xlinker -s -v

# Final image
FROM base

# Copy executable from SwiftPM build directory
COPY --from=builder /app/.build/*/release/Example1 /app/Example1

CMD ["./Example1"]

@Joannis Joannis force-pushed the update/support-linux-native-build-cmake-cpu-backend branch from 94a0d2a to bd63fca Compare January 28, 2026 18:33
@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

By the way, on the topic of:

I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes

Should that be swift format -ri .?

@Joannis Joannis force-pushed the update/support-linux-native-build-cmake-cpu-backend branch from bd63fca to a932421 Compare January 28, 2026 18:48
@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

I was indeed able to remove the MLXFast change now. I can make a follow-up PR for the CI side of things. If you want to use a specific workflow let me know.

@Joannis
Copy link
Contributor Author

Joannis commented Jan 28, 2026

This screenshot is the dockerfile running on my WendyOS Linux device:

image

@davidkoski
Copy link
Collaborator

@davidkoski I don't know what your policy is regarding this, but it would be good to have SwiftPM on Linux in CI

Yes, we can add it in this PR or a followup, but for sure we would want it.

@davidkoski
Copy link
Collaborator

By the way, on the topic of:

I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes

Should that be swift format -ri .?

pre-commit is a command that consumes the .pre-commit-config.yaml -- the swift format piece is probably the one that 99.9% of the people care about so this might be a good change in wording.

@davidkoski
Copy link
Collaborator

davidkoski commented Jan 28, 2026

I was indeed able to remove the MLXFast change now. I can make a follow-up PR for the CI side of things. If you want to use a specific workflow let me know.

Either way is fine -- since adding this doesn't break anything existing (e.g. there is no linux swiftpm build right now) we are safe to merge without CI in place.

Since I think this is complete perhaps we should merge this and have another PR for CI? That keeps us moving forward and keeps this away from merge conflicts.

"framework",
"include-framework",
"metal-cpp",
// Exclude Metal backend files on Linux, but keep no_metal.cpp for stubs
Copy link
Collaborator

Choose a reason for hiding this comment

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

It is too bad swiftpm doesn't have an easier way to do this!

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It's a bigger problem than it seems too, as #if os(macOS) evaluates to true when cross-compiling from macOS -> Linux

Copy link
Collaborator

Choose a reason for hiding this comment

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

Oh right! I ran into that problem elsewhere -- this is not about the target but where it is running.

"mlx/tests",

// special handling for cuda -- we need to keep one file:
// mlx/mlx/backend/cuda/no_cuda.cpp
Copy link
Collaborator

Choose a reason for hiding this comment

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

I guess this will be refactored in a similar way once we have swiftpm + cuda (assuming it is possible)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Correct, it looks it to become possible. We (at the Build & Packaging Workgroup) have discussed this requirement many times over at length. For now, we'll just have to let the swift evolution process go on and happen.

let mlxSwiftExcludes: [String] = [
"GPU+Metal.swift",
"MLXArray+Metal.swift",
"MLXFast.swift",
Copy link
Collaborator

Choose a reason for hiding this comment

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

This one is unfortunate -- did it not build? Or does it just not have CPU implementations?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It doesn't have a CPU implementation, at least as far as I can tell..

Copy link
Collaborator

@davidkoski davidkoski left a comment

Choose a reason for hiding this comment

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

Looks great, thank you!

@davidkoski
Copy link
Collaborator

CI runs and I reviewed and it looks good -- what do you think about merging now and we can get the CI running later?

@Joannis
Copy link
Contributor Author

Joannis commented Jan 29, 2026

Yeah let's merge it now! Happy to have CI in a follow-up

@davidkoski davidkoski merged commit f688d89 into ml-explore:main Jan 29, 2026
7 checks passed
@Joannis Joannis deleted the update/support-linux-native-build-cmake-cpu-backend branch January 29, 2026 10:04
@heckj
Copy link
Contributor

heckj commented Jan 30, 2026

I got all exited about this PR merging and went and cobbled up a little container that loaded the deps and then tried to use the MLX-swift package as a dependency, and I'm afraid the lack of MLXFast blew it up in my face.

I was just reaching for

import MLX
import MLXNN
import MLXOptimizers

print("Hello Swift")

but it looks like trying to reach for MLXNN as a module isn't happy:

89.59 /app/.build/checkouts/mlx-swift/Source/MLXNN/Normalization.swift:107:9: error: cannot find 'MLXFast' in scope
89.59 105 |
89.59 106 |     open func callAsFunction(_ x: MLXArray) -> MLXArray {
89.59 107 |         MLXFast.layerNorm(x, weight: weight, bias: bias, eps: eps)
89.59     |         `- error: cannot find 'MLXFast' in scope
89.59 108 |     }
89.59 109 | }
89.59

@davidkoski
Copy link
Collaborator

Yeah, we may need a compatibility layer -- this is all over in MLXNN as well.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants