Code generation and cartesian iteration

[timholy]

Compared to manual iteration, the new cartesian iterator has a small performance penalty:

function sumcart_manual(A::AbstractMatrix)
    s = 0.0
    @inbounds for j = 1:size(A,2), i = 1:size(A,1)
        s += A[i,j]
    end
    s
end

function sumcart_iter(A)
    s = 0.0
    @inbounds for I in CartesianIndices(size(A))
        s += A[I]
    end
    s
end

A = rand(10^4,10^4);
sumcart_manual(A);
sumcart_iter(A);

julia> @time sumcart_manual(A)
elapsed time: 0.176363505 seconds (96 bytes allocated)
4.999591778337127e7

julia> @time sumcart_iter(A)
elapsed time: 0.250890195 seconds (96 bytes allocated)
4.999591778337127e7

In this gist, I compare the code that gets generated. The comparison reveals that the CartesianIndex is not fully elided by the compiler.

Any ideas what's happening?

EDIT: the core problem is that next contains a branch, and testing done is a 2nd branch, so there are two branches per iteration. See #16035 (comment)

[timholy]

Here's an example that doesn't rely on all the code-generation facilities in the iterators in Base:

module Iter

import Base: start, next, done, getindex

immutable Raw2
    i1::Int
    i2::Int
end
Raw2(dims::(Int,Int)) = Raw2(dims[1], dims[2])

start(sz::Raw2) = Raw2(1,1)
@inline next(sz::Raw2, state::Raw2) = state.i1 < sz.i1 ?
                                     (state, Raw2(state.i1+1,state.i2)) :
                                     (state, Raw2(1, state.i2+1))
done(sz::Raw2, state::Raw2) = state.i2 > sz.i2

getindex(A::AbstractArray, I::Raw2) = A[I.i1,I.i2]

end

function mysum(A)
    s = 0.0
    @inbounds for I in Iter.Raw2(size(A))
        s += A[I]
    end
    s
end

function mysum2(A)
    s = 0.0
    @inbounds for j = 1:size(A,2), i = 1:size(A,1)
        s += A[i,j]
    end
    s
end

Then do

A = reshape(1:12, 3, 4)
@code_llvm mysum(A)
@code_llvm mysum2(A)

You should see that the code for mysum is not nearly as well-optimized as mysum2.

[timholy]

@simonster, while we're on the topic of SubArrays and indexing code generation, if you have time I'd be grateful for any insights you have here.

Comparing mysum and mysum2 for A = rand(10^4, 10^4) reveals that there's a 40% speed gap, too. The critical section of @code_llvm for mysum is

L3:                                               ; preds = %L2, %if1
  %_var0.0 = phi [2 x %Raw2] [ %41, %L2 ], [ %37, %if1 ]
  %42 = extractvalue [2 x %Raw2] %_var0.0, 0, !dbg !796, !julia_type !803
  %43 = extractvalue [2 x %Raw2] %_var0.0, 1, !dbg !796, !julia_type !803
  %44 = extractvalue %Raw2 %42, 0, !dbg !804
  %45 = add i64 %44, -1, !dbg !804
  %46 = extractvalue %Raw2 %42, 1, !dbg !804
  %47 = add i64 %46, -1, !dbg !804
  %48 = mul i64 %47, %13, !dbg !804
  %49 = add i64 %45, %48, !dbg !804
  %50 = getelementptr double* %28, i64 %49, !dbg !804
  %51 = load double* %50, align 8, !dbg !804, !tbaa %jtbaa_user
  %52 = fadd double %s.0, %51, !dbg !804
  %53 = extractvalue %Raw2 %43, 1, !dbg !804
  %54 = icmp slt i64 %26, %53, !dbg !804
  br i1 %54, label %L6, label %L, !dbg !804

compared to

L2:                                               ; preds = %L2.preheader, %L2
  %s.1 = phi double [ %24, %L2 ], [ %s.0, %L2.preheader ]
  %"#s10.0" = phi i64 [ %20, %L2 ], [ 1, %L2.preheader ]
  %20 = add i64 %"#s10.0", 1, !dbg !798
  %21 = add i64 %19, %"#s10.0", !dbg !805
  %22 = getelementptr double* %10, i64 %21, !dbg !805
  %23 = load double* %22, align 8, !dbg !805, !tbaa %jtbaa_user
  %24 = fadd double %s.1, %23, !dbg !805
  %25 = icmp eq i64 %"#s10.0", %15, !dbg !805
  br i1 %25, label %L5, label %L2, !dbg !805

