Using Swift 6.1 traits to conditionally compile for MLX Cuda backend #259

zamderax · 2025-07-29T23:54:59Z

updated the git submodule to the Commit 4ad5341 which is https://github.com/ml-explore/mlx/releases/tag/v0.27.1
then we use Swift Package Manager Traits https://github.com/swiftlang/swift-evolution/blob/main/proposals/0450-swiftpm-package-traits.md with https://github.com/swiftlang/swift-evolution/blob/main/proposals/0135-package-manager-support-for-differentiating-packages-by-swift-version.md to create a [email protected]
Then I updated update-mlx.sh to include steel_gemm_segmented which generated a lot of new file changes in mlx-generated

- Fixed indexing issue in reduce.cpp to correctly calculate output index. - Improved simd_reduce implementations in reduce_utils.cpp to handle integral and non-integral types, including complex numbers. - Added handling for zero input in complex power operation in scan.cpp. - Updated softmax.cpp to use finite minimum for values outside axis size. - Enhanced Conv2D input and weight loaders in steel_conv_general.cpp for better safety and performance. - Introduced segmented GEMM kernel in steel_gemm_segmented.cpp for improved matrix multiplication. - Modified unary and ternary operations to support batch processing in unary.cpp and ternary.cpp. - Added utility functions for complex exponential calculations in unary_ops.cpp. - Updated utils.cpp to define WorkPerThread struct for better thread utilization. - Updated update-mlx.sh to include new segmented GEMM kernel in build process.

davidkoski · 2025-07-30T02:45:02Z

This looks really good! The CMakeLists.txt might not work -- it references mlx-c, which references a different version of mlx. Let's see what CI says.

I will try it out tomorrow as well.

davidkoski · 2025-07-30T13:25:34Z

CI hit the swift-format item first -- can you run that? We can try again

zamderax · 2025-07-30T18:11:52Z

Okay I ran swift format --recursive .

davidkoski · 2025-07-30T19:46:42Z

Source/MLXNN/Module.swift

-            } else {
-                return value!
-            }
+            return value!


I dob't think we want this change -- there is a warning but I think (and I think there is a test) that this does actually do something.

For some reason that swift-format didn't do it, it is still failing in CI. I can try it later today.

davidkoski · 2025-07-30T21:57:47Z

Source/MLX/MLXArray+Bytes.swift

    /// - ``asData(access:)``
    public func asMTLBuffer(device: any MTLDevice, noCopy: Bool = false) -> (any MTLBuffer)? {
-        let data = asData(access: noCopy ? .noCopyIfContiguous : .copy)
+        _ = asData(access: noCopy ? .noCopyIfContiguous : .copy)


This looks like a bug maybe -- perhaps this line should be deleted entirely.

davidkoski · 2025-07-30T21:58:04Z

Source/MLXNN/Module.swift

-            } else {
-                return module!
-            }
+            return module!


This change should also be removed

davidkoski · 2025-07-30T22:04:17Z

Aha, swift-format needs to be run on the new Package.swift!

Look at the two items in Module.swift -- I think you were fixing warnings but we need those, despite the warning it does nothing. In fact testOptionInfos will fail if this is changed.

2. Reverted changes in Module.swift lines 1378 and 1487 3. Fixed the bug in MLXArray+Bytes.swift line 288

zamderax · 2025-07-31T04:30:49Z

Great!

I ran swift-format --in-place [email protected]
I reverted the changes in Module.swift lines 1378 and 1487
Then I think I addressed that bug in MLXArray+Bytes.swift line 288

davidkoski · 2025-07-31T05:53:29Z

OK, seeing a failure here:

Testing failed:
        'cuda_runtime.h' file not found
        Testing cancelled because the build failed.

building this:

