Applied

src/LazyArrays.jl

@@ -28,7 +28,7 @@ import Base: ReinterpretArray, ReshapedArray, AbstractCartesianIndex, Slice,
 
 import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcasted,
                         combine_eltypes, DefaultArrayStyle, instantiate, materialize,
-                        materialize!
+                        materialize!, eltypes
 
 import LinearAlgebra: AbstractTriangular, checksquare, pinv
 
@@ -39,10 +39,12 @@ import FillArrays: AbstractFill
 import StaticArrays: StaticArrayStyle
 
 export Mul, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
-        Hcat, Vcat, Kron, BroadcastArray, cache, Ldiv, Inv, PInv, Diff, Cumsum
+        Hcat, Vcat, Kron, BroadcastArray, cache, Ldiv, Inv, PInv, Diff, Cumsum,
+        applied
 
 include("memorylayout.jl")
 include("cache.jl")
+include("lazyapplying.jl")
 include("lazybroadcasting.jl")
 include("lazyconcat.jl")
 include("linalg/linalg.jl")

src/lazyapplying.jl

@@ -0,0 +1,105 @@
+
+
+abstract type ApplyStyle end
+struct DefaultApplyStyle <: ApplyStyle end
+struct LayoutApplyStyle{Layouts<:Tuple} <: ApplyStyle
+    layouts::Layouts
+end
+
+# Used for when a lazy version should be constructed on materialize
+struct LazyArrayApplyStyle <: ApplyStyle end
+
+ApplyStyle(f, args...) = DefaultApplyStyle()
+
+struct Applied{Style<:ApplyStyle, F, Args<:Tuple}
+    style::Style
+    f::F
+    args::Args
+end
+
+applied(f, args...) = Applied(ApplyStyle(f, args...), f, args)
+_materialize(A::Applied, _) = A.f(materialize.(A.args)...)
+materialize(M::Applied{<:LayoutApplyStyle}) = _materialize(M, axes(M))
+materialize(A::Applied) = A.f(materialize.(A.args)...)
+
+
+
+similar(M::Applied) = similar(M, eltype(M))
+
+struct ApplyBroadcastStyle <: BroadcastStyle end
+
+@inline copyto!(dest::AbstractArray, bc::Broadcasted{ApplyBroadcastStyle}) =
+    _copyto!(MemoryLayout(dest), dest, bc)
+# Use default broacasting in general
+@inline _copyto!(_, dest, bc::Broadcasted) = copyto!(dest, Broadcasted{Nothing}(bc.f, bc.args, bc.axes))
+
+BroadcastStyle(::Type{<:Applied}) = ApplyBroadcastStyle()
+
+
+struct MatrixFunctionStyle{F} <: ApplyStyle end
+
+for f in (:exp, :sin, :cos, :sqrt)
+    @eval ApplyStyle(::typeof($f), ::AbstractMatrix) = MatrixFunctionStyle{typeof($f)}()
+end
+
+materialize(A::Applied{<:MatrixFunctionStyle,<:Any,<:Tuple{<:Any}}) =
+    A.f(materialize(first(A.args)))
+
+axes(A::Applied{<:MatrixFunctionStyle}) = axes(first(A.args))
+size(A::Applied{<:MatrixFunctionStyle}) = size(first(A.args))
+eltype(A::Applied{<:MatrixFunctionStyle}) = eltype(first(A.args))
+
+getindex(A::Applied{<:MatrixFunctionStyle}, k::Int, j::Int) =
+    materialize(A)[k,j]
+
+
+struct ApplyArray{T, N, App<:Applied} <: AbstractArray{T,N}
+    applied::App
+end
+
+const ApplyVector{T, App<:Applied} = ApplyArray{T, 1, App}
+const ApplyMatrix{T, App<:Applied} = ApplyArray{T, 2, App}
+
+
+ApplyArray{T,N}(M::App) where {T,N,App<:Applied} = ApplyArray{T,N,App}(M)
+ApplyArray{T}(M::Applied) where {T} = ApplyArray{T,ndims(M)}(M)
+ApplyArray(M::Applied) = ApplyArray{eltype(M)}(M)
+ApplyVector(M::Applied) = ApplyVector{eltype(M)}(M)
+ApplyMatrix(M::Applied) = ApplyMatrix{eltype(M)}(M)
+
+ApplyArray(f, factors...) = ApplyArray(applied(f, factors...))
+ApplyArray{T}(f, factors...) where T = ApplyArray{T}(applied(f, factors...))
+ApplyArray{T,N}(f, factors...) where {T,N} = ApplyArray{T,N}(applied(f, factors...))
+
+ApplyVector(f, factors...) = ApplyVector(applied(f, factors...))
+ApplyMatrix(f, factors...) = ApplyMatrix(applied(f, factors...))
+
+axes(A::ApplyArray) = axes(A.applied)
+size(A::ApplyArray) = map(length, axes(A))
+
+IndexStyle(::ApplyArray{<:Any,1}) = IndexLinear()
+
+@propagate_inbounds getindex(A::ApplyArray{T,N}, kj::Vararg{Int,N}) where {T,N} =
+    materialize(A.applied)[kj...]
+
+materialize(A::Applied{LazyArrayApplyStyle}) = ApplyArray(A)
+
+@inline copyto!(dest::AbstractArray, M::Applied) = _copyto!(MemoryLayout(dest), dest, M)
+@inline _copyto!(_, dest::AbstractArray, M::Applied) = copyto!(dest, materialize(M))
+
+broadcastable(M::Applied) = M
+
+# adjoint(A::MulArray) = MulArray(reverse(adjoint.(A.applied.args))...)
+# transpose(A::MulArray) = MulArray(reverse(transpose.(A.applied.args))...)
+
+
+struct  ApplyLayout{F, LAY} <: MemoryLayout
+    f::F
+    layouts::LAY
+end
+
+MemoryLayout(M::ApplyArray) = ApplyLayout(M.applied.f, MemoryLayout.(M.applied.args))
+
+
+# _flatten(A::ApplyArray, B...) = _flatten(A.applied.args..., B...)
+# flatten(A::MulArray) = MulArray(Mul(_flatten(A.applied.args...)))

