Flux.jl: Batched Matrix Multiply

[onetonfoot]

It would be nice to able to use batched matrix multiples like those in BatchedRoutines.jl, ortorch.bmm since it seems roughly twice as fast.

A = rand(10,20,1000)
B = rand(20,10,1000)

function bmm(A,B)
    C = Array{eltype(A)}(undef,size(A,1),size(B,2),size(A,3))
    for i in 1:size(A,3)
        C[:,:,i] = A[:,:,i] * B[:,:,i]
    end
    C
end

@btime bmm(A,B) #   974.955 μs (3002 allocations: 5.07 MiB)

@btime batched_gemm('N','N',A,B) #   433.808 μs (2 allocations: 781.33 KiB)

However I can't seem to get it to work with a TrackedArray ?

[avik-pal]

BatchedRoutines.jl ccalls in BLAS, LAPACK, etc for computing the results. So it will not work directly with TrackedArrays. But we can simply define the custom gradients (I dont know if we can handle this in any other way) to handle it. The process to describe custom gradients are here.

I had written a version for batched matrix multiplication for computing gram matrices. So it is essentially the batched version of x * x'. It can be adapted for A * B batched versions by just changing the gradient function.

gram_matrix_fast(x::TrackedArray) = track(gram_matrix_fast, x)

gram_matrix_fast(x::CuArray) = CuArrays.CUBLAS.gemm_strided_batched('T', 'N', x, x)

∇gram_matrix_fast(x::CuArray, Δ::CuArray) =
    CuArrays.CUBLAS.gemm_strided_batched('N', 'N', x, Δ .+ Δ') 

@grad function gram_matrix_fast(x)
    return gram_matrix_fast(data(x)), Δ -> (∇gram_matrix_fast(data(x), Δ), )
end

This was pretty much faster than simple multiplication (as your example). You can swap the functions from BatchedRoutines in place of the CUBLAS functions.

[onetonfoot]

Thanks for the help I got something working.

function bmm(A::TrackedArray, B::TrackedArray)
    track(bmm,A,B)
end

function bmm(A::Array,B::Array)
    batched_gemm('N','N',A,B)
end

function bmm(A::CuArray,B::CuArray)
    batched_gemm('N','N',A,B)
end

@grad function bmm(A,B)
     bmm(data(A),data(B)) ,Δ -> (batched_gemm('N' ,'T', Δ , data(B)), batched_gemm( 'T' ,'N',data(A) , Δ))
end