diff --git a/mlx/distributed/jaccl/lib/examples/CMakeLists.txt b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt index 48b7522d..9efbb513 100644 --- a/mlx/distributed/jaccl/lib/examples/CMakeLists.txt +++ b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt @@ -35,6 +35,8 @@ endfunction() # Examples build_example(minimal_env.cpp) build_example(minimal_cfg.cpp) +build_example(monte_carlo_pi.cpp) +build_example(file_broadcast.cpp) # Benchmarks build_example(allreduce_bench.cpp) diff --git a/mlx/distributed/jaccl/lib/examples/file_broadcast.cpp b/mlx/distributed/jaccl/lib/examples/file_broadcast.cpp new file mode 100644 index 00000000..4106a373 --- /dev/null +++ b/mlx/distributed/jaccl/lib/examples/file_broadcast.cpp @@ -0,0 +1,360 @@ +// Copyright © 2025 Apple Inc. +// +// File Broadcast with JACCL +// +// This example demonstrates distributed file transfer using JACCL's all_sum +// operation to broadcast a file from any rank to all other machines. +// +// The algorithm: +// 1. The sender rank reads the file into memory +// 2. All other ranks allocate zero-filled buffers of the same size +// 3. Use all_sum to broadcast: sender has data, others have zeros +// 4. After all_sum, all ranks have the file data +// 5. All ranks write the file to disk +// +// For large files, the transfer is chunked to manage memory efficiently. +// +// Usage: +// Set environment variables (see README.md), then run: +// +// ./jaccl_file_broadcast -f [-s ] [-o ] +// +// Or with mlx.launch: +// +// mlx.launch --hostfile hosts.json ./jaccl_file_broadcast -f myfile.bin +// +// Example output (4 ranks, sender rank 2): +// Rank 0 of 4: Received 10485760 bytes from rank 2 (982.5 MB/s) +// Rank 1 of 4: Received 10485760 bytes from rank 2 (985.2 MB/s) +// Rank 2 of 4: Sent 10485760 bytes (980.1 MB/s) +// Rank 3 of 4: Received 10485760 bytes from rank 2 (978.9 MB/s) + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static void usage(const char* prog) { + std::cerr + << "Usage: " << prog << " [options]\n" + << " -f File to broadcast (required)\n" + << " -s Sender rank (default: 0)\n" + << " -o Output directory (default: current dir)\n" + << " -c Chunk size in bytes (default: 67108864 = 64MB)\n" + << " -v Verbose output\n" + << " -h Show this help\n"; +} + +static bool file_exists(const std::string& path) { + struct stat buffer; + return (stat(path.c_str(), &buffer) == 0); +} + +static std::int64_t file_size(const std::string& path) { + struct stat buffer; + if (stat(path.c_str(), &buffer) != 0) { + return -1; + } + return static_cast(buffer.st_size); +} + +static bool create_directory(const std::string& path) { + if (path.empty() || path == ".") { + return true; + } + return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; +} + +static std::string basename(const std::string& path) { + size_t pos = path.find_last_of("/\\"); + return (pos == std::string::npos) ? path : path.substr(pos + 1); +} + +struct BroadcastStats { + std::int64_t total_bytes; + std::int64_t chunks_sent; + std::int64_t chunks_received; + double total_time_ms; + int sender_rank; +}; + +int main(int argc, char** argv) { + std::string input_file; + std::string output_dir = "."; + int sender_rank = 0; + std::int64_t chunk_size = 67108864; + bool verbose = 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 == "-f" && i + 1 < argc) { + input_file = argv[++i]; + } else if (arg == "-s" && i + 1 < argc) { + sender_rank = std::atoi(argv[++i]); + } else if (arg == "-o" && i + 1 < argc) { + output_dir = argv[++i]; + } else if (arg == "-c" && i + 1 < argc) { + chunk_size = std::atoll(argv[++i]); + } else if (arg == "-v" || arg == "--verbose") { + verbose = true; + } else { + std::cerr << "Unknown option: " << arg << "\n"; + usage(argv[0]); + return 1; + } + } + + if (input_file.empty()) { + std::cerr << "Error: Input file is required (-f )\n"; + usage(argv[0]); + return 1; + } + + auto group = jaccl::init(); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + int rank = group->rank(); + int nranks = group->size(); + + if (sender_rank < 0 || sender_rank >= nranks) { + std::cerr << "Error: Sender rank " << sender_rank << " is out of range [0, " + << nranks << ")\n"; + return 1; + } + + std::int64_t total_file_size = 0; + if (rank == sender_rank) { + if (!file_exists(input_file)) { + std::cerr << "Error: File not found: " << input_file << "\n"; + return 1; + } + total_file_size = file_size(input_file); + if (total_file_size < 0) { + std::cerr << "Error: Cannot read file size: " << input_file << "\n"; + return 1; + } + } + + group->all_sum( + &total_file_size, &total_file_size, sizeof(int64_t), jaccl::Int64); + + if (!create_directory(output_dir)) { + std::cerr << "Error: Cannot create output directory: " << output_dir + << "\n"; + return 1; + } + + std::string output_file = output_dir == "." + ? basename(input_file) + : output_dir + "/" + basename(input_file); + + if (verbose) { + std::printf( + "Rank %d of %d: Broadcasting '%s' (%ld bytes) from rank %d\n", + rank, + nranks, + input_file.c_str(), + static_cast(total_file_size), + sender_rank); + } + + auto t_start = std::chrono::high_resolution_clock::now(); + + std::int64_t num_chunks = (total_file_size + chunk_size - 1) / chunk_size; + if (num_chunks == 0) { + num_chunks = 1; + } + + const int num_buffers = 4; + std::vector> buffers( + num_buffers, std::vector(chunk_size, 0)); + + std::ifstream infile; + std::ofstream outfile; + + if (rank == sender_rank) { + infile.open(input_file, std::ios::binary); + if (!infile.good()) { + std::cerr << "Error: Cannot open input file: " << input_file << "\n"; + return 1; + } + } + + outfile.open(output_file, std::ios::binary); + if (!outfile.good()) { + std::cerr << "Error: Cannot open output file: " << output_file << "\n"; + return 1; + } + + std::atomic next_read_chunk{0}; + std::atomic next_comm_chunk{0}; + std::atomic next_write_chunk{0}; + std::atomic read_done{false}; + std::atomic comm_done{false}; + + std::vector> buffer_ready(num_buffers); + std::vector> buffer_written(num_buffers); + for (int i = 0; i < num_buffers; i++) { + buffer_ready[i] = false; + buffer_written[i] = false; + } + + std::vector chunk_sizes(num_chunks); + for (std::int64_t i = 0; i < num_chunks; i++) { + chunk_sizes[i] = std::min(chunk_size, total_file_size - i * chunk_size); + } + + std::thread reader_thread; + if (rank == sender_rank) { + reader_thread = std::thread([&]() { + while (true) { + std::int64_t chunk_idx = next_read_chunk.fetch_add(1); + if (chunk_idx >= num_chunks) { + break; + } + std::int64_t offset = chunk_idx * chunk_size; + std::int64_t this_chunk_size = chunk_sizes[chunk_idx]; + int buffer_idx = chunk_idx % num_buffers; + + infile.seekg(offset, std::ios::beg); + infile.read( + reinterpret_cast(buffers[buffer_idx].data()), + this_chunk_size); + + std::fill( + buffers[buffer_idx].begin() + this_chunk_size, + buffers[buffer_idx].end(), + 0); + + buffer_ready[buffer_idx] = true; + } + read_done = true; + }); + } else { + read_done = true; + } + + std::thread writer_thread([&]() { + while (true) { + std::int64_t chunk_idx = next_write_chunk.load(); + if (chunk_idx >= num_chunks && comm_done) { + break; + } + if (chunk_idx >= num_chunks) { + std::this_thread::yield(); + continue; + } + + int buffer_idx = chunk_idx % num_buffers; + if (!buffer_written[buffer_idx]) { + std::this_thread::yield(); + continue; + } + + std::int64_t this_chunk_size = chunk_sizes[chunk_idx]; + outfile.write( + reinterpret_cast(buffers[buffer_idx].data()), + this_chunk_size); + + buffer_written[buffer_idx] = false; + next_write_chunk.fetch_add(1); + } + }); + + for (std::int64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + std::int64_t this_chunk_size = chunk_sizes[chunk_idx]; + int buffer_idx = chunk_idx % num_buffers; + + if (rank == sender_rank) { + while (!buffer_ready[buffer_idx] && !read_done) { + std::this_thread::yield(); + } + } + + std::fill( + buffers[buffer_idx].begin() + this_chunk_size, + buffers[buffer_idx].end(), + 0); + + group->all_sum( + buffers[buffer_idx].data(), + buffers[buffer_idx].data(), + this_chunk_size, + jaccl::UInt8); + + buffer_written[buffer_idx] = true; + next_comm_chunk.fetch_add(1); + + if (verbose) { + double progress = 100.0 * (chunk_idx + 1) / num_chunks; + std::printf( + "Rank %d: Progress %.1f%% (%ld/%ld chunks)\n", + rank, + progress, + static_cast(chunk_idx + 1), + static_cast(num_chunks)); + } + } + + comm_done = true; + + if (reader_thread.joinable()) { + reader_thread.join(); + } + writer_thread.join(); + + infile.close(); + outfile.close(); + + auto t_end = std::chrono::high_resolution_clock::now(); + double elapsed_ms = + std::chrono::duration(t_end - t_start).count(); + double elapsed_sec = elapsed_ms / 1000.0; + double bandwidth_mbps = (total_file_size / (1024.0 * 1024.0)) / elapsed_sec; + + if (rank == sender_rank) { + std::printf( + "Rank %d of %d: Sent %ld bytes from '%s' (%.1f MB/s)\n", + rank, + nranks, + static_cast(total_file_size), + input_file.c_str(), + bandwidth_mbps); + } else { + std::printf( + "Rank %d of %d: Received %ld bytes from rank %d to '%s' (%.1f MB/s)\n", + rank, + nranks, + static_cast(total_file_size), + sender_rank, + output_file.c_str(), + bandwidth_mbps); + } + + if (verbose) { + std::printf( + "Rank %d: Total time: %.2f ms, Chunks: %ld, Chunk size: %ld bytes\n", + rank, + elapsed_ms, + static_cast(num_chunks), + static_cast(chunk_size)); + } + + return 0; +} diff --git a/mlx/distributed/jaccl/lib/examples/monte_carlo_pi.cpp b/mlx/distributed/jaccl/lib/examples/monte_carlo_pi.cpp new file mode 100644 index 00000000..d710b1d0 --- /dev/null +++ b/mlx/distributed/jaccl/lib/examples/monte_carlo_pi.cpp @@ -0,0 +1,152 @@ +// Copyright © 2025 Apple Inc. +// +// Monte Carlo Pi Estimation with JACCL +// +// This example demonstrates distributed Monte Carlo simulation using JACCL +// to estimate the value of π. Each rank generates random points independently +// and uses all-reduce to aggregate the results across all machines. +// +// The algorithm: +// 1. Each rank generates N random points in the unit square [0,1] x [0,1] +// 2. Count how many fall inside the quarter circle (x² + y² ≤ 1) +// 3. Use all_sum to aggregate hits and total points across all ranks +// 4. π ≈ 4 × (hits / total) +// +// Usage: +// Set environment variables (see README.md), then run: +// +// ./jaccl_monte_carlo_pi [-n ] +// +// Or with mlx.launch: +// +// mlx.launch --hostfile hosts.json ./jaccl_monte_carlo_pi -n 10000000 +// +// Example output (4 ranks, 10M points each): +// Rank 2 of 4 +// Local: 7854321 hits out of 10000000 points +// Global: 31416789 hits out of 40000000 points +// Estimated π = 3.141679 (error: 0.000086) + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static void usage(const char* prog) { + std::cerr << "Usage: " << prog << " [options]\n" + << " -n Points per rank (default: 1000000)\n" + << " -s Random seed base (default: 42)\n" + << " -h Show this help\n"; +} + +struct MonteCarloResult { + int64_t hits; + int64_t total; +}; + +MonteCarloResult estimate_pi_local(int64_t num_points, unsigned int seed) { + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(0.0, 1.0); + + int64_t hits = 0; + for (int64_t i = 0; i < num_points; i++) { + double x = dist(rng); + double y = dist(rng); + if (x * x + y * y <= 1.0) { + hits++; + } + } + + return {hits, num_points}; +} + +int main(int argc, char** argv) { + int64_t points_per_rank = 1000000; + unsigned int seed_base = 42; + + 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 == "-n" && i + 1 < argc) { + points_per_rank = std::atoll(argv[++i]); + } else if (arg == "-s" && i + 1 < argc) { + seed_base = static_cast(std::atoi(argv[++i])); + } else { + std::cerr << "Unknown option: " << arg << "\n"; + usage(argv[0]); + return 1; + } + } + + auto group = jaccl::init(); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + int rank = group->rank(); + int nranks = group->size(); + + std::printf("Rank %d of %d\n", rank, nranks); + std::printf( + "Generating %ld random points (seed: %u)...\n", + static_cast(points_per_rank), + seed_base + static_cast(rank)); + + auto t0 = std::chrono::high_resolution_clock::now(); + + MonteCarloResult local = estimate_pi_local( + points_per_rank, seed_base + static_cast(rank)); + + auto t1 = std::chrono::high_resolution_clock::now(); + double local_time = + std::chrono::duration(t1 - t0).count(); + + std::printf( + "Rank %d: %ld hits out of %ld points (%.2f ms)\n", + rank, + static_cast(local.hits), + static_cast(local.total), + local_time); + + MonteCarloResult global = {0, 0}; + + group->all_sum(&local.hits, &global.hits, sizeof(int64_t), jaccl::Int64); + group->all_sum(&local.total, &global.total, sizeof(int64_t), jaccl::Int64); + + if (rank == 0) { + double pi_estimate = 4.0 * static_cast(global.hits) / + static_cast(global.total); + double error = std::abs(pi_estimate - M_PI); + + std::printf("\n=== Results ===\n"); + std::printf( + "Global: %ld hits out of %ld points\n", + static_cast(global.hits), + static_cast(global.total)); + std::printf("Estimated π = %.10f\n", pi_estimate); + std::printf("True π = %.10f\n", M_PI); + std::printf("Error = %.10f (%.6f%%)\n", error, 100.0 * error / M_PI); + + double total_time = + std::chrono::duration(t1 - t0).count(); + std::printf("\nPerformance:\n"); + std::printf("Total points: %ld\n", static_cast(global.total)); + std::printf("Time: %.2f ms\n", total_time); + std::printf( + "Points/sec: %.0f\n", + static_cast(global.total) / (total_time / 1000.0)); + } + + return 0; +}