Skip to content

Commit 29e8d33

Browse files
committed
Merge pull request #13614 from JuliaLang/mb/nolinear
Detangle linear indexing from non-scalar indexing
2 parents 2bb63a9 + 64525aa commit 29e8d33

File tree

1 file changed

+8
-44
lines changed

1 file changed

+8
-44
lines changed

Diff for: base/multidimensional.jl

+8-44
Original file line numberDiff line numberDiff line change
@@ -187,14 +187,14 @@ index_shape_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
187187
checkbounds(A, I...)
188188
_unsafe_getindex(l, A, I...)
189189
end
190-
@generated function _unsafe_getindex(l::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
190+
@generated function _unsafe_getindex(::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
191191
N = length(I)
192192
quote
193193
# This is specifically *not* inlined.
194194
@nexprs $N d->(I_d = to_index(I[d]))
195195
dest = similar(A, @ncall $N index_shape A I)
196196
@ncall $N checksize dest I
197-
@ncall $N _unsafe_getindex! dest l A I
197+
@ncall $N _unsafe_getindex! dest A I
198198
end
199199
end
200200

@@ -221,7 +221,7 @@ end
221221

222222
# Indexing with an array of indices is inherently linear in the source, but
223223
# might be able to be optimized with fast dividing integers
224-
@inline function _unsafe_getindex!(dest::AbstractArray, ::LinearIndexing, src::AbstractArray, I::AbstractArray)
224+
@inline function _unsafe_getindex!(dest::AbstractArray, src::AbstractArray, I::AbstractArray)
225225
D = eachindex(dest)
226226
Ds = start(D)
227227
for idx in I
@@ -231,26 +231,8 @@ end
231231
dest
232232
end
233233

234-
# Fast source - compute the linear index
235-
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearFast, src::AbstractArray, I::Union{Real, AbstractVector, Colon}...)
236-
N = length(I)
237-
quote
238-
$(Expr(:meta, :inline))
239-
stride_1 = 1
240-
@nexprs $N d->(stride_{d+1} = stride_d*size(src, d))
241-
$(symbol(:offset_, N)) = 1
242-
D = eachindex(dest)
243-
Ds = start(D)
244-
@nloops $N i dest d->(offset_{d-1} = offset_d + (unsafe_getindex(I[d], i_d)-1)*stride_d) begin
245-
d, Ds = next(D, Ds)
246-
unsafe_setindex!(dest, unsafe_getindex(src, offset_0), d)
247-
end
248-
dest
249-
end
250-
end
251-
# Slow source - index with the indices provided.
252-
# TODO: this may not be the full dimensionality; that case could be optimized
253-
@generated function _unsafe_getindex!(dest::AbstractArray, ::LinearSlow, src::AbstractArray, I::Union{Real, AbstractVector, Colon}...)
234+
# Always index with the exactly indices provided.
235+
@generated function _unsafe_getindex!(dest::AbstractArray, src::AbstractArray, I::Union{Real, AbstractVector, Colon}...)
254236
N = length(I)
255237
quote
256238
$(Expr(:meta, :inline))
@@ -292,8 +274,8 @@ _iterable(v) = repeated(v)
292274
checkbounds(A, J...)
293275
_unsafe_setindex!(l, A, x, J...)
294276
end
295-
@inline function _unsafe_setindex!(l::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
296-
_unsafe_batchsetindex!(l, A, _iterable(x), to_indexes(J...)...)
277+
@inline function _unsafe_setindex!(::LinearIndexing, A::AbstractArray, x, J::Union{Real,AbstractArray,Colon}...)
278+
_unsafe_batchsetindex!(A, _iterable(x), to_indexes(J...)...)
297279
end
298280

299281
# 1-d logical indexing: override the above to avoid calling find (in to_index)
@@ -315,25 +297,7 @@ function _unsafe_setindex!(::LinearIndexing, A::AbstractArray, x, I::AbstractArr
315297
A
316298
end
317299

318-
# Use iteration over X so we don't need to worry about its storage
319-
@generated function _unsafe_batchsetindex!(::LinearFast, A::AbstractArray, X, I::Union{Real,AbstractArray,Colon}...)
320-
N = length(I)
321-
quote
322-
@nexprs $N d->(I_d = I[d])
323-
idxlens = @ncall $N index_lengths A I
324-
@ncall $N setindex_shape_check X (d->idxlens[d])
325-
Xs = start(X)
326-
stride_1 = 1
327-
@nexprs $N d->(stride_{d+1} = stride_d*size(A,d))
328-
$(symbol(:offset_, N)) = 1
329-
@nloops $N i d->(1:idxlens[d]) d->(offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
330-
v, Xs = next(X, Xs)
331-
unsafe_setindex!(A, v, offset_0)
332-
end
333-
A
334-
end
335-
end
336-
@generated function _unsafe_batchsetindex!(::LinearSlow, A::AbstractArray, X, I::Union{Real,AbstractArray,Colon}...)
300+
@generated function _unsafe_batchsetindex!(A::AbstractArray, X, I::Union{Real,AbstractArray,Colon}...)
337301
N = length(I)
338302
quote
339303
@nexprs $N d->(I_d = I[d])

0 commit comments

Comments
 (0)