Automatic handling of batching

[THargreaves]

There are many places where our code could be simplified by writing generic code that automatically handled batching (rather than requiring special code for batched filtering). This is especially important when we consider multiple levels of batching, e.g. a Gaussian sum filter on multiple independent targets (which has two batch dimensions).

It should be possible to handle this all automatically with the right sort of dispatching (since the algorithms are exactly the same). This is exactly how you would do it in Python (TF/JAX/PyTorch) where you would write the KF algorithm in a single matrix setting and broadcasting would make it work automatically in batch settings.

[THargreaves]

Here are some notes discussing this issue in more detail.

Attached as PDF since LaTeX isn't rendering properly on GitHub.

Batched Interface Notes.pdf

GeneralisedFilters.jl Batched Interface Notes

Introduction

This memo discusses the design of a batched interface for filtering with
GeneralisedFilters.jl.

I will start by providing a few motivating examples for why we specifically need
a batching interface in GeneralisedFilters.jl, though it should be noted that
batched linear algebra has far wider applications in the Julia ecosystem.

I will then discuss existing tooling in the Julia ecosystem for batched linear
algebra as well as approaches taken by others, in particular, the Python deep
learning frameworks.

Finally, I will share some rough ideas for what a Julian batched interface could
look like. It would be great to discuss these ideas further in our next meeting.

Motivation

At a high-level, a well-designed batched interface will allow us to create a
clean and intuitive API for filtering with GeneralisedFilters.jl that reduces
code duplication whilst providing support for:

• Batched filtering algorithms/Rao-Blackwellisation
• Optimal use of StaticArrays.jl
• GPU acceleration via CUDA
• Multiple-target tracking (MTT) with:
  • Multiple independent filters
  • A single dependent filter, that outputs multiple observations

I'll expand upon these tasks in the rest of the section.

Only one of these (multiple independent filters) is supported by Stone Soup.

Note

These features are related to a few other problems that we should consider
together:

• "Element type"-getters for initialisation storage containers
• Initialisation of filter states
• Settings with multiple measurements for one target
• Multiple batching dimensions (e.g. MTT, each with a RBPF)
• Implementation of a Gaussian-sum filter
• Automatic use of struct-of-array layouts for our CUDA code

From a philosophical perspective, a batched interface should allow us to
maintain the core ideal of SSMProblems.jl/GeneralisedFilters.jl—that the user
should be able to define a model/algorithm once and have it work in various
settings (e.g. parallel with CUDA, MTT, etc.) with no/minimal redefinitions
required.

Batched Filtering/RBPF/Independent MTT

There are many settings in which we would like to run multiple independent
filters in parallel. For example:

• As the inner step of a RBPF
• In MTT, where each target has its own filter

Note, the two settings above differ in that the first has the same observation
across the batch (homogeneous), while the second has different observations for
each filter (heterogeneous).

A user may also want to run parallel filtering using multiple models with
different dynamics, which is a related but slightly more involved problem, still
worth being considered.

Currently, we support single KF on the CPU and batched KF using CUDA, but these
have completely separate implementations, despite the maths being identical.

Ideally, we should have one generic Kalman Filter implementation that can be
used for:

• Single KF
• Batched KF using CUDA
• Batched KF using other backends (e.g. contiguous static arrays)

Important

This bloat matters more when we start implementing EKF, UKF, SR-EKF, SR-UKF,
forward algorithm, information filter (IF), EIF, etc. plus all of the related
backward filters and smoothers. Or, when we consider the various permutations of
homogeneous/heterogeneous models, observations, controls, which is would be
restrictive not to support.

Roughly, the user should be able to run a filtering update step where the
inputs are:

• A single LGSSM or a batch of LGSSMs
• A single observation or a batch of observations
• A single control or a batch of controls
• A single state or a batch of states

and the results should be "broadcasted" appropriately.

Optimal use of StaticArrays.jl

Not only does a batched interface provide code quality benefits, it is crucial
for performance. Clearly, batching is required to take advantage of GPU
acceleration but it is actually essential for low-dimensional CPU performance
too.

Tip

StaticArrays.jl speeds up linear algebra for low dimensions by
stack-allocating the elements of the array (essentially storing them as a
tuple). The "static" refers to static sizing (i.e. we know the size of the
arrays at compile time); the arrays can be mutable if needed.

