Skip to content

Reduce allocations for multipling LazyTensor of sparse and dense #80

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.

Sign up for GitHub

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

Merged
merged 17 commits into from
Mar 27, 2023
Merged

Reduce allocations for multipling LazyTensor of sparse and dense #80

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
671fae8
make shape and strides a tuple
AmitRotem Feb 23, 2023
2a19fd9
make _strides type stable for tuple
AmitRotem Feb 23, 2023
0a5e5b5
reduce calls to _comp_size
AmitRotem Feb 23, 2023
e051a0a
remove type constraint from _get_shape_and_srtides
AmitRotem Feb 23, 2023
52c0551
fix typo
AmitRotem Feb 24, 2023
6070884
allow _gemm_recursive_lazy_dense matrix vector mul
AmitRotem Feb 24, 2023
5b473d6
type fix
AmitRotem Feb 25, 2023
9911d6e
Merge branch 'qojulia:master' into master
AmitRotem Mar 17, 2023
cca18de
relax type in _gemm_puresparse\!, _gemm_recursive_lazy_dense, and _ge…
Mar 21, 2023
4d7d294
Merge branch 'qojulia:master' into master
AmitRotem Mar 21, 2023
b7bdb94
fix size(AbstractOperator)
Mar 22, 2023
7e64f54
add dim check for _gemm_puresparse
Mar 22, 2023
a616377
test LazyTensor with sparse dimension mismatch
Mar 22, 2023
316ba34
change error thrown by size to ErrorException
AmitRotem Mar 22, 2023
6984ff0
test size of AbstractOperator
AmitRotem Mar 22, 2023
8aa2ab9
use lazy string
Mar 24, 2023
f91c5a0
split check_mul\!_compatibility and specialize for vectors
Mar 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/operators_dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,10 @@ function _strides(shape)
return S
end

function _strides(shape::Ty)::Ty where Ty <: Tuple
accumulate(*, (1,Base.front(shape)...))
end

# Dense operator version
@generated function _ptrace(::Type{Val{RANK}}, a,
shape_l, shape_r,
Expand Down
27 changes: 12 additions & 15 deletions src/operators_lazytensor.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -609,7 +609,7 @@ end
function _gemm_recursive_lazy_dense(i_k, N_k, K, J, val,
shape, strides_k, strides_j,
indices, h::LazyTensor,
op::Matrix, result::Matrix)
op::VecOrMat, result::VecOrMat)
if i_k > N_k
for I=1:size(op, 2)
result[J, I] += val*op[K, I]
Expand Down Expand Up @@ -648,34 +648,31 @@ function _gemm_puresparse(alpha, op::Matrix, h::LazyTensor{B1,B2,F,I,T}, beta, r
rmul!(result, beta)
end
N_k = length(h.basis_r.bases)
shape = [min(h.basis_l.shape[i], h.basis_r.shape[i]) for i=1:length(h.basis_l.shape)]
strides_j = _strides(h.basis_l.shape)
strides_k = _strides(h.basis_r.shape)
shape, strides_j, strides_k = _get_shape_and_strides(h)
_gemm_recursive_dense_lazy(1, N_k, 1, 1, alpha*h.factor, shape, strides_k, strides_j, h.indices, h, op, result)
end

function _gemm_puresparse(alpha, h::LazyTensor{B1,B2,F,I,T}, op::Matrix, beta, result::Matrix) where {B1,B2,F,I,T<:Tuple{Vararg{SparseOpPureType}}}
function _gemm_puresparse(alpha, h::LazyTensor{B1,B2,F,I,T}, op::VecOrMat, beta, result::VecOrMat) where {B1,B2,F,I,T<:Tuple{Vararg{SparseOpPureType}}}
if iszero(beta)
fill!(result, beta)
elseif !isone(beta)
rmul!(result, beta)
end
N_k = length(h.basis_l.bases)
shape = [min(h.basis_l.shape[i], h.basis_r.shape[i]) for i=1:length(h.basis_l.shape)]
strides_j = _strides(h.basis_l.shape)
strides_k = _strides(h.basis_r.shape)
shape, strides_j, strides_k = _get_shape_and_strides(h)
_gemm_recursive_lazy_dense(1, N_k, 1, 1, alpha*h.factor, shape, strides_k, strides_j, h.indices, h, op, result)
end

function _get_shape_and_strides(h)
shape_l, shape_r = _comp_size(h.basis_l), _comp_size(h.basis_r)
shape = min.(shape_l, shape_r)
strides_j, strides_k = _strides(shape_l), _strides(shape_r)
return shape, strides_j, strides_k
end

_mul_puresparse!(result::DenseOpType{B1,B3},h::LazyTensor{B1,B2,F,I,T},op::DenseOpType{B2,B3},alpha,beta) where {B1,B2,B3,F,I,T<:Tuple{Vararg{SparseOpPureType}}} = (_gemm_puresparse(alpha, h, op.data, beta, result.data); result)
_mul_puresparse!(result::DenseOpType{B1,B3},op::DenseOpType{B1,B2},h::LazyTensor{B2,B3,F,I,T},alpha,beta) where {B1,B2,B3,F,I,T<:Tuple{Vararg{SparseOpPureType}}} = (_gemm_puresparse(alpha, op.data, h, beta, result.data); result)

function _mul_puresparse!(result::Ket{B1},a::LazyTensor{B1,B2,F,I,T},b::Ket{B2},alpha,beta) where {B1,B2,F,I,T<:Tuple{Vararg{SparseOpPureType}}}
b_data = reshape(b.data, length(b.data), 1)
result_data = reshape(result.data, length(result.data), 1)
_gemm_puresparse(alpha, a, b_data, beta, result_data)
result
end
_mul_puresparse!(result::Ket{B1},a::LazyTensor{B1,B2,F,I,T},b::Ket{B2},alpha,beta) where {B1,B2,F,I,T<:Tuple{Vararg{SparseOpPureType}}} = (_gemm_puresparse(alpha, a, b.data, beta, result.data); result)

function _mul_puresparse!(result::Bra{B2},a::Bra{B1},b::LazyTensor{B1,B2,F,I,T},alpha,beta) where {B1,B2,F,I,T<:Tuple{Vararg{SparseOpPureType}}}
a_data = reshape(a.data, 1, length(a.data))
Expand Down