Merge branch 'main' into ap/updated_no_nan

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.153"
+version = "0.2.155"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -32,6 +32,7 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -54,6 +55,7 @@ ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
 ReactantDLFP8TypesExt = "DLFP8Types"
+ReactantFillArraysExt = "FillArrays"
 ReactantFloat8sExt = "Float8s"
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
@@ -77,6 +79,7 @@ Downloads = "1.6"
 EnumX = "1"
 Enzyme = "0.13.49"
 EnzymeCore = "0.8.11"
+FillArrays = "1.13"
 Float8s = "0.1"
 Functors = "0.5"
 GPUArraysCore = "0.2"
@@ -97,15 +100,15 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.232"
+Reactant_jll = "0.0.233"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
-unzip_jll = "6"
 YaoBlocks = "0.13, 0.14"
 julia = "1.10"
+unzip_jll = "6"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

deps/ReactantExtra/WORKSPACE

@@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"
 
 NSYNC_SHA256 = ""
 
-ENZYMEXLA_COMMIT = "ad1e0392bcea0afaf75d12445608a08550055ea6"
+ENZYMEXLA_COMMIT = "45f1d0d5a47ee706500eff8841b710b9da112ec6"
 
 ENZYMEXLA_SHA256 = ""
 

docs/src/api/api.md

@@ -21,6 +21,13 @@ within_compile
 @trace
 ```
 
+## Reactant data types
+
+```@docs
+ConcreteRArray
+ConcreteRNumber
+```
+
 ## Inspect Generated HLO
 
 ```@docs

docs/src/introduction/index.md

@@ -13,8 +13,8 @@ Pkg.add("Reactant")
 
 ## Quick Start
 