CompileC /Users/distiller/Library/Developer/Xcode/DerivedData/project-focxxzqvztridmeptmirzegwtlnj/Build/Intermediates.noindex/mlx-swift.build/Debug/Cmlx.build/Objects-normal/arm64/worker.o /Users/distiller/project/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp normal arm64 c++ com.apple.compilers.llvm.clang.1_0.compiler (in target 'Cmlx' from project 'mlx-swift')

So I think the problem may be that the Package.swift (old one) needs to exclude the cuda files from the build -- it picks them up by default.

It also makes me think we need a code 16.3 (and maybe linux) builder. The Xcode side of things can be set up like this:

but we want 15.3.0 and 16.3.0.

…pported CUDA backend files in the older Package.swift

zamderax · 2025-08-03T06:50:26Z

Alright let's give this another try.

davidkoski · 2025-08-04T15:33:47Z

Looks like Package.swift is now failing on swift-format. You can try these locally with:

pre-commit run --all-files

The next step after that is:

xcodebuild test -scheme mlx-swift-Package -destination 'platform=OS X'
xcodebuild test -scheme mlx-swift-Package -destination 'platform=iOS' -sdk iphoneos

I found I had to change both Package.swift files:

-                "mlx/mlx/backend/no_metal",
+                "mlx/mlx/backend/no_gpu",

Now I am hitting this issue: https://developer.apple.com/forums/thread/779744 (though I suspect it is not related to this PR). Per that thread I may try Xcode 16.4.

Finally this runs in CI:

brew install cmake ninja
rm -rf build
mkdir build
cd build
cmake .. -G Ninja
ninja

Surprisingly the latter build for me -- this picks up the mlx dependency transitively via mlx-c and I thought it might fail since the versions do not match, but it built ok.

Anyway, I think these are the current problems:

formatting in Package.swift
no_metal -> no_gpu

zamderax mentioned this pull request Jul 29, 2025

mlx-swift with CUDA for Linux and Windows? #258

Open

Apply Swift format

7bf851f

davidkoski reviewed Jul 30, 2025

View reviewed changes

Source/MLXNN/Module.swift Outdated

} else {

return module!

}

return module!

Copy link

Collaborator

davidkoski Jul 30, 2025

Choose a reason for hiding this comment

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

This change should also be removed

1. Ran swift-format on [email protected]

b9e588e

2. Reverted changes in Module.swift lines 1378 and 1487 3. Fixed the bug in MLXArray+Bytes.swift line 288

Update CircleCI config to parameterize Xcode version and exclude unsu…

84e487f

…pported CUDA backend files in the older Package.swift

davidkoski mentioned this pull request Sep 24, 2025

adopt mlx 0.29.1 and related mlx-c #273

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Using Swift 6.1 traits to conditionally compile for MLX Cuda backend #259

Using Swift 6.1 traits to conditionally compile for MLX Cuda backend #259

Uh oh!

zamderax commented Jul 29, 2025

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

zamderax commented Jul 30, 2025

Uh oh!

davidkoski Jul 30, 2025

Uh oh!

davidkoski Jul 30, 2025

Uh oh!

davidkoski Jul 30, 2025

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

zamderax commented Jul 31, 2025

Uh oh!

davidkoski commented Jul 31, 2025

Uh oh!

zamderax commented Aug 3, 2025

Uh oh!

davidkoski commented Aug 4, 2025

Uh oh!

Uh oh!

Using Swift 6.1 traits to conditionally compile for MLX Cuda backend #259

Are you sure you want to change the base?

Using Swift 6.1 traits to conditionally compile for MLX Cuda backend #259

Uh oh!

Conversation

zamderax commented Jul 29, 2025

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

zamderax commented Jul 30, 2025

Uh oh!

davidkoski Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

davidkoski Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

davidkoski Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

davidkoski commented Jul 30, 2025

Uh oh!

zamderax commented Jul 31, 2025

Uh oh!

davidkoski commented Jul 31, 2025

Uh oh!

zamderax commented Aug 3, 2025

Uh oh!

davidkoski commented Aug 4, 2025

Uh oh!

Uh oh!