Julep: solving tricky iteration problems

[timholy]

The problem (a brief summary)

In a PR for a blog post, I described a number of challenges for several potential new AbstractArray types. Very briefly, the types and the challenges they exhibit are:

• ReshapedArrays: the only efficient way of indexing some types is by indexing the parent array directly (i.e., with an iterator optimized for the pre-reshaping array). This means that indexing a 2D array as B[i,j], for integer i and j, may not always be efficient.
• TransposedMatrix (for Taking vector transposes seriously LinearAlgebra.jl#42): iteration is most efficient when it produces a cache-friendly sequence of array accesses. Unfortunately, algorithms that mix column-major and row-major arrays imply that we can't naively allow eachindex(A) to always return the most efficient iterator.
• OffsetArrays (sometimes called "Fortran arrays"): these violate the assumptions that array indexes run from 1:size(A,d), and more generally question how one should match elements of two different
  arrays.

Possible APIs

In the blog post, only two concrete algorithms were considered: matrix-vector multiplication and copy!. At the risk of being insufficiently general, let's explore possible APIs in terms of these two functions.

EDIT: note that the thoughts on the API are evolving as a consequence of this discussion. To jump ahead to an updated proposal, see #15648 (comment). There may be further updates even later. The material below is left for reference, but some of it is outdated.

The "raw" API

For the moment, let's worry about correctness and efficiency more than elegance. Correctness trumps all other considerations. Efficiency means that we have to be able to iterate over the matrix in
cache-friendly order. That means that the array B has to be "in charge" of the sequence of values and pick its most efficient iterator.

I should emphasize that I haven't taken even the first step towards attempting to implement this API. So this is crystal-ball development, trying to imagine what might work while hoping to shake
out likely problems in advance via community feedback. That caveat understood, here's a possible implementation of matrix-vector multiplication:

RB = eachindex(B)
Rv = eachindex(RB[*,:], v)
Rdest = eachindex(RB[:,*], dest)
for (idest, iB, iv) in zip(Rdest, RB, Rv)
    dest[idest] += B[iB] * v[iv]
end

eachindex(iter, obj) should perform bounds-checking, to ensure that the indexes of iter align with obj (even handling things like OffsetArrays). The notation RB[*,:] means "corresponding to the
second dimension of RB"; interestingly, this expression parses without error, attempting to dispatch to getindex with argument types Tuple{typeof(RB), Function, Colon}. Hopefully there isn't another pending meaning for indexing with a Function.

Because either idest or iv is constant over the "inner" loop (depending on storage order of B), ideally one would like to hoist the corresponding access dest[idest] or v[iv]. I'd argue that's a
job for the compiler, and not something we should have to bake into the API.

Because zip(X, Y) does not allow Y to see the state of X, it seems quite possible that this could not be made efficient in general (consider the case described for ReshapedArrays, where the
efficiently-parametrized iterator for the first dimension depends on the state of the index for the second dimension). If this proves to be the case, we have to abandon zip and write this more like

RB = eachindex(B)
for (iB, idest, iv) in couple(RB, (RB[:,*], dest), (RB[*,:], v))
    dest[idest] += B[iB] * v[iv]
end

where couple is a new exported function.

One interesting thing about this API is that sparse multiplication might be implemented efficiently just by replacing the first line with

RB = eachstored(B)

so that only "stored" elements of B are visited. (@tkelman may have been the first to propose an eachstored iterator.) It should be pointed out that several recent discussions have addressed some of the complications that arise from Inf and NaN, and it remains to be determined whether a version that doesn't ignore these complications can be written with this (or similar) API.

Likewise, copy! might be implemented

Rdest = eachindex(dest)
for (idest, isrc) in couple(Rdest, (Rdest, src))
    dest[idest] = src[isrc]
end

At least for the array types described here, this API seems plausibly capable of generating both correct and (fairly) efficient code. One might ideally want to write generic "tiled" implementations, but that seems beyond the scope of this initial effort.

Making it prettier with "sufficiently smart iteration"

This API is still more complicated than naive iteration. One might possibly be able to hide a lot of this with a magical @iterate macro,

@iterate B dest[i] += B[i,j]*v[j]

for matrix-vector multiplication and

@iterate dest dest[i] = src[i]

for copy!. The idea is that the @iterate macro would expand these expressions into their longer forms above. Note that the first argument to @iterate specifies which object is "in charge" of the
sequence of operations. @iterate might be too limited for many things, but may be worth exploring. The KernelTools repository may have some useful tidbits.

[meggart]

Thanks a lot for writing the blog post and this issue, I learned a lot from it.

I have a question: You propose

    RB = eachindex(B)
    Rv = eachindex(RB[*,:], v)

If B is an Array RB will simply be 1:length(B), and lose all the shape information of B and makes your life harder interpreting second eachindex call.

In general, if you can come up with a couple function (which I think would be great), do you need the prior eachindex call at all? Why not simply do for (iB, idest, iv) in couple(B, (B[:,*], dest), (B[*,:], v)). where again B is the array that determines the iteration order.

So if possible I would prefer having a single couple function that could be heavily specialized for different argument types and one would need eachindex and zip only internally (called by couple).

I am also not sure if it is always a good idea to put one iterator in charge of the whole iteration order. How would you, for example formulate the loop for map!(dest,f,a,b), where the 3 arrays might have different storage orders? Could one implement a couple function that lets the three arrays dest, a, and b decide via "majority vote" how to iterate?

[timholy]

You're absolutely right; the intent was to keep all the shape information, and I overlooked the problem you pointed out. I was trying to avoid directly indexing arrays with something weird (the function *), since it seems like that could be unpopular. But what I wrote won't work. A possible alternative would be B[?,:], because ? is not defined yet, and so indexing arrays with ? seems like it might be OK.

Your two points also come together very nicely if one thinks about matrix-matrix (rather than matrix-vector) multiplication. First, you pretty much have to have a syntax that addresses the arrays directly, rather than some "master" iterator. Second, you're right that you probably don't want any iterator in charge. Something like this?

for (idest, iA, jdest, jB, kA, kB) in couple((dest[:,?], A[:,?]), (dest[?,:], B[?,:]), (A[?,:], B[:,?]))
    dest[idest,jdest] += A[idest,kA]*B[kB,jB]
end

which could hopefully be simplified to

@iterate dest[i,j] += A[i,k]*B[k,j]

Here, there's not actually much magic in @iterate (other than parsing the pattern of indexes and array-accesses), all the hard work is being done by couple.

My big worry is that couple could be quite hard to write. Still, it seems like the right place to put the complexity: do a bunch of analysis right at the beginning of the loop to figure out the optimal access pattern.

In some situations it may be difficult/impossible to make it type-stable (though all the ?/: indexing is designed to make this more feasible). As a fallback one could do

function foo!(dest, A, B)
    iters = couple(...)
    _foo(dest, A, B, iters)
end

@noinline function _foo!(dest, A, B, iters)
    for (idest, iA, jdest, jB, kA, kB) in iters
        ...
    end
end

However, for the @iterate macro this would need to be

function foo!(dest, A, B)
    iters = couple(...)
    @noinline function _foo!(dest, A, B, iters)
        for (idest, iA, jdest, jB, kA, kB) in iters
            ...
        end
    end
    _foo(dest, A, B, iters)
end

I haven't yet tested whether _foo! could be placed inside foo! and yet solve the type-instability problem. (I suspect so, but worth testing.)

[eschnett]

On Tue, Mar 29, 2016 at 11:18 AM, Tim Holy notifications@github.com wrote:

You're absolutely right; the intent was to keep all the shape information,
and I overlooked the problem you pointed out. I was trying to avoid
directly indexing arrays with something weird (the function *), since it
seems like that could be unpopular. But what I wrote won't work. A possible
alternative would be B[?,:], because ? is not defined yet, and so
indexing arrays with ? seems like it might be OK.

Your two points also come together very nicely if one thinks about
matrix-matrix (rather than matrix-vector) multiplication. First, you pretty
much have to have a syntax that addresses the arrays directly, rather than
some "master" iterator. Second, you're right that you probably don't want
any iterator in charge. Something like this?

for (idest, iA, jdest, jB, kA, kB) in couple((dest[:,?], A[:,?]), (dest[?,:], B[?,:]), (A[?,:], B[:,?]))
dest[idest,jdest] += A[idest,kA]*B[kB,jB]end

Ooooooh. This is now awfully close to an abstract index notation, which
allows generalizing this to more than two dimensions:

for magic in joined(dest[:i,:j], A[;i,:k], B[:k,:j])
dest[magic,:i,:j] += A[magic,:i,:k] * B[magic,:k,:j]
end

Obviously the syntax for the loop body can be improved, maybe with a macro,
maybe with clever types (e.g. turning the indices i,j,k into functions that
take the array as argument).

-erik

which could hopefully be simplified to

@iterate dest[i,j] += A[i,k]*B[k,j]

Here, there's not actually much magic in @iterate (other than parsing the
pattern of indexes and array-accesses).

