-
-
Notifications
You must be signed in to change notification settings - Fork 5.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
draft: symmetric sparse matrix support #22200
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,76 @@ | ||
# This file is a part of Julia. License is MIT: https://julialang.org/license | ||
|
||
function Symmetric(A::SparseMatrixCSC, uplo::Symbol=:U) | ||
checksquare(A) | ||
Symmetric{eltype(A), typeof(A)}(A, Base.LinAlg.char_uplo(uplo)) # preserve A | ||
end | ||
|
||
(*)(A::Symmetric{TA,SparseMatrixCSC{TA,S}}, x::StridedVecOrMat{Tx}) where {TA,S,Tx} = A_mul_B(A, x) | ||
|
||
function A_mul_B!(α::Number, A::Symmetric{TA,SparseMatrixCSC{TA,S}}, B::StridedVecOrMat, β::Number, C::StridedVecOrMat) where {TA,S} | ||
|
||
A.data.n == size(B, 1) || throw(DimensionMismatch()) | ||
A.data.m == size(C, 1) || throw(DimensionMismatch()) | ||
|
||
A.uplo == 'U' ? A_mul_B_U_kernel!(α, A, B, β, C) : A_mul_B_L_kernel!(α, A, B, β, C) | ||
end | ||
|
||
function A_mul_B(A::Symmetric{TA,SparseMatrixCSC{TA,S}}, x::StridedVector{Tx}) where {TA,S,Tx} | ||
T = promote_type(TA, Tx) | ||
A_mul_B!(one(T), A, x, zero(T), similar(x, T, A.data.n)) | ||
end | ||
|
||
function A_mul_B(A::Symmetric{TA,SparseMatrixCSC{TA,S}}, B::StridedMatrix{Tx}) where {TA,S,Tx} | ||
T = promote_type(TA, Tx) | ||
A_mul_B!(one(T), A, B, zero(T), similar(B, T, (A.data.n, size(B, 2)))) | ||
end | ||
|
||
function A_mul_B_U_kernel!(α::Number, A::Symmetric{TA,SparseMatrixCSC{TA,S}}, B::StridedVecOrMat, β::Number, C::StridedVecOrMat) where {TA,S} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There are just 2 lines that are different between the function A_mul_B_UL_kernel!(α::Number, A::Symmetric{TA,SparseMatrixCSC{TA,S}},
B::StridedVecOrMat, β::Number, C::StridedVecOrMat, ::Type{Val{uplo}}) where {TA,S,uplo}
colptr = A.data.colptr
rowval = A.data.rowval
nzval = A.data.nzval
if β != 1
β != 0 ? scale!(C, β) : fill!(C, zero(eltype(C)))
end
@inbounds for k = 1 : size(C, 2)
@inbounds for col = 1 : A.data.n
αxj = α * B[col, k]
tmp = TA(0)
@inbounds for j = (uplo == :U ? (colptr[col] : (colptr[col + 1] - 1)) :
(colptr[col] : (colptr[col + 1] - 1)))
row = rowval[j]
uplo == :U ? (row > col && break) : (row > col && break) # assume indices are sorted
a = nzval[j]
C[row, k] += a * αxj
row == col || (tmp += a * B[row, k])
end
C[col, k] += α * tmp
end
end
C
end There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There's a performance impact for that, which is why I separated them. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In fact, there's a performance impact for allowing There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
That should be compiled away.
If the impact is noticeable it makes sense to have a separate kernel for matvec. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Using the matrix in my gist, here's the benchmark of
Now
and
I didn't know about |
||
|
||
colptr = A.data.colptr | ||
rowval = A.data.rowval | ||
nzval = A.data.nzval | ||
if β != 1 | ||
β != 0 ? scale!(C, β) : fill!(C, zero(eltype(C))) | ||
end | ||
@inbounds for k = 1 : size(C, 2) | ||
@inbounds for col = 1 : A.data.n | ||
αxj = α * B[col, k] | ||
tmp = TA(0) | ||
@inbounds for j = colptr[col] : (colptr[col + 1] - 1) | ||
row = rowval[j] | ||
row > col && break # assume indices are sorted | ||
a = nzval[j] | ||
C[row, k] += a * αxj | ||
row == col || (tmp += a * B[row, k]) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this should have an alpha There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Of course, thanks. I chose to multiply There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ps: should the kernels be exported? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't think so. |
||
end | ||
C[col, k] += tmp | ||
end | ||
end | ||
C | ||
end | ||
|
||
function A_mul_B_L_kernel!(α::Number, A::Symmetric{TA,SparseMatrixCSC{TA,S}}, B::StridedVecOrMat, β::Number, C::StridedVecOrMat) where {TA,S} | ||
|
||
colptr = A.data.colptr | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 4 space indent, and shouldn't start a function body with a blank line There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should be fixed. Besides doing the same job for |
||
rowval = A.data.rowval | ||
nzval = A.data.nzval | ||
if β != 1 | ||
β != 0 ? scale!(C, β) : fill!(C, zero(eltype(C))) | ||
end | ||
@inbounds for k = 1 : size(C, 2) | ||
@inbounds for col = 1 : A.data.n | ||
αxj = α * B[col, k] | ||
tmp = TA(0) | ||
@inbounds for j = (colptr[col + 1] - 1) : -1 : colptr[col] | ||
row = rowval[j] | ||
row < col && break | ||
a = nzval[j] | ||
C[row, k] += a * αxj | ||
row == col || (tmp += a * B[row, k]) | ||
end | ||
C[col, k] += tmp | ||
end | ||
end | ||
C | ||
end |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is not needed since it is equivalent to the default constructor:
julia/base/linalg/symmetric.jl
Line 43 in 063e8f1
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right. Fixed.