feat: sharding group (#1811)

* feat: sharding group

* Update src/Ops.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

src/Compiler.jl

@@ -1580,7 +1580,8 @@ function compile_mlir!(
     args,
     compile_options::CompileOptions,
     callcache=default_callcache(),
-    sdycache=default_sdycache();
+    sdycache=default_sdycache(),
+    sdygroupidcache=default_sdygroupidcache();
     fn_kwargs=(),
     backend="gpu",
     runtime::Union{Val{:PJRT},Val{:IFRT}},
@@ -1597,6 +1598,7 @@ function compile_mlir!(
     MLIR.IR.activate!(MLIR.IR.body(mod))
     activate_callcache!(callcache)
     activate_sdycache!(sdycache)
+    activate_sdygroupidcache!(sdygroupidcache)
 
     # Save in the TLS whether we are raising.  We identify that condition by
     # checking whether the user set an explicit list of passes, or chose
@@ -1623,6 +1625,7 @@ function compile_mlir!(
     finally
         deactivate_raising!(is_raising)
         deactivate_sdycache!(sdycache)
+        deactivate_sdygroupidcache!(sdygroupidcache)
         deactivate_callcache!(callcache)
         MLIR.IR.deactivate!(MLIR.IR.body(mod))
         MLIR.IR.deactivate!(mod)
@@ -3832,7 +3835,7 @@ function register_thunk(
     )
 end
 
-for cache_type in (:callcache, :sdycache)
+for cache_type in (:callcache, :sdycache, :sdygroupidcache)
     activate_fn = Symbol(:activate_, cache_type, :!)
     deactivate_fn = Symbol(:deactivate_, cache_type, :!)
     has_fn = Symbol(:_has_, cache_type)
@@ -3879,6 +3882,14 @@ function default_sdycache()
     }()
 end
 
+mutable struct SdyGroupIDCounter{T}
+    @atomic group_id::T
+end
+
+function default_sdygroupidcache()
+    return SdyGroupIDCounter{Int}(0), Base.IdDict{Union{TracedRArray,TracedRNumber},Int}()
+end
+
 function default_callcache()
     return Dict{
         Vector,

src/Ops.jl

@@ -1,6 +1,6 @@
 # This module reflects the HLO ops defined in the openxla/stablehlo repo (plus some extras).
 # If you want to add some check or test, the StableHLO spec should be taken as the source of truth, not the Julia or Reactant semantics.
-# Julia and Reactant semantics should be considered on the higher abstractions that use these 
+# Julia and Reactant semantics should be considered on the higher abstractions that use these
 module Ops
 using ..MLIR: MLIR
 using ..MLIR.Dialects: stablehlo, chlo, enzyme, enzymexla
@@ -3474,4 +3474,45 @@ end
     )
 end
 
+@noinline function sharding_group(
+    inputs::Union{TracedRArray,TracedRNumber}...;
+    group_id::Union{Integer,Nothing}=nothing,
+    location=mlir_stacktrace("sharding_group", @__FILE__, @__LINE__),
+)
+    @assert length(inputs) > 1 "At least two inputs are required to form a sharding group, \
+                                got $(length(inputs))"
+
+    counter, cache = Reactant.Compiler.sdygroupidcache()
+
+    group_ids = unique([cache[input] for input in inputs if haskey(cache, input)])
+    if length(group_ids) > 1
+        error("All inputs must belong to the same sharding group. Found multiple group \
+               ids: $(group_ids)")
+    end
+
+    if length(group_ids) == 0
+        if group_id === nothing
+            group_id = @atomic counter.group_id
+            @atomic counter.group_id += 1
+        end
+    else
+        found_group_id = only(group_ids)
+        if group_id !== nothing && found_group_id != group_id
+            error("Provided group_id $(group_id) does not match the existing group_id \
+                   $(found_group_id) for the inputs. All inputs must belong to the same \
+                   sharding group.")
+        end
+        group_id = found_group_id
+    end
+
+    for input in inputs
+        if !haskey(cache, input)
+            cache[input] = group_id
+            MLIR.Dialects.sdy.sharding_group(input.mlir_data; group_id=group_id, location)
+        end
+    end
+
+    return nothing
+end
+
 end # module Ops

test/sharding.jl

@@ -509,3 +509,29 @@ end
         @warn "Not enough addressable devices to run sharding tests"
     end
 end
+
+@testset "Sharding Group" begin
+    if length(Reactant.devices()) ≥ 4 && Reactant.XLA.runtime() isa Val{:IFRT}
+        mesh = Sharding.Mesh(reshape(0:3, 2, 2), (:x, :y))
+        sharding = Sharding.NamedSharding(mesh, (:x, :y))
+
+        function shard_groups(x)
+            y = (x' * x)[1:4, :]
+            Reactant.Ops.sharding_group(x, y)
+            z = y .+ x
+            Reactant.Ops.sharding_group(z, y)
+            return z
+        end
+
+        x = Reactant.to_rarray(
+            Reactant.TestUtils.construct_test_array(Float32, 4, 128); sharding
+        )
+
+        hlo = repr(@code_hlo shard_groups(x))
+
+        @test count("sharding_group", hlo) == 3
+        @test count("group_id=0", hlo) == 3
+    else
+        @warn "Not enough addressable devices to run sharding tests"
+    end
+end