FixedSizeArrays support (And Memory Support)

Project.toml

@@ -47,6 +47,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
+FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 
 [sources]
 ReactantCore = {path = "lib/ReactantCore"}
@@ -57,6 +58,7 @@ ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
 ReactantDLFP8TypesExt = "DLFP8Types"
 ReactantFillArraysExt = "FillArrays"
+ReactantFixedSizeArraysExt = "FixedSizeArrays"
 ReactantFloat8sExt = "Float8s"
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
@@ -82,6 +84,7 @@ EnumX = "1"
 Enzyme = "0.13.74"
 EnzymeCore = "0.8.13"
 FillArrays = "1.13"
+FixedSizeArrays = "1.2.0"
 Float8s = "0.1"
 Functors = "0.5"
 GPUArraysCore = "0.2"

ext/ReactantFixedSizeArraysExt.jl

@@ -0,0 +1,36 @@
+module ReactantFixedSizeArraysExt
+
+using FixedSizeArrays
+using Reactant
+using Reactant: TracedRArray, TracedRNumber, Ops
+using ReactantCore: ReactantCore
+
+function Reactant.traced_type_inner(
+    @nospecialize(_::Type{FixedSizeArrays.FixedSizeArrayDefault{T,N}}),
+    seen,
+    @nospecialize(mode::Reactant.TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T,N}
+    T2 = Reactant.TracedRNumber{T}
+    return FixedSizeArrays.FixedSizeArrayDefault{T2,N}
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen,
+    @nospecialize(prev::FixedSizeArrays.FixedSizeArrayDefault{T,N}),
+    @nospecialize(path),
+    mode;
+    kwargs...,
+) where {T,N}
+    shape = size(prev)
+    return reshape(
+        Reactant.make_tracer(
+            seen, parent(prev), (path..., 1), mode; kwargs..., track_numbers=Number
+        ),
+        shape,
+    )
+end
+
+end

src/Tracing.jl

