From 6467f919c9ed29ef94b49d5f0f3fdd6e239a6958 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Sat, 23 May 2026 06:04:58 -0500 Subject: [PATCH] Bundle FM + Workspace as CachedFactoredMatrix A small wrapper type that pairs a `FactoredMatrix` with a pre-allocated `Workspace` so they travel together through APIs that take 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. Constraint documented in the docstring: the bundled `Workspace` has a single `p`, so the same `cfm` is usable for operations that all share the same outer dimension. Mixed outer dimensions need separate bundles. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 2 +- src/FactoredMatrices.jl | 85 ++++++++++++++++++++++++++++++++++++++++- test/runtests.jl | 49 ++++++++++++++++++++++++ 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 021aa0c..6decf3d 100644 --- a/Project.toml +++ b/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 "] [deps] diff --git a/src/FactoredMatrices.jl b/src/FactoredMatrices.jl index 7c955bd..e1c53f8 100644 --- a/src/FactoredMatrices.jl +++ b/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) diff --git a/test/runtests.jl b/test/runtests.jl index 43d6217..6eec1d2 100644 --- a/test/runtests.jl +++ b/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