implementation for batch-wise matrix multiplication

src/NNlib.jl

@@ -13,6 +13,7 @@ end
 
 include("activation.jl")
 include("softmax.jl")
+include("batched/batchedmul.jl")
 include("gemm.jl")
 include("conv.jl")
 include("pooling.jl")

src/batched/batchedadjtrans.jl

@@ -0,0 +1,67 @@
+using LinearAlgebra
+import Base: -
+
+"""
+    BatchedTranspose{T, N, S} <: AbstractBatchedMatrix{T, N}
+Batched transpose. Transpose a batch of matrix.
+"""
+struct BatchedTranspose{T, S} <: AbstractArray{T, 3}
+    parent::S
+    BatchedTranspose{T, S}(X::S) where {T, S} = new{T, S}(X)
+end
+
+"""
+    batched_transpose(A)
+Lazy batched transpose.
+"""
+batched_transpose(A::AbstractArray{T}) where T = BatchedTranspose(A)
+batched_transpose(A::BatchedTranspose) = A.parent
+
+"""
+    BatchedAdjoint{T, N, S} <: AbstractBatchedMatrix{T, N}
+Batched ajoint. Transpose a batch of matrix.
+"""
+struct BatchedAdjoint{T, S} <: AbstractArray{T, 3}
+    parent::S
+    BatchedAdjoint{T, S}(X::S) where {T, S} = new{T, S}(X)
+end
+
+"""
+    batched_adjoint(A)
+Lazy batched adjoint.
+"""
+batched_adjoint(A::AbstractArray{T, 3}) where T = BatchedAdjoint(A)
+batched_adjoint(A::BatchedAdjoint) = A.parent
+
+BatchedAdjoint(A) = BatchedAdjoint{Base.promote_op(adjoint,eltype(A)),typeof(A)}(A)
+BatchedTranspose(A) = BatchedTranspose{Base.promote_op(transpose,eltype(A)),typeof(A)}(A)
+
+
+const BatchedAdjOrTrans{T, S} = Union{BatchedTranspose{T, S}, BatchedAdjoint{T, S}}
+
+LinearAlgebra.wrapperop(A::BatchedAdjoint) = batched_adjoint
+LinearAlgebra.wrapperop(B::BatchedTranspose) = batched_transpose
+
+# AbstractArray Interface
+Base.length(A::BatchedAdjOrTrans) = length(A.parent)
+Base.size(m::BatchedAdjOrTrans) = (size(m.parent, 2), size(m.parent, 1), size(m.parent, 3))
+Base.axes(m::BatchedAdjOrTrans) = (axes(m.parent, 2), axes(m.parent, 1), axes(m.parent, 3))
+
+Base.IndexStyle(::Type{<:BatchedAdjOrTrans}) = IndexCartesian()
+Base.@propagate_inbounds Base.getindex(m::BatchedTranspose, i::Int, j::Int, k::Int) = getindex(m.parent, j, i, k)
+Base.@propagate_inbounds Base.getindex(m::BatchedAdjoint, i::Int, j::Int, k::Int) = adjoint(getindex(m.parent, j, i, k))
+Base.@propagate_inbounds Base.setindex!(m::BatchedAdjOrTrans, v, i::Int, j::Int, k::Int) = setindex!(m.parent, v, j, i, k)
+
+Base.similar(A::BatchedAdjOrTrans, T::Type, dims::Dims) = similar(A.parent, T, dims)
+Base.similar(A::BatchedAdjOrTrans, dims::Dims) = similar(A.parent, dims)
+Base.similar(A::BatchedAdjOrTrans, T::Type) = similar(A.parent, T, size(A))
+Base.similar(A::BatchedAdjOrTrans) = similar(A.parent, size(A))
+
+Base.parent(A::BatchedAdjOrTrans) = A.parent
+
+(-)(A::BatchedAdjoint)   = BatchedAdjoint(  -A.parent)
+(-)(A::BatchedTranspose) = BatchedTranspose(-A.parent)
+
+Base.copy(A::BatchedTranspose) = BatchedTranspose(copy(A.parent))
+Base.copy(A::BatchedAdjoint) = BatchedAdjoint(copy(A.parent))
+

src/batched/batchedmul.jl

@@ -0,0 +1,38 @@
+# batch-wise matrix multiplication
+# wrapper for batched_gemm!
+export batched_mul, batched_transpose, batched_adjoint
+
+
+include("./batchedadjtrans.jl")
+
+function batched_mul(A::AbstractArray{T, 3}, B::AbstractArray{T, 3}) where T
+    size(A, 3) == size(B, 3) || throw(DimensionMismatch("batch size mismatch"))
+    batched_mul!(similar(A, (size(A, 1), size(B, 2), size(A, 3))), A, B)
+end
+
+"""
+    batched_mul!(C, A, B) -> C
+batched `mul!`.
+"""
+function batched_mul! end
+
+_unbatch(A) = A
+_unbatch(A::BatchedAdjOrTrans) = A.parent
+
+# bmm
+const _BATCHED_MATRIX_LIST = [
+        (:(AbstractArray{T, 3}), 'N'),
+        (:(BatchedTranspose{T, <:AbstractArray{T, 3}}), 'T'),
+        (:(BatchedAdjoint{T, <:AbstractArray{T, 3}}), 'C')
+]
+
+for (TA, transA) in _BATCHED_MATRIX_LIST, (TB, transB) in _BATCHED_MATRIX_LIST
+    @eval begin
+        function batched_mul!(C::AbstractArray{T, 3}, A::$TA, B::$TB) where T
+            batched_gemm!($transA, $transB, one(T), _unbatch(A), _unbatch(B), zero(T), C)
+            C
+        end
+
+
+    end
+end

