add fallback batched_mul

Author: Michael Abbott

src/batched/batchedmul.jl

@@ -2,37 +2,62 @@
 # 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"))
+    C = similar(A, (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,9 +22,11 @@ function bmm_adjtest(a,b; adjA = false, adjB = false)
     cat(c...; dims = 3)
 end
 
-@testset "Batched Matrix Multiplication" begin
+
+@testset "Batched Matrix Multiplication" 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)
@@ -34,7 +36,7 @@ end
 
 
     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)
@@ -44,4 +46,15 @@ end
 
     @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_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))
+
 end