From fa6a4e8b4ed98dd08da7f0384f05003541523d87 Mon Sep 17 00:00:00 2001
From: anupsv <6407789+anupsv@users.noreply.github.com>
Date: Wed, 20 May 2026 18:50:29 -0700
Subject: [PATCH] Expose Cmlx product + enable jaccl distributed backend
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Downstream consumers (e.g. d-inference's provider-swift) need direct
access to the underlying mlx-c symbols for distributed-group setup and
custom collective ops. This commit:

1. Adds `Cmlx` to the package's library products so callers can declare
   `.product(name: "Cmlx", package: "mlx-swift")`. The target already
   exists; only the public product entry was missing.

2. Enables the jaccl distributed backend by un-excluding its source
   files (jaccl.cpp, mesh.cpp, ring.cpp, utils.cpp) and excluding the
   `no_jaccl.cpp` stub instead. jaccl is the Apple-Silicon-friendly
   collective backend used for cluster pipeline inference over
   Thunderbolt 5; the other distributed backends (mpi, ring, nccl) remain
   excluded since they're not supported on macOS.

3. Surfaces `mlx-c-distributed_group.h` and `mlx-c-distributed.h` in
   the umbrella `Cmlx.h` so the C distributed-group API is reachable
   from Swift via the Cmlx module.

No behavior change for existing consumers — MLX, MLXNN, MLXRandom etc.
continue to work exactly as before. Build size grows by the jaccl
sources (~600 LOC of C++).
---
 Package.swift                        | 13 +++++--------
 Source/Cmlx/include-framework/Cmlx.h |  2 ++
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/Package.swift b/Package.swift
index 17a4178f..c5e4de42 100644
--- a/Package.swift
+++ b/Package.swift
@@ -111,10 +111,8 @@ let cmlx = Target.target(
         // vendor docs
         "vendor-README.md",
 
-        // example code + mlx-c distributed
+        // example code only (distributed C API now enabled)
         "mlx-c/examples",
-        "mlx-c/mlx/c/distributed.cpp",
-        "mlx-c/mlx/c/distributed_group.cpp",
 
         // vendored library, include header only
         "json",
@@ -190,15 +188,13 @@ let cmlx = Target.target(
         "mlx/mlx/backend/metal/kernels",
         "mlx/mlx/backend/metal/nojit_kernels.cpp",
 
-        // do not build distributed support (yet)
+        // non-jaccl distributed backends: excluded
         "mlx/mlx/distributed/mpi/mpi.cpp",
         "mlx/mlx/distributed/ring/ring.cpp",
         "mlx/mlx/distributed/nccl/nccl.cpp",
         "mlx/mlx/distributed/nccl/nccl_stub",
-        "mlx/mlx/distributed/jaccl/jaccl.cpp",
-        "mlx/mlx/distributed/jaccl/mesh.cpp",
-        "mlx/mlx/distributed/jaccl/ring.cpp",
-        "mlx/mlx/distributed/jaccl/utils.cpp",
+        // jaccl stub excluded now that the real backend is built
+        "mlx/mlx/distributed/jaccl/no_jaccl.cpp",
     ],
     cSettings: [
         .headerSearchPath("mlx"),
@@ -226,6 +222,7 @@ let package = Package(
 
     products: [
         // main targets
+        .library(name: "Cmlx", targets: ["Cmlx"]),
         .library(name: "MLX", targets: ["MLX"]),
         .library(name: "MLXRandom", targets: ["MLXRandom"]),
         .library(name: "MLXNN", targets: ["MLXNN"]),
diff --git a/Source/Cmlx/include-framework/Cmlx.h b/Source/Cmlx/include-framework/Cmlx.h
index 8c007786..b99023a1 100644
--- a/Source/Cmlx/include-framework/Cmlx.h
+++ b/Source/Cmlx/include-framework/Cmlx.h
@@ -11,6 +11,8 @@
 #include <Cmlx/mlx-device.h>
 #include <Cmlx/mlx-distributed-distributed.h>
 #include <Cmlx/mlx-distributed-ops.h>
+#include <Cmlx/mlx-c-distributed_group.h>
+#include <Cmlx/mlx-c-distributed.h>
 #include <Cmlx/mlx-einsum.h>
 #include <Cmlx/mlx-export.h>
 #include <Cmlx/mlx-fast.h>