Tentative JACCL examples

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)

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 <file> [-s <sender_rank>] [-o <output_dir>]
+//
+//   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 <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <sys/stat.h>
+#include <atomic>
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <string>
+#include <thread>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr
+      << "Usage: " << prog << " [options]\n"
+      << "  -f <file>       File to broadcast (required)\n"
+      << "  -s <rank>       Sender rank (default: 0)\n"
+      << "  -o <dir>        Output directory (default: current dir)\n"
+      << "  -c <bytes>      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<std::int64_t>(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 <file>)\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<long>(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<std::vector<std::uint8_t>> buffers(
+      num_buffers, std::vector<std::uint8_t>(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<std::int64_t> next_read_chunk{0};
+  std::atomic<std::int64_t> next_comm_chunk{0};
+  std::atomic<std::int64_t> next_write_chunk{0};
+  std::atomic<bool> read_done{false};
+  std::atomic<bool> comm_done{false};
+
+  std::vector<std::atomic<bool>> buffer_ready(num_buffers);
+  std::vector<std::atomic<bool>> buffer_written(num_buffers);
+  for (int i = 0; i < num_buffers; i++) {
+    buffer_ready[i] = false;
+    buffer_written[i] = false;
+  }
+
+  std::vector<std::int64_t> 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<char*>(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<const char*>(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<long>(chunk_idx + 1),
+          static_cast<long>(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<double, std::milli>(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<long>(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<long>(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<long>(num_chunks),
+        static_cast<long>(chunk_size));
+  }
+
+  return 0;
+}

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 <points_per_rank>]
+//
+//   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 <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <random>
+#include <string>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr << "Usage: " << prog << " [options]\n"
+            << "  -n <points>  Points per rank (default: 1000000)\n"
+            << "  -s <seed>    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<double> 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<unsigned int>(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<long>(points_per_rank),
+      seed_base + static_cast<unsigned int>(rank));
+
+  auto t0 = std::chrono::high_resolution_clock::now();
+
+  MonteCarloResult local = estimate_pi_local(
+      points_per_rank, seed_base + static_cast<unsigned int>(rank));
+
+  auto t1 = std::chrono::high_resolution_clock::now();
+  double local_time =
+      std::chrono::duration<double, std::milli>(t1 - t0).count();
+
+  std::printf(
+      "Rank %d: %ld hits out of %ld points (%.2f ms)\n",
+      rank,
+      static_cast<long>(local.hits),
+      static_cast<long>(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<double>(global.hits) /
+        static_cast<double>(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<long>(global.hits),
+        static_cast<long>(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<double, std::milli>(t1 - t0).count();
+    std::printf("\nPerformance:\n");
+    std::printf("Total points: %ld\n", static_cast<long>(global.total));
+    std::printf("Time: %.2f ms\n", total_time);
+    std::printf(
+        "Points/sec: %.0f\n",
+        static_cast<double>(global.total) / (total_time / 1000.0));
+  }
+
+  return 0;
+}