diff --git a/docs/src/api/api.md b/docs/src/api/api.md index 2c0bf8eddf..1b9fa1e1df 100644 --- a/docs/src/api/api.md +++ b/docs/src/api/api.md @@ -32,6 +32,8 @@ Reactant.to_rarray ```@docs ConcreteRArray ConcreteRNumber +RArraySpec +RNumberSpec ``` ## Inspect Generated HLO diff --git a/src/Compiler.jl b/src/Compiler.jl index 4b62f914b8..9f0b885330 100644 --- a/src/Compiler.jl +++ b/src/Compiler.jl @@ -3947,7 +3947,12 @@ function __resolve_device_and_client(client, seen_args, linear_args, is_sharded) if length(linear_args) > 0 devices_list = [] for (k, v) in seen_args - !(v isa TracedRArray || v isa TracedRNumber) && continue + if ( + !(v isa TracedRArray || v isa TracedRNumber) || + (k isa Reactant.RArraySpec || k isa Reactant.RNumberSpec) + ) + continue + end buffer = k.data isa Tuple ? only(k.data) : k.data push!(devices_list, XLA.device(buffer)) end diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl index 6a57d31ba8..5bb4b96f7a 100644 --- a/src/ConcreteRArray.jl +++ b/src/ConcreteRArray.jl @@ -286,6 +286,23 @@ function Base.show(io::IO, X::Union{ConcretePJRTScalar,ConcreteIFRTScalar}) return nothing end +function Base.showarg(io::IO, ::RArraySpec{T,N}, toplevel) where {T,N} + toplevel || print(io, "::") + print(io, "RArraySpec{$T,$N}") + # TODO: Add sharding info + return nothing +end + +function Base.print_array(io::IO, ::RArraySpec) + print(io, "") + return nothing +end + +function Base.show(io::IO, X::RArraySpec) + print(io, "$(typeof(X))()") + return nothing +end + function Base.print_array(io::IO, X::Union{AnyConcretePJRTArray,AnyConcreteIFRTArray}) if isempty(X) print(io, "") diff --git a/src/Reactant.jl b/src/Reactant.jl index 6d31c41c04..ab0c36e543 100644 --- a/src/Reactant.jl +++ b/src/Reactant.jl @@ -124,6 +124,9 @@ unwrapped_eltype(::TracedRNumber{T}) where {T} = T unwrapped_eltype(::Type{<:AbstractArray{T,N}}) where {T,N} = unwrapped_eltype(T) unwrapped_eltype(::AbstractArray{T,N}) where {T,N} = unwrapped_eltype(T) +unwrapped_eltype(::Type{RArraySpec{T,N}}) where {T,N} = T +unwrapped_eltype(::RArraySpec{T,N}) where {T,N} = T + include("Ops.jl") Base.push!(no_rewrite_ancestor_modules, Ops) @@ -289,6 +292,8 @@ export ConcreteRArray, ConcretePJRTNumber, ConcreteIFRTArray, ConcreteIFRTNumber, + RArraySpec, + RNumberSpec, @compile, @code_hlo, @code_mhlo, diff --git a/src/Tracing.jl b/src/Tracing.jl index ef96f2412a..4fabe2799c 100644 --- a/src/Tracing.jl +++ b/src/Tracing.jl @@ -7,6 +7,7 @@ TracedToTypes = 6 NoStopTracedTrack = 7 TracedToJAX = 8 + ArrayToSpec = 9 end function convert_to_jax_dtype_struct end @@ -69,6 +70,8 @@ Base.@nospecializeinfer function traced_type_inner( else error("Unsupported runtime $runtime") end + elseif mode == ArrayToSpec && T <: track_numbers + return RNumberSpec{T} elseif (mode == NoStopTracedTrack || mode == TracedTrack || mode == TracedSetPath) && T <: track_numbers return TracedRNumber{T} @@ -238,7 +241,7 @@ Base.@nospecializeinfer function traced_type_inner( elseif mode == ArrayToConcrete @assert runtime isa Val{:PJRT} if T0 isa UnionAll - return ConcretePJRTNumbe{T,_unwrap_val(ndevices)} where {T} + return ConcretePJRTNumber{T,_unwrap_val(ndevices)} where {T} else return ConcretePJRTNumber{T,_unwrap_val(ndevices)} end @@ -329,6 +332,45 @@ Base.@nospecializeinfer function traced_type_inner( end end +Base.@nospecializeinfer function traced_type_inner( + @nospecialize(T::Type{<:RArraySpec}), + seen, + @nospecialize(mode::TraceMode), + @nospecialize(track_numbers::Type), + @nospecialize(ndevices), + @nospecialize(runtime) +) + if mode == ConcreteToTraced + T´ = Base.unwrap_unionall(T) + T, N = T´.parameters + T´´ = TracedRArray{T,N} + T_ret = N isa Core.TypeVar ? UnionAll(N, T´´) : T´´ + T_ret2 = T isa Core.TypeVar ? UnionAll(T, T_ret) : T_ret + return T_ret2 + else + throw("Unsupported mode: $mode") + end +end + +Base.@nospecializeinfer function traced_type_inner( + @nospecialize(T::Type{<:RNumberSpec}), + seen, + @nospecialize(mode::TraceMode), + @nospecialize(track_numbers::Type), + @nospecialize(ndevices), + @nospecialize(runtime) +) + if mode == ConcreteToTraced + T´ = Base.unwrap_unionall(T) + T = T´.parameters + T´´ = TracedRNumber{T} + T_ret = T isa Core.TypeVar ? UnionAll(T, T´´) : T´´ + return T_ret + else + throw("Unsupported mode: $mode") + end +end + Base.@nospecializeinfer function traced_type_inner( @nospecialize(T::Type{MissingTracedValue}), seen, @@ -482,6 +524,8 @@ Base.@nospecializeinfer function traced_type_inner( else error("Unsupported runtime $runtime") end + elseif mode == ArrayToSpec && T <: ReactantPrimitive + A_wrapper = RArraySpec end # WARN replacing typevars first is required to construct the UnionAlls correctly @@ -504,6 +548,8 @@ Base.@nospecializeinfer function traced_type_inner( end end error("Unsupported runtime $runtime") + elseif mode == ArrayToSpec && T <: ReactantPrimitive + return RArraySpec{T,N} else return Array{ traced_type_inner(T, seen, mode, track_numbers, ndevices, runtime),N @@ -520,7 +566,7 @@ Base.@nospecializeinfer function traced_type_inner( @nospecialize(ndevices), @nospecialize(runtime) ) - if mode == ArrayToConcrete + if mode == ArrayToConcrete || mode == ArrayToSpec A´ = A isa UnionAll ? Array{Bool} : Array{Bool,ndims(A)} return traced_type_inner(A´, seen, mode, track_numbers, ndevices, runtime) else @@ -617,7 +663,7 @@ Base.@nospecializeinfer function traced_type_inner( @nospecialize(ndevices), @nospecialize(runtime) ) - if mode == ArrayToConcrete + if mode == ArrayToConcrete || mode == ArrayToSpec return ReactantRNG{ traced_type_inner(Array{UInt64,1}, seen, mode, track_numbers, ndevices, runtime) } @@ -1407,6 +1453,38 @@ Base.@nospecializeinfer function make_tracer( return res end +Base.@nospecializeinfer function make_tracer( + seen, @nospecialize(prev::RArraySpec{T,N}), @nospecialize(path), mode; kwargs... +) where {T,N} + if mode != ConcreteToTraced + throw( + ArgumentError( + "RArraySpec can only be used with ConcreteToTraced mode, got $mode" + ), + ) + end + haskey(seen, prev) && return seen[prev]::TracedRArray{T,N} + res = TracedRArray{T,N}((path,), nothing, size(prev)) + seen[prev] = res + return res +end + +Base.@nospecializeinfer function make_tracer( + seen, @nospecialize(prev::RNumberSpec{T}), @nospecialize(path), mode; kwargs... +) where {T} + if mode != ConcreteToTraced + throw( + ArgumentError( + "RNumberSpec can only be used with ConcreteToTraced mode, got $mode" + ), + ) + end + haskey(seen, prev) && return seen[prev]::TracedRNumber{T} + res = TracedRNumber{T}((path,), nothing) + seen[prev] = res + return res +end + Base.@nospecializeinfer function make_tracer( seen, prev::ConcretePJRTNumber{T}, @@ -1688,6 +1766,8 @@ Base.@nospecializeinfer function make_tracer( runtime isa Val{:IFRT} && return ConcreteIFRTNumber(prev; sharding, device, client) error("Unsupported runtime $runtime") + elseif mode == ArrayToSpec + return RNumberSpec{RT}(; sharding) else if mode == TracedTrack || mode == NoStopTracedTrack res = TracedRNumber{RT}((path,), broadcast_to_size(prev, ()).mlir_data) @@ -1780,6 +1860,8 @@ Base.@nospecializeinfer function make_tracer( runtime isa Val{:IFRT} && (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client)) error("Unsupported runtime $runtime") + elseif mode == ArrayToSpec + return seen[prev] = RArraySpec{eltype(RT),ndims(RT)}(size(prev); sharding) elseif mode == TracedToTypes # Original array can get mutated so we store a copy: push!(path, copy(prev)) @@ -1842,7 +1924,7 @@ end Base.@nospecializeinfer function make_tracer( seen, @nospecialize(prev::BitArray), @nospecialize(path), mode; kwargs... ) - if mode == ArrayToConcrete + if mode == ArrayToConcrete || mode == ArrayToSpec return make_tracer(seen, Array(prev), path, mode; kwargs...) else return prev @@ -1877,6 +1959,8 @@ Base.@nospecializeinfer function make_tracer( runtime isa Val{:IFRT} && (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client)) error("Unsupported runtime $runtime") + elseif mode == ArrayToSpec + return seen[prev] = RArraySpec{eltype(prev),ndims(prev)}(size(prev); sharding) elseif mode == TracedToTypes # Original array can get mutated so we store a copy: push!(path, copy(prev)) @@ -2084,7 +2168,7 @@ end Base.@nospecializeinfer function make_tracer( seen, @nospecialize(prev::Random.AbstractRNG), @nospecialize(path), mode; kwargs... ) - if mode == ArrayToConcrete + if mode == ArrayToConcrete || mode == ArrayToSpec TracedRandom.should_warn_if_not_natively_supported(prev) return ReactantRNG( make_tracer(seen, TracedRandom.make_seed(prev), (path..., 1), mode; kwargs...), @@ -2095,7 +2179,14 @@ Base.@nospecializeinfer function make_tracer( end """ - to_rarray(x; track_numbers=false, sharding=NoSharding(), device=nothing, client=nothing, runtime=nothing) + to_rarray( + x; + track_numbers=false, + sharding=NoSharding(), + device=nothing, + client=nothing, + runtime=nothing, + ) Convert a Julia value `x` into its Reactant equivalent by tracing through the structure. Arrays are converted to `ConcreteRArray`, and (optionally) scalar numbers are converted @@ -2114,6 +2205,9 @@ to `ConcreteRNumber`. - `device`: Target device for the resulting array. - `client`: XLA client to use. - `runtime`: Backend runtime to use (`Val(:PJRT)` or `Val(:IFRT)`). +- `convert_to_specification::Union{Val{true},Val{false}}=Val{false}()`: Whether to convert + the arrays to a specification (RArraySpec/RNumberSpec) instead of a concrete array. + This enables AoT compiling reactant functions without allocating the actual arrays. ## Examples @@ -2142,13 +2236,17 @@ become compile-time constants. sharding=Sharding.Sharding.NoSharding(), device=nothing, client=nothing, + convert_to_specification::Union{Val{true},Val{false}}=Val{false}(), ) runtime === nothing && (runtime = XLA.runtime()) track_numbers isa Bool && (track_numbers = track_numbers ? Number : Union{}) - return to_rarray_internal(x, track_numbers, sharding, runtime, device, client) + return to_rarray_internal( + convert_to_specification, x, track_numbers, sharding, runtime, device, client + ) end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2160,7 +2258,7 @@ end OrderedIdDict(), x, (), - ArrayToConcrete; + convert_to_specification isa Val{true} ? ArrayToSpec : ArrayToConcrete; track_numbers, sharding, runtime, @@ -2171,6 +2269,7 @@ end # fast paths avoiding make_tracer function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(::TracedRArray), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2182,6 +2281,7 @@ function to_rarray_internal( end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::ConcretePJRTArray), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2193,6 +2293,7 @@ end end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::ConcreteIFRTArray), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2204,6 +2305,7 @@ end end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2211,12 +2313,15 @@ end @nospecialize(device), @nospecialize(client) ) + convert_to_specification isa Val{true} && + return RArraySpec{eltype(x),ndims(x)}(size(x); sharding) runtime isa Val{:PJRT} && return ConcretePJRTArray(x; sharding, device, client) runtime isa Val{:IFRT} && return ConcreteIFRTArray(x; sharding, device, client) return error("Unsupported runtime $runtime") end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::Array{T}), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2225,6 +2330,8 @@ end @nospecialize(client) ) where {T<:Number} if reactant_primitive(T) !== nothing + convert_to_specification isa Val{true} && + return RArraySpec{to_reactant_primitive(T),ndims(x)}(size(x); sharding) if runtime isa Val{:PJRT} return ConcretePJRTArray(to_reactant_primitive.(x); sharding, device, client) elseif runtime isa Val{:IFRT} @@ -2233,11 +2340,18 @@ end error("Unsupported runtime $runtime") end return @invoke to_rarray_internal( - x::Any, track_numbers::Type, sharding, runtime, device, client + convert_to_specification, + x::Any, + track_numbers::Type, + sharding, + runtime, + device, + client, ) end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::ConcretePJRTNumber), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2249,6 +2363,7 @@ end end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::ConcreteIFRTNumber), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2260,6 +2375,7 @@ end end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::ReactantPrimitive), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2268,6 +2384,7 @@ end @nospecialize(client) ) if typeof(x) <: track_numbers + convert_to_specification isa Val{true} && return RNumberSpec{eltype(x)}(; sharding) runtime isa Val{:PJRT} && return ConcretePJRTNumber(x; sharding, device, client) runtime isa Val{:IFRT} && return ConcreteIFRTNumber(x; sharding, device, client) error("Unsupported runtime $runtime") @@ -2276,6 +2393,7 @@ end end @inline function to_rarray_internal( + @nospecialize(convert_to_specification), @nospecialize(x::Number), @nospecialize(track_numbers::Type), @nospecialize(sharding), @@ -2284,14 +2402,22 @@ end @nospecialize(client) ) if reactant_primitive(typeof(x)) !== nothing + convert_to_specification isa Val{true} && + return RNumberSpec{to_reactant_primitive(eltype(x))}(; sharding) runtime isa Val{:PJRT} && - return ConcretePJRTArray(to_reactant_primitive(x); sharding, device, client) + return ConcretePJRTNumber(to_reactant_primitive(x); sharding, device, client) runtime isa Val{:IFRT} && - return ConcreteIFRTArray(to_reactant_primitive(x); sharding, device, client) + return ConcreteIFRTNumber(to_reactant_primitive(x); sharding, device, client) error("Unsupported runtime $runtime") end return @invoke to_rarray_internal( - x::Any, track_numbers::Type, sharding, runtime, device, client + convert_to_specification, + x::Any, + track_numbers::Type, + sharding, + runtime, + device, + client, ) end diff --git a/src/Types.jl b/src/Types.jl index 409a69d7cc..c2e24e1924 100644 --- a/src/Types.jl +++ b/src/Types.jl @@ -135,6 +135,91 @@ const AnyTracedRVector{T} = AnyTracedRArray{T,1} const AnyTracedRMatrix{T} = AnyTracedRArray{T,2} const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}} +## RArraySpec +""" + RArraySpec{T,N}(shape::NTuple{N,Int}) + RArraySpec{T}(shape::NTuple{N,Int}) + +Lightweight structure that specifies the shape and element type (dtype) of an array +without allocating the actual array data. Similar to JAX's `RArraySpec`. + +This is useful for compiling functions without constructing the full `ConcreteRArray`, +which can save memory and improve compilation performance. + +# Examples +```julia +# Specify shape and dtype for a 2D array +spec = Reactant.RArraySpec{Float32}((10, 20)) + +# Compile a function using just the spec +f(x) = sum(x) +compiled_f = Reactant.compile(f, (spec,)) + +# Execute with actual data +x = Reactant.ConcreteRArray(rand(Float32, 10, 20)) +result = compiled_f(x) +``` + +See also: [`compile`](@ref), [`ConcreteRArray`](@ref) +""" +struct RArraySpec{T,N} <: RArray{T,N} + shape::NTuple{N,Int} + # TODO: Sharding + + function RArraySpec{T,N}( + shape::NTuple{N,Int}; sharding=Sharding.NoShardInfo() + ) where {T,N} + return new{T,N}(shape) + end +end + +function RArraySpec{T}(shape::NTuple{N,Int}; sharding=Sharding.NoShardInfo()) where {T,N} + return RArraySpec{T,N}(shape; sharding) +end + +Base.size(x::RArraySpec) = x.shape +Base.ndims(::RArraySpec{T,N}) where {T,N} = N +Base.eltype(::RArraySpec{T}) where {T} = T + +@leaf RArraySpec + +""" + RNumberSpec{T}() + +Lightweight structure that specifies the element type (dtype) of a number +without allocating the actual number data. Similar to JAX's `RNumberSpec`. + +This is useful for compiling functions without constructing the full `ConcreteRNumber`, +which can save memory and improve compilation performance. + +# Examples +```julia +# Specify dtype for a number +spec = Reactant.RNumberSpec{Float32}() + +# Compile a function using just the spec +f(x) = x + 1 +compiled_f = Reactant.compile(f, (spec,)) + +# Execute with actual data +x = Reactant.ConcreteRNumber(1.0f0) +result = compiled_f(x) +``` + +See also: [`compile`](@ref), [`ConcreteRNumber`](@ref) +""" +struct RNumberSpec{T} <: RNumber{T} + # TODO: Sharding + + function RNumberSpec{T}(; sharding=Sharding.NoShardInfo()) where {T} + return new{T}() + end +end + +Base.eltype(::RNumberSpec{T}) where {T} = T + +@leaf RNumberSpec + # Concrete Types ## ConcretePJRTNumber mutable struct ConcretePJRTNumber{T,D} <: AbstractConcreteNumber{T} diff --git a/test/core/compile.jl b/test/core/compile.jl index a3a07c4300..2e8e6613e3 100644 --- a/test/core/compile.jl +++ b/test/core/compile.jl @@ -640,3 +640,61 @@ end @test Array(y[:Mhalo]) ≈ [1.0f0, 2.0f0] @test Array(y[:x]) ≈ [2.0f0, 3.0f0] end + +@testset "RArraySpec compilation" begin + @testset "Basic compilation with RArraySpec" begin + # Compile using RArraySpec instead of ConcreteRArray + spec = RArraySpec{Float32}((10, 20)) + compiled_f = Reactant.compile(sum, (spec,)) + + # Execute with actual data + x = Reactant.ConcreteRArray(rand(Float32, 10, 20)) + result = compiled_f(x) + + @test result isa Reactant.ConcreteRNumber{Float32} + @test result ≈ sum(Array(x)) + end + + @testset "Multiple arguments with RArraySpec" begin + spec1 = Reactant.RArraySpec((5, 5), Float64) + spec2 = Reactant.RArraySpec((5, 5), Float64) + compiled_f = Reactant.compile(.+, (spec1, spec2)) + + x = Reactant.ConcreteRArray(rand(Float64, 5, 5)) + y = Reactant.ConcreteRArray(rand(Float64, 5, 5)) + result = compiled_f(x, y) + + @test result isa Reactant.ConcreteRArray{Float64,2} + @test result ≈ Array(x) .+ Array(y) + end + + @testset "RArraySpec with different dtypes" begin + for dtype in [Float32, Float64] + spec = Reactant.RArraySpec((10,), dtype) + compiled_f = Reactant.compile(Base.BroadcastFunction(sin), (spec,)) + + x = Reactant.ConcreteRArray(rand(dtype, 10)) + result = compiled_f(x) + + @test result isa Reactant.ConcreteRArray{dtype,1} + @test result ≈ sin.(Array(x)) + end + end + + @testset "RArraySpec constructor variations" begin + # Test different constructor forms + spec1 = Reactant.RArraySpec{Float32,2}((3, 4)) + @test size(spec1) == (3, 4) + @test eltype(spec1) == Float32 + @test ndims(spec1) == 2 + + spec2 = Reactant.RArraySpec((3, 4), Float32) + @test size(spec2) == (3, 4) + @test eltype(spec2) == Float32 + + # Test with integer tuple (not Int tuple) + spec3 = Reactant.RArraySpec((3, 4), Float64) + @test size(spec3) == (3, 4) + @test eltype(spec3) == Float64 + end +end