Only define BLAS optimizations for Float32 and Float64, not their union

Author: devmotion

src/weights.jl

@@ -440,16 +440,16 @@ end
 #     (a) A is a vector: we invoke the vector version wsum above.
 #         The internal function that implements this is _wsum1!
 #
-#     (b) A is a dense matrix with eltype <: BlasReal: we call gemv!
+#     (b) A is a dense matrix with `Float32` or `Float64` `eltype`: we call gemv!
 #         The internal function that implements this is _wsum2_blas!
 #
-#     (c) A is a contiguous array with eltype <: BlasReal:
+#     (c) A is a contiguous array with `Float32` or `Float64` `eltype`:
 #         dim == 1: treat A like a matrix of size (d1, d2 x ... x dN)
 #         dim == N: treat A like a matrix of size (d1 x ... x d(N-1), dN)
 #         otherwise: decompose A into multiple pages, and apply _wsum2!
 #         for each
 #
-#     (d) A is a general dense array with eltype <: BlasReal:
+#     (d) A is a general dense array with `Float32` or `Float64` `eltype`:
 #         dim <= 2: delegate to (a) and (b)
 #         otherwise, decompose A into multiple pages
 
@@ -463,56 +463,6 @@ function _wsum1!(R::AbstractArray, A::AbstractVector, w::AbstractVector, init::B
     return R
 end
 
-function _wsum2_blas!(R::StridedVector{T}, A::StridedMatrix{T}, w::StridedVector{T}, dim::Int, init::Bool) where T<:BlasReal
-    trans = dim == 1 ? 'T' : 'N'
-    BLAS.gemv!(trans, true, A, w, !init, R)
-    return R
-end
-
-function _wsumN!(R::StridedArray{T}, A::StridedArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N}
-    if dim == 1
-        m = size(A, 1)
-        n = div(length(A), m)
-        _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 1, init)
-    elseif dim == N
-        n = size(A, N)
-        m = div(length(A), n)
-        _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 2, init)
-    else # 1 < dim < N
-        m = 1
-        for i = 1:dim-1; m *= size(A, i); end
-        n = size(A, dim)
-        k = 1
-        for i = dim+1:N; k *= size(A, i); end
-        Av = reshape(A, (m, n, k))
-        Rv = reshape(R, (m, k))
-        for i = 1:k
-            _wsum2_blas!(view(Rv,:,i), view(Av,:,:,i), w, 2, init)
-        end
-    end
-    return R
-end
-
-function _wsumN!(R::StridedArray{T}, A::DenseArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N}
-    @assert N >= 3
-    if dim <= 2
-        m = size(A, 1)
-        n = size(A, 2)
-        npages = 1
-        for i = 3:N
-            npages *= size(A, i)
-        end
-        rlen = ifelse(dim == 1, n, m)
-        Rv = reshape(R, (rlen, npages))
-        for i = 1:npages
-            _wsum2_blas!(view(Rv,:,i), view(A,:,:,i), w, dim, init)
-        end
-    else
-        _wsum_general!(R, identity, A, w, dim, init)
-    end
-    return R
-end
-
 ## general Cartesian-based weighted sum across dimensions
 
 @generated function _wsum_general!(R::AbstractArray{RT}, f::supertype(typeof(abs)),
@@ -571,25 +521,10 @@ end
     end
 end
 
-# N = 1
-_wsum!(R::StridedArray{T}, A::DenseArray{T,1}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal} =
-    _wsum1!(R, A, w, init)
-
-# N = 2
-_wsum!(R::StridedArray{T}, A::DenseArray{T,2}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal} =
-    (_wsum2_blas!(view(R,:), A, w, dim, init); R)
-
-# N >= 3
-_wsum!(R::StridedArray{T}, A::DenseArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N} =
-    _wsumN!(R, A, w, dim, init)
-
 _wsum!(R::AbstractArray, A::AbstractArray, w::AbstractVector, dim::Int, init::Bool) =
     _wsum_general!(R, identity, A, w, dim, init)
 
-## wsum! and wsum
-
 wsumtype(::Type{T}, ::Type{W}) where {T,W} = typeof(zero(T) * zero(W) + zero(T) * zero(W))
-wsumtype(::Type{T}, ::Type{T}) where {T<:BlasReal} = T
 
 """
     wsum!(R::AbstractArray, A::AbstractArray,
@@ -662,7 +597,6 @@ _mean!(R::AbstractArray, A::AbstractArray, w::AbstractWeights, dims::Int) =
     rmul!(Base.sum!(R, A, w, dims), inv(sum(w)))
 
 wmeantype(::Type{T}, ::Type{W}) where {T,W} = typeof((zero(T)*zero(W) + zero(T)*zero(W)) / one(W))
-wmeantype(::Type{T}, ::Type{T}) where {T<:BlasReal} = T
 
 """
     mean(A::AbstractArray, w::AbstractWeights[, dims::Int])
@@ -692,6 +626,76 @@ function mean(A::AbstractArray, w::UnitWeights; dims::Union{Colon,Int}=:)
     return mean(A, dims=dims)
 end
 
+##### BLAS optimizations for Float32 and Float64 #####
+
+for T in (:Float32, :Float64)
+    @eval begin
+        function _wsum!(R::StridedArray{$T}, A::DenseArray{$T,N}, w::StridedVector{$T}, dim::Int, init::Bool) where {N}
+            if N === 1
+                _wsum1!(R, A, w, init)
+            elseif N === 2
+                _wsum2_blas!(view(R,:), A, w, dim, init)
+            else # N >= 3
+                _wsumN!(R, A, w, dim, init)
+            end
+            return R
+        end
+
+        wsumtype(::Type{$T}, ::Type{$T}) = $T
+        wmeantype(::Type{$T}, ::Type{$T}) = $T
+
+        function _wsum2_blas!(R::StridedVector{$T}, A::StridedMatrix{$T}, w::StridedVector{$T}, dim::Int, init::Bool)
+            trans = dim == 1 ? 'T' : 'N'
+            BLAS.gemv!(trans, true, A, w, !init, R)
+            return R
+        end
+
+        function _wsumN!(R::StridedArray{$T}, A::StridedArray{$T,N}, w::StridedVector{$T}, dim::Int, init::Bool) where {N}
+            if dim == 1
+                m = size(A, 1)
+                n = div(length(A), m)
+                _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 1, init)
+            elseif dim == N
+                n = size(A, N)
+                m = div(length(A), n)
+                _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 2, init)
+            else # 1 < dim < N
+                m = 1
+                for i = 1:dim-1; m *= size(A, i); end
+                n = size(A, dim)
+                k = 1
+                for i = dim+1:N; k *= size(A, i); end
+                Av = reshape(A, (m, n, k))
+                Rv = reshape(R, (m, k))
+                for i = 1:k
+                    _wsum2_blas!(view(Rv,:,i), view(Av,:,:,i), w, 2, init)
+                end
+            end
+            return R
+        end
+
+        function _wsumN!(R::StridedArray{$T}, A::DenseArray{$T,N}, w::StridedVector{$T}, dim::Int, init::Bool) where {N}
+            @assert N >= 3
+            if dim <= 2
+                m = size(A, 1)
+                n = size(A, 2)
+                npages = 1
+                for i = 3:N
+                    npages *= size(A, i)
+                end
+                rlen = ifelse(dim == 1, n, m)
+                Rv = reshape(R, (rlen, npages))
+                for i = 1:npages
+                    _wsum2_blas!(view(Rv,:,i), view(A,:,:,i), w, dim, init)
+                end
+            else
+                _wsum_general!(R, identity, A, w, dim, init)
+            end
+            return R
+        end
+    end
+end
+
 ##### Weighted quantile #####
 
 """