My big worry is that couple could be quite hard to write. Still, it seems
like the right place to put the complexity: do a bunch of analysis right at
the beginning of the loop to figure out the optimal access pattern.

In some situations it may be difficult/impossible to make it type-stable
(though all the ?/: indexing is designed to make this more feasible). As
a fallback one could do

function foo(dest, A, B)
iters = couple(...)
_foo(dest, A, B, iters)end
@noinline function _foo(dest, A, B, iters)
for (idest, iA, jdest, jB, kA, kB) in iters
...
endend

However, for the @iterate macro this would need to be

function foo(dest, A, B)
iters = couple(...)
@noinline function _foo(dest, A, B, iters)
for (idest, iA, jdest, jB, kA, kB) in iters
...
end
end
_foo(dest, A, B, iters)end

I haven't yet tested whether _foo could be placed inside foo and yet
solve the type-instability problem. (I suspect so, but worth testing.)

—
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub
#15648 (comment)

Erik Schnetter schnetter@gmail.com
http://www.perimeterinstitute.ca/personal/eschnetter/

[timholy]

@eschnett, what you're suggesting is exactly what the @iterate macro does. Symbols don't have different types, so to do this via functions you have to introduce the appropriate types. But macros can parse the symbols, and emit code that is properly type-stable. In other words, the @iterate macro is effectively a pretty wrapper around couple.