-Reactant provides two new array types at its core, a ConcreteRArray and a TracedRArray. A
-ConcreteRArray is an underlying buffer to whatever device data you wish to store and can be
+Reactant provides two new array types at its core, a [`ConcreteRArray`](@ref) and a `TracedRArray`. A
+`ConcreteRArray` is an underlying buffer to whatever device data you wish to store and can be
 created by converting from a regular Julia Array.
 
 ```@example quickstart

ext/ReactantFillArraysExt.jl

@@ -0,0 +1,121 @@
+module ReactantFillArraysExt
+
+using Reactant: Reactant, TracedUtils, TracedRNumber, Ops, Sharding, unwrapped_eltype
+using ReactantCore: ReactantCore
+using FillArrays: FillArrays, AbstractFill, Fill, Ones, Zeros, OneElement
+using GPUArraysCore: @allowscalar
+
+# Tracing
+Reactant._parent_type(T::Type{<:AbstractFill}) = T
+Reactant._parent_type(T::Type{<:OneElement}) = T
+
+for AT in (Fill, Ones, Zeros)
+    @eval Base.@nospecializeinfer function Reactant.traced_type_inner(
+        @nospecialize(FA::Type{$(AT){T,N,Axes}}),
+        seen,
+        mode::Reactant.TraceMode,
+        @nospecialize(track_numbers::Type),
+        @nospecialize(sharding),
+        @nospecialize(runtime)
+    ) where {T,N,Axes}
+        # T will be a number so we need to trace it
+        return $(AT){
+            Reactant.traced_type_inner(T, seen, mode, Number, sharding, runtime),N,Axes
+        }
+    end
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen, @nospecialize(prev::Fill{T,N,Axes}), @nospecialize(path), mode; kwargs...
+) where {T,N,Axes}
+    return Fill(
+        Reactant.make_tracer(
+            seen, prev.value, (path..., 1), mode; kwargs..., track_numbers=Number
+        ),
+        prev.axes,
+    )
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen,
+    @nospecialize(prev::Ones{T,N,Axes}),
+    @nospecialize(path),
+    mode;
+    @nospecialize(sharding = Sharding.NoSharding()),
+    @nospecialize(runtime = nothing),
+    kwargs...,
+) where {T,N,Axes}
+    return Ones(
+        Reactant.traced_type_inner(T, seen, mode, Number, sharding, runtime), prev.axes
+    )
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen,
+    @nospecialize(prev::Zeros{T,N,Axes}),
+    @nospecialize(path),
+    mode;
+    @nospecialize(sharding = Sharding.NoSharding()),
+    @nospecialize(runtime = nothing),
+    kwargs...,
+) where {T,N,Axes}
+    return Zeros(
+        Reactant.traced_type_inner(T, seen, mode, Number, sharding, runtime), prev.axes
+    )
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(FA::Type{OneElement{T,N,I,A}}),
+    seen,
+    mode::Reactant.TraceMode,
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T,N,I,A}
+    # T will be a number so we need to trace it
+    return OneElement{
+        Reactant.traced_type_inner(T, seen, mode, Number, sharding, runtime),N,I,A
+    }
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen, @nospecialize(prev::OneElement{T,N,I,A}), @nospecialize(path), mode; kwargs...
+) where {T,N,I,A}
+    return OneElement(
+        Reactant.make_tracer(
+            seen, prev.val, (path..., 1), mode; kwargs..., track_numbers=Number
+        ),
+        prev.ind,
+        prev.axes,
+    )
+end
+
+# Materialize into a dense array
+function ReactantCore.materialize_traced_array(x::Fill{T}) where {T}
+    return TracedUtils.broadcast_to_size(
+        TracedUtils.promote_to(TracedRNumber{unwrapped_eltype(T)}, x.value), size(x)
+    )
+end
+
+function ReactantCore.materialize_traced_array(x::Ones{T}) where {T}
+    return TracedUtils.broadcast_to_size(unwrapped_eltype(T)(1), size(x))
+end
+
+function ReactantCore.materialize_traced_array(x::Zeros{T}) where {T}
+    return TracedUtils.broadcast_to_size(unwrapped_eltype(T)(0), size(x))
+end
+
+function ReactantCore.materialize_traced_array(x::OneElement{T}) where {T}
+    y = TracedUtils.broadcast_to_size(unwrapped_eltype(T)(0), size(x))
+    @allowscalar setindex!(y, x.val, x.ind...)
+    return y
+end
+
+# some functions to avoid bad performance
+for AT in (Fill, Ones, Zeros, OneElement)
+    @eval function Base.similar(x::$AT{<:TracedRNumber}, ::Type{T}, dims::Dims) where {T}
+        return TracedUtils.broadcast_to_size(unwrapped_eltype(T)(0), dims)
+    end
+end
+
+end

ext/ReactantKernelAbstractionsExt.jl

@@ -10,6 +10,9 @@ using Adapt: Adapt
 
 export ReactantBackend
 
+# ToDo: Include XLA client, device and sharding in ReactantBackend struct, to
+# support more complex applications? If so, need to adapt implementation of
+# `KA.get_backend` and `KA.allocate` accordingly.
 struct ReactantBackend <: KA.GPU end
 
 function Base.getproperty(x::ReactantBackend, sym::Symbol)
@@ -22,16 +25,23 @@ function Base.getproperty(x::ReactantBackend, sym::Symbol)
     end
 end
 
-KA.allocate(n::ReactantBackend, ::Type{T}, dims::Tuple) where {T} = KA.zeros(b, T, dims)
-function KA.zeros(::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
-    return Reactant.to_rarray(zeros(T, dims))
+function KA.allocate(::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
+    return ConcreteRArray{T}(undef, dims)
 end
-function KA.ones(::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
-    return Reactant.to_rarray(ones(T, dims))
+
+function KA.zeros(b::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
+    A = KA.allocate(b, T, dims)
+    isempty(A) || fill!(A, zero(T))
+    return A
+end
+function KA.ones(b::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
+    A = KA.allocate(b, T, dims)
+    isempty(A) || fill!(A, one(T))
+    return A
 end
 
 KA.get_backend(::Reactant.AnyTracedRArray) = ReactantBackend()
-KA.get_backend(::Reactant.AnyConcretePJRTArray) = ReactantBackend()
+KA.get_backend(::Reactant.AnyConcreteRArray) = ReactantBackend()
 function KA.synchronize(::ReactantBackend) end
 
 Adapt.adapt_storage(::ReactantBackend, a::Array) = a

src/Compiler.jl

@@ -288,7 +288,7 @@ function create_result(
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
 
-        @assert haskey(result_stores, path)
+        @assert haskey(result_stores, path) "Expected $(path) in $(keys(result_stores))"
         restore = result_stores[path]
         delete!(result_stores, path)
         if path_to_shard_info !== nothing && haskey(path_to_shard_info, path)

src/ConcreteRArray.jl

@@ -83,6 +83,7 @@ for T in Base.uniontypes(ReactantPrimitive)
 end
 
 function Base.convert(::Type{T}, x::AbstractConcreteNumber) where {T<:Number}
+    T == typeof(x) && return x
     return convert(T, to_number(x))
 end
 
@@ -377,16 +378,8 @@ end
     device::Union{Nothing,XLA.PJRT.Device}=nothing,
     sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
 ) where {S}
-    client = client === nothing ? XLA.default_backend() : client
-
-    if idx isa Int && device === nothing
-        device = XLA.get_device(client, idx)
-    end
-
-    sdata, sharding = sharding(client, device, S, dims)
-
-    return ConcretePJRTArray{S,length(dims),length(sdata),typeof(sharding)}(
-        sdata, dims, sharding
+    return ConcretePJRTArray{S}(
+        undef, dims; client=client, idx=idx, device=device, sharding=sharding
     )
 end
 
@@ -417,7 +410,7 @@ function Base.similar(a::ConcreteIFRTArray{T}, ::Type{S}=T, dims::Dims=size(a))
 end
 Base.similar(a::ConcreteIFRTArray, dims::Dims) = similar(a, eltype(a), dims)
 function Base.similar(::Type{ConcreteIFRTArray{T}}, dims) where {T}
-    return ConcreteIFRTArray(similar(Array{T}, dims))
+    return ConcreteIFRTArray{T}(undef, dims)
 end
 
 # Broadcasting interface

src/Reactant.jl

@@ -42,11 +42,20 @@ function ancestor(T::Type{<:AbstractArray})
         p_T == T && return T
         return ancestor(p_T)
     end
+    if applicable(_parent_type, T)
+        p_T = _parent_type(T)
+        p_T == T && return T
+        return ancestor(p_T)
+    end
     @warn "`Adapt.parent_type` is not implemented for $(T). Assuming $T isn't a wrapped \
            array." maxlog = 1
     return T
 end
 
+# A lot of packages don't define `Adapt.parent_type`. We use `_parent_type` as a way to
+# define the parent type of an array without type-piracy.
+function _parent_type end
+
 include("accelerators/Accelerators.jl")
 
 using .Accelerators.TPU: has_tpu

src/TracedRArray.jl

@@ -568,7 +568,7 @@ function overloaded_mapreduce(
     original_dims = dims
     dims isa Int && (dims = Int64[dims])
     dims isa Colon && (dims = collect(Int64, 1:N))
-    dims isa AbstractVector{<:Integer} || (dims = collect(Int64, dims))
+    dims isa Vector{Int64} || (dims = collect(Int64, dims))
 
     op_in_T = unwrapped_eltype(Core.Compiler.return_type(f, Tuple{T}))
     reduce_init = __default_init(op_in_T, op)