Skip to content

Commit 0db2c4f

Browse files
ranochapiever
andauthored
implement specialization of Base.dataids to speed-up broadcast (#196)
* implement specialization of Base.dataids to speed-up broadcast * dataids should return a tuple of UInts Co-authored-by: Pietro Vertechi <pietro.vertechi@protonmail.com> * add explicit test of dataids Co-authored-by: Pietro Vertechi <pietro.vertechi@protonmail.com>
1 parent 9db8de0 commit 0db2c4f

File tree

2 files changed

+6
-0
lines changed

2 files changed

+6
-0
lines changed

src/structarray.jl

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -452,3 +452,6 @@ BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(S
452452

453453
Base.similar(bc::Broadcasted{StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,N,ElType} =
454454
isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc))
455+
456+
# for aliasing analysis during broadcast
457+
Base.dataids(u::StructArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), values(components(u)), init=())

test/runtests.jl

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -822,6 +822,9 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El
822822
@test isa(@inferred(s .+ r), StructArray)
823823
@test s .+ r == StructArray{ComplexF64}((s.re .+ r, s.im))
824824

825+
# used inside of broadcast but we also test it here explicitly
826+
@test isa(@inferred(Base.dataids(s)), NTuple{N, UInt} where {N})
827+
825828
s = StructArray{ComplexF64}((MyArray(rand(2,2)), MyArray(rand(2,2))))
826829
@test_throws MethodError s .+ s
827830
end

0 commit comments

Comments
 (0)