src/lazybroadcasting.jl

@@ -94,14 +94,3 @@ broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::AbstractArray{T,N}, b::Zeros{V,
     broadcast(DefaultArrayStyle{N}(), *, a, b)
 broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::AbstractArray{V,N}) where {T,V,N} =
     broadcast(DefaultArrayStyle{N}(), *, a, b)
-
-const Add = BroadcastArray{<:Any, <:Any, <:Broadcasted{<:Any, <:Any, typeof(+)}}
-const AddVector = Add{<:Any, 1}
-const AddMatrix = Add{<:Any, 2}
-
-"""
-    Add(A1, A2, …, AN)
-
-A lazy representation of `A1 .+ A2 .+ … .+ AN`; i.e., a shorthand for `BroadcastArray(+, A1, A2, …, AN)`.
-"""
-Add(As...) = BroadcastArray(+, As...)

src/linalg/add.jl

@@ -0,0 +1,53 @@
+####
+# These are special routines to make operations involving +
+# more efficient
+####
+
+
+const Add{Factors<:Tuple} = Applied{<:Any, typeof(+), Factors}
+
+size(M::Add, p::Int) = size(M)[p]
+axes(M::Add, p::Int) = axes(M)[p]
+ndims(M::Add) = ndims(first(M.args))
+
+length(M::Add) = prod(size(M))
+size(M::Add) = length.(axes(M))
+axes(M::Add) = axes(first(M.args))
+
+
+eltype(M::Add) = Base._return_type(+, eltype.(M.args))
+
+const AddArray{T,N,Factors<:Tuple} = ApplyArray{T,N,<:Add{Factors}}
+const AddVector{T,Factors<:Tuple} = AddArray{T,1,Factors}
+const AddMatrix{T,Factors<:Tuple} = AddArray{T,2,Factors}
+
+AddArray(factors...) = ApplyArray(+, factors...)
+
+"""
+    Add(A1, A2, …, AN)
+
+A lazy representation of `A1 + A2 + … + AN`; i.e., a shorthand for `applied(+, A1, A2, …, AN)`.
+"""
+Add(As...) = applied(+, As...)
+
+
+getindex(M::Add, k::Integer) = sum(getindex.(M.args, k))
+getindex(M::Add, k::Integer, j::Integer) = sum(getindex.(M.args, k, j))
+getindex(M::Add, k::CartesianIndex{1}) = M[convert(Int, k)]
+getindex(M::Add, kj::CartesianIndex{2}) = M[kj[1], kj[2]]
+
+for MulAdd_ in [MatMulMatAdd, MatMulVecAdd]
+    # `MulAdd{<:ApplyLayout{typeof(+)}}` cannot "win" against
+    # `MatMulMatAdd` and `MatMulVecAdd` hence `@eval`:
+    @eval function materialize!(M::$MulAdd_{<:ApplyLayout{typeof(+)}})
+        α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
+        if C ≡ B
+            B = copy(B)
+        end
+        lmul!(β, C)
+        for A in A.applied.args
+            C .= α .* Mul(A, B) .+ C
+        end
+        C
+    end
+end

