Support for dicts (#748)

* Support for dicts

* fix

* Apply suggestions from code review

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: wsmoses

src/Compiler.jl

@@ -22,6 +22,10 @@ import ..Reactant:
 
 import ..ReactantCore: correct_maybe_bcast_call
 
+@inline function traced_getfield(@nospecialize(obj::Dict), field)
+    return Base.getindex(obj, field)
+end
+
 @inline function traced_getfield(@nospecialize(obj), field)
     return Base.getfield(obj, field)
 end
@@ -40,6 +44,10 @@ end
     return Base.setindex!(obj, val, field)
 end
 
+@inline function traced_setfield!(@nospecialize(obj::Dict), field, val)
+    return Base.setindex!(obj, field, val)
+end
+
 function create_result(
     tocopy::T, path, result_stores, path_to_shard_info, sharding_mesh
 ) where {T}

src/Sharding.jl

@@ -58,6 +58,7 @@ struct NoSharding <: AbstractSharding end
 
 # This allows us to mark entire branches as NoSharding
 Base.getproperty(::NoSharding, x) = NoSharding()
+Base.getproperty(::NoSharding, x::Symbol) = NoSharding()
 
 function (::NoSharding)(client::XLA.Client, device, x::Union{AbstractArray,Number})
     buffer = XLA.AsyncBuffer(XLA.ArrayFromHostBuffer(client, x, device), nothing)

src/Tracing.jl

@@ -1232,6 +1232,68 @@ function make_tracer(
     return newa
 end
 
+function make_tracer(
+    seen,
+    @nospecialize(prev::Dict{Key,Value}),
+    @nospecialize(path),
+    mode;
+    @nospecialize(track_numbers::Type = Union{}),
+    @nospecialize(sharding = Sharding.NoSharding()),
+    kwargs...,
+) where {Key,Value}
+    RT = Core.Typeof(prev)
+    # XXX: If someone wants to shard the same array with different shardings, we need to
+    #      somehow handle this correctly... Right now we just use the first sharding.
+    if mode != NoStopTracedTrack && haskey(seen, prev)
+        if mode == TracedToTypes
+            visited = seen[prev]
+            push!(path, visited)
+            return nothing
+        end
+        return seen[prev]
+    end
+    if eltype(RT) <: ReactantPrimitive
+        if mode == ArrayToConcrete && return seen[prev] = ConcreteRArray(prev; sharding)
+        elseif mode == TracedToTypes
+            # Original array can get mutated so we store a copy:
+            push!(path, copy(prev))
+            seen[prev] = VisitedObject(length(seen) + 1)
+            return nothing
+        end
+    elseif mode == TracedToTypes
+        push!(path, RT)
+        for (k, v) in prev
+            make_tracer(seen, k, path, mode; track_numbers, sharding, kwargs...)
+            make_tracer(seen, v, path, mode; track_numbers, sharding, kwargs...)
+        end
+        return nothing
+    end
+    Value2 = traced_type(Value, Val(mode), track_numbers, sharding)
+    newa = Dict{Key,Value2}()
+    seen[prev] = newa
+    same = true
+    for (k, v) in prev
+        nv = make_tracer(
+            seen,
+            v,
+            append_path(path, k),
+            mode;
+            track_numbers,
+            sharding=Base.getproperty(sharding, k),
+            kwargs...,
+        )
+        if v !== nv
+            same = false
+        end
+        newa[k] = nv
+    end
+    if same
+        seen[prev] = prev
+        return prev
+    end
+    return newa
+end
+
 function make_tracer(
     seen,
     @nospecialize(prev::Tuple),

test/basic.jl

@@ -901,3 +901,18 @@ end
     @test contains(hlo, "stablehlo.add")
     @test Array(@jit(fn(Base.OneTo(10000)))) ≈ collect(Base.OneTo(10000))
 end
+
+function dip!(x)
+    x[:a] = x[:a] .* x[:b]
+    return nothing
+end
+
+@testset "Dict" begin
+    x = Dict{Symbol,Vector{Float32}}()
+    x[:a] = 2.7 * ones(4)
+    x[:b] = 3.1 * ones(4)
+
+    ra = Reactant.to_rarray(x)
+    Reactant.@jit dip!(ra)
+    ra[:a] ≈ (2.7 * 2) * ones(4)
+end