FixedSizeArray -> ConcreteArray

Author: Qfl3x

ext/ReactantFixedSizeArraysExt.jl

@@ -6,28 +6,25 @@ using Reactant: TracedRArray, TracedRNumber, Ops
 using ReactantCore: ReactantCore
 
 function Reactant.traced_type_inner(
-    @nospecialize(_::Type{FixedSizeArrays.FixedSizeArray{T, N, Memory{I}}}),
+    @nospecialize(_::Type{FixedSizeArrays.FixedSizeArrayDefault{T, N}}),
     seen,
     @nospecialize(mode::Reactant.TraceMode),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
     @nospecialize(runtime)
-) where {T, N, I}
+) where {T, N}
     T2 = Reactant.TracedRNumber{T}
-    I2 = Reactant.TracedRNumber{I}
-    return FixedSizeArrays.FixedSizeArray{T2, N, Memory{I2}}
+    return FixedSizeArrays.FixedSizeArrayDefault{T2, N}
 end
 
 Base.@nospecializeinfer function Reactant.make_tracer(
     seen,
-    @nospecialize(prev::FixedSizeArrays.FixedSizeArray{T, N, Memory{I}}),
+    @nospecialize(prev::FixedSizeArrays.FixedSizeArrayDefault{T, N}),
     @nospecialize(path),
     mode; kwargs...
-) where {T, N, I}
-    return FixedSizeArrays.FixedSizeArray(
-        Reactant.make_tracer(
-            seen, parent(prev), (path..., 1), mode; kwargs..., track_numbers=Number
-        )
+) where {T, N}
+    return Reactant.make_tracer(
+        seen, parent(prev), (path..., 1), mode; kwargs..., track_numbers=Number
     )
 end
 

src/Tracing.jl

@@ -1811,6 +1811,95 @@ Base.@nospecializeinfer function make_tracer(
     return res
 end
 
+
+Base.@nospecializeinfer function make_tracer(
+    seen,
+    @nospecialize(prev::Memory),
+    @nospecialize(path),
+    mode;
+    @nospecialize(track_numbers::Type = Union{}),
+    @nospecialize(sharding = Sharding.NoSharding()),
+    @nospecialize(runtime = nothing),
+    @nospecialize(device = nothing),
+    @nospecialize(client = nothing),
+    kwargs...,
+)
+    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
+            runtime isa Val{:PJRT} &&
+                (return seen[prev] = ConcretePJRTArray(prev; sharding, device, client))
+            runtime isa Val{:IFRT} &&
+                (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client))
+            error("Unsupported runtime $runtime")
+        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 I in eachindex(prev)
+            if isassigned(prev, I)
+                pv = prev[I]
+                make_tracer(
+                    seen,
+                    pv,
+                    path,
+                    mode;
+                    track_numbers,
+                    sharding,
+                    runtime,
+                    device,
+                    client,
+                    kwargs...,
+                )
+            end
+        end
+        return nothing
+    end
+    TT = traced_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
+    newa = Array{TT,ndims(RT)}(undef, size(prev))
+    seen[prev] = newa
+    same = true
+    for I in eachindex(prev)
+        if isassigned(prev, I)
+            pv = prev[I]
+            nv = make_tracer(
+                seen,
+                pv,
+                append_path(path, I),
+                mode;
+                track_numbers,
+                sharding=Base.getproperty(sharding, I),
+                runtime,
+                device,
+                client,
+                kwargs...,
+            )
+            if pv !== nv
+                same = false
+            end
+            @inbounds newa[I] = nv
+        end
+    end
+    if same
+        seen[prev] = prev
+        return prev
+    end
+    return newa
+end
 Base.@nospecializeinfer function make_tracer(
     seen,
     @nospecialize(prev::Sharding.Mesh),

src/Types.jl

@@ -192,6 +192,20 @@ function ConcretePJRTArray(
     return ConcretePJRTArray{T,N,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
 end
 
+function ConcretePJRTArray(
+    data::Memory{T};
+    client::Union{Nothing,XLA.PJRT.Client}=nothing,
+    idx::Union{Int,Nothing}=nothing,
+    device::Union{Nothing,XLA.PJRT.Device}=nothing,
+    sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
+) where {T}
+    theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
+    sharded_data, shardinfo = sharding(theclient, thedevice, data)
+    shape = size(data)
+    nsharded = length(sharded_data)
+    return ConcretePJRTArray{T,1,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+end
+
 Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data)
 XLA.client(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = XLA.client(x.data)
 function XLA.device(x::Union{ConcretePJRTArray,ConcretePJRTNumber})
@@ -320,6 +334,19 @@ function ConcreteIFRTArray(
     return ConcreteIFRTArray{T,N,typeof(shardinfo)}(sharded_data, shape, shardinfo, padding)
 end
 
+function ConcreteIFRTArray(
+    data::Memory{T};
+    client::Union{Nothing,XLA.IFRT.Client}=nothing,
+    idx::Union{Int,Nothing}=nothing,
+    device::Union{Nothing,XLA.IFRT.Device}=nothing,
+    sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
+) where {T}
+    theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
+    sharded_data, shardinfo, padding = sharding(theclient, nothing, data)
+    shape = size(data)
+    return ConcreteIFRTArray{T,1,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+end
+
 # Assemble data from multiple arrays. Needed in distributed setting where each process wont
 # have enough host memory to hold all the arrays. We assume that the data is only provided
 # for all of the addressable devices.

src/xla/PJRT/Buffer.jl

@@ -21,6 +21,22 @@ function Buffer(client::Client, array::Array{T,N}, device::Device) where {T,N}
     return Buffer(buffer)
 end
 
+
+function Buffer(client::Client, memory::Memory{T}, device::Device) where {T}
+    sizear = collect(Int64, reverse(size(memory)))
+    buffer = GC.@preserve memory sizear begin
+        @ccall MLIR.API.mlir_c.ArrayFromHostBuffer(
+            client.client::Ptr{Cvoid},
+            pointer(memory)::Ptr{T},
+            XLA.primitive_type(T)::UInt64,
+            1::Csize_t,
+            pointer(sizear)::Ptr{Int64},
+            device.device::Ptr{Cvoid},
+        )::Ptr{Cvoid}
+    end
+    return Buffer(buffer)
+end
+
 function Base.similar(a::Buffer)
     buffer = GC.@preserve a begin
         @ccall MLIR.API.mlir_c.UninitPJRTBuffer(