Support building MLX-Swift on Linux through SwiftPM (CPU only)#304
Conversation
|
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. |
878fc10 to
e8dd62b
Compare
|
@Joannis you may rebase this PR now since the other PR got merged yesterday. Happy to help with the review once it's rebased. |
e8dd62b to
76494af
Compare
|
Rebased @incertum |
Source/MLXNN/Normalization.swift
Outdated
| } | ||
|
|
||
| open func callAsFunction(_ x: MLXArray) -> MLXArray { | ||
| #if canImport(MLXFast) |
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.layerNormakaMLXFast.layerNorm(the new home, you should not import MLXFast and it is not here)MLX.layerNorm(just forwards toMLX.MLXFast.layerNormfor 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 | |||
There was a problem hiding this comment.
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])) |
There was a problem hiding this comment.
I think these references to MLXFast are:
- stale
- 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",
]
),
|
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:
Ideally this would be fixed but we can do a followup PR:
|
|
@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. |
|
@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. |
|
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. |
|
@davidkoski I don't know what your policy is regarding this, but it would be good to have SwiftPM on Linux in CI |
7590346 to
94a0d2a
Compare
|
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"] |
94a0d2a to
bd63fca
Compare
|
By the way, on the topic of:
Should that be |
bd63fca to
a932421
Compare
|
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. |
Yes, we can add it in this PR or a followup, but for sure we would want it. |
|
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 |
There was a problem hiding this comment.
It is too bad swiftpm doesn't have an easier way to do this!
There was a problem hiding this comment.
It's a bigger problem than it seems too, as #if os(macOS) evaluates to true when cross-compiling from macOS -> Linux
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
I guess this will be refactored in a similar way once we have swiftpm + cuda (assuming it is possible)
There was a problem hiding this comment.
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", |
There was a problem hiding this comment.
This one is unfortunate -- did it not build? Or does it just not have CPU implementations?
There was a problem hiding this comment.
It doesn't have a CPU implementation, at least as far as I can tell..
davidkoski
left a comment
There was a problem hiding this comment.
Looks great, thank you!
|
CI runs and I reviewed and it looks good -- what do you think about merging now and we can get the CI running later? |
|
Yeah let's merge it now! Happy to have CI in a follow-up |
|
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: |
|
Yeah, we may need a compatibility layer -- this is all over in MLXNN as well. |

Based on #293
Proposed changes
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes