Collaboration, transfer, and adding iterators

[timholy]

JuliaLang/julia#31563 (comment), posted by @AriMKatz, drew my attention to this package. I agree that there appear to be a subset of shared goals between this package and https://github.com/timholy/ArrayIteration.jl. This is a proposal to collaborate so we don't duplicate effort.

To summarize ArrayIteration (because the README is ancient and talks about some stuff that long ago made it into Base): if you had to boil it down to a single goal, it's to provide an efficient fallback for A*B for matrices A and B with arbitrary sparsity patterns and representation. The general problem is that if you have to specialize on both typeof(A) and typeof(B), you have to write O(N^2) methods where N is the number of array types. (And wrappers like Adjoint make the problem worse because it's essentially O(4N^2) or O(9N^2).) ArrayIteration takes the stance that you can solve the problem with O(N) methods if you instead create a fancy iterator for each matrix that allows synchronization with another iterator. You sync the row iterator of A to the column iterator of B and thereby mostly just visit entries that are nonzero in both.

ArrayIteration is basically waiting for JuliaLang/julia#34126 (so we don't pay the allocation penalty for wrappers) and probably JuliaLang/julia#34847 before I trying bringing it back to life. The state of the art is the teh/stored2 branch which was almost functional, before the performance problems due to wrapper creation put it on hold. (The README, however, was not updated to reflect the state of that branch.)

I'd propose we all get together in a package hosted at JuliaArrays, which seems like the natural home for generic array stuff. We could start by transferring this one, since I'm guessing it's working whereas ArrayIteration is stalled out. If anyone wants to just steal the code in ArrayIteration that's fine too, but I don't think keeping it under the JuliaDiffEq umbrella makes sense if it's supposed to be generic array handling. We could also keep them two packages, but in that case we'd want to make sure they're independent but interoperable.

Also CCing @vchuravy whose https://github.com/JuliaGPU/KernelAbstractions.jl was also mentioned. I am less certain there are true shared goals there, but just in case.

[ChrisRackauckas]

I'd propose we all get together in a package hosted at JuliaArrays, which seems like the natural home for generic array stuff. We could start by transferring this one, since I'm guessing it's working whereas ArrayIteration is stalled out. If anyone wants to just steal the code in ArrayIteration that's fine too, but I don't think keeping it under the JuliaDiffEq umbrella makes sense if it's supposed to be generic array handling. We could also keep them two packages, but in that case we'd want to make sure they're independent but interoperable.

Oh I definitely agree that it shouldn't live in JuliaDiffEq at all. But if we're going to take the time to really finish this, I think it should go directly to Base. This whole system of "lets define some traits, pirate what we need, and then use Requires to affix it to packages" is really just a hacky system to get the traits that I need "widely" defined throughout Julia, since without these traits I couldn't do my work. But in reality, these should go into Base, and then each package should be defining its traits. So essentially, a low level "going rouge" traits package for if some random DiffEq user will define a Jacobian as a AbstractBandedToplitzMatrix and I'll have to go patch it really quickly and tada we have tons of cool features with automatic sparse AD, but it'll probably take years for people to accept some of these traits more broadly.

I think culturally, the steps are to build a sparse AD system, then get big packages like DiffEq, Optim, NLsolve, etc. to all adopt language-wide sparsity handling, then everyone will want to join. We're almost at that "main package adoption level" (DiffEq is already doing it, and Optim + NLsolve are by this summer), so then it's going to be more like "oh, your package doesn't accept a color vector? It really should...". Once that's the case, it'll be simple to say "can we please move the matrix color trait to Base since every AbstractMatrix should be defining it (or using the default)?"

[timholy]

Moving to Base actually seems like the right solution. Or rather, a mixture of Base (for core traits) and LinearAlgebra (which will presumably become a package eventually).

[ChrisRackauckas]

In fact, @dlfivefifty should be in this discussion. Most of the difficult iterators are due to BlockBandedMatrices.jl. But with JuliaLinearAlgebra/BlockBandedMatrices.jl#62 I feel like we're really missing something. Somehow we need to nest structure, since we had BandedMatrix, then BlockBandedMatrix, then BandedBlockBandedMatrix, and now this issue is noticing that for a system of 2D PDEs we need a block matrix of block banded matrices. I think specializing each case is going to be insane 🤣 (though in reality, I "think" this is as complex as actual applications will get? Famous last words). Somehow we need Block{Banded{Block{Banded{T}}}} and then specialize the iteration on this, auto-parallelize it, allow it to GPU, etc.