For such sizes, operations using StaticArrays.jl are roughly 10x faster for
most linear algebra operations in microbenchmarks. This is actually a
pessimistic estimate and in the GeneralisedFilters.jl KF benchmark, we've seen
a speed-up of around 1000x for $D\leq 4$.

The speed-up is particularly pronounced when performing multiple operations in
series, as we are memory-bound for small arrays, and StaticArrays.jl avoids us
having to write to the heap after each operation.

StaticArrays.jl recommends only using static arrays for arrays with up to 100
elements at which point the performance benefits start to diminish as we
become compute-bound.

Vector of static arrays are substantially less performant than the theoretical
bound when the arrays are small. This is because each array is padded to 64-byte alignment, so
memory access is wasteful (there may be other reasons on top of this—the
Array type has quite a bit of bloat).

vs = [rand(1, 1) for _ in 1:2];
Int(pointer(vs[2])) - Int(pointer(vs[1]))
#> 64

Ideally, we want to store batches of matrices as $D\times D \times B$ contiguous
arrays and reinterpret them as static arrays before performing computations.

out = Array{Float64}(undef, D, D, B)
Σs_static = reinterpret(reshape, SMatrix{D, D, Float64, D^2}, reshape(Σs, D^2, B))
for b in 1:B
    out[:, :, b] .= det(Σs_static[b])
end

Note

There might be a better way of implementing the above idea.

For $D=2$, Kalman filter times are:

Vectors of arrays: 955.545 ms 
Vectors of static arrays: 31.600 ms
Reinterpreted static arrays: 3.342 ms

with the last matching the CPU theoretical optimum.

Multiple-Target Tracking with Dependent Dynamics

As far as the official documentation is concerned (though there may be hacks to
achieve this) Stone Soup does provide the ability to have dependent dynamics in
a MTT setting. Instead, each target must have its own independent posterior
state and be updated independently.

As an example, suppose we are tracking 3 objects in a 2D space each having
position/velocity states (i.e. 3 x 4 = 12 total dimensions).

These dynamics may often be dependent (in fact, Simon is particularly interested in
this setting for a fish tracking problem),
requiring us to represent the joint posterior as a Gaussian distribution in $\mathbb{R}^{12}$.

If we represent the current predicted joint latent state as a Gaussian
distribution,

$$ \bar{G}^{X}_{t|t-1} = \mathcal{N}_{12}(\bar{\mu}^{X}_{t|t-1}, \bar{\Sigma}^{X}_{t|t-1}), $$

our projection step (maps latent $X_t$ onto predicted observation
$\hat{Y}_t$), needs to map this to a 3-vector of 2-dimensional Gaussians,

$$ G^{\hat{Y_i}}_{t|t-1} = \mathcal{N}_{2}(\bar{\mu}^{\hat{Y_i}}_{t|t-1}, \bar{\Sigma}^{\hat{Y_i}}_{t|t-1}). $$

These could be represented as simply a vector-of-Gaussians but more likely some
special struct for batching. This is because the vector-of-XYZ approach is
inefficient for StaticArrays.jl as shown above, and is infeasible for CUDA
arrays which should be large and contiguous.

One way to think about this is for our projection matrix $H$ to be a batch of
$2\times 12$ sparse matrices of batch-size 3:

$$ \begin{align*} H_1 &= \begin{bmatrix} 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \end{bmatrix}\\ H_2 &= \begin{bmatrix} 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 \end{bmatrix}\\ H_3 &= \begin{bmatrix} 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 \end{bmatrix} \end{align*} $$

In a "tensor"-style language (e.g. TensorFlow, JAX), this can be viewed as an
array in $\mathbb{R}^{2\times 12 \times 3}$, though we consider it more as $\left(\mathbb{R}^{2\times 12}\right)^3$.

Multiplying this by $\bar{G}^{X}_{t|t-1}$, correctly gives us a 3-batch of
two-dimensional Gaussians. In terms of dimension sizes we have:

$$ \left(\mathbb{R}^{2\times 12}\right)^3 \times \mathbb{R}^{12} = \left(\mathbb{R}^2\right)^3 $$

Important

