initial StructArrays extension

Author: jumerckx

Project.toml

@@ -26,6 +26,7 @@ Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 p7zip_jll = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
 
 [weakdeps]
@@ -71,6 +72,7 @@ ReactantRandom123Ext = "Random123"
 ReactantSparseArraysExt = "SparseArrays"
 ReactantSpecialFunctionsExt = "SpecialFunctions"
 ReactantStatisticsExt = "Statistics"
+ReactantStructArraysExt = "StructArrays"
 ReactantYaoBlocksExt = "YaoBlocks"
 ReactantZygoteExt = "Zygote"
 
@@ -115,6 +117,7 @@ Sockets = "1.10"
 SparseArrays = "1.10"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
+StructArrays = "0.7.2"
 YaoBlocks = "0.13, 0.14"
 Zygote = "0.7"
 julia = "1.10"

ext/ReactantStructArraysExt.jl

@@ -0,0 +1,74 @@
+module ReactantStructArraysExt
+
+import Reactant
+import StructArrays
+
+import StructArrays: StructArrayStyle, StructArray, StructVector, index_type
+import Reactant: TraceMode, TracedToTypes, traced_type_inner, append_path, make_tracer, traced_type
+import Reactant.TracedRArrayOverrides: AbstractReactantArrayStyle, _copy
+import Base.Broadcast: Broadcasted
+
+StructArrays.always_struct_broadcast(::AbstractReactantArrayStyle) = true
+
+function Base.copy(bc::Broadcasted{StructArrays.StructArrayStyle{S, N}}) where {S<:AbstractReactantArrayStyle, N}
+    return _copy(bc)
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(prev::Type{<:StructVector{NT}}),
+    seen,
+    @nospecialize(mode::TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {NT <: NamedTuple}
+    T, N, C, I = prev.parameters
+    C_traced = traced_type_inner(
+        C,
+        seen,
+        mode,
+        track_numbers,
+        sharding,
+        runtime,
+    )
+    T_traced = traced_type_inner(
+        T,
+        seen,
+        mode,
+        # The elements in the NamedTuple are backed by vectors,
+        # these vectors are converted to RArrays so we need to track numbers:
+        Number #= track_numbers =#,
+        sharding,
+        runtime,
+    )
+    return StructVector{T_traced, C_traced, index_type(fieldtypes(C_traced))}
+end
+
+function Reactant.make_tracer(
+    seen, @nospecialize(prev::StructVector{NT}), @nospecialize(path), mode; track_numbers=false, sharding=Reactant.Sharding.Sharding.NoSharding(), runtime=nothing, kwargs...
+) where {NT <: NamedTuple}
+    track_numbers isa Bool && (track_numbers = track_numbers ? Number : Union{})
+    components = getfield(prev, :components)
+    if mode == TracedToTypes
+        push!(path, typeof(prev))
+        for c in components
+            make_tracer(seen, c, path, mode; track_numbers, sharding, runtime, kwargs...)
+        end
+        return nothing
+    end
+    traced_components = make_tracer(seen, components, append_path(path, :components), mode; track_numbers, sharding, runtime, kwargs...)
+    T_traced = traced_type(
+        typeof(prev),
+        Val(mode),
+        track_numbers,
+        sharding,
+        runtime,
+    )
+    return StructVector{first(T_traced.parameters)}(traced_components)
+end
+
+@inline function Reactant.traced_getfield(@nospecialize(obj::StructArray), field)
+    return Base.getfield(obj, field)
+end
+
+end

src/TracedRArray.jl

@@ -459,7 +459,7 @@ first_scalar(x) = @allowscalar first(x)
 
 # we need to override the outer copy method to make sure we never fall back to scalar
 # iteration (see, e.g., CUDA.jl#145)
-function Base.copy(bc::Broadcasted{<:AbstractReactantArrayStyle})
+function _copy(bc)
     fn = if bc.f isa Type && bc.f <: Reactant.ReactantPrimitive
         TracedUtils.TypeCast{bc.f}()
     else
@@ -477,6 +477,10 @@ function Base.copy(bc::Broadcasted{<:AbstractReactantArrayStyle})
     return copyto!(sim, bc)
 end
 
+@noinline function Base.copy(bc::Broadcasted{<:AbstractReactantArrayStyle})
+    return _copy(bc)
+end
+
 function Base.materialize!(
     ::Style, dest, bc::Broadcasted
 ) where {Style<:AbstractReactantArrayStyle}