@peterahrens worked on a Swizzles idea in this direction, but it kind of died down. So my only hesitiation to moving it to Base right now is I feel like we might have an epiphony on how to do this in a more composible way in a few months and then suddenly all of this code needs to change. But essentially the issue is that somehow we need to directly tie iteration and operations to structure, but structure itself is something that can be combinatoric.

[AriMKatz]

I think you meant to tag @peterahrens

[dlfivefifty]

To be honest it feels unnatural working with block-arrays on the "index level": one would usually try to work directly with blocks. In particular I don't think the generic fallback Tim proposes would be very fast for these types of arrays.

That said, perhaps there is a block-wise way to think about iterating over the non-zeros: we have a block-iterator that iterates over the non-zero blocks, which then calls the iterator for each block.

[AriMKatz]

The other potential consideration, is should this be paused for the new traits (traitor+extensible unions etc) infrastructure that was discussed on slack, or will the discussed design be futureproof?

[timholy]

My conclusion from that discussion is that the best interface may be the one we have now, where you manually write out the trait depot.

[dlfivefifty]

To further add to the noise: also relevant is https://github.com/JuliaMatrices/ArrayLayouts.jl which I use to specify traits for the underlying storage of arrays, giving much more general implementations of multiplication algorithms (in there is muladd!). This has proven pretty successful: I get fast multiplication in banded/block-banded matrices for subviews for "free".

Though perhaps this is orthogonal to what we are discussing here, which is how to implement a general sparse multiplication, not how to dispatch

[AriMKatz]

Also, @ChrisRackauckas, how does work on zygote compatibility semantics play into this? Things like value type arrays, @MikeInnes 's https://github.com/FluxML/IArrays.jl and @Keno 's heap allocated array stuff JuliaLang/julia#31630

[timholy]

Re IArrays: it seems like it would be better to add whatever is needed to StaticArrays. They have a good abstract hierarchy and you can get specific features by subclassing it. And you'd get a ton of functionality for free: https://github.com/JuliaArrays/StaticArrays.jl/tree/master/src

[MikeInnes]

The goal is a bit different from StaticArrays since we don't expect them to be specialised on size; they should behave exactly like normal Arrays that happen to be immutable. If / when we have that working well it would likely belong in Base.

[Tokazama]

Might be worth mentioning that part of this may be accomplished with https://github.com/Tokazama/AxisIndices.jl because the whole package is just about mapping values from the user facing input to the underlying memory of the parent array.

[nalimilan]

Has it been considered to stop using Requires.jl and instead move definitions to the packages that define the corresponding types (StaticArrays, CUDA...)? Requires.jl increases the load time significantly, so it's generally considered as a showstopper for lightweight essential packages which could use traits (such as StatsFuns see JuliaStats/StatsFuns.jl#97 (comment)). Also in general it sounds like a better model to have each package define traits than to maintain code for all possible array packages centrally.

[ChrisRackauckas]

Yes. Requires.jl is a solution for today, but as these traits become more standard I hope that some move to Base and the packages then adopt the right overloads. But traits are only useful if they are supported by the ecosystem, so for now we add the support for others. We should probably start the process of upstreaming some of these.

[timholy]

I very much agree with this. It would be best, I think, to have a basal traits package (probably in JuliaArrays) consisting of just the fallback definitions. It would be very lightweight and make it an easy dependency for lots of specific array-type packages.

It's also worth noting that on 1.6 the overhead of Requires is now about 60ms on my machine, enough to be noticeable but not as bad as it was:

julia> @time using ImageCore   # default package definition
  0.431669 seconds (883.65 k allocations: 65.192 MiB)

# Now comment out `using Requires` and the `__init__` function and try again
julia> @time using ImageCore
  0.372618 seconds (846.88 k allocations: 62.834 MiB)

[Tokazama]

What's the next step in moving forward with having other packages directly depend on ArrayInterface? Do we need to transfer to JuliaArrays or should we just start submitting PRs to other packages?