Package-Independent BLAS / LAPACK Rule System

[willtebbutt]

Both Enzyme.jl and Tapir.jl can in principle support low-ish level calls to BLAS / LAPACK functionality. However, both really do need rules for them (certainly Tapir.jl does because BLAS / LAPACK sit behind C-call barriers, and I believe I'm correct in saying that Enzyme.jl also needs rules for them). To this end, I would like to consider implementing a restricted set of rules which is

1. independent of both packages, and
2. usable by both packages.

In this sense, it should be more like DiffRules.jl than ChainRules.jl, albeit it won't involve symbolic rules.

In order for this to happen, it would need to support (in some sense) the union of the interfaces for the rule systems for the two packages, and be just flexible enough to cover everything in BLAS / LAPACK.

Key Point: Limited Ambition

The main idea is to avoid having to write BLAS / LAPACK rules twice, once for Enzyme.jl and once for Tapir.jl.
Moreover, I want a system which is entirely focused on BLAS / LAPACK, with no regard for anything more general than this.
This lets us agree on how to handle each of the possible types we might encounter, rather than having to support an open-ended collection -- this is a much easier task.

BLAS / LAPACK Interfaces

To the best of my knowledge, there are basically 3 or 4 types which appear as arguments in BLAS / LAPACK implementations in Julia:

1. Ptr{T}
2. T
3. Char
4. Int
   where T<:IEEEFloat. It might be that Char and Int are actually pointers to such types in some cases, but the important features is that they're non-differentiable arguments.

The Idea

Implement rules for linear algebra with these types.
They would use shadow memory for pointers, but not for scalars.
They look a lot like both Tapir.jl rules and Enzyme.jl rules.

For example, the rule for gemm would be something like

function linalg_forwards_pass(
    ::typeof(gemm!),
    tA::Char,
    tB::Char,
    m::Int,
    n::Int,
    k::Int,
    alpha::T,
    A::Tuple{Ptr{T}, Ptr{T}}
    lda::Int,
    B::Tuple{Ptr{T}, Ptr{T}},
    ldb::Int,
    beta::T,
    C::Tuple{Ptr{T}, Ptr{T}},
    ldc::Int,
)
    # run forwards pass + return intermediate state needed on reverse-pass
end

function linalg_reverse_pass(
    state,
    ::typeof(gemm!),
    tA::Char,
    tB::Char,
    m::Int,
    n::Int,
    k::Int,
    alpha::T,
    A::Tuple{Ptr{T}, Ptr{T}}
    lda::Int,
    B::Tuple{Ptr{T}, Ptr{T}},
    ldb::Int,
    beta::T,
    C::Tuple{Ptr{T}, Ptr{T}},
    ldc::Int,
)
    # compute derivative of adjoints for isbits types
    # increment / zero shadow memory in the usual manner.
    # return adjoint of derivative for each of the differentiable isbits arguments that are active.
end

We would need some data structures to represent activity / whether or not certain memory needs to be restored on the reverse-pass etc, and decide what happens in those situations. Largely I imagine that this will be the same as in Enzyme.jl.
Types which are identified with their address should have shadow memory if they're active, and not if they're not.

How would Tapir.jl make use of such rules?

Tapir.jl targets LinearAlgebra.BLAS.gemm!, rather than the C-call to gemm, basically because it's a little bit cleaner.
You could use the above to write an rrule!! along the following lines:

function rrule!!(
    ::CoDual{typeof(BLAS.gemm!)},
    transA::CoDual{Char},
    transB::CoDual{Char},
    alpha::CoDual{T},
    A::CoDual{<:MatrixOrView{T}},
    B::CoDual{<:MatrixOrView{T}},
    beta::CoDual{T},
    C::CoDual{Matrix{T}},
) where {T<:Union{Float32, Float64}}
    state = linalg_forwards_pass(
        primal(transA),
        primal(transB),
        m, n, k, 
        primal(alpha),
        (primal(A), tangent(A)),
        ...
    )
    function gemm!_pb!!(::NoRData)
        _, _, _, _, ..., dalpha, ..., = linalg_reverse_pass(state, primal(transA), ...)
        return # organise output of reverse pass with `NoRData` to give the right thing
    end
    return C, gemm!_pb!!
end

Testing

Testing can be done using FiniteDifferences.jl. Due to the heavily restricted interface, it will be easy to write some simple standardised testing functionality. We would not make it available to other packages, as the intention would not be for other packages to extend this collection of rules -- all of the things we wish to support are already known to us.

[yebai]

@willtebbutt, would it be possible to make ChainRules use this new low-level rule package, i.e. Enzyme/Tapir/ChainRules all depend on a shared rule package for BLAS/LAPACK? That would allow us to minimise duplicate code in the ecosystem.

EDIT: can we have a ChainRulesCore.jl V2 to specify these rules?

[willtebbutt]

It would definitely require API changes to ChainRules to make it possible to implement these kinds of rules in ChainRules, because ChainRules v1 doesn't do mutation.

If ChainRules v2 ever happens, we will need to have worked out the design for rules which support mutation well in advance. The best way to do this is in a separate package where we have the freedom to experiment.

To be honest, I really think a standalone package is the way to go regardless. Supporting BLAS / LAPACK is a well-defined problem, because 1) they use a restricted set of types and 2) they are stable on decadal time scales.

In any case, a future version of ChainRules could just wrap the rules from this standalone package.