@@ -1812,6 +1812,97 @@ Base.@nospecializeinfer function make_tracer(
     return res
 end
 
+if isdefined(Base, :Memory)
+    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
+end
+
 Base.@nospecializeinfer function make_tracer(
     seen,
     @nospecialize(prev::Sharding.Mesh),

src/Types.jl

@@ -228,6 +228,24 @@ function ConcretePJRTArray(
     return ConcretePJRTArray{T,N,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
 end
 
+if isdefined(Base, :Memory)
+    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
+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})
@@ -356,6 +374,21 @@ function ConcreteIFRTArray(
     return ConcreteIFRTArray{T,N,typeof(shardinfo)}(sharded_data, shape, shardinfo, padding)
 end
 
+if isdefined(Base, :Memory)
+    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
+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.
@@ -472,8 +505,11 @@ elseif XLA.REACTANT_XLA_RUNTIME == "IFRT"
     ConcreteIFRTArray
 end
 
-@inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} =
-    ConcreteRArray{T}(undef, Dims(shape); kwargs...)
+@inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} = ConcreteRArray{
+    T
+}(
+    undef, Dims(shape); kwargs...
+)
 
 """
     ConcreteRNumber(

src/xla/IFRT/Array.jl

@@ -38,6 +38,30 @@ function Array(
     return Array(buffer)
 end
 
+if isdefined(Base, :Memory)
+    function Array(
+        client::Client,
+        memory::Base.Memory{T},
+        device::Device=XLA.default_device(client),
+        memory_kind::AbstractString=string(convert(MemoryKind, XLA.default_memory(device))),
+    ) where {T<:Reactant.ReactantPrimitive}
+        sizear = collect(Int64, reverse(size(memory)))
+        buffer = GC.@preserve memory sizear begin
+            @ccall MLIR.API.mlir_c.ifrt_client_make_single_shard_array_from_host_buffer(
+                client.client::Ptr{Cvoid},
+                pointer(memory)::Ptr{T},
+                XLA.primitive_type(T)::UInt64,
+                1::Csize_t,
+                sizear::Ptr{Int64},
+                0::Cint, # kAlwaysCopy
+                device.device::Ptr{Cvoid},
+                string(memory_kind)::Cstring,
+            )::Ptr{Cvoid}
+        end
+        return Array(buffer)
+    end
+end
+
 function Array(
     client::Client, array::Base.Array{T,N}, sharding::Sharding
 ) where {T<:Reactant.ReactantPrimitive,N}
@@ -143,6 +167,47 @@ function Array(
     return Array(buffer)
 end
 
+if isdefined(Base, :Memory)
+    function Array(
+        client::Client, memory::Base.Memory{T}, sharding::Sharding
+    ) where {T<:Reactant.ReactantPrimitive}
+        all_devices = XLA.devices(sharding)
+        all_logical_device_ids = collect(Int64, 0:(length(all_devices) - 1))
+        hlo_sharding = convert(XLA.HloSharding, sharding)
+
+        slices, _ = XLA.sharding_to_concrete_array_indices(
+            hlo_sharding, size(memory), all_logical_device_ids
+        )
+
+        seen_slice = Dict{NTuple{N,UnitRange{Int64}},Int}()
+        host_buffers = Base.Array{T,1}[]
+        addressable_shard_indices = Vector{Int64}[]
+
+        cur_shard = 0
+        for (slice, device) in zip(slices, all_devices)
+            XLA.is_addressable(device) || continue
+
+            if haskey(seen_slice, slice)
+                idx = seen_slice[slice]
+                push!(addressable_shard_indices[idx], cur_shard)
+            else
+                host_buffer = let slice = memory[slice...]
+                    slice isa Number ? collect(slice) : slice
+                end
+                push!(host_buffers, host_buffer)
+                push!(addressable_shard_indices, Int64[cur_shard])
+                seen_slice[slice] = length(host_buffers)
+            end
+
+            cur_shard += 1
+        end
+
+        return Array(
+            client, host_buffers, addressable_shard_indices, size(memory), sharding
+        )
+    end
+end
+
 function Array(
     client::Client, array::Base.Array{T,N}, sharding
 ) where {T<:Reactant.ReactantPrimitive,N}
@@ -158,6 +223,23 @@ function Array(
     return Array(client, array, ifrt_sharding)
 end
 
+if isdefined(Base, :Memory)
+    function Array(
+        client::Client, memory::Base.Memory{T}, sharding
+    ) where {T<:Reactant.ReactantPrimitive}
+        @assert sharding isa Reactant.Sharding.AbstractSharding
+        if !(sharding isa Reactant.Sharding.HloSharding)
+            sharding = Reactant.Sharding.HloSharding(sharding, size(memory))
+        end
+
+        (; hlo_sharding, mesh) = sharding
+        devices = XLA.get_device.((client,), mesh.device_ids)
+        ifrt_sharding = Sharding([devices...], hlo_sharding)
+
+        return Array(client, memory, ifrt_sharding)
+    end
+end
+
 @inline function XLA.free_buffer(buffer::Array)
     if buffer.buffer != C_NULL
         @ccall MLIR.API.mlir_c.ifrt_free_array(buffer.buffer::Ptr{Cvoid})::Cvoid

src/xla/PJRT/Buffer.jl

@@ -21,6 +21,23 @@ function Buffer(client::Client, array::Array{T,N}, device::Device) where {T,N}
     return Buffer(buffer)
 end
 
+if isdefined(Base, :Memory)
+    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
+end
+
 function Base.similar(a::Buffer)
     buffer = GC.@preserve a begin
         @ccall MLIR.API.mlir_c.UninitPJRTBuffer(

test/Project.toml

@@ -9,6 +9,7 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"

test/integration/fixedsizearrays.jl

@@ -0,0 +1,17 @@
+
+using Reactant, Test, FixedSizeArrays
+
+fn(x, y) = (2 .* x .- 3) * y'
+
+@testset "FixedSizeArrays" begin
+    @testset "1D" begin
+        x = FixedSizeArray(fill(3.0f0, 100))
+        rx = Reactant.to_rarray(x)
+        @test @jit(fn(rx, rx)) ≈ fn(x, x)
+    end
+    @testset "2D" begin
+        x = FixedSizeArray(fill(3.0f0, (4, 5)))
+        rx = Reactant.to_rarray(x)
+        @test @jit(fn(rx, rx)) ≈ fn(x, x)
+    end
+end

test/memory.jl

@@ -0,0 +1,10 @@
+using Reactant, Test
+
+fn(x, y) = sin.(x) .+ cos.(y)
+
+@testset "Memory test" begin
+    x = Memory{Float32}(fill(2.0f0, 10))
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(fn(x_ra, x_ra)) ≈ fn(x, x)
+end