Skip to content

[CUDA] Build the CUDA argument lists in steps - #442

Merged
davidkoski merged 1 commit into
ml-explore:mainfrom
GoodOlClint:pr/encuda-split-argument-expression
Aug 17, 2026
Merged

davidkoski merged 1 commit into
ml-explore:mainfrom
GoodOlClint:pr/encuda-split-argument-expression

Conversation

@GoodOlClint

@GoodOlClint GoodOlClint commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

A readability/compile-cost cleanup in the CUDA build, and a follow-up to #413 ("Allowing SPM to compile on Linux with CUDA") — no behaviour change.

Three argument lists are assembled as long + chains of arrays. They are hard to read, awkward to extend, and expensive for the type checker, which has to solve the whole chain as one expression:

  • Source/Encuda/encuda-compile.swift — the nvcc argument list (7 concatenations).
  • Plugins/CudaBuild/plugin.swift — the encuda compile and encuda link argument lists, built inline inside createBuildCommands.

Each becomes stepwise appends into a local array, which is what the compiler's own diagnostic suggests when a chain does get too large. The plugin's two move into compileArguments and linkArguments helpers so createBuildCommands reads as a sequence of build commands rather than argument plumbing, and adding one more flag to either list is now a one-line change that can't tip the type checker over.

The concatenation order — and therefore the emitted command lines — is unchanged.

On Swift 6.2: these chains do exceed the 6.2 type checker's budget outright (error: the compiler is unable to type-check this expression in reasonable time), so this change is a prerequisite for building on that toolchain. It is not sufficient, and this PR should not be read as lowering the toolchain floor: SwiftPM 6.2 discards non-Swift files produced by a build-tool plugin (warning: Only Swift is supported for generated plugin source files at this time), so all ~95 generated .cpp files are dropped and every CUDA symbol goes undefined at link time. The swift-tools-version: 6.3 in #413 is load-bearing. Type-check cost and readability are the whole of the rationale here.

Verification

  • pre-commit run --all-files — clean.
  • scripts/verify-docs.sh — passes.
  • xcodebuild build-for-testing -scheme mlx-swift-Package -destination 'platform=macOS' — builds; CmlxTests and MLXTests pass (Xcode 26.6, Swift 6.3.3).
  • Exercised in a real CUDA build: swift build on an sm_121 device (DGX Spark / GB10, CUDA 13.0.88, Swift 6.3.3, aarch64) compiled all 95 .cu files and decoded Llama-3.2-1B-Instruct-4bit at 196.4 tok/s and gemma-4-12B-it-4bit at 23.0 tok/s.

Note CI does not cover the SwiftPM CUDA lane — linux_build_cmake_cuda builds via CMake — so the CUDA-side verification above is on-device rather than in CI.

Related

Two other follow-ups to #413 are open alongside this one, and all three touch Plugins/CudaBuild/plugin.swift. They are functionally independent, but whichever merges first will leave the others needing a trivial rebase — happy to reorder or stack them however you prefer:

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 — no test added: this is a behaviour-preserving refactor of argument-list construction; the existing build is the check, and the emitted command lines are unchanged.
  • I have updated the necessary documentation (if needed) — no API or configuration change.

Three argument lists in the CUDA build are assembled as long `+` chains that the
type checker struggles with. On a Swift 6.2 toolchain they exceed its budget
outright:

    encuda-compile.swift:66:17: error: the compiler is unable to type-check this
    expression in reasonable time; try breaking up the expression into distinct
    sub-expressions
    plugin.swift:172:13: error: the compiler is unable to type-check this
    expression in reasonable time; ...

Appends into local arrays instead, which is behaviourally identical and is what
the diagnostic suggests. The plugin's two chains move into `compileArguments` and
`linkArguments` helpers to keep `createBuildCommands` readable, and adding a
further argument to either list no longer risks tipping the type checker over.

Note this does *not* make the package buildable on Swift 6.2: SwiftPM 6.2 ignores
non-Swift files produced by a build-tool plugin ("Only Swift is supported for
generated plugin source files at this time"), so every generated .cpp is dropped
and the CUDA symbols go undefined at link time. The 6.3 floor is load-bearing.
This is purely a type-check-cost and readability change.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

@davidkoski davidkoski left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Change looks good to me. Thank you!

@davidkoski
davidkoski merged commit 97cf19e into ml-explore:main Aug 17, 2026
7 checks passed
jeethu added a commit to numen-tech/mlx-swift that referenced this pull request Aug 26, 2026
Brings the fork up to upstream main as of 2026-08-26: the eval deadlock fix (ml-explore#461),
the MLXArray(rawPointer:) finalizer leak fix (ml-explore#448), nested initializers (ml-explore#456),
CudaBuild plugin gated to Linux (ml-explore#447), CUDA build fixes (ml-explore#442/ml-explore#451) and the
update-mlx.sh tweak (ml-explore#445). The vendored core submodule stays on
numen-tech/mlx prism-0.31.1-fixes (b2b8a3d); no 0.32 mlx-c bindings exist yet.
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.

2 participants