PR proposal: `AutoBatch`-batch/vectorized AD backend

[mattsignorelli]

I would like to propose an implementation I have that makes it easy to compute the Jacobians of many different (e.g. 100,000) primals in parallel on the CPU/GPU. The only requirement is that the function to evaluate Jacobians through must be "vectorized", in that:

f(x) = ... # (GPU) vectorized function accepting SoA memory layout
a = f([1 2 3 4 5 6])  # gives some values
b = f([7 8 9 10 11 12]) # gives some other values

# But instead faster to do it in one evaluation thanks to SIMD/SIMT:
f([1 2 3 4 5 6; 7 8 9 10 11 12]) == vcat(a, b)

In this case, each "batch element" is a row. I will refer to this case as the batchdim=1 case

Another option is the AoS memory layout, which is not really vectorized but can still be much, much faster in certain cases than evaluating the same function for each primal separately (if there are layers of type instability in the function, for example):

g(x) = ... # AoS memory layout 
a = g([1, 2, 3, 4, 5, 6])  # gives some values
b = g([7, 8, 9, 10, 11, 12]) # gives some other values

# Many cases where this is still faster:
f([1 7; 2 8; 3 9; 4 10; 5 11; 6 12]) == hcat(a, b)

In this case, each "batch element" is a column, so we can let this case be the batchdim=2 case.

For both of these cases, I can get the Jacobians for all "batch elements"/primals together by just reshaping the input/output arrays to be 1-dimensional, and the using AutoSparse. For batchdim=2 case, the sparse Jacobian matrix will be block-diagonal, and for the batchdim=1 case, the sparse Jacobian matrix will have diagonal bands. This works well, and in theory you could just use AutoSparse on this problem and be done. However, there are a couple of important remarks that make me believe a separate type which I call AutoBatch, different from AutoSparse is truly warranted for this special case:

1. I found the automatic sparsity detector and coloring algorithm can be really slow in cases where there are e.g. 10,000-100,000 batch elements/primals . Furthermore, for batchdim=1 and batchdim=2, we know the pattern and coloring of the sparse Jacobian beforehand, and so we can just construct the pattern and use a ConstantColoringAlgorithm instead.
2. Backend-specific optimizations can be made. Immediately, I am thinking about the massive speed that finite differences could achieve if the function is known to be vectorized or if it is faster to evaluate a batch all at once than to re-evaluate the function many times for each tangent. So, e.g. FiniteDiff could provide special overrides for AutoBatch{AutoFiniteDiff} type. This would be a more general way to solve Feature request: vectorized finite difference evaluation FiniteDiff.jl#210

I have currently a working implementation of AutoBatch in my BatchSolve.jl package. I would be happy to work with you on merging this functionality into DI if there is interest! Feedback is also very welcome.

I use these Jacobians that I compute with cuBLAS's library to do batched newton solves, for example, on 100,000 6x6 systems. It is really fast

Let me know what you think!

[mattsignorelli]

The PR would also include the AutoBatch{AutoFiniteDiff} bindings which I'm working on now

[mattsignorelli]

The code is done, including the AutoFiniteDiff bindings. I included an example to show you how much of a speedup you get, even for only 1000 1D systems:

Beyond 1,000 1D systems, the preparation step of just naively using AutoSparse is far too long for me to wait for. however with AutoBatch, we can compute jacobians for e.g. 100,000 1D systems no problem:

It's peanuts for a GPU to handle.

[mattsignorelli]

Also of course, because AutoBatch handles the sparse preparation step manually by default for all backends, we could still have used AutoForwardFromPrimitive(AutoForwardDiff()) and not had a very long waiting time during the preparation step.

[gdalle]

Thanks a lot! I have seen the issue, I'll give you a detailed answer tomorrow 😊

[gdalle]

See also:

• New package: BatchSolve v0.1.0 JuliaRegistries/General#153094

[gdalle]

Thank you for your suggestion! In general I find the idea really interesting and I'd like to pursue it. What do you imagine as the API?

I found the automatic sparsity detector and coloring algorithm can be really slow in cases where there are e.g. 10,000-100,000 batch elements/primals .

Based on your CUDA MWE, I don't think you're using AutoSparse correctly: if you don't specify the coloring algorithm and sparsity detection method, it defaults to what is essentially a dense Jacobian!
Can you try it again with the settings from this tutorial? Sparsity detection itself won't work on CuArrays but you can either

1. perform sparsity detection on the CPU yourself beforehand
2. use the experimental GPU-friendly detector from feat: GPU-based sparsity detection adrhill/SparseConnectivityTracer.jl#297

To clarify, I don't doubt that vectorizing all of this would give insane speedups, just not quite as insane as the ones you show ;)

Furthermore, for batchdim=1 and batchdim=2, we know the pattern and coloring of the sparse Jacobian beforehand, and so we can just construct the pattern and use a ConstantColoringAlgorithm instead.

This reminds me a lot of ExaModels.jl, perhaps we could borrow some inspiration from there? Ping @frapac @amontoison.

So, e.g. FiniteDiff could provide special overrides for AutoBatch{AutoFiniteDiff} type.

Where would these live, here or in FiniteDiff?