diff --git a/Package.swift b/Package.swift index 2fbab7f02..aa299aced 100644 --- a/Package.swift +++ b/Package.swift @@ -141,6 +141,7 @@ let noCudaCmlxExcludes = [ // the wrappers there so they must not also be built directly "mlx/mlx/backend/cpu/compiled.cpp", "mlx/mlx/backend/cpu/jit_compiler.cpp", + "mlx/mlx/distributed/jaccl/no_jaccl.cpp", // opt-out of these backends (using metal) "mlx/mlx/backend/no_gpu", @@ -154,6 +155,7 @@ let noCudaCmlxExcludes = [ let cxxSettings: [CXXSetting] = [ .headerSearchPath("metal-cpp"), + .headerSearchPath("mlx/mlx/distributed/jaccl/lib"), .define("MLX_USE_ACCELERATE"), .define("ACCELERATE_NEW_LAPACK"), @@ -224,7 +226,8 @@ let cmlx = Target.target( "mlx/mlx/backend/metal/kernels", "mlx/mlx/backend/metal/nojit_kernels.cpp", - // distributed backends: enable ring, disable MPI + NCCL + JACCL + // distributed backends: enable ring, disable MPI + NCCL. JACCL is + // built through mlx-conditional on macOS, so its sources are excluded here. "mlx/mlx/distributed/mpi/mpi.cpp", "mlx/mlx/distributed/ring/no_ring.cpp", "mlx/mlx/distributed/nccl/nccl.cpp", diff --git a/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp new file mode 100644 index 000000000..8e8afea5f --- /dev/null +++ b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp @@ -0,0 +1,18 @@ +// Copyright © 2026 Apple Inc. + +#include + +// JACCL (RDMA over Thunderbolt 5) needs , which ships with +// the macOS 26.2 SDK and later. Build the backend and its library when the +// header is present, and the stub otherwise. librdma is loaded at runtime, so +// on an older macOS MLXDistributed.isAvailable(.jaccl) reports false. +#if TARGET_OS_OSX && __has_include() +#include "../mlx/mlx/distributed/jaccl/jaccl.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/rdma.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/ring.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/tcp.cpp" +#else +#include "../mlx/mlx/distributed/jaccl/no_jaccl.cpp" +#endif diff --git a/Source/MLX/MLXArray+Bytes.swift b/Source/MLX/MLXArray+Bytes.swift index b4d1d62dc..666f8a36e 100644 --- a/Source/MLX/MLXArray+Bytes.swift +++ b/Source/MLX/MLXArray+Bytes.swift @@ -31,6 +31,13 @@ extension MLXArray { /// Return the physical size of the backing (assuming it is evaluated) in elements. This should /// only be used when accessing the backing directly, e.g. via `mlx_array_data_uint8()` var physicalSize: Int { + // an empty array (e.g. an empty slice like `a[128...]`) has no backing + // data pointer, so there is nothing to read even if some strides are + // non-zero + if self.size == 0 { + return 0 + } + // nbytes is the logical size of the input, not the physical size return zip(self.shape, self.internalStrides) .map { Swift.abs($0.0 * $0.1) } @@ -43,7 +50,11 @@ extension MLXArray { let shape = self.shape let strides = self.internalStrides - if contiguousDimension == 0 { + if output.isEmpty { + // an empty array has nothing to copy (and no backing pointer) + return + + } else if contiguousDimension == 0 { // entire backing is contiguous from.copyBytes(to: output) @@ -218,10 +229,11 @@ extension MLXArray { if self.contiguousToDimension() == 0 { // the backing is contiguous, we can provide a wrapper // for the contents without a copy - let source = UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx))! - let data = Data( - bytesNoCopy: source, count: size * itemSize, - deallocator: .none) + // an empty array has no backing pointer + let data = + UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx)).map { + Data(bytesNoCopy: $0, count: size * itemSize, deallocator: .none) + } ?? Data() return MLXArrayData( data: data, @@ -234,10 +246,11 @@ extension MLXArray { } case .noCopy: - let source = UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx))! - let data = Data( - bytesNoCopy: source, count: nbytes, - deallocator: .none) + // an empty array has no backing pointer + let data = + UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx)).map { + Data(bytesNoCopy: $0, count: nbytes, deallocator: .none) + } ?? Data() let strides: [Int] if ndim == 0 { diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift index 8570cbacd..bde1dd91a 100644 --- a/Source/MLX/MLXArray.swift +++ b/Source/MLX/MLXArray.swift @@ -562,7 +562,15 @@ public final class MLXArray { /// Note: this is an implementation detail and only visible because of the need to call it from /// other `mlx-swift` modules. public func _updateInternal(_ array: MLXArray) { - mlx_array_set(&self.ctx, array.ctx) + // Free the old array rather than assigning over it: `mlx_array_set` is a C++ copy assignment, + // which skips `array::~array()` and with it the breaking of sibling cycles. An unevaluated + // output of a multi-output primitive (e.g. the `Quantize` behind a freshly built + // `QuantizedLinear`) replaced through `Module.update(parameters:)` would otherwise leak its + // whole graph. Copy before freeing so `a._updateInternal(a)` keeps its contents. + var new = mlx_array_new() + mlx_array_set(&new, array.ctx) + mlx_array_free(self.ctx) + self.ctx = new } /// Internal function for copying the backing `mlx::core::array` context. diff --git a/Tests/MLXTests/DistributedTests.swift b/Tests/MLXTests/DistributedTests.swift index f9defc304..4946ea6bb 100644 --- a/Tests/MLXTests/DistributedTests.swift +++ b/Tests/MLXTests/DistributedTests.swift @@ -42,7 +42,9 @@ class DistributedTests: XCTestCase { /// Backends that are not compiled in report `false` rather than trapping. func testUnavailableBackendsReportFalse() { XCTAssertFalse(MLXDistributed.isAvailable(.mpi)) - XCTAssertFalse(MLXDistributed.isAvailable(.jaccl)) + #if !os(macOS) + XCTAssertFalse(MLXDistributed.isAvailable(.jaccl)) + #endif // nccl is CUDA-only and is never built on Apple platforms #if !os(Linux) diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift index 1a74fdba2..2d0dc4273 100644 --- a/Tests/MLXTests/MLXArrayTests.swift +++ b/Tests/MLXTests/MLXArrayTests.swift @@ -36,6 +36,46 @@ class MLXArrayTests: XCTestCase { XCTAssertEqual(b, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) } + func testAsArrayEmpty() { + // an empty slice has size 0 and no backing pointer; it should return + // an empty array rather than crashing (see issue #480) + let a = MLXArray.ones([1, 200]) + let b = a[128...] + XCTAssertEqual(b.shape, [0, 200]) + XCTAssertEqual(b.size, 0) + XCTAssertEqual(b.asArray(Int.self), []) + XCTAssertEqual(b.asArray(Float.self), []) + + // a directly-constructed empty array behaves the same + let empty = MLXArray([Int32](), [0]) + XCTAssertEqual(empty.asArray(Int32.self), []) + + // a non-contiguous empty array goes through the strided copy path + let t = b.T + XCTAssertEqual(t.shape, [200, 0]) + XCTAssertEqual(t.asArray(Float.self), []) + } + + func testAsDataEmpty() { + // asData shares physicalSize with asArray, so it also returns empty + // data for an empty array (see issue #480) + let b = MLXArray.ones([1, 200])[128...] + for access in [MLXArray.AccessMethod.copy, .noCopyIfContiguous, .noCopy] { + let data = b.asData(access: access) + XCTAssertEqual(data.data.count, 0) + XCTAssertEqual(data.shape, [0, 200]) + } + XCTAssertEqual(b.T.asData(access: .copy).data.count, 0) + } + + func testAsArrayScalar() { + // a scalar (0-dimensional) array yields a single-element array + let a = MLXArray(7) + XCTAssertEqual(a.shape, []) + XCTAssertEqual(a.size, 1) + XCTAssertEqual(a.asArray(Int32.self), [7]) + } + func testAsArrayNonContiguous1() { // skipping elements via slicing let a = MLXArray(0 ..< 9, [3, 3]) @@ -169,4 +209,11 @@ class MLXArrayTests: XCTestCase { } } + func testUpdateInternalWithItself() { + // `Module.update(parameters:)` can hand an array its own value back + let a = MLXArray([1, 2, 3] as [Float]) + a._updateInternal(a) + XCTAssertEqual(a.asArray(Float.self), [1, 2, 3]) + } + } diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 3a62a5c99..1926a191e 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -1113,6 +1113,7 @@ mlx/distributed/jaccl/lib/jaccl/threadpool.h, mlx/distributed/jaccl/lib/jaccl/types.h, mlx/distributed/jaccl/lib/README.md, + mlx/distributed/jaccl/no_jaccl.cpp, mlx/distributed/mpi/CMakeLists.txt, mlx/distributed/mpi/mpi_declarations.h, mlx/distributed/mpi/mpi.cpp, diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index c6b2b5f4f..57d2b9115 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -5,7 +5,7 @@ PRODUCT_BUNDLE_IDENTIFIER = com.apple.mlx.Cmlx DEFINES_MODULE = YES BUILD_LIBRARY_FOR_DISTRIBUTION = YES -HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann +HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann $(SRCROOT)/../Source/Cmlx/mlx/mlx/distributed/jaccl/lib OTHER_CFLAGS = -isysroot $(SDKROOT)