-
-
Notifications
You must be signed in to change notification settings - Fork 5.5k
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
add in-place kron #31069
add in-place kron #31069
Conversation
b12ea18
to
62b45de
Compare
end | ||
|
||
function kron(A::AbstractMatrix{T}, B::Diagonal{S}) where {T<:Number, S<:Number} | ||
@inline function kron!(C::AbstractMatrix, A::AbstractMatrix, B::Diagonal) |
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 should also have a @propagate_inbounds
, I suppose?
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.
Ahh, yes!
stdlib/LinearAlgebra/src/dense.jl
Outdated
|
||
Like [`kron`](@ref), but stores the results in `C`. | ||
|
||
!!! tip |
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.
Is this a recognized Documenter thing? I might go with !!! info
to be safe.
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.
Yes, tip will be a green block, but info is good as well.
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.
Oh nice, I didn't know that. Either is fine then.
It would be nice to get this finished. @Roger-luo Could you please rebase this onto current master and resolve the merge conflicts? One thing I noticed is that there are occurrences of |
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.
Just a few comments that I noticed. Will take a more careful look later.
stdlib/SparseArrays/src/linalg.jl
Outdated
ptr_range = (1:lB) .+ (colptrC[col]-1) | ||
colptrC[col+1] = colptrC[col] + lA*lB | ||
ptr_range = (1:lB) .+ (C.colptr[col]-1) | ||
C.colptr[col+1] = C.colptr[col] + lA*lB |
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.
Wait, there is no C
defined in this function...? Or did you intend to rename this to kron!
and add a first C
argument here? Independent from that, I think one is no longer supposed to refer to colptr
via the dot notation, but rather via getcolptr(C)
, in case you'll need it. Otherwise, CI will throw and fail.
stdlib/SparseArrays/src/linalg.jl
Outdated
return z | ||
end | ||
|
||
kron!(C::SparseMatrixCSC, A::SparseMatrixCSC, x::SparseVector) = kron!(C, A, SparseMatrixCSC(x)) |
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.
Here and below, I think all the SparseMatrixCSC
can be relaxed to AbstractSparseMatrixCSC
.
stdlib/LinearAlgebra/src/dense.jl
Outdated
Bounds checking can be disabled by `@inbounds`, but you need to take care of the shape | ||
of `C`, `A`, `B` yourself. | ||
""" | ||
@inline Base.@propagate_inbounds function kron!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) |
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.
@inline Base.@propagate_inbounds function kron!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) | |
Base.@propagate_inbounds function kron!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) |
should inline already, same below.
Hi sorry, I found the I also polished the |
This is impressive work! I've been looking for ways to repeatedly compute large Kronecker sums with a fixed structure. Below are some benchmarks showing the speedup Below are some benchmarks for calculating a single term in a Kronecker sum. Denseusing BenchmarkTools
using LinearAlgebra
using SparseArrays
for N in 2 .^ [0,1,2,3,4,5,6,7]
A = rand(N,N); B = Diagonal(oneunit(A)); C = kron(A,B);
println("Out-of-place N=$N")
@btime kron($A,$B)
println("In-place N=$N")
@btime kron!($C,$A,$B)
println("")
end Click to expand timings
Sparsep = 1.0
for N in 2 .^ [0,1,2,3,4,5,6,7]
A = sprand(N,N,p); B = oneunit(A); C = kron(A,B);
println("Out-of-place N=$N")
@btime kron($A,$B)
println("In-place N=$N")
@btime kron!($C,$A,$B)
println("")
end Click to expand timings
|
Hi, @elisno I'm glad this helps. Regarding large Kronecker sum, I'm not sure about your problem, but if your matrices are more special (e.g general permutation matrices) we have an even faster implementation (as a summation of quantum circuits) in Yao. You may want to check that out. |
First this needs to get merged, and then maybe someone could add those parts of this PR that are a new feature to |
bump. I think I've fixed all the tests. any chance to review this? @dkarrasch |
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 very nice work, indeed, and consistent with all the *
-mul!
business. I only have a couple of non-functional comments, so this should be quick. I'd only like to ask not to overwrite the PR by push-forced, so that I can quickly review afterwards.
function kron! 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.
What is this good for? For stack traces or something like that?
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 because kron
is defined in Base
(https://github.com/JuliaLang/julia/blob/master/base/operators.jl#L533), thus all other stdlibs need to import Base: kron
, I think it'd be more consistent to define the generic function in Base
instead of LinearAlgebra
.
Some changes in |
the failed test seems unrelated to this PR. I've reverted the extra commit from this PR. |
Let's try again. |
looks good now, all tests pass. |
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 LGTM. I think we may want to announce this in NEWS.md, since this is of interest to external packages as well. I'll leave this open for a few days, hoping that another pair of eyes takes a look.
@Roger-luo Could you please add an item to NEWS.md, announcing this feature? You may need to rebase the PR, but I'm not sure, to add to the Julia v1.6 version of NEWS.md. |
5d08842
to
23bfba1
Compare
add sparse arrays
The failed test on MacOS seems unrelated. |
bump, anyone else could take a look at this PR maybe? |
This is useful to avoid extra allocation when
A
,B
is large inkron(A, B)
. Since the implementation itself is exactly the same with originalkron
, I guess it's better to have it directly instdlib
s