Turn off struct broadcast by default. (#260)

* Turn off struct broadcast by default.

* add doc and rename

* Add more unalias test.

* Also overload `instantiate` and `_axes`

This makes sure `Tuple`'s axes static during struct static broadcast.

* Update wording

---------

Co-authored-by: Pietro Vertechi <[email protected]>

Author: N5N3

src/StructArrays.jl

@@ -38,5 +38,6 @@ function GPUArraysCore.backend(::Type{T}) where {T<:StructArray}
     isconsistent || throw(ArgumentError("all component arrays must have the same GPU backend"))
     return backend
 end
+always_struct_broadcast(::GPUArraysCore.AbstractGPUArrayStyle) = true
 
 end # module

src/staticarrays_support.jl

@@ -38,35 +38,10 @@ end
 # This looks costly, but the compiler should be able to optimize them away
 Broadcast._axes(bc::Broadcasted{<:StructStaticArrayStyle}, ::Nothing) = axes(replace_structarray(bc))
 
-to_staticstyle(@nospecialize(x::Type)) = x
-to_staticstyle(::Type{StructStaticArrayStyle{N}}) where {N} = StaticArrayStyle{N}
-
-"""
-    replace_structarray(bc::Broadcasted)
-
-An internal function transforms the `Broadcasted` with `StructArray` into
-an equivalent one without it. This is not a must if the root `BroadcastStyle`
-supports `AbstractArray`. But some `BroadcastStyle` limits the input array types, 
-e.g. `StaticArrayStyle`, thus we have to omit all `StructArray`.
-"""
-function replace_structarray(bc::Broadcasted{Style}) where {Style}
-    args = replace_structarray_args(bc.args)
-    return Broadcasted{to_staticstyle(Style)}(bc.f, args, nothing)
-end
-function replace_structarray(A::StructArray)
-    f = Instantiator(eltype(A))
-    args = Tuple(components(A))
-    return Broadcasted{StaticArrayStyle{ndims(A)}}(f, args, nothing)
-end
-replace_structarray(@nospecialize(A)) = A
-
-replace_structarray_args(args::Tuple) = (replace_structarray(args[1]), replace_structarray_args(tail(args))...)
-replace_structarray_args(::Tuple{}) = ()
-
 # StaticArrayStyle has no similar defined.
 # Overload `Base.copy` instead.
-@inline function Base.copy(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M}
-    sa = copy(convert(Broadcasted{StaticArrayStyle{M}}, bc))
+@inline function try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M}
+    sa = copy(bc)
     ET = eltype(sa)
     isnonemptystructtype(ET) || return sa
     elements = Tuple(sa)

src/structarray.jl

@@ -494,7 +494,8 @@ function Base.showarg(io::IO, s::StructArray{T}, toplevel) where T
 end
 
 # broadcast
-import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, Unknown
+import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, Unknown, ArrayConflict
+using Base.Broadcast: combine_styles
 
 struct StructArrayStyle{S, N} <: AbstractArrayStyle{N} end
 
