Memory efficient way to implement masked matrix multiplication

WuYunfan: <p>I want to multiply two dense matrices A(N, d) and B(d, N). The resulting matrix will be of size N x N, where N is very large and I can’t use normal matrix multiplication function to do this.<br>
However, I only want values at a few positions which are specified in a sparse matrix C(N, N). For example, I only want <strong>E</strong> non-zero entries in C.<br>
My problem is the same as discussed in <a href="https://discuss.pytorch.org/t/matrix-multiplication-with-values-only-at-masked-entries/11430" class="inline-onebox">Matrix multiplication with values only at masked entries</a>.<br>
However, I don’t find a memory efficient way to solve this problem.<br>
The solution provided by <strong>tom</strong> take a memory complexity of <strong>O(Ed)</strong>.<br>
Spliting it into chunks does not help as it still occupies memory for <strong>gradient ackpropagation</strong>.</p>

googlebot: <aside class="quote no-group" data-username="WuYunfan" data-post="1" data-topic="127543">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wuyunfan/48/40577_2.png" class="avatar"> WuYunfan:</div>
<blockquote>
<p>Spliting it into chunks does not help as it still occupies memory for <strong>gradient ackpropagation</strong>.</p>
</blockquote>
</aside>
<p>matmul doesn’t store output AB, but your further indexing does. so, you can try writing custom autograd.Function (A,B) (returning (A@B)[indexes])  that would use sparse gradient of AB in backward (then do sparse-dense matmuls to compute gradients of A and B).</p>
<p>PS: I mean using chunked matmuls in forward, can’t think of a better way atm</p>

WuYunfan: <aside class="quote no-group" data-username="googlebot" data-post="2" data-topic="127543">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>Function (A,B) (returning (A@B)[indexes])</p>
</blockquote>
</aside>
<p>I can’t do matmuls A@B there as it takes a time complexity of <strong>O(N^2 d)</strong>.<br>
Specifically, I have a gpraph with <strong>N nodes</strong> and <strong>E</strong> edges.<br>
N is several millions and E is dozens of millions.</p>

googlebot: <p>if you have memory problem with this:</p>
<blockquote>
<p>vals = (A[rows]*B[cols]).sum(1)</p>
</blockquote>
<p>maybe make it a function and wrap it with utils.checkpoint</p>

WuYunfan: <aside class="quote no-group" data-username="googlebot" data-post="4" data-topic="127543">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>maybe make it a function and wrap it with utils.checkpoint</p>
</blockquote>
</aside>
<p>It works well with my code, thanks a lot.</p>

shenshanf: <p>see this link:<a href="https://github.com/facebookresearch/xformers/blob/main/xformers/csrc/attention/cuda/matmul.cu" rel="noopener nofollow ugc">matmul_with_mask</a> may help you.</p>

shenshanf: <p>detail code see: <a href="https://github.com/shenshanf/matmul_with_mask" class="inline-onebox" rel="noopener nofollow ugc">GitHub - shenshanf/matmul_with_mask: matrix multiplication with sparse mask that indicate the operation position</a></p>

shenshanf: <p>dgl lib: <a href="https://docs.dgl.ai/generated/dgl.sparse.bsddmm.html" class="inline-onebox" rel="noopener nofollow ugc">dgl.sparse.bsddmm — DGL 1.1.2 documentation</a></p>