Bundle FM + Workspace as CachedFactoredMatrix

Project.toml

@@ -1,6 +1,6 @@
 name = "FactoredMatrices"
 uuid = "13d07ad6-c0d4-11e8-35b1-91a532e61cc8"
-version = "1.0.0"
+version = "1.1.0"
 authors = ["Tim Holy <tim.holy@gmail.com>"]
 
 [deps]

src/FactoredMatrices.jl

@@ -12,7 +12,7 @@ import Base: *, Array, Matrix, size, length, similar, ==, hash, adjoint, transpo
 import LinearAlgebra: mul!, dot
 
 export FactoredMatrix
-VERSION >= v"1.11" && eval(Meta.parse("public Workspace"))
+VERSION >= v"1.11" && eval(Meta.parse("public Workspace, CachedFactoredMatrix"))
 
 """
     FactoredMatrix(U::AbstractMatrix{T}, V::AbstractMatrix{T}) where T
@@ -21,6 +21,9 @@ Store the matrix `M = U*V` in factored form for efficient computation.
 Multiplications exploit the factored structure without forming the full dense
 matrix.
 
+`FactoredMatrix <: Factorization`, meaning it does not support indexing or iteration.
+The only supported interface is multiplication.
+
 For allocation-free repeated multiplications, create a
 [`FactoredMatrices.Workspace`](@ref) and pass it as `mul!(C, A, B; cache=ws)`.
 """