for mysum2.

Anything obvious to change to narrow the gap between these?

[simonster]

This generates better code for me:

start(sz::Raw2) = Raw2(0,0)
@inline function next(sz::Raw2, state::Raw2)
    i1 = state.i1+1
    i2 = state.i2
    if i1 == sz.i1
        i1 = 0
        i2 = state.i2+1
    end
    return (state, Raw2(i1, i2))
end
done(sz::Raw2, state::Raw2) = state.i2 == sz.i2

getindex(A::AbstractArray, I::Raw2) = A[I.i1+1,I.i2+1]

This is ~5% slower on my system instead of ~20%. But both the original code and this change are still suboptimal because the inner loop has a branch:

L:                                                ; preds = %L.preheader, %L2
  %s.0 = phi double [ %24, %L2 ], [ 0.000000e+00, %L.preheader ]
  %"#s2307.0" = phi %Raw2 [ %19, %L2 ], [ zeroinitializer, %L.preheader ]
  %9 = ptrtoint %jl_value_t* %8 to i64
  %10 = extractvalue %Raw2 %"#s2307.0", 0, !dbg !1614
  %11 = add i64 %10, 1, !dbg !1614
  %12 = extractvalue %Raw2 %"#s2307.0", 1, !dbg !1614
  %13 = icmp eq i64 %11, %9, !dbg !1614
  br i1 %13, label %if1, label %L2, !dbg !1614

if1:                                              ; preds = %L
  %14 = add i64 %12, 1, !dbg !1614
  br label %L2, !dbg !1614

L2:                                               ; preds = %L, %if1
  %_var1.0 = phi i64 [ %11, %L ], [ 0, %if1 ]
  %_var2.0 = phi i64 [ %12, %L ], [ %14, %if1 ]
  %15 = ptrtoint %jl_value_t* %2 to i64
  %16 = ptrtoint %jl_value_t* %8 to i64
  %17 = bitcast i8* %6 to double*
  %18 = insertvalue %Raw2 undef, i64 %_var1.0, 0, !dbg !1614
  %19 = insertvalue %Raw2 %18, i64 %_var2.0, 1, !dbg !1614, !julia_type !1621
  %20 = mul i64 %12, %16, !dbg !1622
  %21 = add i64 %20, %10, !dbg !1622
  %22 = getelementptr double* %17, i64 %21, !dbg !1622
  %23 = load double* %22, align 8, !dbg !1622, !tbaa %jtbaa_user
  %24 = fadd double %s.0, %23, !dbg !1622
  %25 = icmp eq i64 %_var2.0, %15, !dbg !1622
  br i1 %25, label %L5, label %L, !dbg !1622

This is pretty much what you'd expect given the iteration protocol. In Julia, it would look something like:

@label start
newi1 = i1 + 1
if newi1 == sz1
    newi1 = 0
    newi2 += 1
else
    newi2 = i2
end
ind = i2*sz1+i1
@inbounds s += A[ind+1]
newi2 == sz2 && @goto done
i1 = newi1
i2 = newi2
@goto start
@label done

(This generates exactly the same assembly for the inner loop). Rearranging this a little allows LLVM to optimize it to match the performance of mysum2:

@label start
ind = i2*sz1+i1
@inbounds s += A[ind+1]
newi1 = i1 + 1
if newi1 == sz1
    newi1 = 0
    newi2 += 1
    newi2 == sz2 && @goto done
else
    newi2 = i2
end
i1 = newi1
i2 = newi2
@goto start
@label done

But I can't see how to express this in the current iteration framework.

edit: Ref iteration protocol issues #6125, #8149, #9182

[StefanKarpinski]

Worth keeping in mind wrt #8149 – if we could make it easier to express internal iteration and get a performance bump at the same time, that would be the best.

[timholy]

Thanks a million, @simonster. The impact of these "stylistic" differences is quite surprising to me---I've learned to lean on LLVM to do the right thing, but this is a good counterexample.

I'll look into adopting your version for our code generation for next.

[timholy]

Hmm, on my system, my original version is faster than your modification. This is weird.

[simonster]

