Skip to content

Commit 5ffa937

Browse files
committed
Refactor mul! to use @interface
1 parent 3fcb544 commit 5ffa937

File tree

1 file changed

+5
-26
lines changed

1 file changed

+5
-26
lines changed

src/abstractsparsearrayinterface.jl

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -325,40 +325,19 @@ function mul_indices(I1::CartesianIndex{2}, I2::CartesianIndex{2})
325325
return CartesianIndex(I1[1], I2[2])
326326
end
327327

328-
using LinearAlgebra: mul!
329-
function default_mul!!(
330-
a_dest::AbstractMatrix,
331-
a1::AbstractMatrix,
332-
a2::AbstractMatrix,
333-
α::Number=true,
334-
β::Number=false,
335-
)
336-
mul!(a_dest, a1, a2, α, β)
337-
return a_dest
338-
end
339-
340-
function default_mul!!(
341-
a_dest::Number, a1::Number, a2::Number, α::Number=true, β::Number=false
342-
)
343-
return a1 * a2 * α + a_dest * β
344-
end
328+
using LinearAlgebra: LinearAlgebra, mul!
345329

346330
# a1 * a2 * α + a_dest * β
347-
function sparse_mul!(
348-
a_dest::AbstractArray,
349-
a1::AbstractArray,
350-
a2::AbstractArray,
351-
α::Number=true,
352-
β::Number=false;
353-
(mul!!)=(default_mul!!),
331+
@interface ::AbstractSparseArrayInterface function LinearAlgebra.mul!(
332+
C::AbstractArray, A::AbstractArray, B::AbstractArray, α::Number, β::Number
354333
)
355334
a_dest .*= β
356335
β′ = one(Bool)
357336
for I1 in eachstoredindex(a1)
358337
for I2 in eachstoredindex(a2)
359338
I_dest = mul_indices(I1, I2)
360339
if !isnothing(I_dest)
361-
a_dest[I_dest] = mul!!(a_dest[I_dest], a1[I1], a2[I2], α, β′)
340+
a_dest[I_dest] = mul!(a_dest[I_dest], a1[I1], a2[I2], α, β′)
362341
end
363342
end
364343
end
@@ -368,7 +347,7 @@ end
368347
function ArrayLayouts.materialize!(
369348
m::MatMulMatAdd{<:AbstractSparseLayout,<:AbstractSparseLayout,<:AbstractSparseLayout}
370349
)
371-
sparse_mul!(m.C, m.A, m.B, m.α, m.β)
350+
@interface SparseArrayInterface() mul!(m.C, m.A, m.B, m.α, m.β)
372351
return m.C
373352
end
374353

0 commit comments

Comments
 (0)