# Poor performance with custom multidimensional iterator
**URL:** https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529
**Category:** Performance
**Created:** [September 12, 2020, 9:21pm UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529 "2020-09-12T21:21:29Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@benninkrs](https://discourse.julialang.org/u/benninkrs)
#### Post date: [September 12, 2020, 9:21pm UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529/1 "2020-09-12T21:21:29Z")
I am trying to implement a custom iterator for multidimensional arrays, which among other things, allows one to specify the order in which the axes are iterated. Sadly, I am having trouble getting even close to the performance of `CartesianIndices` even though almost all the code is identical. In particular, I found that the performance depends greatly on whether the index calculation is done within the iterator or in the caller.
Below is a MWE. `Iter1` is basically a stripped-down version of `CartesianIndices` and achieves good performance. It returns a tuple which is converted to an appropriate linear index in the caller. `Iter2` is nearly identical to `Iter1`; the only difference is that the conversion to a linear index is done within the `iterate` function instead of in the caller. Surprisingly to me, this results in nearly 8x slowdown. (This is on Julia 1.5.)
```julia
struct Iter1{N}
strides::Dims{N}
stop::Dims{N}
function Iter1(sz::Dims{N}, p::Dims{N}) where {N}
strides = cumprod(ntuple(i -> i>1 ? sz[i-1] : 1, Val(N)))
new{N}(ntuple(i->strides[p[i]], Val(N)), sz)
end
end
# (Exactly the same as Iter1)
struct Iter2{N}
strides::Dims{N}
stop::Dims{N}
function Iter2(sz::Dims{N}, p::Dims{N}) where {N}
strides = cumprod(ntuple(i -> i>1 ? sz[i-1] : 1, Val(N)))
new{N}(ntuple(i->strides[p[i]], Val(N)), sz)
end
end
# Iterating Iter1 returns a Tuple
@inline function Base.iterate(it::Iter1{N}) where {N}
if any(map(s -> s<1, it.stop))
return nothing
end
I = ntuple(i->1, Val(N))
return I, I
end
@inline function Base.iterate(it::Iter1{N}, state::Dims{N}) where {N}
valid, I = __inc(state, it.stop)
valid || return nothing
return I, I
end
# Iterating Iter2 returns a linear index
@inline function Base.iterate(it::Iter2{N}) where {N}
if any(map(s -> s<1, it.stop))
return nothing
end
I = ntuple(i->1, Val(N))
return 1, I
end
@inline function Base.iterate(it::Iter2, state::Dims{N}) where {N}
valid, I = __inc(state, it.stop)
valid || return nothing
return calc_index(I, it.strides), I
end
# helper functions
@inline __inc(::Tuple{}, ::Tuple{}) = false, ()
@inline function __inc(state::Dims{N}, stop::Dims{N}) where {N}
if state[1] < stop[1]
return true, (state[1]+1, Base.tail(state)...)
end
valid, I = __inc(Base.tail(state), Base.tail(stop))
return valid, (1, I...)
end
@inline function calc_index(I::Dims{N}, strides::Dims{N}) where {N}
1 + sum(ntuple(i -> (I[i]-1)*strides[i], Val(N)))
end
# Benchmarks
function test1(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
it = Iter1(size(A), p)
for I in it
iA = calc_index(I, it.strides)
A[iA] += 1.0
end
end
function test2(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
it = Iter2(size(A), p)
for iA in it
A[iA] += 1.0
end
end
function test2(A::AbstractArray{Float64,N}) where {N}
it = Iter2(size(A), ntuple(i->1+N-i, Val(N)))
for iA in it
A[iA] += 1.0
end
end
using BenchmarkTools
A = rand(4,4,4,4,4)
@btime test1($A, (5,4,3,2,1)) # 710 ns
@btime test2($A, (5,4,3,2,1)) # 5.2 μs
@btime test2($A) # 1.6 μs
```
Note that the only difference between `test1` and `test2` is where `calc_index` is called. Evidently the compiler is utilizing information about `it.strides` in `test1` that it is not using in `iterate(::Iter2)`. Hard-coding the permutation `p` (the second version of `test2`) helps somewhat, but is still 2x slower than `test1`.
I’d greatly appreciate any insight anyone has to offer. I’ve been pulling my hair out for two days trying to figure out why my iterator is performing so poorly.
---
### Author:  [@tim.holy](https://discourse.julialang.org/u/tim.holy)
#### Post date: [September 15, 2020, 1:16am UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529/2 "2020-09-15T01:16:39Z")
Without having looked at your code, let me start with a dumb question: can you just combine `PermutedDimsArray` with ordinary `CartesianIndices` iteration?
---
### Author:  [@benninkrs](https://discourse.julialang.org/u/benninkrs)
#### Post date: [September 16, 2020, 2:57am UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529/3 "2020-09-16T02:57:06Z")
It’s not a dumb question at all. I originally tried multiple variations on permuting `CartiesianIndex`’s, but they didn’t perform nearly as well as without permutation. Then I thought I might directly iterate the subscripts in a custom order, so that no permutation would be required. Unfortunately that performed poorly as well. I finally determined that calculating a linear index with custom strides was fastest, but only if the calculation was done in the caller. At that point I threw my hands up and posted my problem here 🙂.
For the record, here’s a benchmark with PermutedDimsArray:
```julia
function test0(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
AP = PermutedDimsArray(A, p)
for ci in CartesianIndices(AP)
A[ci] += 1.0
end
end
@btime test0($A, (5,4,3,2,1)) # 2.3 μs
```
Incidentally, my actual goal is more complicated than just changing the iteration sequence of a single array – it is to create an iterator for tensor operations, as a sort of alternative to existing tensor packages. But the issue posted here lies at the heart of it. And as a secondary goal I am hoping to get a better understanding of the idiosyncracies of compiler optimization (or the lack thereof).
---
### Author:  [@tim.holy](https://discourse.julialang.org/u/tim.holy)
#### Post date: [September 16, 2020, 7:40am UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529/4 "2020-09-16T07:40:20Z")
That code only works if `A` is of isotropic size. Here’s what I meant:
```julia
function inc_cartesian!(B::AbstractArray)
for I in CartesianIndices(B)
B[I] += 1
end
return B
end
A = rand(4, 5, 6, 7, 8) # notice that this is a bigger array
AP = PermutedDimsArray(copy(A), (5,4,3,2,1))
julia> @btime inc_cartesian!($A);
5.718 μs (0 allocations: 0 bytes)
julia> @btime inc_cartesian!($AP);
4.308 μs (0 allocations: 0 bytes)
```
There’s a crucial difference here: in my version, Julia’s compiler knows the permutation, but in yours it doesn’t. Why? You passed the permutation as a value, but I passed it as a type parameter:
```julia
julia> typeof(AP)
PermutedDimsArray{Float64, 5, (5, 4, 3, 2, 1), (5, 4, 3, 2, 1), Array{Float64, 5}}
```
The permutation itself (as well as the inverse permutation, which in this case happens to be the same) is encoded in the type.
If you doubt this achieves what you were asking, let’s try a different implementation:
```julia
julia> function address_order!(B)
k = 0
for I in CartesianIndices(B)
B[I] = k += 1
end
end
address_order! (generic function with 1 method)
julia> A = zeros(Int, 3, 4);
julia> AP = PermutedDimsArray(copy(A), (2, 1));
julia> address_order!(A);
julia> address_order!(AP);
julia> A
3×4 Matrix{Int64}:
1 4 7 10
2 5 8 11
3 6 9 12
julia> AP.parent
3×4 Matrix{Int64}:
1 2 3 4
5 6 7 8
9 10 11 12
```
You see that `AP.parent` was visited in row-major order.
---
### Author:  [@benninkrs](https://discourse.julialang.org/u/benninkrs)
#### Post date: [September 16, 2020, 5:48pm UTC](https://discourse.julialang.org/t/poor-performance-with-custom-multidimensional-iterator/46529/5 "2020-09-16T17:48:22Z")
> That code only works if `A` is of isotropic size.
Oops. That’s because I had a typo. It should’ve been
```julia
function test0(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
AP = PermutedDimsArray(A, p)
for ci in CartesianIndices(AP)
AP[ci] += 1.0
end
end
```
which clocks in at a whopping 77 microseconds (!). To put your suggestion in my framework above:
```julia
function test0_TH(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
AP = PermutedDimsArray(A, p)
inc_cartesian!(AP)
end
julia> @btime test0_TH($A, (5,4,3,2,1));
1.010 μs (4 allocations: 208 bytes)
```
which i think is pretty good – it’s not too far off the 700ns for the unpermuted iteration. However, this approach is a bit less convenient than I was hoping for: It requires the user to create a permuted array and define a separate function to do the loop, whereas I am aiming for an iterator that can be used in a transparent way. But I bet a macro could give me something close to the kind of API I would like.
Incidentally, one of my earliest attempts was to incorporate index permutation based on a type parameter (a la `PermutedDimsArray`) into an iterator, along these lines:
```julia
struct Iter3{N, P}
stop::Dims{N}
function Iter3(sz::Dims{N}, p::Dims{N}) where {N}
new{N,p}(sz)
end
end
@inline function Base.iterate(it::Iter3{N,P}) where {N,P}
if any(map(s -> s<1, it.stop))
return nothing
end
I = ntuple(i->1, Val(N))
return CartesianIndex{N}(map(j->I[j], P)), I
end
@inline function Base.iterate(it::Iter3{N,P}, state::Dims{N}) where {N,P}
valid, I = __inc(state, it.stop)
valid || return nothing
return CartesianIndex{N}(map(j->I[j], P)), I
end
function test3(A::AbstractArray{Float64,N}, p::Dims{N}) where {N}
for iA in Iter3(size(A), p)
A[iA] += 1.0
end
end
```
Although the permutation is a type parameter as it was in your case, the performance is quite poor:
```julia
julia> @btime test3($A, (5,4,3,2,1));
44.399 μs (2050 allocations: 144.05 KiB)
```
I suspect this is because there is no function barrier between the type-unstable calling function and the iterating function.