Hi, I am optimizing some old non-performance critical code of mine and run into a lot of allocations when using slicing operations with sparse and dense matrices.
I distilled it to this MWE
# -----------------------------------
# Initialize everything. This comes from somewhere else in the actual code.
using SparseArrays
p = 6; n = 14; Tf = 248;Tfn = n*Tf;
AzeroAp = ones(n,n*(p+1)); # collects [A0 -A1 -A2 ... -Ap]
H_Int = zeros(Int,Tfn,Tfn) # initializes the H matrix to get the linear indices
# fill the block-diagonals and off-diagonals with p+1 (nxn) matrices
for ij = 0:p
for ii = 1:div(Tfn,n)-ij
H_Int[ ij*n + (ii-1)*n + 1 : n + (ii-1)*n + ij*n, (ii-1)*n + 1 : n + (ii-1)*n] = ones(Int,n,n)
end
end
H_sp = sparse(H_Int) # set everything else to 0 (density around 2%, sparsity 98%)
AzeroAp_LinInd = LinearIndices(AzeroAp) # get the linear indices of [A0 -A1 -A2 ... -Ap]
# ----------------------------------
# THIS IS THE CRITICAL PART
@time for ij = 0:p # loop over the block-diagonal and off diagonals
for ii = 1:div(Tfn,n)-ij
H_sp[ ij*n + (ii-1)*n + 1 : n + (ii-1)*n + ij*n, (ii-1)*n + 1 : n + (ii-1)*n] = AzeroAp_LinInd[ :, 1 + (ij-0)*n : n + (ij-0)*n]
end
end
What the code does
I am creating the the following big block-banded matrix[1]:
It has \mathbf{A}_0 on the main block diagonal, \mathbf{A}_1 on the diagonal below that and so on until \mathbf{A}_p. p is always a small number, max is 6, n is also less than 20, while we have Tf block-rows, where Tf could be 200, so total Tf*n \times Tf*n.
In the code above, the variable AzeroAp := \left[ \mathbf{A}_0, -\mathbf{A}_1, \dots, -\mathbf{A}_p \right].
The objective of the code is to map the linear indices of AzeroAp to the corresponding linear indices of these matrices in \mathbf{H}.[2]. For a small example with p=1 and n=2
The indices corresponding to \mathbf{A_0} are
or in vector form [1, 2, 5, 6, 11, 12, 15,16]', while the indices corresponding to \mathbf{A_1} are [3, 4, 7, 8]' in vector form. So if I call H[3,4,7,8] = AzeroAp[5, 6, 7, 8] I will update the Hmatrix. The last loop essentially gives me these linear indices for that updating.
The question
When we come to the section with the crucial part, where I map the indices, this allocates like hell. For the above example it uses 8GB of memory.
julia> @time for ij = 0:p
for ii = 1:div(Tfn,n)-ij
H_sp[ ij*n + (ii-1)*n + 1 : n + (ii-1)*n + ij*n, (ii-1)*n + 1 : n + (ii-1)*n] = AzeroAp_LinInd[ :, 1 + (ij-0)*n : n + (ij-0)*n]
end
end
6.088194 seconds (51.16 k allocations: 8.645 GiB, 38.78% gc time)
Before coming here, as it is normal nowadays, I asked AI and the explanation was that slicing sparse and dense matrices allocates and that I should use setindex [3]. Then, the AI suggested this “fix”
for ij = 0:p
for ii = 1:div(Tfn, n) - ij
row_start = ij * n + (ii - 1) * n + 1
row_end = n + (ii - 1) * n + ij * n
col_start = (ii - 1) * n + 1
col_end = n + (ii - 1) * n
# Directly update the sparse matrix
for i = row_start:row_end
for j = col_start:col_end
H_sp[i, j] = AzeroAp_LinInd[i - row_start + 1, j - col_start + 1]
end
end
end
end
It effing works and I don’t understand why. What am I doing wrong? Why is it that when it is split across rows and columns works?
julia> 0.108702 seconds (1.35 M allocations: 36.757 MiB)
TLDR: why does the last for loop in the MWE allocate a lot, I thought I was just looking into the matrices and mappign their values from the AzeroAp to H_sp
Additional info for the curious
reveal
While this is unncessary for the question at hand you might be wondering what the hell am I doing.
The main problem is that in every iteration (say from 20000), I get new \mathbf{A} matrices, so I need to update the \mathbf{H} matrix with the new [\mathbf{A_0} \dots \mathbf{A_p}]. So I thought its easiest to save the linear indices, so that I can do
H_sp[Lind] = AzeroAp[Rind]
So before the main loop I get the linear indices vectors, Lind and Rind, respectively, once with the above code. In that sense it is not performance critical, because it is only ran once but still I thought I am doing something wrong with that many allocations. Other than that the code and the updating work very well.
Maybe there is a much easier way generate Lindand Rind but I didn’t know how and that is how I came up with the above solution. I also tried all the block-banded packages, as well as the .nzval property of sparse matrices but none worked in the way I wanted them to.
Yes, I am aware and did all try the
BlockBandedand sparseBlockBandedpackages in the Julia ecosystem ↩︎I will discuss this design choice later as to not make the exposition too large, but there are definitely other ways to achieve the same result. ↩︎
honestly this is what I thought I was doing ↩︎