Fix pathing (#1337)

* Fix pathing

* 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

@@ -1378,8 +1378,10 @@ function compile_mlir!(
         seen_args,
         ret,
         linear_args,
+        skipped_args,
         in_tys,
         linear_results,
+        skipped_results,
         is_sharded,
     ) = mlir_fn_res
     compiled_f = mlir_fn_res.f
@@ -2107,8 +2109,10 @@ function compile_mlir!(
         seen_args,
         ret,
         linear_args,
+        skipped_args,
         in_tys,
         linear_results2,
+        skipped_results,
         mlir_fn_res.num_partitions,
         mlir_fn_res.num_replicas,
         mlir_fn_res.is_sharded,
@@ -2399,7 +2403,6 @@ The _linearized arguments_ do not directly refer to the  are the arguments that
 function codegen_flatten!(
     linear_args,
     seen_args,
-    result_stores,
     is_sharded::Bool,
     linear_parameter_shardings,
     client,
@@ -2804,14 +2807,16 @@ function codegen_unflatten!(
                 if length(path) > 0
                     needs_cache_dict = true
                     # XXX: we might need to handle sharding here
-                    unflatcode = :(traced_setfield_buffer!(
-                        $(runtime),
-                        $(cache_dict),
-                        $(concrete_res_name_final),
-                        $(unflatcode),
-                        $(Meta.quot(path[end])),
-                        $(path),
-                    ))
+                    unflatcode = quote
+                        traced_setfield_buffer!(
+                            $(runtime),
+                            $(cache_dict),
+                            $(concrete_res_name_final),
+                            $(unflatcode),
+                            $(Meta.quot(path[end])),
+                            $(path),
+                        )
+                    end
                 else
                     unflatcode = :(traced_setfield!(
                         $(unflatcode), :data, $(concrete_res_name_final), $(path)
@@ -3280,7 +3285,6 @@ function compile(f, args; sync=false, kwargs...)
     flatten_arg_names, flatten_code, resharded_inputs = codegen_flatten!(
         linear_args,
         seen_args,
-        result_stores,
         mlir_fn_res.is_sharded,
         XLA.get_parameter_shardings(exec), # TODO: use the same workflow as output shardings to parse the tensor sharding attributes directly if possible
         client,

src/Ops.jl

@@ -2329,7 +2329,7 @@ end
 end
 
 @noinline function call(f, args...)
-    seen = Dict()
+    seen = Reactant.OrderedIdDict()
     cache_key = []
     Reactant.make_tracer(seen, (f, args...), cache_key, Reactant.TracedToTypes)
     cache = Reactant.Compiler.callcache()

src/TracedRArray.jl

@@ -1,6 +1,7 @@
 module TracedRArrayOverrides
 
 using Adapt: WrappedArray
+using Adapt: Adapt
 using Base.Broadcast
 using Base.Broadcast: BroadcastStyle, Broadcasted, AbstractArrayStyle, instantiate
 
@@ -506,6 +507,14 @@ function Base.similar(::TracedRArray, ::Type{T}, dims::Dims{N}) where {T,N}
     return Ops.fill(zero(unwrapped_eltype(T)), dims)
 end
 
+function Base.show(io::IOty, X::AnyTracedRArray) where {IOty<:Union{IO,IOContext}}
+    print(io, Core.Typeof(X), "(")
+    if Adapt.parent(X) !== X
+        Base.show(io, Adapt.parent(X))
+    end
+    return print(io, ")")
+end
+
 function Base.show(io::IOty, X::TracedRArray{T,N}) where {T,N,IOty<:Union{IO,IOContext}}
     return print(io, "TracedRArray{", T, ",", N, "N}(", X.paths, ", size=", size(X), ")")
     # TODO this line segfaults if MLIR IR has not correctly been generated

src/TracedUtils.jl

@@ -228,8 +228,10 @@ mutable struct CompiledMlirFnResult{F,TR,Re,Rt,LA,LR,PA,CR,M,MA,RS,GD,DA}
     seen_args::OrderedIdDict
     ret::Rt
     linear_args::Vector{LA}
+    skipped_args::Vector{LA}
     in_tys::Vector{MLIR.IR.Type}
     linear_results::Vector{LR}
+    skipped_results::Vector{LR}
     num_partitions::Int
     num_replicas::Int
     is_sharded::Bool
@@ -333,7 +335,7 @@ function make_mlir_fn(
         end
     end
 
-    (func2, traced_result, ret, linear_args, in_tys, linear_results, num_partitions, is_sharded, unique_meshes, mutated_args, global_device_ids) = finalize_mlir_fn(
+    (func2, traced_result, ret, linear_args, in_tys, linear_results, skipped_results, num_partitions, is_sharded, unique_meshes, mutated_args, global_device_ids) = finalize_mlir_fn(
         result,
         traced_args,
         linear_args,
@@ -373,8 +375,10 @@ function make_mlir_fn(
         seen_args,
         ret,
         linear_args,
+        skipped_args,
         in_tys,
         linear_results,
+        skipped_results,
         num_partitions,
         num_replicas,
         is_sharded,
@@ -628,9 +632,49 @@ function finalize_mlir_fn(
     end
 
     linear_results = Reactant.TracedType[]
+    skipped_results = Reactant.TracedType[]
     for (k, v) in seen_results
         v isa Reactant.TracedType || continue
         if any(Base.Fix1(===, k), skipped_args)
+            push!(skipped_results, v)
+
+            _, argpath = get_argidx(v, argprefix)
+
+            @assert has_idx(v, argprefix)
+
+            newpaths = Tuple[]
+            for path in v.paths
+                if length(path) == 0
+                    continue
+                end
+                if path[1] == argprefix
+                    continue
+                end
+                if path[1] == resargprefix
+                    original_arg = args[path[2]]
+                    for p in path[3:end]
+                        original_arg = Reactant.Compiler.traced_getfield(original_arg, p)
+                    end
+                    if !(
+                        original_arg isa Union{
+                            Reactant.ConcreteRNumber,
+                            Reactant.ConcreteRArray,
+                            Reactant.TracedType,
+                        }
+                    )
+                        continue
+                    end
+                    push!(newpaths, path)
+                end
+                if path[1] == resprefix
+                    push!(newpaths, path)
+                end
+            end
+
+            if length(newpaths) != 0
+                push!(linear_results, Reactant.repath(v, (newpaths...,)))
+            end
+
             continue
         end
         if args_in_result != :all
@@ -924,6 +968,7 @@ function finalize_mlir_fn(
         linear_args,
         in_tys,
         linear_results,
+        skipped_results,
         num_partitions,
         is_sharded,
         unique_meshes,

src/Tracing.jl

@@ -895,7 +895,7 @@ const traced_type_cache = Dict{Tuple{TraceMode,Type,Any},Dict{Type,Type}}()
 Base.@assume_effects :total @inline function traced_type(
     T::Type, ::Val{mode}, track_numbers::Type, sharding, runtime
 ) where {mode}
-    if mode == TracedSetPath || mode == TracedTrack
+    if mode == TracedSetPath || mode == TracedTrack || mode == TracedToTypes
         return T
     end
 

src/Types.jl

@@ -52,6 +52,10 @@ mutable struct TracedRNumber{T} <: RNumber{T}
     end
 end
 
+function repath(x::TracedRNumber{T}, paths) where {T}
+    return TracedRNumber{T}(paths, x.mlir_data)
+end
+
 @leaf TracedRNumber
 
 ## TracedRArray
@@ -71,6 +75,10 @@ mutable struct TracedRArray{T,N} <: RArray{TracedRNumber{T},N}
     end
 end
 
+function repath(x::TracedRArray{T,N}, paths) where {T,N}
+    return TracedRArray{T,N}(paths, x.mlir_data, x.shape)
+end
+
 @leaf TracedRArray
 Adapt.parent_type(::Type{TracedRArray{T,N}}) where {T,N} = TracedRArray{T,N}
 

test/constructor.jl

@@ -0,0 +1,139 @@
+using Reactant, Test, Adapt
+
+struct MyGrid{FT,AT} <: AbstractVector{FT}
+    data::AT
+    radius::FT
+end
+
+Adapt.parent(x::MyGrid) = x.data
+
+Base.getindex(x::MyGrid, args...) = Base.getindex(x.data, args...)
+
+Base.size(x::MyGrid) = Base.size(x.data)
+
+function Base.show(io::IOty, X::MyGrid) where {IOty<:Union{IO,IOContext}}
+    print(io, Core.Typeof(X), "(")
+    if Adapt.parent(X) !== X
+        Base.show(io, Adapt.parent(X))
+    end
+    return print(io, ")")
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(OA::Type{MyGrid{FT,AT}}),
+    seen,
+    mode::Reactant.TraceMode,
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {FT,AT}
+    FT2 = Reactant.traced_type_inner(FT, seen, mode, track_numbers, sharding, runtime)
+    AT2 = Reactant.traced_type_inner(AT, seen, mode, track_numbers, sharding, runtime)
+
+    for NF in (AT2,)
+        FT2 = Reactant.promote_traced_type(FT2, eltype(NF))
+    end
+
+    res = MyGrid{FT2,AT2}
+    return res
+end
+
+@inline Reactant.make_tracer(seen, @nospecialize(prev::MyGrid), args...; kwargs...) =
+    Reactant.make_tracer_via_immutable_constructor(seen, prev, args...; kwargs...)
+
+struct MyGrid2{FT,AT} <: AbstractVector{FT}
+    data::AT
+    radius::FT
+    bar::FT
+end
+
+Adapt.parent(x::MyGrid2) = x.data
+
+Base.getindex(x::MyGrid2, args...) = Base.getindex(x.data, args...)
+
+Base.size(x::MyGrid2) = Base.size(x.data)
+
+function Base.show(io::IOty, X::MyGrid2) where {IOty<:Union{IO,IOContext}}
+    print(io, Core.Typeof(X), "(")
+    if Adapt.parent(X) !== X
+        Base.show(io, Adapt.parent(X))
+    end
+    return print(io, ")")
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(OA::Type{MyGrid2{FT,AT}}),
+    seen,
+    mode::Reactant.TraceMode,
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {FT,AT}
+    FT2 = Reactant.traced_type_inner(FT, seen, mode, track_numbers, sharding, runtime)
+    AT2 = Reactant.traced_type_inner(AT, seen, mode, track_numbers, sharding, runtime)
+
+    for NF in (AT2,)
+        FT2 = Reactant.promote_traced_type(FT2, eltype(NF))
+    end
+
+    res = MyGrid2{FT2,AT2}
+    return res
+end
+
+@inline Reactant.make_tracer(seen, @nospecialize(prev::MyGrid2), args...; kwargs...) =
+    Reactant.make_tracer_via_immutable_constructor(seen, prev, args...; kwargs...)
+
+function update!(g)
+    @allowscalar g.data[1] = g.radius
+    return nothing
+end
+
+function selfreturn(g)
+    return g
+end
+
+function call_update!(g)
+    @trace update!(g)
+end
+
+function call_selfreturn(g)
+    @trace selfreturn(g)
+end
+
+@testset "Custom construction" begin
+    g = MyGrid([3.14, 1.59], 2.7)
+    rg = Reactant.to_rarray(g)
+
+    @jit update!(rg)
+    @test convert(Array, rg.data) == [2.7, 1.59]
+
+    rg = Reactant.to_rarray(g)
+    res = @jit selfreturn(rg)
+    @test convert(Array, res.data) == [3.14, 1.59]
+    @test res.radius == 2.7
+    @show typeof(res)
+    @test typeof(res.radius) <: ConcreteRNumber
+
+    rg = Reactant.to_rarray(g)
+
+    @jit call_update!(rg)
+    @test convert(Array, rg.data) == [2.7, 1.59]
+
+    rg = Reactant.to_rarray(g)
+    res = @jit call_selfreturn(rg)
+    @test convert(Array, res.data) == [3.14, 1.59]
+    @test res.radius == 2.7
+    @show typeof(res)
+    @test typeof(res.radius) <: ConcreteRNumber
+end
+
+@testset "Custom construction2 " begin
+    g = Ref(MyGrid([3.14, 1.59], 2.7))
+    g = (g, g)
+
+    rg = Reactant.to_rarray(g)
+    res = @jit selfreturn(rg)
+    @test convert(Array, res[1][].data) == [3.14, 1.59]
+    @test convert(Array, res[2][].data) == [3.14, 1.59]
+    @test res[1][].data == res[2][].data
+end

test/runtests.jl

@@ -15,6 +15,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Layout" include("layout.jl")
         @safetestset "Tracing" include("tracing.jl")
         @safetestset "Basic" include("basic.jl")
+        @safetestset "Constructor" include("constructor.jl")
         @safetestset "Autodiff" include("autodiff.jl")
         @safetestset "Complex" include("complex.jl")
         @safetestset "Broadcast" include("bcast.jl")