@@ -524,6 +525,82 @@ Base.@pure cst(::Type{SA}) where {SA} = combine_style_types(array_types(SA).para
 
 BroadcastStyle(::Type{SA}) where {SA<:StructArray} = StructArrayStyle{typeof(cst(SA)), ndims(SA)}()
 
+"""
+    always_struct_broadcast(style::BroadcastStyle)
+
+Check if `style` supports struct-broadcast natively, which means:
+1) `Base.copy` is not overloaded.
+2) `Base.similar` is defined.
+3) `Base.copyto!` supports `StructArray`s as broadcasted arguments.
+
+If any of the above conditions are not met, then this function should
+not be overloaded.
+In that case, try to overload [`try_struct_copy`](@ref) to support out-of-place
+struct-broadcast.
+"""
+always_struct_broadcast(::Any) = false
+always_struct_broadcast(::DefaultArrayStyle) = true
+always_struct_broadcast(::ArrayConflict) = true
+
+"""
+    try_struct_copy(bc::Broadcasted)
+
+Entry for non-native outplace struct-broadcast.
+
+See also [`always_struct_broadcast`](@ref).
+"""
+try_struct_copy(bc::Broadcasted) = copy(bc)
+
+function Base.copy(bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N}
+    if always_struct_broadcast(S())
+        return invoke(copy, Tuple{Broadcasted}, bc)
+    else
+        return try_struct_copy(replace_structarray(bc))
+    end
+end
+
+"""
+    replace_structarray(bc::Broadcasted)
+
+An internal function transforms the `Broadcasted` with `StructArray` into
+an equivalent one without it. This is not a must if the root `BroadcastStyle`
+supports `AbstractArray`. But some `BroadcastStyle` limits the input array types, 
+e.g. `StaticArrayStyle`, thus we have to omit all `StructArray`.
+"""
+function replace_structarray(bc::Broadcasted{Style}) where {Style}
+    args = replace_structarray_args(bc.args)
+    Style′ = parent_style(Style())
+    return Broadcasted{Style′}(bc.f, args, bc.axes)
+end
+function replace_structarray(A::StructArray)
+    f = Instantiator(eltype(A))
+    args = Tuple(components(A))
+    Style = typeof(combine_styles(args...))
+    return Broadcasted{Style}(f, args, axes(A))
+end
+replace_structarray(@nospecialize(A)) = A
+
+replace_structarray_args(args::Tuple) = (replace_structarray(args[1]), replace_structarray_args(tail(args))...)
+replace_structarray_args(::Tuple{}) = ()
+
+parent_style(@nospecialize(x)) = typeof(x)
+parent_style(::StructArrayStyle{S, N}) where {S, N} = S
+parent_style(::StructArrayStyle{S, N}) where {N, S<:AbstractArrayStyle{N}} = S
+parent_style(::StructArrayStyle{S, N}) where {S<:AbstractArrayStyle{Any}, N} = S
+parent_style(::StructArrayStyle{S, N}) where {S<:AbstractArrayStyle, N} = typeof(S(Val(N)))
+
+# `instantiate` and `_axes` might be overloaded for static axes.
+function Broadcast.instantiate(bc::Broadcasted{Style}) where {Style <: StructArrayStyle}
+    Style′ = parent_style(Style())
+    bc′ = Broadcast.instantiate(convert(Broadcasted{Style′}, bc))
+    return convert(Broadcasted{Style}, bc′)
+end
+
+function Broadcast._axes(bc::Broadcasted{Style}, ::Nothing) where {Style <: StructArrayStyle}
+    Style′ = parent_style(Style())
+    return Broadcast._axes(convert(Broadcasted{Style′}, bc), nothing)
+end
+
 # Here we use `similar` defined for `S` to build the dest Array.
 function Base.similar(bc::Broadcasted{StructArrayStyle{S, N}}, ::Type{ElType}) where {S, N, ElType}
     bc′ = convert(Broadcasted{S}, bc)
@@ -532,12 +609,22 @@ end
 
 # Unwrapper to recover the behaviour defined by parent style.
 @inline function Base.copyto!(dest::AbstractArray, bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N}
-    return copyto!(dest, convert(Broadcasted{S}, bc))
+    bc′ = always_struct_broadcast(S()) ? convert(Broadcasted{S}, bc) : replace_structarray(bc)
+    return copyto!(dest, bc′)
 end
 
 @inline function Broadcast.materialize!(::StructArrayStyle{S}, dest, bc::Broadcasted) where {S}
-    return Broadcast.materialize!(S(), dest, bc)
+    bc′ = always_struct_broadcast(S()) ? bc : replace_structarray(bc)
+    return Broadcast.materialize!(S(), dest, bc′)
 end
 
 # for aliasing analysis during broadcast
+function Broadcast.broadcast_unalias(dest::StructArray, src::AbstractArray)
+    if dest === src || any(Base.Fix2(===, src), components(dest))
+        return src
+    else
+        return Base.unalias(dest, src)
+    end
+end
+
 Base.dataids(u::StructArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), values(components(u)), init=())

test/runtests.jl

@@ -1183,6 +1183,7 @@ for S in (1, 2, 3)
         Base.setindex!(A::$MyArray, val, i::Int) = A.A[i] = val
         Base.size(A::$MyArray) = Base.size(A.A)
         Base.BroadcastStyle(::Type{<:$MyArray}) = Broadcast.ArrayStyle{$MyArray}()
+        StructArrays.always_struct_broadcast(::Broadcast.ArrayStyle{$MyArray}) = true
     end
 end
 Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray1}}, ::Type{ElType}) where ElType =
@@ -1247,6 +1248,16 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
     @test Base.broadcasted(+, reshape(1:2,1,1,2), s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{3}}
     @test Base.broadcasted(+, s, MyArray1(rand(2))) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{Any}}
 
+    #parent_style
+    @test StructArrays.parent_style(StructArrayStyle{Broadcast.DefaultArrayStyle{0},2}()) == Broadcast.DefaultArrayStyle{2}
+    @test StructArrays.parent_style(StructArrayStyle{Broadcast.Style{Tuple},2}()) == Broadcast.Style{Tuple}
+
+    # allocation test for overloaded `broadcast_unalias`
+    StructArrays.always_struct_broadcast(::Broadcast.ArrayStyle{MyArray1}) = false
+    f(s) = s .+= 1
+    f(s)
+    @test (@allocated f(s)) == 0
+    
     # issue #185
     A = StructArray(randn(ComplexF64, 3, 3))
     B = randn(ComplexF64, 3, 3)
@@ -1288,6 +1299,8 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         a = StructArray{ComplexF64}(undef, 1)
         allocated(a) = @allocated  a .+ 1
         @test allocated(a) == 2allocated(a.re)
+        allocated2(a) = @allocated a .= complex.(a.im, a.re)
+        @test allocated2(a) == 0
     end
 
     @testset "StructStaticArray" begin
@@ -1299,7 +1312,7 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         @test (@inferred bclog(s)) isa typeof(s)
         test_allocated(bclog, s)
         @test abs.(s) .+ ((1,) .+ (1,2,3,4,5,6,7,8,9,10)) isa SMatrix
-        bc = Base.broadcasted(+, s, s);
+        bc = Base.broadcasted(+, s, s, ntuple(identity, 10));
         bc = Base.broadcasted(+, bc, bc, s);
         @test @inferred(Broadcast.axes(bc)) === axes(s)
     end
@@ -1317,6 +1330,14 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         @test backend(bcmul2(sa)) === backend(sa)
         @test (sa .+= 1) === sa
     end
+
+    @testset "StructSparseArray" begin
+        a = sprand(10, 10, 0.5)
+        b = sprand(10, 10, 0.5)
+        c = StructArray{ComplexF64}((a, b))
+        d = identity.(c)
+        @test d isa SparseMatrixCSC
+    end
 end
 
 @testset "map" begin