Skip to content

Commit e53284b

Browse files
ViralBShahclaude
andauthored
mul! into dense outputs runs in GraphBLAS; fix Matrix(A') (#166)
* Run mul! into dense outputs in GraphBLAS (#125) mul!(C, A::GBArray, B) with dense B and C fell back to LinearAlgebra's generic single-threaded matmul, so Arpack/KrylovKit never used GraphBLAS. The 5-arg mul! now multiplies in GraphBLAS and writes alpha*A*B + beta*C into C. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fix Matrix/Vector of a transposed GBArray Matrix(A') threw a MethodError on wait(::Transpose); materialize first. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent ced917d commit e53284b

4 files changed

Lines changed: 58 additions & 2 deletions

File tree

‎src/abstractgbarray.jl‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ function Base.copy(A::M) where {M<:AbstractGBArray}
177177
M(_copyGrBMat(A.p), A.fill)
178178
end
179179

180-
function Base.Matrix(A::GBArrayOrTranspose)
180+
function Base.Matrix(A::AbstractGBArray)
181181
format = sparsitystatus(A)
182182
if format === Dense()
183183
T = unsafeunpack!(A, Dense())
@@ -195,7 +195,7 @@ function Base.Matrix(A::GBArrayOrTranspose)
195195
return M
196196
end
197197

198-
function Base.Vector(A::GBVectorOrTranspose)
198+
function Base.Vector(A::AbstractGBVector)
199199
format = sparsitystatus(A)
200200
if format === Dense()
201201
T = unsafeunpack!(A, Dense())
@@ -213,6 +213,10 @@ function Base.Vector(A::GBVectorOrTranspose)
213213
return M
214214
end
215215

216+
# The unpacking above needs an AbstractGBArray, so materialize the transpose first.
217+
Base.Matrix(A::Transpose{<:Any, <:AbstractGBArray}) = Matrix(copy(A))
218+
Base.Vector(A::Transpose{<:Any, <:AbstractGBVector}) = Vector(parent(A))
219+
216220
function SparseArrays.SparseMatrixCSC(A::GBArrayOrTranspose)
217221
T = copy(A) # avoid changing sparsity of A and destroying it.
218222
return unsafeunpack!(T, SparseMatrixCSC; attachfinalizer = true)

‎src/operations/mul.jl‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,44 @@ function Base.:*(
165165
return @_densepack B (*(A, B, op; mask, accum, desc))
166166
end
167167

168+
# Dense outputs: C = α * A * B + β * C, so that mul! callers (Arpack, KrylovKit, ...) run in
169+
# GraphBLAS rather than the generic LinearAlgebra fallback. See #125.
170+
# pack needs a DenseArray: alias contiguous views (e.g. Arpack's views into workd), copy others.
171+
_asdense(B::DenseVecOrMat) = B
172+
_asdense(B::Transpose) = transpose(_asdense(parent(B)))
173+
function _asdense(B::StridedVecOrMat)
174+
contiguous = stride(B, 1) == 1 && (B isa AbstractVector || stride(B, 2) == size(B, 1))
175+
return contiguous ? unsafe_wrap(Array, pointer(B), size(B)) : copy(B)
176+
end
177+
178+
function _densemul!(C, A, B, α, β)
179+
size(A, 2) == size(B, 1) || throw(DimensionMismatch("size(A, 2) != size(B, 1)"))
180+
size(A, 1) == size(C, 1) || throw(DimensionMismatch("size(A, 1) != size(C, 1)"))
181+
size(B, 2) == size(C, 2) || throw(DimensionMismatch("size(B, 2) != size(C, 2)"))
182+
T = GC.@preserve B A * pack(_asdense(B))
183+
R = unsafeunpack!(T, Dense())
184+
try
185+
R′ = reshape(R, size(C))
186+
iszero(β) ? (C .= α .* R′) : (C .= α .* R′ .+ β .* C)
187+
finally
188+
_jlfree(R)
189+
end
190+
return C
191+
end
192+
193+
const _GBVecMatOrTrans = Union{GBMatrixOrTranspose, GBVectorOrTranspose}
194+
195+
for (TC, TB) in (
196+
(:StridedVector, :(StridedVector{<:valid_union})),
197+
(:StridedMatrix, :(Union{StridedVecOrMat{<:valid_union}, Transpose{<:Any, <:StridedVecOrMat{<:valid_union}}})),
198+
)
199+
@eval begin
200+
LinearAlgebra.mul!(C::$TC, A::_GBVecMatOrTrans, B::$TB, α::Number, β::Number) =
201+
_densemul!(C, A, B, α, β)
202+
LinearAlgebra.mul!(C::$TC, A::_GBVecMatOrTrans, B::$TB) = _densemul!(C, A, B, true, false)
203+
end
204+
end
205+
168206
function Base.:*((⊕)::Monoid, (⊗)::Function)
169207
return function(A::GBArrayOrTranspose, B::GBArrayOrTranspose; mask=nothing, accum=nothing, desc=nothing)
170208
*(A, B, (⊕, ⊗); mask, accum, desc)

‎test/gbarray.jl‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@
3838
# #104: non-default fill shows up in show
3939
@test !occursin("fill value", sprint(show, MIME"text/plain"(), A))
4040
@test occursin("fill value: 17", sprint(show, MIME"text/plain"(), setfill(A, 17)))
41+
A = GBMatrix(sprand(ComplexF64, 5, 4, 0.5))
42+
@test Matrix(A') == Matrix(A)'
43+
@test Matrix(transpose(A)) == transpose(Matrix(A))
4144
end
4245
@testset "sparse" begin
4346
#Construction and indexing correct for sparse matrices

‎test/operations/mul.jl‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,15 @@
2121
@test m' * n == Matrix(m)' * Vector(n)
2222
@test n' * m == Vector(n)' * Matrix(m)
2323
@test n' * n == GBVector([Vector(n)' * Vector(n)])
24+
25+
# mul! into dense outputs runs in GraphBLAS (#125)
26+
S = sprand(ComplexF64, 20, 15, 0.3)
27+
A = GBMatrix(S)
28+
x, X, y, Y = rand(ComplexF64, 15), rand(ComplexF64, 15, 3), rand(ComplexF64, 20), rand(ComplexF64, 20, 3)
29+
@test mul!(copy(y), A, x) ≈ S * x
30+
@test mul!(copy(y), A, x, 2, 3) ≈ 2 * S * x + 3 * y
31+
@test mul!(copy(Y), A, X) ≈ S * X
32+
@test mul!(copy(Y), A, X, 2, 3) ≈ 2 * S * X + 3 * Y
33+
@test mul!(zeros(ComplexF64, 15), A', y) ≈ S' * y
34+
@test mul!(view(zeros(ComplexF64, 25), 3:22), A, view([x; x], 16:30)) ≈ S * x
2435
end

0 commit comments

Comments
 (0)