Skip to content

Commit 00e970f

Browse files
vtjnashKristofferC
authored andcommitted
fix performance of ProductIterator for non-concrete iterators (#59711)
This uses inferred element types to make the iterator state type-stable, enabling inlining, which can enable the runtime to see the real types. This is essentially the transform that we want inference to infer and codegen to implement, except made explicit so that we can rely on it actually happening. Fix #56607 ``` julia> @Btime f() 574.500 ns (2 allocations: 272 bytes) # non-const op 253.878 ns (2 allocations: 272 bytes) # const op ``` (cherry picked from commit 498d982)
1 parent 8d99e85 commit 00e970f

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

base/iterators.jl

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ using .Base:
1616
(:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing,
1717
any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex,
1818
tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString,
19-
afoldl
19+
afoldl, @default_eltype
2020
using Core
2121
using Core: @doc
2222

@@ -1140,6 +1140,8 @@ end
11401140
next === nothing && return nothing
11411141
restnext = _piterate(rest...)
11421142
restnext === nothing && return nothing
1143+
VS = @default_eltype(iter1)
1144+
next = Pair{VS, typeof(next[2])}(next[1], next[2])
11431145
return (next, restnext...)
11441146
end
11451147
@inline function iterate(P::ProductIterator)
@@ -1152,15 +1154,16 @@ end
11521154
@inline _piterate1(::Tuple{}, ::Tuple{}) = nothing
11531155
@inline function _piterate1(iters, states)
11541156
iter1 = first(iters)
1155-
next = iterate(iter1, first(states)[2])
1156-
restnext = tail(states)
1157+
state1, restnext... = states
1158+
next = iterate(iter1, state1[2])
11571159
if next === nothing
11581160
isdone(iter1) === true && return nothing
11591161
restnext = _piterate1(tail(iters), restnext)
11601162
restnext === nothing && return nothing
11611163
next = iterate(iter1)
11621164
next === nothing && return nothing
11631165
end
1166+
next = Pair{fieldtype(typeof(state1), 1), typeof(next[2])}(next[1], next[2])
11641167
return (next, restnext...)
11651168
end
11661169
@inline function iterate(P::ProductIterator, states)

0 commit comments

Comments
 (0)