src/gemm.jl

@@ -56,3 +56,50 @@ for (gemm, elt) in gemm_datatype_mappings
         end
     end
 end
+
+for (gemm, elt) in gemm_datatype_mappings
+    @eval begin
+        @inline function batched_gemm!(transA::AbstractChar,
+                               transB::AbstractChar,
+                               alpha::($elt),
+                               A::AbstractArray{$elt, 3},
+                               B::AbstractArray{$elt, 3},
+                               beta::($elt),
+                               C::AbstractArray{$elt, 3})
+            @assert !Base.has_offset_axes(A, B, C)
+            @assert size(A, 3) == size(B, 3) == size(C, 3) "batch size mismatch"
+            m = size(A, transA == 'N' ? 1 : 2)
+            ka = size(A, transA == 'N' ? 2 : 1)
+            kb = size(B, transB == 'N' ? 1 : 2)
+            n = size(B, transB == 'N' ? 2 : 1)
+            if ka != kb || m != size(C,1) || n != size(C,2)
+                throw(DimensionMismatch("A has size ($m,$ka), B has size ($kb,$n), C has size $(size(C))"))
+            end
+            LinearAlgebra.BLAS.chkstride1(A)
+            LinearAlgebra.BLAS.chkstride1(B)
+            LinearAlgebra.BLAS.chkstride1(C)
+
+            ptrA = Base.unsafe_convert(Ptr{$elt}, A)
+            ptrB = Base.unsafe_convert(Ptr{$elt}, B)
+            ptrC = Base.unsafe_convert(Ptr{$elt}, C)
+
+            for k in 1:size(A, 3)
+                ccall((@blasfunc($(gemm)), libblas), Nothing,
+                      (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
+                       Ref{BlasInt}, Ref{$elt}, Ptr{$elt}, Ref{BlasInt},
+                       Ptr{$elt}, Ref{BlasInt}, Ref{$elt}, Ptr{$elt},
+                       Ref{BlasInt}),
+                      transA, transB, m, n,
+                      ka, alpha, ptrA, max(1,Base.stride(A,2)),
+                      ptrB, max(1,Base.stride(B,2)), beta, ptrC,
+                      max(1,Base.stride(C,2)))
+
+                ptrA += size(A, 1) * size(A, 2) * sizeof($elt)
+                ptrB += size(B, 1) * size(B, 2) * sizeof($elt)
+                ptrC += size(C, 1) * size(C, 2) * sizeof($elt)
+            end
+
+            C
+        end
+    end
+end

test/batchedmul.jl

@@ -0,0 +1,47 @@
+function bmm_test(a,b; transA = false, transB = false)
+    bs = size(a,3)
+    transA && (a = permutedims(a, [2,1,3]))
+    transB && (b = permutedims(b, [2,1,3]))
+    c = []
+    for i = 1:bs
+        push!(c, a[:,:,i]*b[:,:,i])
+    end
+
+    cat(c...; dims = 3)
+end
+
+function bmm_adjtest(a,b; adjA = false, adjB = false)
+    bs = size(a,3)
+    c = []
+    for i = 1:bs
+        ai = adjA ? adjoint(a[:,:,i]) : a[:,:,i]
+        bi = adjB ? adjoint(b[:,:,i]) : b[:,:,i]
+        push!(c, ai*bi)
+    end
+
+    cat(c...; dims = 3)
+end
+
+@testset "Batched Matrix Multiplication" begin
+    A = randn(7,5,3)
+    B = randn(5,7,3)
+    C = randn(7,6,3)
+
+    @test batched_mul(A, B) == bmm_test(A, B)
+    @test batched_mul(batched_transpose(A), batched_transpose(B)) == bmm_test(A, B; transA = true, transB = true)
+    @test batched_mul(batched_transpose(A), C) == bmm_test(A, C; transA = true)
+    @test batched_mul(A, batched_transpose(A)) == bmm_test(A, A; transB = true)
+
+
+    cA = randn(Complex{Float64}, 7,5,3)
+    cB = randn(Complex{Float64}, 5,7,3)
+    cC = randn(Complex{Float64}, 7,6,3)
+
+    @test batched_mul(cA, cB) == bmm_adjtest(cA, cB)
+    @test batched_mul(batched_adjoint(cA), batched_adjoint(cB)) == bmm_adjtest(cA, cB; adjA = true, adjB = true)
+    @test batched_mul(batched_adjoint(cA), cC) == bmm_adjtest(cA, cC; adjA = true)
+    @test batched_mul(cA, batched_adjoint(cA)) == bmm_adjtest(cA, cA; adjB = true)
+
+    @test batched_transpose(batched_transpose(A)) == A
+    @test batched_adjoint(batched_adjoint(cA)) == cA
+end

test/runtests.jl

@@ -2,5 +2,6 @@ using NNlib, Test
 
 include("activation.jl")
 include("conv.jl")
+include("batchedmul.jl")
 include("pooling.jl")
 include("inference.jl")