@@ -71,6 +74,39 @@ Workspace{T}(k::Int, p::Int) where {T} = Workspace{T}(Matrix{T}(undef, k, p), Ma
 
 Workspace(A::FactoredMatrix{T}, p::Int) where {T} = Workspace{T}(size(A.U, 2), p)
 
+"""
+    CachedFactoredMatrix(M::FactoredMatrix, ws::Workspace)
+
+Bundle a [`FactoredMatrix`](@ref) with a pre-allocated [`Workspace`](@ref) so
+they travel together through an API that takes a single matrix-like argument.
+Multiplications against a `CachedFactoredMatrix` automatically use the bundled
+`Workspace` as their cache, sparing callers from threading a separate cache
+argument through every call site.
+
+`CachedFactoredMatrix <: Factorization`, meaning it does not support indexing or
+iteration. The only supported interface is multiplication.
+
+The bundled `Workspace` carries a single `p`, which must equal both `size(B, 2)`
+(for `cfm * B`) and `size(A, 1)` (for `A * cfm`).  Use a single `cfm` only when
+all such operations share the same outer dimension; otherwise build separate
+`cfm`s with their own `Workspace`s.
+
+Reuse one `Workspace` per task when running concurrently — see the
+[`Workspace`](@ref) docs.
+
+# Example
+```julia
+ws  = Workspace(M, max(size(B,2), size(A,1)))
+cfm = CachedFactoredMatrix(M, ws)
+cfm * B    # uses ws as cache
+A * cfm    # uses ws as cache
+```
+"""
+struct CachedFactoredMatrix{T, FT <: FactoredMatrix{T}, WT <: Workspace{T}} <: Factorization{T}
+    M::FT
+    ws::WT
+end
+
 size(A::FactoredMatrix) = size(A.U, 1), size(A.V, 2)
 size(A::FactoredMatrix, d::Integer) = d == 1 ? size(A.U, 1) : (d == 2 ? size(A.V, 2) : 1)
 similar(A::FactoredMatrix, dims) = Array{eltype(A.U)}(undef, dims)
@@ -264,6 +300,43 @@ end
 *(A::FactoredMatrix, B) = A.U * (A.V * B)
 *(A, B::FactoredMatrix) = (A * B.U) * B.V
 
+# --- CachedFactoredMatrix ---
+
+size(C::CachedFactoredMatrix) = size(C.M)
+size(C::CachedFactoredMatrix, d::Integer) = size(C.M, d)
+length(C::CachedFactoredMatrix) = length(C.M)
+
+# Delegate mul! to the underlying FactoredMatrix dispatch, threading the
+# bundled Workspace through as `cache`.
+mul!(C::AbstractMatrix, A::CachedFactoredMatrix, B::AbstractMatrix) = mul!(C, A.M, B; cache = A.ws)
+mul!(C::AbstractVecOrMat, A::CachedFactoredMatrix, b::AbstractVector) = mul!(C, A.M, b; cache = A.ws)
+mul!(C::AbstractMatrix, A::AbstractMatrix, B::CachedFactoredMatrix) = mul!(C, A, B.M; cache = B.ws)
+
+mul!(C::AbstractMatrix, A::CachedFactoredMatrix, B::AbstractMatrix, α::Number, β::Number) =
+    mul!(C, A.M, B, α, β; cache = A.ws)
+mul!(C::AbstractVector, A::CachedFactoredMatrix, b::AbstractVector, α::Number, β::Number) =
+    mul!(C, A.M, b, α, β; cache = A.ws)
+mul!(C::AbstractMatrix, A::AbstractMatrix, B::CachedFactoredMatrix, α::Number, β::Number) =
+    mul!(C, A, B.M, α, β; cache = B.ws)
+
+# Allocating `*` forms — delegate to the FactoredMatrix dispatch directly so the
+# bundled Workspace propagates to the inner `mul!`s that produce small
+# intermediates.  The intermediate matrices `(A * B.M.U)` etc. are themselves
+# fresh allocations; the Workspace's job is to eliminate the small `templeft` /
+# `tempright` allocations inside those `mul!`s.
+function *(A::CachedFactoredMatrix{T}, B::AbstractMatrix) where {T}
+    out = Matrix{T}(undef, size(A, 1), size(B, 2))
+    return mul!(out, A.M, B; cache = A.ws)
+end
+function *(A::CachedFactoredMatrix{T}, b::AbstractVector) where {T}
+    out = Vector{T}(undef, size(A, 1))
+    return mul!(out, A.M, b; cache = A.ws)
+end
+function *(A::AbstractMatrix, B::CachedFactoredMatrix{T}) where {T}
+    out = Matrix{T}(undef, size(A, 1), size(B, 2))
+    return mul!(out, A, B.M; cache = B.ws)
+end
+
 # --- misc ---
 
 function dot(A::FactoredMatrix, B::FactoredMatrix)
@@ -272,6 +345,16 @@ function dot(A::FactoredMatrix, B::FactoredMatrix)
     return sum(M1 .* conj(M2))
 end
 
+# Closed-form `||UV||²_F = tr((U'U) (VV'))` so callers can write `sum(abs2, A)`
+# without materializing the dense product.  Subtyping `Factorization` (rather
+# than `AbstractMatrix`) means the default `sum(f, iter)` would otherwise
+# MethodError on iteration — this override gives a fast path.  The `conj` on
+# the right factor is needed for complex `T`; for real `T` it's a no-op.  The
+# outer `real` discards the negligible imaginary noise that floating-point
+# roundoff in the elementwise product can leave behind.
+Base.sum(::typeof(abs2), A::FactoredMatrix) = real(sum((A.U' * A.U) .* conj.(A.V * A.V')))
+Base.sum(::typeof(abs2), C::CachedFactoredMatrix) = sum(abs2, C.M)
+
 """
     ssd(A::FactoredMatrix, B::FactoredMatrix)
 

test/runtests.jl

@@ -283,4 +283,53 @@ end
             R = copy(c5); mul!(R, Mf, b5, α, β; cache = ws5); @test R ≈ α * M * b5 + β * c5
         end
     end
+
+    @testset "CachedFactoredMatrix" begin
+        for T in (Float32, Float64, ComplexF64)
+            U  = myrand(T, 10, 3)
+            V  = myrand(T, 3, 15)
+            Mf = FactoredMatrix(U, V)
+            M  = Array(Mf)
+            # A single bundled `Workspace` carries a single `p`, so the same
+            # `p` must equal both `size(B, 2)` (for `Mf * B`) and
+            # `size(A, 1)` (for `A * Mf`).  Pick `p = 5` and shape both that
+            # way.
+            B  = myrand(T, 15, 5)   # right multiplicand (used as `Mf * B`)
+            A  = myrand(T, 5, 10)   # left multiplicand  (used as `A * Mf`)
+            b  = vec(myrand(T, 15, 1))  # right vector
+
+            ws  = FactoredMatrices.Workspace(Mf, 5)
+            cfm = FactoredMatrices.CachedFactoredMatrix(Mf, ws)
+
+            @test cfm isa LinearAlgebra.Factorization{T}
+            @test !(cfm isa AbstractMatrix)    # fail-fast on getindex/iterate
+            @test size(cfm) == size(M)
+            @test size(cfm, 1) == 10 && size(cfm, 2) == 15
+            @test length(cfm) == length(M)
+            @test eltype(cfm) === T
+
+            # `*` returns the same result as the bare FactoredMatrix.
+            @test cfm * B ≈ M * B
+            @test cfm * b ≈ M * b
+            @test A * cfm ≈ A * M
+
+            # `mul!` 3-arg and 5-arg.
+            R = myrand(T, 10, 5);      mul!(R, cfm, B); @test R ≈ M * B
+            R = myrand(T, 5, 15);      mul!(R, A, cfm); @test R ≈ A * M
+            v = vec(myrand(T, 10, 1)); mul!(v, cfm, b); @test v ≈ M * b
+            α, β = T(2), T(3)
+            R0 = myrand(T, 10, 5);      R = copy(R0); mul!(R, cfm, B, α, β); @test R ≈ α * (M * B) + β * R0
+            R0 = myrand(T, 5, 15);      R = copy(R0); mul!(R, A, cfm, α, β); @test R ≈ α * (A * M) + β * R0
+            v0 = vec(myrand(T, 10, 1)); v = copy(v0); mul!(v, cfm, b, α, β); @test v ≈ α * (M * b) + β * v0
+
+            # Element access falls through to a MethodError (deliberate — the
+            # type promises only matrix-product semantics, not indexing).
+            @test_throws MethodError cfm[1, 1]
+
+            # `sum(abs2, ·)` works on both the bare `FactoredMatrix` and the
+            # cached bundle via the closed-form override.
+            @test sum(abs2, Mf)  ≈ sum(abs2, M)
+            @test sum(abs2, cfm) ≈ sum(abs2, M)
+        end
+    end
 end