Julia might actually allow us to produce a cleaner system than the Python deep
learning frameworks. Specifically, in e.g. TensorFlow, a 2D array could either
be a matrix or a batched vector. Likewise a 3D array could be a batch of
matrices or vectors with two batching dimensions.

We can then compare our projected means in $\left(\mathbb{R}^{2}\right)^3$ to a
batched vector of observations $\left(\mathbb{R}^{2}\right)^{N_{\text{obs}}}$ to
perform data association.

Warning

TODO: There are certainly more cases like this that need to be considered to
double-check they can all be represented through batching. This is especially
the case when we start nesting batch dimensions (e.g. MTT with RBPF using
Gaussian-sum filter as inner filter—3 levels of nesting)

Additional Considerations

I'll briefly make a point about the first two additional considerations, which
are concerned with the initialisation of the filter and initialisation of
the storage type.

We currently implement a latent dynamics type as a subtype of a
LatentDynamics{T, ET} where T gives the arithmetic type (e.g. Float64) and
ET gives the element type of the non-batched CPU code. This introduces many
issues when we try to expand to batched settings.

For example, let's say we want to perform batch forward simulation using CUDA.
In the non-batched CPU case, we know that the output trajectory will be a
Vector{ET}, but there's no guaranteed way to know this in a batched setting.

In the batched CUDA setting, this is likely to be a Vector{CuArray{T}} but
that is not guaranteed.

A more pressing issue is how to initialise the filter state. Following our
philosophy that the same model should work in all of a non-batched, batched
CUDA, and batched StaticArrays.jl setting, the model itself shouldn't contain
information on the types of the states.

I propose instead that we take a more hands off approach to these types
concerns, and aim to decouple the model dynamics from the filter states by
writing the filtering steps generically for non-batched and batched cases.

Important

I'm not sure what impacts this would have on type stability and whether this
would cause problems for auto-diff. This is something worth discussing.

For example, suppose we define some (deterministic for now) latent dynamics by

simulate(rng::AbstractRNG, dyn::ExampleDynamics, prev_state, step::Integer; u, kwargs...)
    state = dyn.A * prev_state + u
    return state
end

This says nothing about the types of the state or control input u. It is a
generic operation.

If we were to pass in a Vector{Float64} for prev_state, we would get a
Vector{Float64} back. If we were to pass in a batched vector however, and had
properly defined * and + for this case, we would correctly get a batch of
predictions.

Importantly, with the correct broadcasting we could use the same A and u
(non-batched) but still consider batched states.

This becomes a bit more complicated when we consider random dynamics. I don't
know what the best form of this could look like, but it could be something like:

simulate(rng::AbstractRNG, dyn::ExampleDynamics, prev_state, step::Integer; u, kwargs...)
    state = dyn.A * prev_state + u
    return rand(Normal(state, dyn.Q))  # with special dispatch for the batched case
end

To make this idea work, I think the user would need to specify either a starting
type or a starting posterior for the filter. The user could therefore use the
same model and depending on whether they run it with a Vector{Float64} or
CuArray{Float64, 3} would get different batching behaviours.

This is similar to the behaviour Kalman.jl, where you pass in an initial
Gaussian distribution.

I appreciate that this might not be ideal for all cases, but we could make it
simpler for new users by:

• Defaulting to the non-batched CPU case if no type/initial posterior is provided
• Creating some constructors that help initialise the filter state

Existing Tooling

Julia Ecosystem

Batched Linear Algebra Libraries

CUDA.jl and MAGMA.jl wrap a considerable number of batched linear algebra
operations that run on CUDA GPUs. They come in nonstrided (vector of pointers
to 2D CuArrays) and strided (3D CuArray) forms.

These methods are faffy to use and require a lot of boilerplate code to set up

julia> methods(CUDA.CUBLAS.getrs_strided_batched)
# 2 methods for generic function "getrs_strided_batched" from CUDA.CUBLAS:
 [1] getrs_strided_batched(trans::Char, A::CuArray{<:Any, 3}, B::CuArray{<:Any, 3}, pivotArray::Union{Nothing, CuMatrix{Int32}, DenseCuMatrix{Int32}})
     @ ~/.julia/packages/CUDA/1kIOw/lib/cublas/wrappers.jl:2078
 [2] getrs_strided_batched(trans::Char, A::CuArray{<:Any, 3}, B::CuArray{<:Any, 3})
     @ ~/.julia/packages/CUDA/1kIOw/lib/cublas/wrappers.jl:2078

