fix: non-iota mesh execution

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -786,7 +786,7 @@ xla::CompileOptions GenerateCompileOptions(int64_t device_id, bool is_sharded,
     for (int64_t i = 0; i < num_mesh_ids; ++i) {
       int64_t mesh_id = mesh_ids[i];
       assert(mesh_id >= 0);
-      device_assignment(0, mesh_id) = i;
+      device_assignment(0, i) = mesh_id;
     }
     options.executable_build_options.set_device_assignment(device_assignment);
 
@@ -945,31 +945,20 @@ void PrintPjRtBuffer(PjRtBuffer *buffer) {
 }
 
 extern "C" void XLAExecute(xla::PjRtLoadedExecutable *exec, int op_args_len,
-                           PjRtBuffer **op_args, const int64_t *mesh_ids,
-                           int64_t num_mesh_ids, uint8_t *is_arg_donatable,
-                           int num_results, PjRtBuffer **op_results,
-                           uint8_t *futures, FutureType **future_results) {
-  // Ensure argument_handles is structured as num_mesh_ids x num_args
-  std::vector<std::vector<PjRtBuffer *>> argument_handles(num_mesh_ids);
-  int num_args = op_args_len / num_mesh_ids;
+                           PjRtBuffer **op_args, int64_t num_devices,
+                           uint8_t *is_arg_donatable, int num_results,
+                           PjRtBuffer **op_results, uint8_t *futures,
+                           FutureType **future_results) {
+  // Ensure argument_handles is structured as num_devices x num_args
+  std::vector<std::vector<PjRtBuffer *>> argument_handles(num_devices);
+  int num_args = op_args_len / num_devices;
 
   // Distribute arguments across devices
-  for (int device_idx = 0; device_idx < num_mesh_ids; ++device_idx) {
-    int64_t mesh_id = mesh_ids[device_idx];
-
-    // Validate mesh_id
-    if (mesh_id < 0 || mesh_id >= num_mesh_ids) {
-      ReactantThrowError(("Invalid mesh_id " + std::to_string(mesh_id) +
-                          " at device_idx " + std::to_string(device_idx))
-                             .c_str());
-    }
-
+  for (int device_idx = 0; device_idx < num_devices; ++device_idx) {
     argument_handles[device_idx].reserve(num_args);
     for (int arg_idx = 0; arg_idx < num_args; ++arg_idx) {
-      // Assuming op_args is a flat array of size num_devices * num_args
-      // where arguments for each device are contiguous
       argument_handles[device_idx].push_back(
-          op_args[mesh_id * num_args + arg_idx]);
+          op_args[device_idx * num_args + arg_idx]);
     }
   }
 
@@ -989,40 +978,40 @@ extern "C" void XLAExecute(xla::PjRtLoadedExecutable *exec, int op_args_len,
               argument_handles),
           options, returned_futures));
 
-  if (results.size() != num_mesh_ids) {
+  if (results.size() != num_devices) {
     ReactantThrowError((" results.size()=" + std::to_string(results.size()) +
-                        " num_mesh_ids=" + std::to_string(num_mesh_ids) + "\n")
+                        " num_devices=" + std::to_string(num_devices) + "\n")
                            .c_str());
   }
 