I tried on another system, which I probably should have done earlier. On my Sandy Bridge system on Linux, mysum is 37% slower than mysum2 with your original version and 35% slower than mysum2 with my modifications. On my Haswell system on OS X, mysum is 22% slower than mysum2 with your original version and 7% slower than mysum2 with my modifications. Both systems are running LLVM 3.3 and the generated assembly is identical, so the performance difference seems to be in the architecture. I'm just benchmarking with:

function bench(A)
    t1 = 0.0
    t2 = 0.0
    for i = 1:50
        t1 += @elapsed mysum(A)
        t2 += @elapsed mysum2(A)
    end
    t1/t2
end

[timholy]

From reading a bit more, it seems that the cost of branches is a big difference among architectures (mine's a Sandy Bridge).

Just for fun I tried #9182, but by itself it didn't change the performance. I wondered if it did the nextstate at the end of the loop, rather than right after nextval, if it might fare better---you wouldn't need to essentially have two state variables (old and new).

To test this, I followed your lead and wrote out some explicit versions using while, and got pretty interesting results:

function mysum_while1(A)
    s = 0.0
    i = j = 1
    inew = jnew = 1
    while j <= size(A,2)
        if inew < size(A,1)
            inew += 1
        else
            inew = 1
            jnew += 1
        end
        @inbounds s += A[i,j]
        i = inew
        j = jnew
    end
    s
end

function mysum_while2(A)
    s = 0.0
    i = j = 1
    while j <= size(A,2)
        @inbounds s += A[i,j]
        if i < size(A,1)
            i += 1
        else
            i = 1
            j += 1
        end
    end
    s
end

function mysum_while3(A)
    s = 0.0
    j = 1
    while j <= size(A,2)
        i = 1
        while i <= size(A,1)
            @inbounds s += A[i,j]
            i += 1
        end
        j += 1
    end
    s
end

function mysum_while4(A)
    s = 0.0
    i = j = 1
    inew = jnew = 1
    notdone = true
    while notdone
        inew += 1
        if inew > size(A, 1)
            inew = 1
            jnew += 1
            notdone = jnew <= size(A,2)
        end
        @inbounds s += A[i,j]
        i = inew
        j = jnew
    end
    s
end


function mysum_while5(A)
    s = 0.0
    i = j = 1
    notdone = true
    while notdone
        @inbounds s += A[i,j]
        i += 1
        if i > size(A, 1)
            i = 1
            j += 1
            notdone = j <= size(A,2)
        end
    end
    s
end

function mysum_while6(A)
    s = 0.0
    i = j = 1
    while j <= size(A,2)
        @inbounds s += A[i,j]
        i += 1
        if i > size(A, 1)
            i = 1
            j += 1
        end
    end
    s
end

Timing results after warmup:

julia> @time mysum_while1(A)
elapsed time: 0.258024361 seconds (96 bytes allocated)
4.9997446791483015e7

julia> @time mysum_while2(A)
elapsed time: 0.216438899 seconds (96 bytes allocated)
4.9997446791483015e7

julia> @time mysum_while3(A)
elapsed time: 0.157289985 seconds (96 bytes allocated)
4.9997446791483015e7

julia> @time mysum_while4(A)
elapsed time: 0.227041972 seconds (96 bytes allocated)
4.9997446791483015e7

julia> @time mysum_while5(A)
elapsed time: 0.140768946 seconds (96 bytes allocated)
4.9997446791483015e7

julia> @time mysum_while6(A)
elapsed time: 0.190532756 seconds (96 bytes allocated)
4.9997446791483015e7

So, at least on my machine

• the placement of the increment matters (end of the loop is best)
• using a boolean variable for while testing seems to help

Interestingly, the iterators in multidimensional already use the boolean version, so it might just be a case of rewriting the way next works. But first, can you check on your system that mysum_while5 is among the best performers?

[timholy]

One more:

function mysum_while7(A)
    s = 0.0
    i = j = 1
    notdone = true
    while notdone
        @inbounds s += A[i,j]
        if i < size(A,1)
            i += 1
        else
            i = 1
            j += 1
            notdone = j <= size(A,2)
        end
    end
    s
end

is just as fast as the fastest ones above. So it really seems to come down to (1) putting the increment at the end, and (2) using a boolean for the while testing (and updating that boolean at the deepest level of the increment testing).