L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Jaccl refactor (#3412)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-04-14 23:52:21 -07:00
committed by GitHub
parent 1fa764fbec
commit 4400504ad5
26 changed files with 2279 additions and 414 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+9
View File
@@ -321,6 +321,15 @@ FetchContent_MakeAvailable(json)
target_include_directories(
mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
if(MLX_BUILD_CPU
AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
AND DEFINED MACOS_SDK_VERSION
AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
${CMAKE_BINARY_DIR}/jaccl)
endif()
add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
target_include_directories(
mlx/distributed/jaccl/.gitignore
+1
View File
@@ -0,0 +1 @@
!lib
mlx/distributed/jaccl/CMakeLists.txt
+2 -6
View File
@@ -1,12 +1,8 @@
if(MLX_BUILD_CPU
AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
target_sources(
mlx
PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp
${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp
${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp)
target_link_libraries(mlx PRIVATE jaccl)
else()
target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp)
endif()
mlx/distributed/jaccl/jaccl.cpp
+138 -143
View File
@@ -1,178 +1,173 @@
// Copyright © 2025 Apple Inc.
#include <fstream>
#include <sstream>
#include <json.hpp>
#include "mlx/distributed/jaccl/jaccl.h"
#include "mlx/backend/cpu/encoder.h"
#include "mlx/distributed/distributed_impl.h"
#include "mlx/distributed/jaccl/mesh.h"
#include "mlx/distributed/jaccl/ring.h"
#include "mlx/distributed/jaccl/utils.h"
#include "mlx/dtype_utils.h"
#include <jaccl/group.h>
#include <jaccl/jaccl.h>
using GroupImpl = mlx::core::distributed::detail::GroupImpl;
using json = nlohmann::json;
namespace mlx::core::distributed::jaccl {
namespace {
struct DeviceFile {
DeviceFile(const char* dev_file) {
std::ifstream f(dev_file);
json devices = json::parse(f);
if (!devices.is_array()) {
throw std::runtime_error(
"[jaccl] The device file should start with an array");
}
/**
* Map MLX Dtype to JACCL Dtype for dispatch.
*/
int dtype_to_jaccl_dtype(Dtype dt) {
switch (dt) {
case bool_:
return ::jaccl::Dtype::Bool;
case int8:
return ::jaccl::Dtype::Int8;
case int16:
return ::jaccl::Dtype::Int16;
case int32:
return ::jaccl::Dtype::Int32;
case int64:
return ::jaccl::Dtype::Int64;
case uint8:
return ::jaccl::Dtype::UInt8;
case uint16:
return ::jaccl::Dtype::UInt16;
case uint32:
return ::jaccl::Dtype::UInt32;
case uint64:
return ::jaccl::Dtype::UInt64;
case float16:
return ::jaccl::Dtype::Float16;
case bfloat16:
return ::jaccl::Dtype::BFloat16;
case float32:
return ::jaccl::Dtype::Float32;
case float64:
return ::jaccl::Dtype::Float64;
case complex64:
return ::jaccl::Dtype::Complex64;
default:
throw std::runtime_error("[jaccl] Unsupported dtype for JACCL operation");
}
}
devices_.resize(devices.size());
for (int rank = 0; rank < devices.size(); rank++) {
auto conn = devices[rank];
if (!conn.is_array()) {
throw std::runtime_error(
"[jaccl] The device file should have an array of arrays");
}
if (conn.size() != devices_.size()) {
std::ostringstream msg;
msg << "[jaccl] The device file should contain the connectivity of each rank to "
<< "all other ranks but rank " << rank << " contains only "
<< conn.size() << " entries.";
throw std::runtime_error(msg.str());
}
/**
* Adapter that wraps a standalone jaccl::Group to implement
* MLX's distributed::detail::GroupImpl interface.
*
* This bridges mlx::core::array to raw pointers for JACCL operations.
*/
class JACCLGroup : public GroupImpl {
public:
JACCLGroup(std::shared_ptr<::jaccl::Group> group)
: group_(std::move(group)) {}
devices_[rank].resize(conn.size());
for (int dst = 0; dst < conn.size(); dst++) {
auto names = conn[dst];
if (names.is_string()) {
devices_[rank][dst].push_back(names);
} else if (names.is_array()) {
for (auto name_it = names.begin(); name_it != names.end();
name_it++) {
devices_[rank][dst].push_back(*name_it);
}
} else if (!names.is_null()) {
throw std::runtime_error(
"[jaccl] Device names should be null, a string or array of strings.");
}
}
}
Stream communication_stream(StreamOrDevice s) override {
return to_stream(s, Device::cpu);
}
int size() {
return devices_.size();
int rank() override {
return group_->rank();
}
bool is_valid_mesh() {
for (int src = 0; src < size(); src++) {
for (int dst = 0; dst < size(); dst++) {
if (devices_[src][dst].size() != static_cast<size_t>(src != dst)) {
return false;
}
}
}
return true;
int size() override {
return group_->size();
}
bool is_valid_ring() {
int num_connections = devices_[0][1].size();
if (num_connections == 0) {
return false;
}
for (int src = 0; src < size(); src++) {
int left = (src + size() - 1) % size();
int right = (src + 1) % size();
for (int dst = 0; dst < size(); dst++) {
if (dst != left && dst != right) {
if (devices_[src][dst].size() != 0) {
return false;
}
} else {
if (devices_[src][dst].size() != num_connections) {
return false;
}
}
}
}
return true;
void all_sum(const array& input, array& output, Stream stream) override {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
int dtype = dtype_to_jaccl_dtype(output.dtype());
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_sum(in_ptr, out_ptr, n_bytes, dtype);
});
}
std::vector<std::string> extract_mesh_connectivity(int rank) {
std::vector<std::string> devices(size());
for (int dst = 0; dst < size(); dst++) {
if (dst != rank) {
devices[dst] = devices_[rank][dst][0];
}
}
return devices;
void all_max(const array& input, array& output, Stream stream) override {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
int dtype = dtype_to_jaccl_dtype(output.dtype());
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_max(in_ptr, out_ptr, n_bytes, dtype);
});
}
std::pair<std::vector<std::string>, std::vector<std::string>>
extract_ring_connectivity(int rank) {
int left = (rank + size() - 1) % size();
int right = (rank + 1) % size();
return std::make_pair(devices_[rank][left], devices_[rank][right]);
void all_min(const array& input, array& output, Stream stream) override {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
int dtype = dtype_to_jaccl_dtype(output.dtype());
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_min(in_ptr, out_ptr, n_bytes, dtype);
});
}
std::vector<std::vector<std::vector<std::string>>> devices_;
void all_gather(const array& input, array& output, Stream stream) override {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
group_->all_gather(in_ptr, out_ptr, n_bytes);
});
}
void send(const array& input, int dst, Stream stream) override {
auto data = input.data<char>();
size_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch(
[data, n_bytes, dst, this]() { group_->send(data, n_bytes, dst); });
}
void recv(array& out, int src, Stream stream) override {
auto data = out.data<char>();
size_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch(
[data, n_bytes, src, this]() { group_->recv(data, n_bytes, src); });
}
void sum_scatter(const array& input, array& output, Stream stream) override {
throw std::runtime_error("[jaccl] sum_scatter not supported.");
}
std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
throw std::runtime_error("[jaccl] Group split not supported.");
}
private:
std::shared_ptr<::jaccl::Group> group_;
};
} // namespace
namespace mlx::core::distributed::jaccl {
bool is_available() {
return ibv().is_available();
return ::jaccl::is_available();
}
std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
const char* dev_file = std::getenv("MLX_IBV_DEVICES");
const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR");
const char* rank_str = std::getenv("MLX_RANK");
const char* ring = std::getenv("MLX_JACCL_RING");
if (!is_available() || !dev_file || !coordinator || !rank_str) {
if (strict) {
std::ostringstream msg;
msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), "
<< "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
<< "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "")
<< "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "")
<< "\" and MLX_JACCL_COORDINATOR=\""
<< ((coordinator) ? coordinator : "");
throw std::runtime_error(msg.str());
}
auto group = ::jaccl::init(strict);
if (group == nullptr) {
return nullptr;
}
auto rank = std::atoi(rank_str);
bool prefer_ring = ring != nullptr;
DeviceFile devices(dev_file);
if (rank >= devices.size() || rank < 0) {
std::ostringstream msg;
msg << "[jaccl] Invalid rank " << rank << ". It should be between 0 and "
<< devices.size();
throw std::runtime_error(msg.str());
}
if (prefer_ring && devices.is_valid_ring()) {
auto [left, right] = devices.extract_ring_connectivity(rank);
return std::make_shared<RingGroup>(
rank, devices.size(), left, right, coordinator);
} else if (devices.is_valid_mesh()) {
auto device_names = devices.extract_mesh_connectivity(rank);
return std::make_shared<MeshGroup>(rank, device_names, coordinator);
} else if (devices.is_valid_ring()) {
auto [left, right] = devices.extract_ring_connectivity(rank);
return std::make_shared<RingGroup>(
rank, devices.size(), left, right, coordinator);
} else {
throw std::runtime_error(
"[jaccl] The device file should define a valid mesh or a valid ring.");
}
return std::make_shared<JACCLGroup>(std::move(group));
}
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/jaccl.h
+2
View File
@@ -1,5 +1,7 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include "mlx/distributed/distributed.h"
namespace mlx::core::distributed::jaccl {
mlx/distributed/jaccl/lib/CMakeLists.txt
+83
View File
@@ -0,0 +1,83 @@
cmake_minimum_required(VERSION 3.24)
project(jaccl LANGUAGES CXX)
# Default to Release when built standalone. Without this CMake uses an empty
# build type (-O0) which severely impacts the reduction / memcpy hot paths.
if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
set(CMAKE_BUILD_TYPE
Release
CACHE STRING "Build type" FORCE)
endif()
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
include(FetchContent)
# nlohmann/json for device file parsing
message(STATUS "Downloading json for JACCL")
FetchContent_Declare(
json
URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
FetchContent_MakeAvailable(json)
# Check platform and SDK version requirements
if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
message(STATUS "JACCL requires macOS (Darwin). Skipping JACCL build.")
return()
endif()
# Try to determine MACOS_SDK_VERSION if not set
if(NOT DEFINED MACOS_SDK_VERSION)
execute_process(
COMMAND xcrun --sdk macosx --show-sdk-version
OUTPUT_VARIABLE MACOS_SDK_VERSION
OUTPUT_STRIP_TRAILING_WHITESPACE)
if(MACOS_SDK_VERSION)
message(STATUS "Detected macOS SDK version: ${MACOS_SDK_VERSION}")
endif()
endif()
if(DEFINED MACOS_SDK_VERSION AND MACOS_SDK_VERSION VERSION_LESS "26.2")
message(STATUS "JACCL requires macOS SDK >= 26.2. Skipping JACCL build.")
return()
endif()
add_library(
jaccl
${CMAKE_CURRENT_SOURCE_DIR}/jaccl/tcp.cpp
${CMAKE_CURRENT_SOURCE_DIR}/jaccl/rdma.cpp
${CMAKE_CURRENT_SOURCE_DIR}/jaccl/mesh.cpp
${CMAKE_CURRENT_SOURCE_DIR}/jaccl/ring.cpp
${CMAKE_CURRENT_SOURCE_DIR}/jaccl/jaccl.cpp)
target_include_directories(
jaccl PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
$<INSTALL_INTERFACE:include>)
target_include_directories(
jaccl PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
target_compile_features(jaccl PUBLIC cxx_std_20)
# Install targets
install(
TARGETS jaccl
EXPORT jacclTargets
LIBRARY DESTINATION lib
ARCHIVE DESTINATION lib
RUNTIME DESTINATION bin
INCLUDES
DESTINATION include)
install(
DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/
DESTINATION include/jaccl
FILES_MATCHING
PATTERN "*.h")
install(
EXPORT jacclTargets
FILE jacclTargets.cmake
NAMESPACE jaccl::
DESTINATION lib/cmake/jaccl)
mlx/distributed/jaccl/lib/README.md
+328
View File
@@ -0,0 +1,328 @@
# JACCL
**JACCL** is a low-latency distributed communication library designed for macOS
systems with Thunderbolt 5 connectivity.
## Overview
JACCL leverages RDMA (Remote Direct Memory Access) over Thunderbolt to achieve
communication latency an order of magnitude lower than traditional TCP-based
approaches. This makes it ideal for:
- Tensor parallelism in large model inference
- High-performance distributed training
- Low-latency collective operations between Macs
JACCL was made possible by Apple's RDMA over Thunderbolt technology introduced
in macOS 26.2.
## Features
- **Mesh Topology**: Fully connected communication where each node can directly
communicate with any other node
- **Ring Topology**: High-bandwidth ring all-reduce for large messages
- **Collective Operations**:
- `all_sum`: Sum values across all nodes
- `all_max`: Element-wise maximum across all nodes
- `all_min`: Element-wise minimum across all nodes
- `all_gather`: Gather data from all nodes
- **Point-to-Point Operations**:
- `send`: Send data to a specific node
- `recv`: Receive data from a specific node
- **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32,
Float64, Complex64
## Requirements
- macOS SDK >= 26.2
- Thunderbolt 5 connectivity between nodes
- RDMA over Thunderbolt enabled (requires macOS recovery mode setup)
## Enabling RDMA over Thunderbolt
RDMA over Thunderbolt must be enabled in macOS recovery mode:
1. Start your Mac in [recovery mode](https://support.apple.com/en-us/102518)
2. Open Terminal from Utilities -> Terminal
3. Run: `rdma_ctl enable`
4. Reboot
To verify RDMA is enabled, run:
```bash
ibv_devices
```
You should see output like:
```
device node GUID
------ ----------------
rdma_en2 8096a9d9edbaac05
rdma_en3 8196a9d9edbaac05
rdma_en5 8396a9d9edbaac05
```
## Building
JACCL can be built as a standalone library:
```bash
cd mlx/distributed/jaccl/lib
mkdir build && cd build
cmake ..
make
```
You can also include it in your own project via CMake:
```
FetchContent_Declare(
jaccl
GIT_REPOSITORY https://github.com/ml-explore/mlx.git
GIT_TAG main
SOURCE_SUBDIR mlx/distributed/jaccl/lib
)
FetchContent_MakeAvailable(jaccl)
```
## Usage
### Environment Variables
The easiest way to intiialize JACCL is by using the following environment
variables:
- **JACCL_RANK** / **MLX_RANK**: The rank of this process (0-based integer)
- **JACCL_IBV_DEVICES** / **MLX_IBV_DEVICES**: Path to a JSON file describing
device connectivity
- **JACCL_COORDINATOR** / **MLX_JACCL_COORDINATOR**: IP:port of the coordinator
(rank 0 listener)
- **JACCL_RING** / **MLX_JACCL_RING**: (Optional) Prefer ring topology over
mesh
### Device File Format
The device file is a JSON array where each entry describes the RDMA devices
connecting that rank to all other ranks:
```json
[
[null, "rdma_en5", "rdma_en4", "rdma_en3"],
["rdma_en5", null, "rdma_en3", "rdma_en4"],
["rdma_en4", "rdma_en3", null, "rdma_en5"],
["rdma_en3", "rdma_en4", "rdma_en5", null]
]
```
For a valid mesh, `devices[i][j]` should contain the device name connecting
rank `i` to rank `j`, or `null` if `i == j`.
For a valid ring, only adjacent nodes should have device names (all others
should be null).
### Basic Example
```cpp
#include <iostream>
#include <jaccl/jaccl.h>
int main() {
// Initialize JACCL group
auto group = jaccl::init();
if (!group) {
std::cerr << "Failed to initialize JACCL" << std::endl;
return 1;
}
std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
// Perform all-reduce sum
float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
float output[10];
group->all_sum(input, output, sizeof(input), jaccl::Float32);
std::cout << "Result: " << output[0] << std::endl;
return 0;
}
```
You can also manually define the configuration instead of reading it from
environment variables.
```cpp
#include <iostream>
#include <jaccl/jaccl.h>
int main() {
auto cfg = jaccl::Config()
.set_rank(0)
.set_coordinator("192.168.1.1:32132")
.set_devices({
{{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}},
{{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}},
{{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}},
{{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}}
});
auto group = jaccl::init(cfg);
if (!group) {
std::cerr << "Failed to initialize JACCL" << std::endl;
return 1;
}
std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
// Perform all-reduce sum
float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
float output[10];
group->all_sum(input, output, sizeof(input), jaccl::Float32);
std::cout << "Result: " << output[0] << std::endl;
return 0;
}
```
### Using with MLX
JACCL integrates seamlessly with MLX's distributed communication:
```python
import mlx.core as mx
# Initialize with JACCL backend
world = mx.distributed.init(backend="jaccl")
# Perform distributed operations
x = mx.ones((10,))
result = mx.distributed.all_sum(x, group=world)
```
Launch with `mlx.launch`:
```bash
mlx.launch --backend jaccl --hostfile hosts.json my_script.py
```
## Hostfile Example
For use with `mlx.launch`, create a hostfile JSON:
```json
{
"backend": "jaccl",
"hosts": [
{
"ssh": "m3-ultra-1",
"ips": ["192.168.1.1"],
"rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"]
},
{
"ssh": "m3-ultra-2",
"ips": [],
"rdma": ["rdma_en5", null, "rdma_en3", "rdma_en4"]
},
{
"ssh": "m3-ultra-3",
"ips": [],
"rdma": ["rdma_en4", "rdma_en3", null, "rdma_en5"]
},
{
"ssh": "m3-ultra-4",
"ips": [],
"rdma": ["rdma_en3", "rdma_en4", "rdma_en5", null]
}
]
}
```
## Automatic Configuration
MLX provides `mlx.distributed_config` to automatically discover and configure
Thunderbolt connectivity:
```bash
# Visualize connections
mlx.distributed_config --verbose \
--hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
--over thunderbolt --dot | dot -Tpng | open -f -a Preview
# Auto-configure and generate hostfile
mlx.distributed_config --verbose \
--hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
--over thunderbolt --backend jaccl \
--auto-setup --output m3-ultra-jaccl.json
```
## API
The main API of JACCL is the communication group. It provides efficient
high-level collectives.
**Note: JACCL does no memory allocation. All output pointers should point to a
location with sufficient memory allocated to hold the result.**
```cpp
class Group {
public:
virtual ~Group() {}
// Helper functions to know which process we are in the group
virtual int rank() = 0;
virtual int size() = 0;
// All reduce implementations. Input and output of the same size the
// reduction happens according to dtype and across the group.
virtual void all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0;
virtual void all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0;
virtual void all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0;
// All gather implementation. The output is group->size() * n_bytes.
virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0;
// Simple send/recv primitives.
virtual void send(const void* input, size_t n_bytes, int dst) = 0;
virtual void recv(void* output, size_t n_bytes, int src) = 0;
};
```
All that is left to use JACCL (except the communication group) is
```cpp
std::shared_ptr<Group> init(bool strict = false);
std::shared_ptr<Group> init(const Config& cfg, bool strict = false);
```
that create the communication group from environment variables or from the
configuration object. The latter allows one to configure JACCL using means
other than environment variables.
```cpp
class Config {
public:
Config();
Config& set_rank(int rank);
Config& set_coordinator(std::string coordinator);
Config& set_devices(std::vector<std::vector<std::vector<std::string>>> devices);
Config& prefer_ring(bool prefer = true);
bool is_valid_mesh() const;
bool is_valid_ring() const;
}
```
## License
JACCL is part of MLX and is released under the same license.
## Acknowledgments
The name JACCL (pronounced Jackal) stands for Jack and Angelos’ Collective
Communication Library and it is an obvious pun to Nvidia’s NCCL but also
tribute to Jack Beasley who led the development of RDMA over Thunderbolt at
Apple.
mlx/distributed/jaccl/lib/examples/CMakeLists.txt
+40
View File
@@ -0,0 +1,40 @@
cmake_minimum_required(VERSION 3.24)
project(jaccl_examples LANGUAGES CXX)
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
include(FetchContent)
# cmake-format: off
#
# Treating the nearby directory like a fetched dependency to simulate what a
# user would do in their own C++ project. Ideally you would write something like
# the following:
#
# FetchContent_Declare(
# jaccl
# GIT_REPOSITORY https://github.com/ml-explore/mlx.git
# GIT_TAG main
# SOURCE_SUBDIR mlx/distributed/jaccl/lib
# )
#
# cmake-format: on
FetchContent_Declare(jaccl SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../)
FetchContent_MakeAvailable(jaccl)
# Helper function to build examples
function(build_example SRCFILE)
get_filename_component(example_name ${SRCFILE} NAME_WE)
set(target "jaccl_${example_name}")
add_executable(${target} ${SRCFILE})
target_link_libraries(${target} PRIVATE jaccl)
message(STATUS "Building JACCL example: ${target}")
endfunction()
# Examples
build_example(minimal_env.cpp)
build_example(minimal_cfg.cpp)
# Benchmarks
build_example(allreduce_bench.cpp)
mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp
+286
View File
@@ -0,0 +1,286 @@
// Copyright © 2025 Apple Inc.
//
// JACCL All-Reduce Benchmark
//
// Measures bandwidth and latency of all_sum across a sweep of message sizes,
// similar in spirit to the NCCL all-reduce benchmark (nccl-tests).
//
// Usage:
// Set the environment variables described in jaccl.h, then run:
//
// ./jaccl_allreduce_bench [-w <warmup_iters>] [-n <iters>]
// [-b <min_bytes>] [-e <max_bytes>]
// [-f <step_factor>] [-d <datatype>]
// [-c] [-h]
//
// Or use the MLX launcher:
//
// mlx.launch --hostfile hosts.json ./jaccl_allreduce_bench
//
// The arguments are:
//
// -w Warmup iterations per message size (default: 5)
// -n Timed iterations per message size (default: 20)
// -b Minimum message size in bytes (default: 1K)
// -e Maximum message size in bytes (default: 256M)
// -f Multiplicative step factor (default: 2)
// -d Datatype: float32, float16, bfloat16 (default: float32)
// -c Check correctness (default: off)
// -h Print this help message
#include <jaccl/jaccl.h>
#include <jaccl/types.h>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <memory>
#include <numeric>
#include <string>
#include <vector>
static void usage(const char* prog) {
std::cerr
<< "Usage: " << prog << " [options]\n"
<< " -w <warmup> Warmup iterations (default: 5)\n"
<< " -n <iters> Timed iterations (default: 20)\n"
<< " -b <min_bytes> Minimum message size (default: 1K)\n"
<< " -e <max_bytes> Maximum message size (default: 256M)\n"
<< " -f <factor> Multiplicative step factor (default: 2)\n"
<< " -d <dtype> float32|float16|bfloat16 (default: float32)\n"
<< " -c Check correctness\n"
<< " -h Show this help\n";
}
static size_t parse_size(const char* s) {
char* end = nullptr;
double val = std::strtod(s, &end);
if (end && (*end == 'K' || *end == 'k'))
val *= 1024;
else if (end && (*end == 'M' || *end == 'm'))
val *= 1024 * 1024;
else if (end && (*end == 'G' || *end == 'g'))
val *= 1024 * 1024 * 1024;
return static_cast<size_t>(val);
}
static std::string fmt_bytes(size_t bytes) {
const char* units[] = {"B", "KB", "MB", "GB", "TB"};
int idx = 0;
double val = static_cast<double>(bytes);
while (val >= 1024.0 && idx < 4) {
val /= 1024.0;
idx++;
}
char buf[32];
if (val == static_cast<int>(val))
std::snprintf(buf, sizeof(buf), "%d %s", static_cast<int>(val), units[idx]);
else
std::snprintf(buf, sizeof(buf), "%.2f %s", val, units[idx]);
return buf;
}
// Conversion from algorithm bandwidth to bus bandwidth for a ring reduce.
static double bus_factor(int nranks) {
return 2.0 * (nranks - 1) / static_cast<double>(nranks);
}
int main(int argc, char** argv) {
int warmup_iters = 5;
int timed_iters = 20;
size_t min_bytes = 1024;
size_t max_bytes = 256 * 1024 * 1024;
int step_factor = 2;
std::string dtype_str = "float32";
bool check = false;
for (int i = 1; i < argc; i++) {
std::string arg = argv[i];
if (arg == "-h" || arg == "--help") {
usage(argv[0]);
return 0;
} else if (arg == "-w" && i + 1 < argc) {
warmup_iters = std::atoi(argv[++i]);
} else if (arg == "-n" && i + 1 < argc) {
timed_iters = std::atoi(argv[++i]);
} else if (arg == "-b" && i + 1 < argc) {
min_bytes = parse_size(argv[++i]);
} else if (arg == "-e" && i + 1 < argc) {
max_bytes = parse_size(argv[++i]);
} else if (arg == "-f" && i + 1 < argc) {
step_factor = std::atoi(argv[++i]);
} else if (arg == "-d" && i + 1 < argc) {
dtype_str = argv[++i];
} else if (arg == "-c") {
check = true;
} else {
std::cerr << "Unknown option: " << arg << "\n";
usage(argv[0]);
return 1;
}
}
jaccl::Dtype dtype;
size_t elem_size;
if (dtype_str == "float32") {
dtype = jaccl::Float32;
elem_size = 4;
} else if (dtype_str == "float16") {
dtype = jaccl::Float16;
elem_size = 2;
} else if (dtype_str == "bfloat16") {
dtype = jaccl::BFloat16;
elem_size = 2;
} else {
std::cerr << "Unsupported dtype: " << dtype_str << "\n";
return 1;
}
auto group = jaccl::init(/* strict= */ true);
int rank = group->rank();
int nranks = group->size();
if (rank == 0) {
std::cout << "# JACCL All-Reduce Benchmark\n"
<< "# Ranks: " << nranks << "\n"
<< "# Dtype: " << dtype_str << "\n"
<< "# Warmup: " << warmup_iters << " iters\n"
<< "# Timed: " << timed_iters << " iters\n"
<< "# Sizes: " << fmt_bytes(min_bytes) << " .. "
<< fmt_bytes(max_bytes) << " (x" << step_factor << ")\n"
<< "#\n";
// Table header (NCCL-style)
std::cout << std::left << std::setw(14) << "# size" << std::right
<< std::setw(12) << "count" << std::setw(12) << "type"
<< std::setw(14) << "time (us)" << std::setw(14) << "algo BW"
<< std::setw(14) << "bus BW";
if (check)
std::cout << std::setw(10) << "check";
std::cout << "\n";
std::cout << std::left << std::setw(14) << "# (bytes)" << std::right
<< std::setw(12) << "(elems)" << std::setw(12) << ""
<< std::setw(14) << "" << std::setw(14) << "(GB/s)"
<< std::setw(14) << "(GB/s)";
if (check)
std::cout << std::setw(10) << "";
std::cout << "\n";
}
size_t max_elems = max_bytes / elem_size;
max_bytes = max_elems * elem_size;
std::vector<char> sendbuf(max_bytes);
// Fill send buffer with a simple deterministic pattern (rank + 1) casted to
// the target type, so correctness checks are straightforward: after all_sum
// every element should equal sum_{r=0}^{nranks-1} (r + 1) =
// nranks*(nranks+1)/2.
auto fill_buffer = [&](char* buf, size_t n_bytes) {
float val = static_cast<float>(rank + 1);
if (dtype == jaccl::Float32) {
auto* p = reinterpret_cast<float*>(buf);
size_t n = n_bytes / sizeof(float);
for (size_t i = 0; i < n; i++) {
p[i] = val;
}
} else if (dtype == jaccl::Float16) {
// Write via the library's float16_t
auto* p = reinterpret_cast<jaccl::float16_t*>(buf);
size_t n = n_bytes / sizeof(jaccl::float16_t);
for (size_t i = 0; i < n; i++) {
p[i] = jaccl::float16_t(val);
}
} else if (dtype == jaccl::BFloat16) {
auto* p = reinterpret_cast<jaccl::bfloat16_t*>(buf);
size_t n = n_bytes / sizeof(jaccl::bfloat16_t);
for (size_t i = 0; i < n; i++) {
p[i] = jaccl::bfloat16_t(val);
}
}
};
auto check_buffer = [&](const char* buf, size_t n_bytes) -> bool {
float expected = static_cast<float>(nranks) * (nranks + 1) / 2.0f;
float tol = (dtype == jaccl::Float32) ? 1e-5f : 1e-1f;
size_t n = n_bytes / elem_size;
for (size_t i = 0; i < n; i++) {
float val;
if (dtype == jaccl::Float32) {
val = reinterpret_cast<const float*>(buf)[i];
} else if (dtype == jaccl::Float16) {
val = static_cast<float>(
reinterpret_cast<const jaccl::float16_t*>(buf)[i]);
} else {
val = static_cast<float>(
reinterpret_cast<const jaccl::bfloat16_t*>(buf)[i]);
}
if (std::abs(val - expected) > tol) {
return false;
}
}
return true;
};
double bf = bus_factor(nranks);
for (size_t nbytes = min_bytes; nbytes <= max_bytes;
nbytes *= static_cast<size_t>(step_factor)) {
// Round down to element boundary
size_t n = std::max((nbytes / elem_size) * elem_size, elem_size);
size_t count = n / elem_size;
fill_buffer(sendbuf.data(), n);
// Warmup
for (int i = 0; i < warmup_iters; i++) {
group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
}
// Timed iterations
auto t0 = std::chrono::high_resolution_clock::now();
for (int i = 0; i < timed_iters; i++) {
group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
}
auto t1 = std::chrono::high_resolution_clock::now();
double elapsed_us =
std::chrono::duration<double, std::micro>(t1 - t0).count();
double avg_us = elapsed_us / timed_iters;
// Bandwidth in GB/s
double algo_bw = (static_cast<double>(n) / avg_us) / 1e3;
double bus_bw = algo_bw * bf;
// Correctness check
std::string check_result;
if (check) {
fill_buffer(sendbuf.data(), n);
group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
check_result = check_buffer(sendbuf.data(), n) ? "OK" : "FAIL";
}
if (rank == 0) {
std::cout << std::left << std::setw(14) << n << std::right
<< std::setw(12) << count << std::setw(12) << dtype_str
<< std::setw(14) << std::fixed << std::setprecision(1) << avg_us
<< std::setw(14) << std::fixed << std::setprecision(2)
<< algo_bw << std::setw(14) << std::fixed
<< std::setprecision(2) << bus_bw;
if (check)
std::cout << std::setw(10) << check_result;
std::cout << "\n";
}
}
if (rank == 0) {
std::cout << "# Done.\n";
}
return 0;
}
mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp
+36
View File
@@ -0,0 +1,36 @@
#include <jaccl/jaccl.h>
#include <iostream>
int main() {
auto cfg =
jaccl::Config()
.set_rank(0) // should be different per node
.set_coordinator("192.168.1.1:32132") // rank 0 will listen here
.set_devices(
{{{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}},
{{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}},
{{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}},
{{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}}});
auto group = jaccl::init(cfg);
if (!group) {
std::cerr << "Failed to initialize JACCL" << std::endl;
return 1;
}
std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
// Perform all-reduce sum
float input[10] = {
1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
float output[10];
group->all_sum(input, output, sizeof(input), jaccl::Float32);
std::cout << "Result: ";
for (auto o : output) {
std::cout << o << " ";
}
std::cout << std::endl;
return 0;
}
mlx/distributed/jaccl/lib/examples/minimal_env.cpp
+28
View File
@@ -0,0 +1,28 @@
#include <jaccl/jaccl.h>
#include <iostream>
int main() {
// Initialize JACCL group
auto group = jaccl::init();
if (!group) {
std::cerr << "Failed to initialize JACCL" << std::endl;
return 1;
}
std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
// Perform all-reduce sum
float input[10] = {
1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
float output[10];
group->all_sum(input, output, sizeof(input), jaccl::Float32);
std::cout << "Result: ";
for (auto o : output) {
std::cout << o << " ";
}
std::cout << std::endl;
return 0;
}
mlx/distributed/jaccl/lib/jaccl/group.h
+58
View File
@@ -0,0 +1,58 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <cstddef>
#include <memory>
namespace jaccl {
/**
* Abstract base class for a JACCL communication group.
*/
class Group {
public:
virtual ~Group() {}
virtual int rank() = 0;
virtual int size() = 0;
virtual void
all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0;
virtual void
all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0;
virtual void
all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0;
virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0;
virtual void send(const void* input, size_t n_bytes, int dst) = 0;
virtual void recv(void* output, size_t n_bytes, int src) = 0;
};
/**
* Type IDs for dispatch in the standalone JACCL library.
*
* Users pass one of these to all_sum/all_max/all_min so JACCL knows how to
* interpret the data for typed reduction operations.
*/
enum Dtype {
Bool = 0,
Int8,
Int16,
Int32,
Int64,
UInt8,
UInt16,
UInt32,
UInt64,
Float16,
BFloat16,
Float32,
Float64,
Complex64,
};
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/jaccl.cpp
+228
View File
@@ -0,0 +1,228 @@
// Copyright © 2025 Apple Inc.
#include <fstream>
#include <sstream>
#include <json.hpp>
#include "jaccl/jaccl.h"
#include "jaccl/mesh.h"
#include "jaccl/rdma.h"
#include "jaccl/ring.h"
using json = nlohmann::json;
namespace {
std::vector<std::vector<std::vector<std::string>>> parse_devices_json(
const char* dev_file) {
std::ifstream f(dev_file);
json devices = json::parse(f);
if (!devices.is_array()) {
throw std::runtime_error(
"[jaccl] The device file should start with an array");
}
std::vector<std::vector<std::vector<std::string>>> result(devices.size());
for (int rank = 0; rank < devices.size(); rank++) {
auto conn = devices[rank];
if (!conn.is_array()) {
throw std::runtime_error(
"[jaccl] The device file should have an array of arrays");
}
if (conn.size() != devices.size()) {
std::ostringstream msg;
msg << "[jaccl] The device file should contain the connectivity of each rank to "
<< "all other ranks but rank " << rank << " contains only "
<< conn.size() << " entries.";
throw std::runtime_error(msg.str());
}
result[rank].resize(conn.size());
for (int dst = 0; dst < conn.size(); dst++) {
auto names = conn[dst];
if (names.is_string()) {
result[rank][dst].push_back(names);
} else if (names.is_array()) {
for (auto name_it = names.begin(); name_it != names.end(); name_it++) {
result[rank][dst].push_back(*name_it);
}
} else if (!names.is_null()) {
throw std::runtime_error(
"[jaccl] Device names should be null, a string or array of strings.");
}
}
}
return result;
}
template <typename First, typename... Rest>
const char* getenv(First first, Rest... rest) {
const char* rs = std::getenv(first);
if (rs != nullptr) {
return rs;
}
if constexpr (sizeof...(rest) > 0) {
return getenv(rest...);
}
return rs;
}
} // namespace
namespace jaccl {
Config::Config() : rank_(0), size_(0) {}
Config& Config::set_rank(int rank) {
rank_ = rank;
return *this;
}
Config& Config::set_coordinator(std::string coordinator) {
coordinator_ = std::move(coordinator);
return *this;
}
Config& Config::set_devices(
std::vector<std::vector<std::vector<std::string>>> devices) {
devices_ = std::move(devices);
size_ = devices_.size();
for (int r = 0; r < size_; r++) {
if (size_ != devices_[r].size()) {
std::ostringstream msg;
msg << "[jaccl] The full connectivity matrix should be provided but we have "
<< size_ << " rows and row " << r << " has " << devices_[r].size()
<< " columns.";
throw std::invalid_argument(msg.str());
}
}
return *this;
}
Config& Config::prefer_ring(bool prefer /* = true */) {
prefer_ring_ = prefer;
return *this;
}
bool Config::is_valid_mesh() const {
if (size_ < 2) {
return false;
}
for (int src = 0; src < size_; src++) {
for (int dst = 0; dst < size_; dst++) {
if ((src == dst && devices_[src][dst].size() != 0) ||
(src != dst && devices_[src][dst].size() == 0)) {
return false;
}
}
}
return true;
}
bool Config::is_valid_ring() const {
if (size_ < 2) {
return false;
}
int num_connections = devices_[0][1].size();
for (int src = 0; src < size_; src++) {
int left = (src + size_ - 1) % size_;
int right = (src + 1) % size_;
for (int dst = 0; dst < size_; dst++) {
if (dst == left || dst == right) {
if (devices_[src][dst].size() != num_connections) {
return false;
}
}
}
}
return true;
}
std::vector<std::string> Config::get_mesh_connectivity() const {
if (!is_valid_mesh()) {
throw std::runtime_error("[jaccl] The devices do not form a valid mesh.");
}
std::vector<std::string> devices(size_);
for (int dst = 0; dst < size_; dst++) {
if (dst != rank_) {
devices[dst] = devices_[rank_][dst][0];
}
}
return devices;
}
std::pair<std::vector<std::string>, std::vector<std::string>>
Config::get_ring_connectivity() const {
if (!is_valid_ring()) {
throw std::runtime_error("[jaccl] The devices do not form a valid ring.");
}
int left = (rank_ + size_ - 1) % size_;
int right = (rank_ + 1) % size_;
return std::make_pair(devices_[rank_][left], devices_[rank_][right]);
}
std::optional<Config> Config::from_env() {
const char* dev_file = getenv("JACCL_IBV_DEVICES", "MLX_IBV_DEVICES");
const char* coordinator =
getenv("JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR");
const char* rank_str = getenv("JACCL_RANK", "MLX_RANK");
const char* ring = getenv("JACCL_RING", "MLX_JACCL_RING");
if (!dev_file || !coordinator || !rank_str) {
return std::nullopt;
}
return Config()
.set_rank(std::atoi(rank_str))
.set_coordinator(coordinator)
.set_devices(parse_devices_json(dev_file))
.prefer_ring(ring != nullptr);
}
bool is_available() {
return ibv().is_available();
}
std::shared_ptr<Group> init(bool strict /* = false */) {
auto cfg = Config::from_env();
if (!cfg.has_value() && strict) {
std::ostringstream msg;
msg << "[jaccl] You need to provide via environment variables a rank "
<< "(JACCL_RANK/MLX_RANK), a device file (JACCL_IBV_DEVICES/"
<< "MLX_IBV_DEVICES) and a coordinator ip/port (JACCL_COORDINATOR/"
<< "MLX_JACCL_COORDINATOR).";
throw std::runtime_error(msg.str());
}
return init(*cfg, strict);
}
std::shared_ptr<Group> init(const Config& cfg, bool strict /* = false */) {
if (cfg.get_prefer_ring() && cfg.is_valid_ring()) {
auto [left, right] = cfg.get_ring_connectivity();
return std::make_shared<RingGroup>(
cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator());
} else if (cfg.is_valid_mesh()) {
auto mesh = cfg.get_mesh_connectivity();
return std::make_shared<MeshGroup>(
cfg.get_rank(), mesh, cfg.get_coordinator());
} else if (cfg.is_valid_ring()) {
auto [left, right] = cfg.get_ring_connectivity();
return std::make_shared<RingGroup>(
cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator());
} else {
if (!strict) {
return nullptr;
}
throw std::runtime_error(
"[jaccl] The configuration should define a valid mesh or a valid ring.");
}
}
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/jaccl.h
+85
View File
@@ -0,0 +1,85 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <memory>
#include <vector>
#include "jaccl/group.h"
namespace jaccl {
class Config {
public:
Config();
Config& set_rank(int rank);
Config& set_coordinator(std::string coordinator);
Config& set_devices(
std::vector<std::vector<std::vector<std::string>>> devices);
Config& prefer_ring(bool prefer = true);
bool is_valid_mesh() const;
bool is_valid_ring() const;
int get_rank() const {
return rank_;
}
int get_size() const {
return size_;
}
std::string get_coordinator() const {
return coordinator_;
}
bool get_prefer_ring() const {
return prefer_ring_;
}
static std::optional<Config> from_env();
friend std::shared_ptr<Group> init(const Config& cfg, bool strict);
private:
std::vector<std::string> get_mesh_connectivity() const;
std::pair<std::vector<std::string>, std::vector<std::string>>
get_ring_connectivity() const;
int rank_;
int size_;
std::string coordinator_;
std::vector<std::vector<std::vector<std::string>>> devices_;
bool prefer_ring_;
};
/**
* Check if JACCL (RDMA over Thunderbolt) is available on this system.
*/
bool is_available();
/**
* Initialize a JACCL communication group from environment variables.
*
* Reads configuration from environment variables:
* - JACCL_RANK / MLX_RANK: The rank of this process
* - JACCL_IBV_DEVICES / MLX_IBV_DEVICES: Path to the device connectivity
* JSON file
* - JACCL_COORDINATOR / MLX_JACCL_COORDINATOR: IP:port of the coordinator
* - JACCL_RING / MLX_JACCL_RING: If set, prefer ring topology
*
* Args:
* strict: If true, throw on failure. If false, return nullptr.
*
* Returns:
* A shared_ptr to the Group, or nullptr on failure.
*/
std::shared_ptr<Group> init(bool strict = false);
/**
* Initialize a JACCL communication group from an explicit Config object.
*/
std::shared_ptr<Group> init(const Config& cfg, bool strict = false);
} // namespace jaccl
mlx/distributed/jaccl/mesh.cpp → mlx/distributed/jaccl/lib/jaccl/mesh.cpp
+61 -74
View File
@@ -1,19 +1,18 @@
// Copyright © 2026 Apple Inc.
#include "mlx/distributed/jaccl/mesh.h"
#include "mlx/backend/cpu/encoder.h"
#include "mlx/distributed/reduction_ops.h"
#include "mlx/dtype_utils.h"
#include "jaccl/mesh.h"
#include "jaccl/reduction_ops.h"
#include "jaccl/types.h"
namespace mlx::core::distributed::jaccl {
namespace jaccl {
MeshGroup::MeshGroup(
int rank,
const std::vector<std::string>& device_names,
const char* coordinator_addr)
const std::string& coordinator_addr)
: rank_(rank),
size_(device_names.size()),
side_channel_(rank_, size_, coordinator_addr),
side_channel_(rank_, size_, coordinator_addr.c_str()),
connections_(create_connections(device_names)) {
if (size_ > MESH_MAX_PEERS) {
std::ostringstream msg;
@@ -26,7 +25,7 @@ MeshGroup::MeshGroup(
initialize();
// Make sure every node has reached here before continuing
side_channel_.all_gather<int>(0);
side_channel_.barrier();
// Create the mesh implementation object
mesh_ = MeshImpl(rank_, size_, connections_, buffers_);
@@ -61,9 +60,9 @@ void MeshGroup::initialize() {
connections_[peer].queue_pair_init();
}
// Gather the information to be exchanged, this also serves as a barrier so
// that all peers have initialized their connections before attempting to
// transition to RTS.
// Gather the information to be exchanged, this also serves as a barrier
// so that all peers have initialized their connections before attempting
// to transition to RTS.
std::vector<Destination> info;
for (auto& conn : connections_) {
info.emplace_back(conn.info());
@@ -106,29 +105,27 @@ void MeshGroup::allocate_buffers() {
for (int i = 0; i < NUM_BUFFERS; i++) {
// Mesh buffers
for (int j = 0; j < size_; j++) {
// This is our send buffer so register it with all pds so we can send
// it to all connected devices.
if (j == rank_) {
// This is our send buffer so register it with all pds so we can
// send it to all connected devices.
for (auto& conn : connections_) {
if (conn.ctx != nullptr) {
buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
.register_to_protection_domain(conn.protection_domain);
}
}
}
// This is the recv buffer from rank j so register it to rank j's
// protection domain.
else {
} else {
// This is the recv buffer from rank j so register it to rank j's
// protection domain.
buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
.register_to_protection_domain(connections_[j].protection_domain);
}
}
// Ring buffers (see ring group for the logic below)
// We register send buffers to both the right and the left.
int left = (rank_ + size_ - 1) % size_;
int right = (rank_ + 1) % size_;
// We register send buffers to both the right and the left.
ring_send_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0]
.register_to_protection_domain(connections_[right].protection_domain);
ring_recv_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0]
@@ -141,78 +138,68 @@ void MeshGroup::allocate_buffers() {
}
}
void MeshGroup::all_sum(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::SumOp<T>{});
void MeshGroup::all_sum(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, SumOp<T>{});
});
}
void MeshGroup::all_max(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
void MeshGroup::all_max(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, MaxOp<T>{});
});
}
void MeshGroup::all_min(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MinOp<T>{});
void MeshGroup::all_min(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, MinOp<T>{});
});
}
void MeshGroup::all_gather(const array& input, array& output, Stream stream) {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
mesh_.all_gather(in_ptr, out_ptr, n_bytes);
});
void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) {
mesh_.all_gather(
static_cast<const char*>(input), static_cast<char*>(output), n_bytes);
}
void MeshGroup::send(const array& input, int dst, Stream stream) {
auto data = input.data<char>();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch(
[data, n_bytes, dst, this]() { mesh_.send(data, n_bytes, dst); });
void MeshGroup::send(const void* input, size_t n_bytes, int dst) {
mesh_.send(static_cast<const char*>(input), n_bytes, dst);
}
void MeshGroup::recv(array& out, int src, Stream stream) {
auto data = out.data<char>();
int64_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch(
[data, n_bytes, src, this]() { mesh_.recv(data, n_bytes, src); });
void MeshGroup::recv(void* output, size_t n_bytes, int src) {
mesh_.recv(static_cast<char*>(output), n_bytes, src);
}
template <typename T, typename ReduceOp>
void MeshGroup::all_reduce(
const array& input,
array& output,
Stream stream,
const void* input,
void* output,
size_t n_bytes,
ReduceOp reduce_op) {
auto in_ptr = input.data<T>();
auto out_ptr = output.data<T>();
int64_t size = input.size();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, size, this, reduce_op]() {
if (size_ > 2 &&
((std::is_same_v<T, bfloat16_t> && size > 65536) ||
size >= 8 * 1024 * 1024 / sizeof(T))) {
ring_.all_reduce<2>(in_ptr, out_ptr, size, 1, reduce_op);
} else {
mesh_.all_reduce(in_ptr, out_ptr, size, reduce_op);
}
});
auto in_ptr = static_cast<const T*>(input);
auto out_ptr = static_cast<T*>(output);
int64_t count = n_bytes / sizeof(T);
if (size_ > 2 &&
((std::is_same_v<T, bfloat16_t> && count > 256 * 1024) ||
count >= 8 * 1024 * 1024 / static_cast<int64_t>(sizeof(T)))) {
ring_.all_reduce<2>(in_ptr, out_ptr, count, 1, reduce_op);
} else {
mesh_.all_reduce(in_ptr, out_ptr, count, reduce_op);
}
}
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/mesh.h → mlx/distributed/jaccl/lib/jaccl/mesh.h
+22 -29
View File
@@ -2,14 +2,12 @@
#pragma once
#include "mlx/distributed/distributed_impl.h"
#include "mlx/distributed/jaccl/mesh_impl.h"
#include "mlx/distributed/jaccl/ring_impl.h"
#include "mlx/distributed/jaccl/utils.h"
#include "jaccl/group.h"
#include "jaccl/mesh_impl.h"
#include "jaccl/rdma.h"
#include "jaccl/ring_impl.h"
using GroupImpl = mlx::core::distributed::detail::GroupImpl;
namespace mlx::core::distributed::jaccl {
namespace jaccl {
/**
* The JACCL communication group for a fully connected mesh. We expect one
@@ -20,16 +18,12 @@ namespace mlx::core::distributed::jaccl {
* information and then configure the connections to be ready for RDMA
* operations.
*/
class MeshGroup : public GroupImpl {
class MeshGroup : public Group {
public:
MeshGroup(
int rank,
const std::vector<std::string>& device_names,
const char* coordinator_addr);
Stream communication_stream(StreamOrDevice s) override {
return to_stream(s, Device::cpu);
}
const std::string& coordinator_addr);
int rank() override {
return rank_;
@@ -39,27 +33,26 @@ class MeshGroup : public GroupImpl {
return size_;
}
void all_sum(const array& input, array& output, Stream stream) override;
void all_max(const array& input, array& output, Stream stream) override;
void all_min(const array& input, array& output, Stream stream) override;
void all_gather(const array& input, array& output, Stream stream) override;
void send(const array& input, int dst, Stream stream) override;
void recv(array& out, int src, Stream stream) override;
void all_sum(const void* input, void* output, size_t n_bytes, int dtype)
override;
void sum_scatter(const array& input, array& output, Stream stream) override {
throw std::runtime_error("[jaccl] sum_scatter not supported.");
}
void all_max(const void* input, void* output, size_t n_bytes, int dtype)
override;
std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
throw std::runtime_error("[jaccl] Group split not supported.");
}
void all_min(const void* input, void* output, size_t n_bytes, int dtype)
override;
void all_gather(const void* input, void* output, size_t n_bytes) override;
void send(const void* input, size_t n_bytes, int dst) override;
void recv(void* output, size_t n_bytes, int src) override;
private:
template <typename T, typename ReduceOp>
void all_reduce(
const array& input,
array& output,
Stream stream,
const void* input,
void* output,
size_t n_bytes,
ReduceOp reduce_op);
/**
@@ -86,4 +79,4 @@ class MeshGroup : public GroupImpl {
RingImpl ring_;
};
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/mesh_impl.h → mlx/distributed/jaccl/lib/jaccl/mesh_impl.h
+106 -44
View File
@@ -2,13 +2,16 @@
#pragma once
#include <memory>
#include <span>
#include "mlx/distributed/jaccl/utils.h"
#include "jaccl/rdma.h"
constexpr int MESH_MAX_PEERS = 8;
constexpr int MESH_PIPELINE = 2;
constexpr int64_t MAX_BUFFER_SIZE = FRAME_SIZE * (1 << (BUFFER_SIZES - 1));
namespace mlx::core::distributed::jaccl {
namespace jaccl {
class MeshImpl {
public:
@@ -17,20 +20,25 @@ class MeshImpl {
int size,
std::vector<Connection>& conns,
std::vector<SharedBuffer>& buffers)
: rank_(rank), size_(size), connections_(conns), buffers_(buffers) {}
: rank_(rank),
size_(size),
connections_(conns),
buffers_(buffers),
staging_mem_(
std::make_unique<char[]>(MESH_PIPELINE * MAX_BUFFER_SIZE)) {}
MeshImpl() : rank_(0), size_(1) {}
template <typename T, typename ReduceOp>
void
all_reduce(const T* in_ptr, T* out_ptr, int64_t size, ReduceOp reduce_op) {
// If not inplace all reduce then copy the input to the output first
if (in_ptr != out_ptr) {
std::memcpy(out_ptr, in_ptr, size * sizeof(T));
}
void all_reduce(const T* in, T* out, int64_t size, ReduceOp reduce_op) {
// Fully connected all reduce with deterministic reduction order.
//
// We copy rank 0's data to the output buffer and then we reduce every
// subsequent rank in-place in the output.
//
// Our own data is copied to a staging buffer to ensure we can reduce it in
// the output when needed.
// Fully connected all reduce
T* data = out_ptr;
auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T));
int64_t N = buffer_size / sizeof(T);
constexpr int PIPELINE = 2;
@@ -38,21 +46,36 @@ class MeshImpl {
int64_t total = static_cast<int64_t>(size);
int num_peers = size_ - 1;
// A helper for convenient access to the staging buffer.
auto local_staging = [&](int buff) -> T* {
return reinterpret_cast<T*>(staging_mem_.get() + buff * MAX_BUFFER_SIZE);
};
// Counters to maintain the state of transfers
int in_flight = 0;
int64_t read_offset = 0;
int completed_send_count[PIPELINE] = {0};
int completed_recv_begin[MESH_MAX_PEERS] = {0};
int completed_recv_end[MESH_MAX_PEERS] = {0};
int recv_end[MESH_MAX_PEERS] = {0};
int reduce_chunk = 0;
int reduce_rank = 0;
// Total number of chunks
int64_t total_chunks = (total + N - 1) / N;
// Prefill the pipeline
int buff = 0;
while (read_offset < total && buff < PIPELINE) {
post_recv_all(sz, buff);
// Copy the local data to send buffer and staging buffer
int64_t elems = std::min(N, total - read_offset);
std::copy(
data + read_offset,
data + std::min(read_offset + N, total),
in + read_offset, in + read_offset + elems, local_staging(buff));
std::copy(
in + read_offset,
in + read_offset + elems,
send_buffer(sz, buff).begin<T>());
recv_end[rank_]++;
post_send_all(sz, buff);
buff++;
@@ -61,14 +84,14 @@ class MeshImpl {
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
while (reduce_chunk < total_chunks) {
// Poll the hardware for completions.
//
// If a send was completed mark how many completions we have received
// for that buffer. If we have sent the buffer to all peers we can
// reuse the buffer so copy the next chunk of data and send it to all.
// Also copy the next chunk into the staging area and advance our
// completed "receives".
//
// If a receive is completed then advance the pointer of completed
// receives.
@@ -84,10 +107,16 @@ class MeshImpl {
if (work_type == SEND_WR && read_offset < total) {
completed_send_count[buff]++;
if (completed_send_count[buff] == num_peers) {
int64_t elems = std::min(N, total - read_offset);
std::copy(
data + read_offset,
data + std::min(read_offset + N, total),
in + read_offset,
in + read_offset + elems,
local_staging(buff));
std::copy(
in + read_offset,
in + read_offset + elems,
send_buffer(sz, buff).begin<T>());
recv_end[rank_]++;
post_send_all(sz, buff);
completed_send_count[buff] = 0;
@@ -97,38 +126,70 @@ class MeshImpl {
}
else if (work_type == RECV_WR) {
completed_recv_end[rank]++;
recv_end[rank]++;
}
}
// Process the completed recv
// Process the received chunks in order.
//
// For each rank we have a range of completed recv defined by a begin
// and end inclusive and exlusive in standard C++ fashion.
//
// When there is an unprocessed receive we first check if we have
// finished sending the write location. If so then we reduce in-place
// and then check if there is more to be received and post a recv.
for (int r = 0; r < size_; r++) {
int s = completed_recv_begin[r];
int e = completed_recv_end[r];
int w = s * N;
while (w < read_offset && e - s > 0) {
int buff = s % PIPELINE;
reduce_op(
recv_buffer(sz, buff, r).begin<T>(),
data + w,
std::min(N, total - w));
w += N;
s++;
if (w + (PIPELINE - 1) * N < total) {
recv_from(sz, r, buff);
// Rank 0 is always copied as is. Our rank is always read from the
// staging area.
while (reduce_chunk < total_chunks) {
// w is our write location so break if it is ahead of the read location.
int64_t w = static_cast<int64_t>(reduce_chunk) * N;
if (w >= read_offset) {
break;
}
// We want to reduce the 'reduce_chunk' chunk but it hasn't arrived
// yet.
if (recv_end[reduce_rank] <= reduce_chunk) {
break;
}
int b = reduce_chunk % PIPELINE;
int64_t elems = std::min(N, total - w);
// Data is read from the staging area
if (reduce_rank == rank_) {
if (reduce_rank == 0) {
std::copy_n(local_staging(b), elems, out + w);
} else {
reduce_op(local_staging(b), out + w, elems);
}
}
// Data is read from the recv buffers
else {
if (reduce_rank == 0) {
std::copy_n(
recv_buffer(sz, b, reduce_rank).begin<T>(), elems, out + w);
} else {
reduce_op(
recv_buffer(sz, b, reduce_rank).begin<T>(), out + w, elems);
}
// Check if we need to post another receive
int64_t next_chunk = static_cast<int64_t>(reduce_chunk) + PIPELINE;
if (next_chunk < total_chunks) {
recv_from(sz, reduce_rank, b);
in_flight++;
}
}
completed_recv_begin[r] = s;
// Means we processed that chunk so move to the next one
reduce_rank++;
if (reduce_rank >= size_) {
reduce_rank = 0;
reduce_chunk++;
}
}
}
// Drain remaining in-flight completions (outstanding sends).
while (in_flight > 0) {
ibv_wc wc[WC_NUM];
int n = poll(connections_, WC_NUM, wc);
in_flight -= n;
}
}
void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) {
@@ -353,6 +414,7 @@ class MeshImpl {
int size_;
std::span<Connection> connections_;
std::span<SharedBuffer> buffers_;
std::unique_ptr<char[]> staging_mem_;
};
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/utils.cpp → mlx/distributed/jaccl/lib/jaccl/rdma.cpp
+17 -8
View File
@@ -5,7 +5,7 @@
#include <iostream>
#include <sstream>
#include "mlx/distributed/jaccl/utils.h"
#include "jaccl/rdma.h"
#define LOAD_SYMBOL(symbol, variable) \
{ \
@@ -31,7 +31,7 @@ void* page_aligned_alloc(size_t num_bytes) {
} // namespace
namespace mlx::core::distributed::jaccl {
namespace jaccl {
IBVWrapper::IBVWrapper() {
librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL);
@@ -182,11 +182,20 @@ const Destination& Connection::info() {
ibv_port_attr port_attr;
ibv().query_port(ctx, 1, &port_attr);
ibv_gid gid;
ibv().query_gid(ctx, 1, 1, &gid);
for (int i = 0; i < port_attr.gid_tbl_len; i++) {
ibv_gid tmp;
if (ibv().query_gid(ctx, 1, i, &tmp) == 0) {
if (*(uint64_t*)&tmp.raw[0] == 0 && *(uint16_t*)&tmp.raw[8] == 0 &&
*(uint16_t*)&tmp.raw[10] == 0xffff) {
gid = tmp;
break;
}
}
}
src.local_id = port_attr.lid;
src.queue_pair_number = queue_pair->qp_num;
src.packet_sequence_number = 7; // TODO: Change to sth random
src.packet_sequence_number = 7;
src.global_identifier = gid;
return src;
@@ -287,10 +296,10 @@ std::vector<Connection> create_connections(
SideChannel::SideChannel(int rank, int size, const char* addr)
: rank_(rank), size_(size) {
auto address = detail::parse_address(addr);
auto address = parse_address(addr);
if (rank_ == 0) {
detail::TCPSocket server(IBV_TAG);
TCPSocket server(IBV_TAG);
server.listen(IBV_TAG, address);
for (int i = 0; i < size - 1; i++) {
@@ -311,7 +320,7 @@ SideChannel::SideChannel(int rank, int size, const char* addr)
}
} else {
sockets_.push_back(
detail::TCPSocket::connect(
TCPSocket::connect(
IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
std::cerr << IBV_TAG << " Connection attempt " << attempt
<< " waiting " << wait << " ms" << std::endl;
@@ -326,4 +335,4 @@ SideChannel::SideChannel(SideChannel&& sc)
sc.size_ = -1;
}
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/utils.h → mlx/distributed/jaccl/lib/jaccl/rdma.h
+11 -6
View File
@@ -5,10 +5,11 @@
#include <infiniband/verbs.h>
#include <span>
#include <sstream>
#include <unordered_map>
#include <vector>
#include "mlx/distributed/utils.h"
#include "jaccl/tcp.h"
constexpr const char* IBV_TAG = "[jaccl]";
constexpr int SEND_WR = 1;
@@ -19,8 +20,6 @@ constexpr int BUFFER_SIZES = 8;
constexpr int NUM_BUFFERS = 2;
constexpr int FRAME_SIZE = 4096;
namespace detail = mlx::core::distributed::detail;
namespace {
template <typename T, typename = void>
@@ -45,7 +44,7 @@ inline std::pair<int, int64_t> buffer_size_from_message(int64_t msg) {
} // namespace
namespace mlx::core::distributed::jaccl {
namespace jaccl {
/**
* Wrapper for the ibverbs API.
@@ -334,10 +333,16 @@ class SideChannel {
return result;
}
void barrier() {
// Twice has proven to be more robust to initialization issues.
all_gather<int>(0);
all_gather<int>(0);
}
private:
int rank_;
int size_;
std::vector<detail::TCPSocket> sockets_;
std::vector<TCPSocket> sockets_;
};
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/reduction_ops.h
+117
View File
@@ -0,0 +1,117 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <cstddef>
#include "jaccl/types.h"
namespace jaccl {
template <typename T>
struct SumOp {
void operator()(const T* input, T* output, size_t N) const {
for (size_t i = 0; i < N; i++) {
output[i] = output[i] + input[i];
}
}
};
template <typename T>
struct MaxOp {
void operator()(const T* input, T* output, size_t N) const {
for (size_t i = 0; i < N; i++) {
output[i] = (output[i] > input[i]) ? output[i] : input[i];
}
}
};
template <typename T>
struct MinOp {
void operator()(const T* input, T* output, size_t N) const {
for (size_t i = 0; i < N; i++) {
output[i] = (output[i] < input[i]) ? output[i] : input[i];
}
}
};
//
// The last piece of the puzzle to use the native bf16 while compiling a single
// binary for all Macs is to compile these functions with
// target("arch=armv8.6-a").
//
// Now we can simply check in runtime and call them only when they are
// supported.
//
#if defined(__aarch64__)
__attribute__((target("arch=armv8.6-a"))) inline void
native_bf16_sum(const void* input, void* output, size_t N) {
auto in = reinterpret_cast<const __bf16*>(input);
auto out = reinterpret_cast<__bf16*>(output);
for (size_t i = 0; i < N; i++) {
out[i] = out[i] + in[i];
}
}
__attribute__((target("arch=armv8.6-a"))) inline void
native_bf16_max(const void* input, void* output, size_t N) {
auto in = reinterpret_cast<const __bf16*>(input);
auto out = reinterpret_cast<__bf16*>(output);
for (size_t i = 0; i < N; i++) {
out[i] = (out[i] > in[i]) ? out[i] : in[i];
}
}
__attribute__((target("arch=armv8.6-a"))) inline void
native_bf16_min(const void* input, void* output, size_t N) {
auto in = reinterpret_cast<const __bf16*>(input);
auto out = reinterpret_cast<__bf16*>(output);
for (size_t i = 0; i < N; i++) {
out[i] = (out[i] < in[i]) ? out[i] : in[i];
}
}
template <>
struct SumOp<bfloat16_t> {
void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
if (has_native_bf16_support()) {
native_bf16_sum(input, output, N);
} else {
for (size_t i = 0; i < N; i++) {
output[i] = output[i] + input[i];
}
}
}
};
template <>
struct MaxOp<bfloat16_t> {
void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
if (has_native_bf16_support()) {
native_bf16_max(input, output, N);
} else {
for (size_t i = 0; i < N; i++) {
output[i] = (output[i] > input[i]) ? output[i] : input[i];
}
}
}
};
template <>
struct MinOp<bfloat16_t> {
void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
if (has_native_bf16_support()) {
native_bf16_min(input, output, N);
} else {
for (size_t i = 0; i < N; i++) {
output[i] = (output[i] < input[i]) ? output[i] : input[i];
}
}
}
};
#endif // defined(__aarch64__)
} // namespace jaccl
mlx/distributed/jaccl/ring.cpp → mlx/distributed/jaccl/lib/jaccl/ring.cpp
+61 -73
View File
@@ -1,22 +1,21 @@
// Copyright © 2026 Apple Inc.
#include "mlx/distributed/jaccl/ring.h"
#include "mlx/backend/cpu/encoder.h"
#include "mlx/distributed/reduction_ops.h"
#include "mlx/dtype_utils.h"
#include "jaccl/ring.h"
#include "jaccl/reduction_ops.h"
#include "jaccl/types.h"
namespace mlx::core::distributed::jaccl {
namespace jaccl {
RingGroup::RingGroup(
int rank,
int size,
const std::vector<std::string>& left_devices,
const std::vector<std::string>& right_devices,
const char* coordinator_addr)
const std::string& coordinator_addr)
: rank_(rank),
size_(size),
n_conns_(left_devices.size()),
side_channel_(rank_, size_, coordinator_addr),
side_channel_(rank_, size_, coordinator_addr.c_str()),
left_(create_connections(left_devices)),
right_(create_connections(right_devices)) {
if (left_.size() > RING_MAX_CONNS || right_.size() > RING_MAX_CONNS) {
@@ -29,8 +28,8 @@ RingGroup::RingGroup(
// Initialize all the connections and allocate buffers
initialize();
// Make sure every node has reached here before continuing
side_channel_.all_gather<int>(0);
// Make sure every node has reached here before continuing.
side_channel_.barrier();
// Create the ring implementation object
ring_ = RingImpl(rank_, size_, left_, right_, send_buffers_, recv_buffers_);
@@ -60,9 +59,9 @@ void RingGroup::initialize() {
conn.queue_pair_init();
}
// Gather the information to be exchanged, this also serves as a barrier so
// that all peers have initialized their connections before attempting to
// transition to RTS.
// Gather the information to be exchanged, this also serves as a barrier
// so that all peers have initialized their connections before attempting
// to transition to RTS.
std::vector<Destination> left_info;
for (auto& conn : left_) {
left_info.emplace_back(conn.info());
@@ -126,40 +125,48 @@ void RingGroup::allocate_buffers() {
}
}
void RingGroup::all_sum(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::SumOp<T>{});
void RingGroup::all_sum(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, SumOp<T>{});
});
}
void RingGroup::all_max(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
void RingGroup::all_max(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, MaxOp<T>{});
});
}
void RingGroup::all_min(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MinOp<T>{});
void RingGroup::all_min(
const void* input,
void* output,
size_t n_bytes,
int dtype) {
dispatch_all_types(dtype, [&](auto type_tag) {
using T = JACCL_GET_TYPE(type_tag);
all_reduce<T>(input, output, n_bytes, MinOp<T>{});
});
}
void RingGroup::all_gather(const array& input, array& output, Stream stream) {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
ring_.all_gather(in_ptr, out_ptr, n_bytes, n_conns_);
});
void RingGroup::all_gather(const void* input, void* output, size_t n_bytes) {
ring_.all_gather(
static_cast<const char*>(input),
static_cast<char*>(output),
n_bytes,
n_conns_);
}
void RingGroup::send(const array& input, int dst, Stream stream) {
void RingGroup::send(const void* input, size_t n_bytes, int dst) {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (dst != right && dst != left) {
@@ -168,16 +175,10 @@ void RingGroup::send(const array& input, int dst, Stream stream) {
<< "but tried to send to " << dst << " from " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
auto data = input.data<char>();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch([data, n_bytes, dst, this]() {
ring_.send(data, n_bytes, dst, n_conns_);
});
ring_.send(static_cast<const char*>(input), n_bytes, dst, n_conns_);
}
void RingGroup::recv(array& out, int src, Stream stream) {
void RingGroup::recv(void* output, size_t n_bytes, int src) {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (src != right && src != left) {
@@ -186,42 +187,29 @@ void RingGroup::recv(array& out, int src, Stream stream) {
<< "but tried to recv from " << src << " to " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
auto data = out.data<char>();
int64_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch([data, n_bytes, src, this]() {
ring_.recv(data, n_bytes, src, n_conns_);
});
ring_.recv(static_cast<char*>(output), n_bytes, src, n_conns_);
}
template <typename T, typename ReduceOp>
void RingGroup::all_reduce(
const array& input,
array& output,
Stream stream,
const void* input,
void* output,
size_t n_bytes,
ReduceOp reduce_op) {
auto in_ptr = input.data<T>();
auto out_ptr = output.data<T>();
int64_t size = input.size();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, size, n_bytes, this, reduce_op]() {
if (size < size_ * 2 * n_conns_) {
ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
return;
}
auto in_ptr = static_cast<const T*>(input);
auto out_ptr = static_cast<T*>(output);
int64_t count = n_bytes / sizeof(T);
if (count < size_ * 2 * n_conns_) {
ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op);
return;
}
if (n_bytes <= 65536) {
ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
return;
}
if (n_bytes <= 65536) {
ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op);
return;
}
ring_.all_reduce<2, T, ReduceOp>(
in_ptr, out_ptr, size, n_conns_, reduce_op);
});
ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, n_conns_, reduce_op);
}
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/ring.h → mlx/distributed/jaccl/lib/jaccl/ring.h
+21 -28
View File
@@ -2,13 +2,11 @@
#pragma once
#include "mlx/distributed/distributed_impl.h"
#include "mlx/distributed/jaccl/ring_impl.h"
#include "mlx/distributed/jaccl/utils.h"
#include "jaccl/group.h"
#include "jaccl/rdma.h"
#include "jaccl/ring_impl.h"
using GroupImpl = mlx::core::distributed::detail::GroupImpl;
namespace mlx::core::distributed::jaccl {
namespace jaccl {
/**
* The JACCL communication group for a ring where each node is connected to its
@@ -19,18 +17,14 @@ namespace mlx::core::distributed::jaccl {
* information and then configure the connections to be ready for RDMA
* operations.
*/
class RingGroup : public GroupImpl {
class RingGroup : public Group {
public:
RingGroup(
int rank,
int size,
const std::vector<std::string>& left_devices,
const std::vector<std::string>& right_devices,
const char* coordinator_addr);
Stream communication_stream(StreamOrDevice s) override {
return to_stream(s, Device::cpu);
}
const std::string& coordinator_addr);
int rank() override {
return rank_;
@@ -40,27 +34,26 @@ class RingGroup : public GroupImpl {
return size_;
}
void all_sum(const array& input, array& output, Stream stream) override;
void all_max(const array& input, array& output, Stream stream) override;
void all_min(const array& input, array& output, Stream stream) override;
void all_gather(const array& input, array& output, Stream stream) override;
void send(const array& input, int dst, Stream stream) override;
void recv(array& out, int src, Stream stream) override;
void all_sum(const void* input, void* output, size_t n_bytes, int dtype)
override;
void sum_scatter(const array& input, array& output, Stream stream) override {
throw std::runtime_error("[jaccl] sum_scatter not supported.");
}
void all_max(const void* input, void* output, size_t n_bytes, int dtype)
override;
std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
throw std::runtime_error("[jaccl] Group split not supported.");
}
void all_min(const void* input, void* output, size_t n_bytes, int dtype)
override;
void all_gather(const void* input, void* output, size_t n_bytes) override;
void send(const void* input, size_t n_bytes, int dst) override;
void recv(void* output, size_t n_bytes, int src) override;
private:
template <typename T, typename ReduceOp>
void all_reduce(
const array& input,
array& output,
Stream stream,
const void* input,
void* output,
size_t n_bytes,
ReduceOp reduce_op);
/**
@@ -86,4 +79,4 @@ class RingGroup : public GroupImpl {
RingImpl ring_;
};
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/ring_impl.h → mlx/distributed/jaccl/lib/jaccl/ring_impl.h
+3 -3
View File
@@ -4,11 +4,11 @@
#include <span>
#include "mlx/distributed/jaccl/utils.h"
#include "jaccl/rdma.h"
constexpr int RING_MAX_CONNS = 4;
namespace mlx::core::distributed::jaccl {
namespace jaccl {
class RingImpl {
public:
@@ -628,4 +628,4 @@ class RingImpl {
std::span<SharedBuffer> recv_buffers_;
};
} // namespace mlx::core::distributed::jaccl
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/tcp.cpp
+206
View File
@@ -0,0 +1,206 @@
// Copyright © 2025 Apple Inc.
#include <netdb.h>
#include <unistd.h>
#include <cstring>
#include <sstream>
#include <thread>
#include "jaccl/tcp.h"
namespace jaccl {
/**
* Parse a sockaddr from an ip and port provided as strings.
*/
address_t parse_address(const std::string& ip, const std::string& port) {
struct addrinfo hints, *res;
std::memset(&hints, 0, sizeof(hints));
hints.ai_family = AF_UNSPEC;
hints.ai_socktype = SOCK_STREAM;
int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res);
if (status != 0) {
std::ostringstream msg;
msg << "Can't parse address " << ip << ":" << port;
throw std::runtime_error(msg.str());
}
address_t result;
memcpy(&result.addr, res->ai_addr, res->ai_addrlen);
result.len = res->ai_addrlen;
freeaddrinfo(res);
return result;
}
/**
* Parse a sockaddr provided as an <ip>:<port> string.
*/
address_t parse_address(const std::string& ip_port) {
auto colon = ip_port.find(":");
if (colon == std::string::npos) {
std::ostringstream msg;
msg << "Can't parse address " << ip_port;
throw std::runtime_error(msg.str());
}
std::string ip(ip_port.begin(), ip_port.begin() + colon);
std::string port(ip_port.begin() + colon + 1, ip_port.end());
return parse_address(ip, port);
}
TCPSocket::TCPSocket(const char* tag) {
sock_ = socket(AF_INET, SOCK_STREAM, 0);
if (sock_ < 0) {
std::ostringstream msg;
msg << tag << " Couldn't create socket (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
}
TCPSocket::TCPSocket(TCPSocket&& s) {
sock_ = s.sock_;
s.sock_ = -1;
}
TCPSocket& TCPSocket::operator=(TCPSocket&& s) {
if (this != &s) {
sock_ = s.sock_;
s.sock_ = -1;
}
return *this;
}
TCPSocket::TCPSocket(int s) : sock_(s) {}
TCPSocket::~TCPSocket() {
if (sock_ > 0) {
shutdown(sock_, 2);
close(sock_);
}
}
int TCPSocket::detach() {
int s = sock_;
sock_ = -1;
return s;
}
void TCPSocket::listen(const char* tag, const address_t& addr) {
int success;
// Make sure we can launch immediately after shutdown by setting the
// reuseaddr option so that we don't get address already in use errors
int enable = 1;
success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
if (success < 0) {
std::ostringstream msg;
msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int));
if (success < 0) {
std::ostringstream msg;
msg << tag << " Couldn't enable reuseport (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
// Bind the socket to the address and port
success = bind(sock_, addr.get(), addr.len);
if (success < 0) {
std::ostringstream msg;
msg << tag << " Couldn't bind socket (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
// Prepare waiting for connections
success = ::listen(sock_, 0);
if (success < 0) {
std::ostringstream msg;
msg << tag << " Couldn't listen (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
}
TCPSocket TCPSocket::accept(const char* tag) {
int peer = ::accept(sock_, nullptr, nullptr);
if (peer < 0) {
std::ostringstream msg;
msg << tag << " Accept failed (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
return TCPSocket(peer);
}
void TCPSocket::send(const char* tag, const void* data, size_t len) {
while (len > 0) {
auto n = ::send(sock_, data, len, 0);
if (n <= 0) {
std::ostringstream msg;
msg << tag << " Send failed with errno=" << errno;
throw std::runtime_error(msg.str());
}
len -= n;
data = static_cast<const char*>(data) + n;
}
}
void TCPSocket::recv(const char* tag, void* data, size_t len) {
while (len > 0) {
auto n = ::recv(sock_, data, len, 0);
if (n <= 0) {
std::ostringstream msg;
msg << tag << " Recv failed with errno=" << errno;
throw std::runtime_error(msg.str());
}
len -= n;
data = static_cast<char*>(data) + n;
}
}
TCPSocket TCPSocket::connect(
const char* tag,
const address_t& addr,
int num_retries,
int wait,
std::function<void(int, int)> cb) {
int sock, success;
// Attempt to connect `num_retries` times with exponential backoff.
for (int attempt = 0; attempt < num_retries; attempt++) {
// Create the socket
sock = socket(AF_INET, SOCK_STREAM, 0);
if (sock < 0) {
std::ostringstream msg;
msg << tag << " Couldn't create socket to connect (error: " << errno
<< ")";
throw std::runtime_error(msg.str());
}
success = ::connect(sock, addr.get(), addr.len);
if (success == 0) {
break;
}
if (cb != nullptr) {
cb(attempt, wait);
}
if (wait > 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(wait));
}
wait <<= 1;
}
if (success < 0) {
std::ostringstream msg;
msg << tag << " Couldn't connect (error: " << errno << ")";
throw std::runtime_error(msg.str());
}
return TCPSocket(sock);
}
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/tcp.h
+67
View File
@@ -0,0 +1,67 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <sys/socket.h>
#include <functional>
#include <string>
namespace jaccl {
struct address_t {
sockaddr_storage addr;
socklen_t len;
const sockaddr* get() const {
return (struct sockaddr*)&addr;
}
};
/**
* Parse a sockaddr from an ip and port provided as strings.
*/
address_t parse_address(const std::string& ip, const std::string& port);
/**
* Parse a sockaddr provided as an <ip>:<port> string.
*/
address_t parse_address(const std::string& ip_port);
/**
* Small wrapper over a TCP socket to simplify initiating connections.
*/
class TCPSocket {
public:
TCPSocket(const char* tag);
TCPSocket(const TCPSocket&) = delete;
TCPSocket& operator=(const TCPSocket&) = delete;
TCPSocket(TCPSocket&& s);
TCPSocket& operator=(TCPSocket&&);
~TCPSocket();
void listen(const char* tag, const address_t& addr);
TCPSocket accept(const char* tag);
void send(const char* tag, const void* data, size_t len);
void recv(const char* tag, void* data, size_t len);
int detach();
operator int() const {
return sock_;
}
static TCPSocket connect(
const char* tag,
const address_t& addr,
int num_retries = 1,
int wait = 0,
std::function<void(int, int)> cb = nullptr);
private:
TCPSocket(int sock);
int sock_;
};
} // namespace jaccl
mlx/distributed/jaccl/lib/jaccl/types.h
+263
View File
@@ -0,0 +1,263 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <complex>
#include <cstdint>
#include <type_traits>
#if defined(__aarch64__) && defined(__APPLE__)
#include <sys/sysctl.h>
#endif
namespace jaccl {
namespace {
union float_bits {
float f;
uint32_t u;
};
} // namespace
#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
#include <arm_fp16.h>
#else
#define __JACCL_HALF_NAN__ 0x7D00
//
// The MLX float16 compatibility fallback. Redefined here to keep JACCL
// standalone. Also it doesn't really need all the ops MXL defines.
//
struct float16_t {
uint16_t bits_;
float16_t(const float& x) : bits_(0) {
float_bits in;
in.f = x;
// Extract sign
uint32_t x_sign_32 = in.u & uint32_t(0x80000000);
uint16_t x_sign_16 = (x_sign_32 >> 16);
if (std::isnan(x)) {
bits_ = x_sign_16 | uint16_t(__JACCL_HALF_NAN__);
} else {
// Union
float_bits inf_scale, zero_scale, magic_bits;
// Find exponent bits and take the max supported by half
uint32_t x_expo_32 = in.u & uint32_t(0x7f800000);
uint32_t max_expo_32 = uint32_t(0x38800000);
x_expo_32 = x_expo_32 < max_expo_32 ? max_expo_32 : x_expo_32;
x_expo_32 += uint32_t(15) << 23;
// Handle scaling to inf as needed
inf_scale.u = uint32_t(0x77800000);
zero_scale.u = uint32_t(0x08800000);
// Combine with magic and let addition do rounding
magic_bits.u = x_expo_32;
magic_bits.f += (std::abs(x) * inf_scale.f) * zero_scale.f;
// Take the lower 5 bits of the exponent
uint32_t x_expo_16 = ((magic_bits.u >> 13) & uint32_t(0x7c00));
// Collect the lower 12 bits which have the mantissa
uint32_t x_mant_16 = magic_bits.u & uint32_t(0x0fff);
// Combine sign, exp and mantissa
bits_ = (x_sign_16 | uint16_t(x_expo_16 + x_mant_16));
}
}
operator float() const {
float_bits out;
uint32_t x_sign_32 = (bits_ << 16) & uint32_t(0x80000000);
uint32_t base = (bits_ << 16);
uint32_t two_base = base + base;
uint32_t denorm_max = 1u << 27;
if (two_base < denorm_max) {
out.u = uint32_t(126) << 23; // magic mask
out.u |= (two_base >> 17); // Bits from fp16
out.f -= 0.5f; // magic bias
} else {
out.u = uint32_t(0xE0) << 23; // exponent offset
out.u += (two_base >> 4); // Bits from fp16
float out_unscaled = out.f; // Store value
out.u = uint32_t(0x7800000); // exponent scale
out.f *= out_unscaled;
}
// Add sign
out.u |= x_sign_32;
return out.f;
}
bool operator<(float16_t x) {
return static_cast<float>(*this) < static_cast<float>(x);
}
bool operator>(float16_t x) {
return static_cast<float>(*this) > static_cast<float>(x);
}
float16_t operator+(float16_t x) {
return static_cast<float>(*this) + static_cast<float>(x);
}
float16_t& operator+=(float16_t x) {
*this = *this + x;
return *this;
}
};
#endif
//
// Check at runtime if the CPU supports native bfloat16 (FEAT_BF16).
//
// This allows us to compile once for all Macs but still enable the feature if
// it is supported.
//
inline bool has_native_bf16_support() {
#if defined(__aarch64__) && defined(__APPLE__)
static bool has_support = []() {
int value = 0;
size_t value_size = sizeof(value);
int success = sysctlbyname(
"hw.optional.arm.FEAT_BF16", &value, &value_size, nullptr, 0);
return success == 0 & value != 0;
}();
return has_support;
#else
return false;
#endif
}
//
// The MLX bfloat16 compatibility fallback.
//
// Contrary to float16 we always define it and we 'll use
// has_native_bf16_support to decide if we are going to use __bf16 instead
// during runtime.
//
#define __JACCL_BFLOAT_NAN__ 0x7FC0
struct bfloat16_t {
uint16_t bits_;
bfloat16_t(const float& x) {
if (std::isnan(x)) {
bits_ = __JACCL_BFLOAT_NAN__;
} else {
float_bits in;
in.f = x;
in.u += (in.u >> 16 & 1) + uint32_t(0x7FFF);
bits_ = in.u >> 16;
}
}
operator float() const {
float_bits out;
out.u = ((uint32_t)bits_) << 16;
return out.f;
}
bool operator<(bfloat16_t x) {
return static_cast<float>(*this) < static_cast<float>(x);
}
bool operator>(bfloat16_t x) {
return static_cast<float>(*this) > static_cast<float>(x);
}
bfloat16_t operator+(bfloat16_t x) {
return static_cast<float>(*this) + static_cast<float>(x);
}
bfloat16_t& operator+=(bfloat16_t x) {
*this = *this + x;
return *this;
}
};
using complex64_t = std::complex<float>;
inline bool operator<(complex64_t lhs, complex64_t rhs) {
return lhs.real() < rhs.real() ||
(lhs.real() == rhs.real() && lhs.imag() < rhs.imag());
}
inline bool operator>(complex64_t lhs, complex64_t rhs) {
return lhs.real() > rhs.real() ||
(lhs.real() == rhs.real() && lhs.imag() > rhs.imag());
}
template <typename T>
struct type_identity {
using type = T;
};
#define JACCL_GET_TYPE(x) typename decltype(x)::type
/**
* Dispatch a function for all supported types based on a Dtype.
*/
template <typename F>
void dispatch_all_types(int dtype, F&& f) {
switch (dtype) {
case Dtype::Bool:
f(type_identity<bool>{});
break;
case Dtype::Int8:
f(type_identity<int8_t>{});
break;
case Dtype::Int16:
f(type_identity<int16_t>{});
break;
case Dtype::Int32:
f(type_identity<int32_t>{});
break;
case Dtype::Int64:
f(type_identity<int64_t>{});
break;
case Dtype::UInt8:
f(type_identity<uint8_t>{});
break;
case Dtype::UInt16:
f(type_identity<uint16_t>{});
break;
case Dtype::UInt32:
f(type_identity<uint32_t>{});
break;
case Dtype::UInt64:
f(type_identity<uint64_t>{});
break;
case Dtype::Float16:
f(type_identity<float16_t>{});
break;
case Dtype::BFloat16:
f(type_identity<bfloat16_t>{});
break;
case Dtype::Float32:
f(type_identity<float>{});
break;
case Dtype::Float64:
f(type_identity<double>{});
break;
case Dtype::Complex64:
f(type_identity<complex64_t>{});
break;
}
}
} // namespace jaccl