[eschnett]

I was trying to generalize the notation for : and ?, and then got side-tracked by how the loop body should look like. So instead of writing

couple((dest[:,?], A[:,?]), (dest[?,:], B[?,:]), (A[?,:], B[:,?]))

one could write

coupled(dest[:i,:j,:k], A[:i,:j,:l], B[:j,:k,:m], C[:m])

which allows for more than two indices that can be used in arbitrary order, allowing transposition.

[mbauman]

Let's look at it from the other side. It seems there are three things @iterate needs to do for an expression like dest[i,j] += A[i,k]*B[k,j]:

• Determine the optimal loop order. This is hard… and an endless task if we consider BLAS-like partial loop striding. I say we punt, and require the user to specify the order.
• Determine how many times to run each loop. This is straightforward.
• Request for an efficient index iterator from each array for the given access pattern. Instead of worrying about coupling these iterators together, what if we just asked each array for an iterator across a given cartesian permutation order? E.g.,

@iterate (i, j, k) begin # specify iteration order
   D[i,j] += A[i,k]*B[k,j]
end

# would expand to:
(sz1, sz2, sz3) = … # get loop lengths and validate array sizes
Ditr = eachindex(D, (1, 2)) # will want Val{(1,2)} or some such
Aitr = eachindex(A, (1, 2))
Bitr = eachindex(B, (2, 1))
Dstate = start(iD)
Astate = start(iA)
for _i=1:sz1
    Bstate = start(iB) # B isn't dependent upon i, @iterate just restarts it every loop
    for _j=1:sz2
        (iD, Dstate) = next(Ditr, Dstate) # D isn't dependent on k, so it hoists it
        Astate′ = Astate # A isn't dependent upon j, so we jump back to where we were
        for _k=1:sz3
            (iB, Bstate) = next(Bitr, Bstate)
            (iA, Astate′) = next(Aitr, Astate′)
            D[iD] += A[iA]*B[iB]
        end
    end
    Astate = Astate′
end

I believe this is sufficiently general that @iterate will always be able to work out where to start, hoist, and hold iteration states. The magic, of course, now happens within eachindex(A, permutation). And while I think that will also be rather challenging to write for more than two dimensions, I think it'd be much simpler than trying to couple iterators together.

