fix *, optimize vector case

Author: alyst

src/dual.jl

@@ -776,53 +776,67 @@ end
 # Efficient left multiplication/division of #
 # Dual array by a constant matrix           #
 #-------------------------------------------#
-
-# creates the copy of x and applies fvalue!(values) to its values, and fpartial!(partial, ix) to its partials
-function _map_dual_components(fvalue!, fpartial!, x::AbstractArray{DT}) where DT <: Dual{<:Any, T} where T
+# creates the copy of x and applies fvalue!(values(y), values(x)) to its values,
+# and fpartial!(partial(y, i), partial(y, i), i) to its partials
+function _map_dual_components!(fvalue!, fpartial!, y::AbstractArray{DT}, x::AbstractArray{DT}) where DT <: Dual{<:Any, T} where T
     N = npartials(DT)
-    res = similar(x) # result
-    t = similar(x, T) # temporary Array{T} for fvalue!/fpartial! application
+    tx = similar(x, T) 
+    ty = similar(y, T) # temporary Array{T} for fvalue!/fpartial! application
     # y allows res to be accessed as Array{T}
-    y = reinterpret(reshape, T, res)
-    @assert size(y) == (N + 1, length(res))
+    yarr = reinterpret(reshape, T, y)
+    @assert size(yarr) == (N + 1, size(y)...)
     ystride = size(y, 1)
 
     # calculate res values
     @inbounds for (j, v) in enumerate(x)
-        t[j] = value(v)
+        tx[j] = value(v)
     end
-    fvalue!(t)
+    fvalue!(ty, tx)
     k = 1
-    @inbounds for tt in t
-        y[k] = tt
+    @inbounds for tt in ty
+        yarr[k] = tt
         k += ystride
     end
 
     # calculate each res partial
     for i in 1:N
         @inbounds for (j, v) in enumerate(x)
-            t[j] = partials(v, i)
+            tx[j] = partials(v, i)
         end
-        fpartial!(t, i)
+        fpartial!(ty, tx, i)
         k = i + 1
-        @inbounds for tt in t
-            y[k] = tt
+        @inbounds for tt in ty
+            yarr[k] = tt
             k += ystride
         end
     end
 
-    return res
+    return y
 end
 
 for MT in (StridedMatrix{<:LinearAlgebra.BlasFloat},
            LowerTriangular{<:LinearAlgebra.BlasFloat},
-           UpperTriangular{<:LinearAlgebra.BlasFloat}),
-    XT in (StridedMatrix{<:Dual}, StridedVector{<:Dual})
-    @eval Base.:\(m::$MT, x::$XT) =
-        _map_dual_components(Base.Fix1(ldiv!, m), (x, _) -> ldiv!(m, x), x)
+           UpperTriangular{<:LinearAlgebra.BlasFloat})
+
+    @eval function Base.:\(m::$MT, x::StridedVector{<:Dual})
+        T = valtype(eltype(x))
+        ldiv!(m', reinterpret(reshape, T, res))
+        return res
+    end
+
+    @eval Base.:\(m::$MT, x::StridedMatrix{<:Dual}) =
+        _map_dual_components!((x, _) -> ldiv!(m, x), (x, _, _) -> ldiv!(m, x), similar(x), x)
+
+    @eval function Base.:*(m::$MT, x::StridedVector{<:Dual})
+        T = valtype(eltype(x))
+        res = similar(x, (size(m, 1),))
+        mul!(reinterpret(reshape, T, res), reinterpret(reshape, T, x), m')
+        return res
+    end
 
-    @eval Base.:*(m::$MT, x::$XT) =
-        _map_dual_components(Base.Fix1(lmul!, m), (x, _) -> lmul!(m, x), x)
+    @eval Base.:*(m::$MT, x::StridedMatrix{<:Dual}) =
+        _map_dual_components!((y, x) -> mul!(y, m, x), (y, x, _) -> mul!(y, m, x),
+                              similar(x, (size(m, 1), size(x, 2))), x)
 end
 
 ###################