Add TurboQuant Swift support primitives

.gitmodules

@@ -3,4 +3,4 @@
 	url = https://github.com/ml-explore/mlx
 [submodule "submodules/mlx-c"]
 	path = Source/Cmlx/mlx-c
-	url = https://github.com/ml-explore/mlx-c
+	url = https://github.com/RNT56/mlx-c

MAINTENANCE.md

@@ -126,13 +126,13 @@ git submodules to include the `mlx` and `mlx-c` repositories.
 When a new version of `mlx` and its equivalent `mlx-c` are to be used, there is a
 process to go through to update `mlx-swift`.
 
-Additionally, SwiftPM supports plugins that can produce derived source for
-building, but this can only produce new swift source. It is possible to use
-plugins to generate new source `.cpp` files and even compile them, but at
-best the `.o` is copied into the output as a resource, not linked.
-This is important because `mlx` has some build-time source generation
-(e.g. `make_compiled_preamble.sh`). This is handled in `mlx-swift` by
-pre-generating the source when updating the `mlx` version.
+Additionally, SwiftPM supports plugins that can produce derived source and
+resources for building. It is possible to use plugins to generate new source
+`.cpp` files and even compile them, but at best the `.o` is copied into the
+output as a resource, not linked. This is important because `mlx` has some
+build-time source generation (e.g. `make_compiled_preamble.sh`). This is
+handled in `mlx-swift` by pre-generating the source when updating the `mlx`
+version, while the SwiftPM Metal library is generated as a build resource.
 
 1. Update the `mlx` and `mlx-c` submodules via `git pull` or `git checkout ...`
    - `Source/Cmlx/mlx`
@@ -143,6 +143,9 @@ pre-generating the source when updating the `mlx` version.
     - this updates headers in Source/Cmlx/include
     - this updates headers in Source/Cmlx/include-framework
     - this generates various files in Source/Cmlx/mlx-generated
+    - SwiftPM builds generate `default.metallib` through the
+      `BuildSwiftPMMetalLibrary` plugin; do not check in copied Metal sources or
+      a concatenated embedded fallback.
 
 4. Fix any build issues with SwiftPM build (opening Package.swift)
 5. Fix any build issues with xcodeproj build (opening xcode/MLX.codeproj), see also [README.xcodeproj.md]
@@ -163,7 +166,9 @@ After updating the mlx/mlx-c version the xcodeproj needs to be brought up to dat
 - no other headers in the project should be included as resources (public/private/project)
     - the easiest way to adjust is look at Project -> Cmlx -> Build Phases and then look at the Headers task
 - similarly there should be _no_ Copy Bundle Resources from the same section
-- compilation issues in .metal files typically mean they are new to the project and need to be removed from Cmlx target membership
+- compilation issues in `.metal` files usually mean the SwiftPM Metal plugin's
+  kernel list or include dependencies need to be updated, or the files need to
+  remain excluded from normal Cmlx target membership
 
 ### Cmlx
 
@@ -181,4 +186,3 @@ Settings, including header search paths are in xcode/xcconfig.
 ### MLX, etc.
 
 These are just normal frameworks that link to Cmlx and others as needed.  The source files are all swift and there are no special settings needed.
-

Package.swift

@@ -71,6 +71,8 @@ import PackageDescription
         "MLXFast.swift",
         "MLXFastKernel.swift",
     ]
+
+    let cmlxPlugins: [Target.PluginUsage]? = nil
 #else
     let platformExcludes: [String] = [
         "mlx/mlx/backend/cpu/compiled.cpp",
@@ -102,6 +104,10 @@ import PackageDescription
     ]
 
     let mlxSwiftExcludes: [String] = []
+
+    let cmlxPlugins: [Target.PluginUsage]? = [
+        "BuildSwiftPMMetalLibrary"
+    ]
 #endif
 
 let cmlx = Target.target(
@@ -211,7 +217,8 @@ let cmlx = Target.target(
         .headerSearchPath("fmt/include"),
         .define("MLX_VERSION", to: "\"0.31.1\""),
     ],
-    linkerSettings: linkerSettings
+    linkerSettings: linkerSettings,
+    plugins: cmlxPlugins
 )
 
 let package = Package(
@@ -240,6 +247,10 @@ let package = Package(
     ],
     targets: [
         cmlx,
+        .plugin(
+            name: "BuildSwiftPMMetalLibrary",
+            capability: .buildTool()
+        ),
         .testTarget(
             name: "CmlxTests",
             dependencies: ["Cmlx"]

Plugins/BuildSwiftPMMetalLibrary/plugin.swift

@@ -0,0 +1,55 @@
+import Foundation
+import PackagePlugin
+
+@main
+struct BuildSwiftPMMetalLibrary: BuildToolPlugin {
+    func createBuildCommands(context: PluginContext, target: any Target) async throws -> [Command] {
+        #if os(Linux)
+            return []
+        #else
+            let packageRoot = context.package.directory
+            let script = packageRoot.appending("tools", "build-swiftpm-metallib.sh")
+            let output = context.pluginWorkDirectory.appending("default.metallib")
+
+            return [
+                .buildCommand(
+                    displayName: "Build SwiftPM default.metallib",
+                    executable: Path("/bin/bash"),
+                    arguments: [script, output],
+                    inputFiles: inputFiles(packageRoot: packageRoot, script: script),
+                    outputFiles: [output]
+                )
+            ]
+        #endif
+    }
+
+    #if !os(Linux)
+        private func inputFiles(packageRoot: Path, script: Path) -> [Path] {
+            let kernelsDirectory = packageRoot.appending(
+                "Source",
+                "Cmlx",
+                "mlx",
+                "mlx",
+                "backend",
+                "metal",
+                "kernels"
+            )
+            var files = [script]
+            files.append(contentsOf: recursivelyCollectedMetalInputs(in: kernelsDirectory))
+            return files
+        }
+
+        private func recursivelyCollectedMetalInputs(in directory: Path) -> [Path] {
+            let fileManager = FileManager.default
+            guard let enumerator = fileManager.enumerator(atPath: directory.string) else {
+                return []
+            }
+
+            return enumerator.compactMap { entry -> Path? in
+                guard let entry = entry as? String else { return nil }
+                guard entry.hasSuffix(".metal") || entry.hasSuffix(".h") else { return nil }
+                return directory.appending(subpath: entry)
+            }.sorted { $0.string < $1.string }
+        }
+    #endif
+}

Source/Cmlx/include/mlx/c/fast.h

@@ -63,6 +63,10 @@ int mlx_fast_cuda_kernel_config_add_template_arg_int(
     mlx_fast_cuda_kernel_config cls,
     const char* name,
     int value);
+int mlx_fast_cuda_kernel_config_add_template_arg_uint32(
+    mlx_fast_cuda_kernel_config cls,
+    const char* name,
+    uint32_t value);
 int mlx_fast_cuda_kernel_config_add_template_arg_bool(
     mlx_fast_cuda_kernel_config cls,
     const char* name,
@@ -133,6 +137,10 @@ int mlx_fast_metal_kernel_config_add_template_arg_int(
     mlx_fast_metal_kernel_config cls,
     const char* name,
     int value);
+int mlx_fast_metal_kernel_config_add_template_arg_uint32(
+    mlx_fast_metal_kernel_config cls,
+    const char* name,
+    uint32_t value);
 int mlx_fast_metal_kernel_config_add_template_arg_bool(
     mlx_fast_metal_kernel_config cls,
     const char* name,