-  for (int device_idx = 0; device_idx < num_mesh_ids; ++device_idx) {
-    int64_t mesh_id = mesh_ids[device_idx];
-    if (results[mesh_id].size() != num_results) {
-      ReactantThrowError((" results[" + std::to_string(mesh_id) + "].size()=" +
-                          std::to_string(results[mesh_id].size()) +
-                          " num_results=" + std::to_string(num_results) + "\n")
-                             .c_str());
+  for (int device_idx = 0; device_idx < num_devices; ++device_idx) {
+    // Remove mesh_id lookup since we're using device_idx ordering
+    if (results[device_idx].size() != num_results) {
+      ReactantThrowError(
+          (" results[" + std::to_string(device_idx) +
+           "].size()=" + std::to_string(results[device_idx].size()) +
+           " num_results=" + std::to_string(num_results) + "\n")
+              .c_str());
     }
   }
 
   // Handle returned futures
   auto future_val = returned_futures.has_value();
   *futures = future_val;
   if (future_val) {
-    if (returned_futures->size() != num_mesh_ids) {
+    if (returned_futures->size() != num_devices) {
       ReactantThrowError((" returned_futures->size()=" +
                           std::to_string(returned_futures->size()) +
-                          " num_mesh_ids=" + std::to_string(num_mesh_ids) +
+                          " num_devices=" + std::to_string(num_devices) +
                           "\n")
                              .c_str());
     }
   }
 
   // Copy results into the output buffers
-  for (int device_idx = 0; device_idx < num_mesh_ids; ++device_idx) {
-    int64_t mesh_id = mesh_ids[device_idx];
+  for (int device_idx = 0; device_idx < num_devices; ++device_idx) {
     for (int result_idx = 0; result_idx < num_results; ++result_idx) {
-      int flat_index = mesh_id * num_results + result_idx;
+      int flat_index = device_idx * num_results + result_idx;
       op_results[flat_index] = results[device_idx][result_idx].release();
       if (future_val) {
         future_results[flat_index] =

src/Compiler.jl

@@ -1307,12 +1307,17 @@ Generate Julia code to call the XLA executable.
 - `nresults`: The number of results to expect.
 """
 function codegen_xla_call(
-    exec, device, flatten_names, donated_args_mask, nresults, is_sharded::Bool, mesh_ids
+    exec,
+    device,
+    flatten_names,
+    donated_args_mask,
+    nresults,
+    is_sharded::Bool,
+    ndevices::Int,
 )
     flatten_buffer_refs = map(n -> :($n.buffer), flatten_names)
 
-    base_symbol_name =
-        is_sharded ? Symbol(:result_buffer_m, length(mesh_ids), :_) : :result_buffer_
+    base_symbol_name = is_sharded ? Symbol(:result_buffer_m, ndevices, :_) : :result_buffer_
     concretized_res_names = Symbol[Symbol(base_symbol_name, i) for i in 1:nresults]
     concretized_res_code = map(enumerate(concretized_res_names)) do (i, varname)
         :($varname = linearized_results[$i])
@@ -1326,11 +1331,10 @@ function codegen_xla_call(
                 GC.@preserve $(flatten_names...) begin
                     linearized_results = XLA.execute(
                         $exec,
-                        $(mesh_ids),
                         ($(flatten_buffer_refs...),),
                         $(Tuple(donated_args_mask)),
                         Val($nresults),
-                        Val($(length(mesh_ids))),
+                        Val($ndevices),
                     )
                 end
                 $(concretized_res_code...)
@@ -1509,7 +1513,7 @@ function compile(f, args; sync=false, kwargs...)
         donated_args_mask,
         length(linear_results),
         mlir_fn_res.is_sharded,
-        mlir_fn_res.is_sharded ? vec(mlir_fn_res.sharding_mesh) : Int64[],
+        mlir_fn_res.is_sharded ? length(mlir_fn_res.sharding_mesh) : 1,
     )
 
     linear_result_shard_info = if mlir_fn_res.is_sharded

src/xla/PJRT/LoadedExecutable.jl

@@ -241,25 +241,22 @@ end
 
 @inline function XLA.execute(
     exec::LoadedExecutable,
-    mesh_ids::Vector{Int64},
     inputs::NTuple{N,Ptr{Cvoid}},
     donated_args::NTuple{M,UInt8},
     ::Val{n_outs},
     ::Val{K},
 ) where {N,M,n_outs,K}
-    @assert length(mesh_ids) == K
     outputs = Ref{NTuple{n_outs * K,Ptr{Cvoid}}}()
     future_res = Ref{NTuple{n_outs * K,Ptr{Cvoid}}}()
     futures = Ref{UInt8}(0)
 
     inputs = Base.RefValue(inputs)
     donated_args = Base.RefValue(donated_args)
-    GC.@preserve inputs donated_args mesh_ids outputs futures future_res begin
+    GC.@preserve inputs donated_args outputs futures future_res begin
         @ccall MLIR.API.mlir_c.XLAExecute(
             exec.exec::Ptr{Cvoid},
             N::Cint,
             inputs::Ptr{Cvoid},
-            mesh_ids::Ptr{Clong},
             K::Clong,
             donated_args::Ptr{UInt8},
             n_outs::Cint,

test/sharding.jl

@@ -104,10 +104,9 @@ end
         mesh = Sharding.Mesh(reshape([4, 6, 0, 2, 7, 3, 1, 5], 4, 2), ("data", "model"))
         x = reshape(collect(Float32, 1:16), 4, 4)
         x_ra = Reactant.to_rarray(
-            x; sharding=Sharding.NamedSharding(mesh, ("data", nothing))
+            x; sharding=Sharding.NamedSharding(mesh, ("data", "model"))
         )
-        # XXX: This needs to be fixed
-        @test_broken Array(@jit sum(x_ra)) ≈ sum(x)
+        @test Array(@jit fn_test2(x_ra)) ≈ fn_test2(x)
     else
         @warn "Not enough addressable devices to run sharding tests"
     end