---

I'm sure you're aware, Tim, but I'd like to also point to TensorOperations.jl, which I think does some of @iterate (just not the fast iterated non-cartesian part). And although it's reason for being is completely different, numpy.einsum is also somewhat similar. Lots of people seem to use it, but I've never been able to sink my teeth into it — I find it quite unreadable.

[timholy]

@eschnett, the reason I think you need to set this up in a type-stable manner is that some array types aren't efficiently indexed by an NTuple of integers. For example, if A is a SparseMatrixCSC, A[i,j] (for integer i, j) does not have impressive performance; a much better strategy is to iterate over an index into rowindex/nzval and keep track of the corresponding i, j as you go. ReshapedArrays are another good example, and the one that drove me to think about this. For both of these cases, you'd really like to use a custom iterator specialized to extract the last bit of performance. (Particularly for the sparse case, if you only need to visit filled values the savings are huge.)

@mbauman, I suspect you're right about some of the challenges. What you're describing is basically how the macros in KernelTools work currently; I also wrote a nifty little macro (@time on steroids) that tests all possible loop orderings and reports the time for each, thus giving the user the chance to hard-code the one that works best. I should take a look at TensorOperations, I haven't looked in a long time so probably don't remember all the tools (or they may have changed).

But for generic code (that might have to handle Matrix and MatrixTranspose), it sure would be nice to be able to be flexible without having to write tons of cases out by hand. May be a pipe dream (I am sure that many have tried to do this before...e.g., Polly).

[timholy]

@mbauman, in your eachindex(A, (d1, d2)), am I reading this right that you're specifying that dimension d1 is the "inner" dimension and d2 is the "outer" dimension? Seems like a promising idea!

[mbauman]

Yes, precisely. The core idea is that each array only needs to concern itself about the order of iteration, and the @iterate macro itself handles starting, hoisting, and holding the iteration state at particular points in the expansion to ensure things stay in sync.

[timholy]

I guess the reason I didn't go this way is because I was attracted by the hope of solving the loop-ordering problem. Since the macro can't see the types, I was looking for a way for it to set things up in a form that would be amenable to type analysis.

[eschnett]

@timholy I don't understand the connection between the order in which the indices are traversed and the syntax to describe which indices need to be coupled.

I agree that sparse (or reshaped, or symmetric where only part of the elements are stored, etc.) arrays require special care.

Let's take this loop as example:

for i,j,k in coupled(r[:i,:j,:k], a[:i,:j,:k], b[:j,:k,:i], c[:k,:i,:j])
   r[i,j,k] = a[i,j,k] + b[j,k,i] + c[k,i,j]
end

How would this look like in the ? / : notation?

[timholy]

Something like couple(iblock, jblock, kblock) where

iblock = (r[:,?,?], a[:,?,?], b[?,?,:], c[?,:,?])
jblock = (r[?,:,?], a[?,:,?], b[:,?,?], c[?,?,:])
kblock = (r[?,?,:], a[?,?,:], b[?,:,?], c[:,?,?])

I.e., you match the indexes with the :s.

One of the points perhaps worth making is that I was also looking for an analog of for i = 2:n..., and the ?/: syntax might generalize reasonably well there. But bad things happen if there are no ?, because then it would dispatch just to regular getindex.

All this is to simply explain my reasoning; while I'm attracted by solving the loop ordering problem, I recognize that this may not be possible/practical. (I don't know the literature here.)

[eschnett]

@timholy Okay; that's a straightforward transformation from the (:i,:j,:k) syntax I like. There's essentially one block per index, specifying where the index appears for each array.

[meggart]

Determine the optimal loop order. This is hard… and an endless task if we consider BLAS-like partial loop striding. I say we punt, and require the user to specify the order.

I just want to support @timholy 's point that I think (at least partially) solving the loop ordering problem would be very nice.Ootherwise one would have to accept that generic library functions like map! and copy! might be slow for transposed matrices, since the user would not be able to influence the loop order.

[timholy]

Perhaps I'm most interested in a stepwise approach: if possible, pick an API that doesn't prevent us from tackling this someday, but punt on loop-ordering for now. Revamping how we do array iteration is a big problem no matter how we slice it, and we might make faster progress by first choosing the smallest feasible unit and getting that working.