src/linalg/blasbroadcasting.jl

@@ -12,7 +12,6 @@ struct ArrayMulArrayStyle{StyleA, StyleB, p, q} <: BroadcastStyle end
 @inline copyto!(dest::AbstractArray, bc::Broadcasted{<:ArrayMulArrayStyle}) =
     _copyto!(MemoryLayout(dest), dest, bc)
 # Use default broacasting in general
-@inline _copyto!(_, dest, bc::Broadcasted) = copyto!(dest, Broadcasted{Nothing}(bc.f, bc.args, bc.axes))
 
 const BArrayMulArray{styleA, styleB, p, q, T, V} =
     Broadcasted{ArrayMulArrayStyle{styleA,styleB,p,q}, <:Any, typeof(identity),
@@ -55,13 +54,11 @@ similar(M::Broadcasted{<:ArrayMulArrayStyle}, ::Type{ElType}) where ElType =
     Array{ElType}(undef,size(M.args[1]))
 
 
-
-broadcastable(M::ArrayMulArray) = M
 instantiate(bc::Broadcasted{<:ArrayMulArrayStyle}) = bc
 
 
 @inline function _copyto!(_, dest::AbstractArray{T}, M::ArrayMulArray) where T
-    A,B = M.factors
+    A,B = M.args
     materialize!(MulAdd(one(T), A, B, zero(T), dest))
 end
 
@@ -73,34 +70,34 @@ end
 
 @inline function _copyto!(_, dest::AbstractArray{T}, bc::BConstArrayMulArray) where T
     α,M = bc.args
-    A,B = M.factors
+    A,B = M.args
     materialize!(MulAdd(α, A, B, zero(T), dest))
 end
 
 @inline function _copyto!(_, dest::AbstractArray{T}, bc::BArrayMulArrayPlusArray) where T
     M,C = bc.args
-    A,B = M.factors
+    A,B = M.args
     copyto!(dest, MulAdd(one(T), A, B, one(T), C))
 end
 
 @inline function _copyto!(_, dest::AbstractArray{T}, bc::BArrayMulArrayPlusConstArray) where T
     M,βc = bc.args
     β,C = βc.args
-    A,B = M.factors
+    A,B = M.args
     copyto!(dest, MulAdd(one(T), A, B, β, C))
 end
 
 @inline function _copyto!(_, dest::AbstractArray{T}, bc::BConstArrayMulArrayPlusArray) where T
     αM,C = bc.args
     α,M = αM.args
-    A,B = M.factors
+    A,B = M.args
     copyto!(dest, MulAdd(α, A, B, one(T), C))
 end
 
 @inline function _copyto!(_, dest::AbstractArray, bc::BConstArrayMulArrayPlusConstArray)
     αM,βc = bc.args
     α,M = αM.args
-    A,B = M.factors
+    A,B = M.args
     β,C = βc.args
     copyto!(dest, MulAdd(α, A, B, β, C))
 end

src/linalg/blasmul.jl

@@ -42,8 +42,7 @@ BroadcastStyle(::Type{<:MatMulVecAdd{StyleA,StyleB,StyleC}}) where {StyleA,Style
 BroadcastStyle(::Type{<:MatMulMatAdd{StyleA,StyleB,StyleC}}) where {StyleA,StyleB,StyleC} =
     ArrayMulArrayStyle{StyleA,StyleB,2,2}()
 
-broadcastable(M::MatMulMatAdd) = M
-broadcastable(M::MatMulVecAdd) = M
+broadcastable(M::MulAdd) = M
 
 const BlasMatMulVec{StyleA,StyleB,StyleC,T} = MulAdd{StyleA,StyleB,StyleC,T,<:AbstractMatrix{T},<:AbstractVector{T},<:AbstractVector{T}}
 const BlasMatMulMat{StyleA,StyleB,StyleC,T} = MulAdd{StyleA,StyleB,StyleC,T,<:AbstractMatrix{T},<:AbstractMatrix{T},<:AbstractMatrix{T}}
@@ -184,6 +183,9 @@ end
 
 function materialize!(M::MatMulMatAdd)
     α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
+    if C ≡ B
+        B = copy(B)
+    end
     ts = tile_size(eltype(A), eltype(B), eltype(C))
     if iszero(β) # false is a "strong" zero to wipe out NaNs
         ts == 0 ? default_blasmul!(α, A, B, false, C) : tiled_blasmul!(ts, α, A, B, false, C)
@@ -194,23 +196,10 @@ end
 
 function materialize!(M::MatMulVecAdd)
     α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
-    default_blasmul!(α, A, B, iszero(β) ? false : β, C)
-end
-
-for MulAdd_ in [MatMulMatAdd, MatMulVecAdd]
-    # `MulAdd{<:BroadcastLayout{typeof(+)}}` cannot "win" against
-    # `MatMulMatAdd` and `MatMulVecAdd` hence `@eval`:
-    @eval function materialize!(M::$MulAdd_{<:BroadcastLayout{typeof(+)}})
-        α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
-        if C ≡ B
-            B = copy(B)
-        end
-        lmul!(β, C)
-        for A in A.broadcasted.args
-            C .= α .* Mul(A, B) .+ C
-        end
-        C
+    if C ≡ B
+        B = copy(B)
     end
+    default_blasmul!(α, A, B, iszero(β) ? false : β, C)
 end
 
 # make copy to make sure always works
@@ -329,15 +318,15 @@ materialize!(M::BlasMatMulVec{<:HermitianLayout{<:AbstractRowMajor},<:AbstractSt
 @inline function _copyto!(::AbstractStridedLayout, dest::AbstractVector{T},
          M::MatMulVec{<:TriangularLayout{UPLO,UNIT,<:AbstractColumnMajor},
                                    <:AbstractStridedLayout, T, T}) where {UPLO,UNIT,T <: BlasFloat}
-    A,x = M.factors
+    A,x = M.args
     x ≡ dest || copyto!(dest, x)
     BLAS.trmv!(UPLO, 'N', UNIT, triangulardata(A), dest)
 end
 
 @inline function _copyto!(::AbstractStridedLayout, dest::AbstractVector{T},
          M::MatMulVec{<:TriangularLayout{UPLO,UNIT,<:AbstractRowMajor},
                                    <:AbstractStridedLayout, T, T}) where {UPLO,UNIT,T <: BlasFloat}
-    A,x = M.factors
+    A,x = M.args
     x ≡ dest || copyto!(dest, x)
     BLAS.trmv!(UPLO, 'T', UNIT, transpose(triangulardata(A)), dest)
 end
@@ -346,7 +335,7 @@ end
 @inline function _copyto!(::AbstractStridedLayout, dest::AbstractVector{T},
          M::MatMulVec{<:TriangularLayout{UPLO,UNIT,<:ConjLayout{<:AbstractRowMajor}},
                                    <:AbstractStridedLayout, T, T}) where {UPLO,UNIT,T <: BlasFloat}
-    A,x = M.factors
+    A,x = M.args
     x ≡ dest || copyto!(dest, x)
     BLAS.trmv!(UPLO, 'C', UNIT, triangulardata(A)', dest)
 end
@@ -356,7 +345,7 @@ end
 
 
 function _copyto!(_, dest::AbstractMatrix, M::MatMulMat{<:TriangularLayout})
-    A,X = M.factors
+    A,X = M.args
     size(dest,2) == size(X,2) || thow(DimensionMismatch("Dimensions must match"))
     @views for j in axes(dest,2)
         dest[:,j] .= Mul(A, X[:,j])