Skip to content

Commit 498d982

Browse files
authored
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 ```
1 parent 9c8886c commit 498d982

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, mod1
19+
afoldl, mod1, @default_eltype
2020
using .Core
2121
using Core: @doc
2222

@@ -1176,6 +1176,8 @@ end
11761176
next === nothing && return nothing
11771177
restnext = _piterate(rest...)
11781178
restnext === nothing && return nothing
1179+
VS = @default_eltype(iter1)
1180+
next = Pair{VS, typeof(next[2])}(next[1], next[2])
11791181
return (next, restnext...)
11801182
end
11811183
@inline function iterate(P::ProductIterator)
@@ -1188,15 +1190,16 @@ end
11881190
@inline _piterate1(::Tuple{}, ::Tuple{}) = nothing
11891191
@inline function _piterate1(iters, states)
11901192
iter1 = first(iters)
1191-
next = iterate(iter1, first(states)[2])
1192-
restnext = tail(states)
1193+
state1, restnext... = states
1194+
next = iterate(iter1, state1[2])
11931195
if next === nothing
11941196
isdone(iter1) === true && return nothing
11951197
restnext = _piterate1(tail(iters), restnext)
11961198
restnext === nothing && return nothing
11971199
next = iterate(iter1)
11981200
next === nothing && return nothing
11991201
end
1202+
next = Pair{fieldtype(typeof(state1), 1), typeof(next[2])}(next[1], next[2])
12001203
return (next, restnext...)
12011204
end
12021205
@inline function iterate(P::ProductIterator, states)

0 commit comments

Comments
 (0)