Merge pull request #181 from mcabbott/batch

Add fallback batched_mul!

Author: CarloLucibello

src/batched/batchedadjtrans.jl

@@ -1,35 +1,38 @@
 using LinearAlgebra
 import Base: -
 
-"""
+_batched_doc = """
+    batched_transpose(A::AbstractArray{T,3})
+    batched_adjoint(A)
+
+Equivalent to applying `transpose` or `adjoint` to each matrix `A[:,:,k]`.
+
+These exist to control how `batched_mul` behaves,
+as it operated on such matrix slices of an array with `ndims(A)==3`.
+
     BatchedTranspose{T, N, S} <: AbstractBatchedMatrix{T, N}
-Batched transpose. Transpose a batch of matrix.
+    BatchedAdjoint{T, N, S}
+
+Lazy wrappers analogous to `Transpose` and `Adjoint`, returned by `batched_transpose`
 """
+
+@doc _batched_doc
 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.
-"""
+@doc _batched_doc
 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.
-"""
+@doc _batched_doc
 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.
-"""
+@doc _batched_doc
 batched_adjoint(A::AbstractArray{T, 3}) where T = BatchedAdjoint(A)
 batched_adjoint(A::BatchedAdjoint) = A.parent
 

src/batched/batchedmul.jl

@@ -2,37 +2,63 @@
 # 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)
+"""
+    batched_mul(A, B) -> C
+
+Batched matrix multiplication. Result has `C[:,:,k] == A[:,:,k] * B[:,:,k]` for all `k`.
+"""
+function batched_mul(A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T1, T2}
+    axes(A, 3) == axes(B, 3) || throw(DimensionMismatch("batch size mismatch"))
+    T = promote_type(T1, T2)
+    C = similar(A, T, (axes(A, 1), axes(B, 2), axes(A, 3)))
+    batched_mul!(C, A, B)
 end
 
 """
     batched_mul!(C, A, B) -> C
-batched `mul!`.
+
+In-place batched matrix multiplication,
+equivalent to `mul!(C[:,:,k], A[:,:,k], B[:,:,k])` for all `k`.
 """
 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')
+# batched_gemm!
+
+const _GemmFloat = Union{Float64, Float32, ComplexF64, ComplexF32}
+
+_BATCHED_GEMM_LIST = [
+    (:(StridedArray{T, 3}), 'N'),
+    (:(BatchedTranspose{T, <:StridedArray{T, 3}}), 'T'),
+    (:(BatchedAdjoint{T, <:StridedArray{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
+for (TA, transA) in _BATCHED_GEMM_LIST, (TB, transB) in _BATCHED_GEMM_LIST
+    @eval function batched_mul!(C::StridedArray{T, 3}, A::$TA, B::$TB) where {T<:_GemmFloat}
+        batched_gemm!($transA, $transB, one(T), _unbatch(A), _unbatch(B), zero(T), C)
+        C
+    end
+end
 
+# fallback
 
+_BATCHED_LIST = [
+    (:(AbstractArray{<:Any, 3}), :identity),
+    (:(BatchedTranspose{<:Any, <:AbstractArray{<:Any, 3}}), :transpose),
+    (:(BatchedAdjoint{<:Any, <:AbstractArray{<:Any, 3}}), :adjoint)
+]
+for (TA, fA) in _BATCHED_LIST, (TB, fB) in _BATCHED_LIST
+    @eval function batched_mul!(C::AbstractArray{<:Any, 3}, A::$TA, B::$TB)
+        axes(A, 3) == axes(B, 3) == axes(C, 3) || throw(DimensionMismatch("batch size mismatch"))
+        @debug "calling fallback method for batched_mul!" typeof(A) typeof(B) typeof(C)
+        A′, B′ = _unbatch(A), _unbatch(B)
+        @inbounds for k in axes(C, 3)
+            @views mul!(C[:,:,k], $fA(A′[:,:,k]), $fB(B′[:,:,k]))
+        end
+        C
     end
 end

test/batchedmul.jl

@@ -22,26 +22,40 @@ function bmm_adjtest(a,b; adjA = false, adjB = false)
     cat(c...; dims = 3)
 end
 
-@testset "Batched Matrix Multiplication" begin
+
+@testset "batched_mul: Float64 * $TB" for TB in [Float64, Float32]
+
     A = randn(7,5,3)
-    B = randn(5,7,3)
+    B = randn(TB, 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)
+    @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)
+    cB = randn(Complex{TB}, 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_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
+
+    TBi = TB==Float64 ? Int64 : Int32
+    iA = rand(1:99, 7,5,3)
+    iB = TB.(rand(1:99, 5,7,3))
+    iC = zeros(Int, 7,6,3)
+    @test batched_mul(iA, iB) == bmm_adjtest(iA, iB)
+    @test batched_mul(cA, iB) ≈ bmm_adjtest(cA, iB)
+
+    @test_throws DimensionMismatch batched_mul(rand(2,2,2), rand(TB, 2,2,10))
+    @test_throws DimensionMismatch batched_mul(rand(2,2,2), rand(TB, 10,2,2))
+    @test_throws Exception batched_mul!(zeros(2,2,10), rand(2,2,2), rand(TB, 2,2,2))
 
-    @test batched_transpose(batched_transpose(A)) == A
-    @test batched_adjoint(batched_adjoint(cA)) == cA
 end