No wrappers are provided to make their use easier (e.g. just calling * rather
than gemm_strided_batched).

Another limitation of these wrappers is that they don't support broadcasting. If
I want to perform $AB_i$ over many $i$, I will need to manually create a vector of
pointers all pointing to $A$ to reuse the same $A$ for each $i$.

There are a few dormant projects implementing additional batched linear algebra
operations using CUDA.jl or KA.jl but none of these provided interfaces, just
low-level operations. See:

• BatchedBLAS.jl
• BatchedRoutines.jl

The closest we have to a batched interface for CUDA is NNlib.jl which
provides some
batching helper functions for multiplication and transposition only. It doesn't
define any special types for batching and so still has its own API rather than
dispatching on * and '. It does handle batch-singleton broadcasting in some
settings though.

Turing Ecosystem

Tor has some old work prototyping a batching system for Bijectors.jl,
here. I'm not
sure it really achieves what we're after, but I need to think more about it.

Python Ecosystem

The Python deep learning libraries support batching by implicitly expanding
dimensions, allowing operations between batches and single matrices/vectors
without explicit reshaping.

batch_size = 32
A = torch.randn(batch_size, 10, 20) 
b = torch.randn(20)
result = torch.matmul(A, b)

This is quite a powerful and flexible feature that should likely be the
foundation of our design. Perhaps however we can take advantage of Julia's
broadcasting and dispatch capabilities to provide a more powerful interface.

Note

The main limitations of these Python libraries are:

• They do not distinguish between array and batch dimensions which could cause ambiguity
• They use a separate API to NumPy are so require redefinitions for GPU code
• They do not support stack-allocated batched operations like we could achieve
  with StaticArrays.jl
• Their fusion capabilities are limited (JAX) or non-existent (PyTorch, TensorFlow)

These libraries almost entirely rely on the CUDA and MAGMA libraries to
implement the core CUDA kernels; they just provide a clean wrapper. These
wrappers involve creating additional pointers to the same data to handle broadcasting.

JAX goes one step further through the use of vmap which allows the user to
take a function written for singleton matrices/vectors and apply it in a batched
setting.

def matrix_op(mat):
    return jnp.linalg.inv(mat @ mat.T + jnp.eye(mat.shape[0]) * 0.1)

batch_result = jax.vmap(matrix_op)(batch_mats)

Hypothetically, this should have the benefit of being able to fuse together
kernels to reduce global memory access, but benchmarks seem to suggest this is
only effective for simple operations such as matmuls.

I'm not sure how feasible it would be to implement something similar in Julia,
but it's very cool.

Implementation Suggestions

In this section, I'll discuss some rough ideas that I have considered, some that
show promise and others that have flaws.

Batched Types

A first-pass approach (that seems to work fairly well) is to take inspiration from
the Python deep learning libraries but take advantage of Julia's dispatch system
by defining batched types.

We could start with abstract types AbstractBatchedVectors and
AbstractBatchedMatrices and then define concrete subtypes for each backend.
E.g.

struct CuBatchedMatrix{T, B} <: AbstractBatchedMatrix{T}
    data::CuArray{T, 3}
end

Tip

We could possibly have AbstractBatchedVectors <: AbstractVector since it
behaves in the same way. We would just need to define the expected methods.

More generally, if the matrices have variable sizes we could represent this as

struct CuVarBatchedMatrix{T, B} <: AbstractBatchedMatrix{T}
    data::Vector{CuMatrix{T}}
    ptrs::CuVector{CuPtr}
end

which allows us to cache the pointers to the data (which have overhead in
computing from a vector of CuArrays).

Note

We could possibly have AbstractBatchedVectors <: AbstractVector since it
behaves in the same way. We would just need to define the expected methods.

We would then have to define wrappers for these batched types that call the
appropriate CUDA or MAGMA functions. Importantly, we would be dispatching on the
regular functions used for singleton matrices/vectors, e.g. *, cholesky.

This would allow us to use the same API for both batched and non-batched
operations. For mixed singleton/batched operations, we would need to create a wrapper that
duplicates pointers to the same data, but that shouldn't be tricky.

