performance - Julia - CartesianIndices 的性能

Question

我正在研究 Base.IteratorsMD 的源代码，我确信 CartesianIndex 的 Base.nextind 的实现(在 julia/base/multidimensional.jl 中，参见此处 https://github.com/JuliaLang/julia/blob/55e36cc308b66d3472990a06b2797f9f9154ea0a/base/multidimensional.jl#L142-L150 ) 必须非常低效:

function Base.nextind(a::AbstractArray{<:Any,N}, i::CartesianIndex{N}) where {N}
    iter = CartesianIndices(axes(a))
    return CartesianIndex(inc(i.I, first(iter).I, last(iter).I))
end

我假设在每次调用中创建 CartesianIndices(...) 的新实例一定非常昂贵，因为在我的机器上调用 CartesianIndices(...) 通过 REPL 似乎创建了所有索引。所以我为替代实现编写了以下基准脚本 mynextind:

using Base.IteratorsMD

function mynextind(a::AbstractArray{<:Any,N}, i::CartesianIndex{N}) where {N}
    dims = (x.stop for x in axes(a))
    return CartesianIndex(Base.IteratorsMD.inc(i.I, CartesianIndex(ones(Int64, length(dims))...).I, CartesianIndex(dims...).I))
end

function f(func, M)
    c = CartesianIndex(1,1)
    while true
        (c = func(M, c)) == CartesianIndex(size(M)...) && break
    end
end

A = rand(100,100)

@btime f(Base.nextind, A)
@btime f(mynextind, A)

Base.nextind 比 mynextind 快几个数量级(7 微秒对 12 毫秒)。此外，Base.nextind 似乎没有进行任何内存分配。现在我试图理解为什么会这样。 CartesianIndices 是否真的创建了所有索引？

Answer 1

I assumed that creating a new instance of CartesianIndices(...) in every call must be very expensive, since on my machine a call to CartesianIndices(...) via REPL seems to create all indices.

这是无效的。

当您在 REPL 中打印 CartesianIndices(...) 时，将打印所有元素。但这并不意味着它将创建所有索引。

正如您在此处的 CartesianIndices 结构的源代码中所见: https://github.com/JuliaLang/julia/blob/64d8ca49122609d1c10c72a96d1711b95346980a/base/multidimensional.jl#L248

当您创建类似 CartesianIndices(axes(A)) 的实例时，只会存储轴。您看到所有索引都被打印出来的原因是 CartesianIndices 是一个 AbstractArray。它还实现了 getindex 方法。因此，当您键入 CartesianIndices(axes(A)) 时，所有索引均按需计算。

你可以在这里看到更多关于 CartesianIndices 的讨论:https://discourse.julialang.org/t/psa-replacement-of-ind2sub-sub2ind-in-julia-0-7/14666