I think this approach gets us most of the way to where we want to be. At the
very least it seems to solve our code reuse problem, although it is not
optimally performant as follows.

Fusion

The above strategy works well if you are only planning to apply one individual
batched operation at a time, but is not sub-optimal if we wish to perform
multiple batched operations in series.

This is because such operations are memory bound and so benefit greatly from
having all operations fused into one so global memory access can be minimised.
This does not just apply to CUDA, but also CPU operations using StaticArrays.jl
as discussed above.

Julia's multiple dispatch might offer us a powerful way of cleanly achieving
this which would help but our filtering algorithms outperform alternatives
written in Python.

I'm not entirely sure what the best way to achieve this is, and this is probably
the area where we could benefit most from discussions.

Limitations of Broadcasting

Suppose we have defined,

function kalman_predict(μ, Σ, A, b, Q)
    μ_pred = A * μ + b
    Σ_pred = A * Σ * A' + Q
    
    return μ_pred, Σ_pred
end

After defining appropriate methods for linear algebra operations using
BatchedCuMatrix, BatchedStaticMatrix, etc. , we could then call this
function with batched inputs. This however would not be fused.

It would be nice if we could instead write

kalman_predict.(μs, Σs, A, b, Q)

where μs and Σs are batched inputs, and take advantage of Julia's broadcast
system by defining special broadcast
styles
for BatchedCuMatrix and BatchedStaticMatrix.

The issue is, that if we call this function with all regular inputs, Julia will
try to broadcast the operation over the matrix/vector input elements and error.

I think the only way to get around this but be to wrap the inputs and parameters
in special types to stop the default broadcasting from being triggered.
Alternatively, we could define a special maybe_broadcast method but this feels
clunky.

A Possible Approach

A more manual approach would require the user to manually specify for what
batched types and for what functions, fusion should be applied. This is a bit
more clunky but is easier to implement and is less likely to lead to unexpected
behaviour. This could be added by the use of certain helpers.

For example we could have a macro @static_batched that takes a function and
batches it for static arrays. For example the Kalman predict function above could
be defined as:

function kalman_predict(
    μ::BatchedStaticVector{T},
    Σ::BatchedStaticMatrix{T},
    A::BatchedStaticMatrix{T},
    b::BatchedStaticVector{T},
    Q::BatchedStaticMatrix{T}
)
    B = length(μ)
    D = length(μ[1])

    μ_ri = reinterpret(reshape, SVector{T, D}, μ)
    Σ_ri = reinterpret(reshape, SMatrix{D, D, T, D^2}, reshape(Σ, D^2, B))
    A_ri = reinterpret(reshape, SMatrix{D, D, T, D^2}, reshape(A, D^2, B))
    b_ri = reinterpret(reshape, SVector{T, D}, b)
    Q_ri = reinterpret(reshape, SMatrix{D, D, T, D^2}, reshape(Q, D^2, B))

    μ_pred = Array{T}(undef, D, B)
    Σ_pred = Array{T}(undef, D, D, B)    

    for b in 1:B
        μ_pred[:, b], Σ_pred[:, :, b] .= kalman_predict(
            μ_ri[b], Σ_ri[b], A_ri[b], b_ri[b], Q_ri[b]
        )
    end

    return μ_pred, Σ_pred
end

and this transformation could be defined automatically.

Likewise, for CUDA kernel fusion, we could use the automatic fusion methods that
I have been working on with Joseph.

Distributions.jl

As a final point, it is worth considering how a batched interface would interact
with Distributions.jl (or at the very least how it could be used for batched
sampling).

We would like to replicate the behaviour of Python deep learning libraries where
you can generate a batch of samples from a distribution by passing in a batch
size:

import torch
import torch.distributions as D

batch_size = 32
dist = D.MultivariateNormal(torch.zeros(10), torch.eye(10))
samples = dist.sample((batch_size,))

Conclusion

It is clear that Julia, and specifically the GeneralisedFilters.jl ecosystem,
would greatly benefit from a more powerful interface for batched linear algebra.

Replicating the capabilities of PyTorch and TensorFlow is quite an approachable
task, and just requires a bit of grunt work in defining what CUBLAS/MAGMA
kernels to call in each case.

Going beyond this to take advantage of function fusion is a more involved task
but comes with significant performance benefits that would put our library in a
league of its own.