Matmul broadcasting makes copies?

thuyenvn: <p>Based on the <a href="https://pytorch.org/docs/stable/torch.html#torch.matmul" rel="nofollow noopener">docs</a>, matmul will broadcast to make inputs compatible.<br>
From this <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L268" rel="nofollow noopener">line</a>, it seems the following script will expand and make copies (1000 copies) of the second tensor (to due <code>contiguous()</code>).</p>
<pre><code class="lang-python">x = torch.rand(1000, 10, 3, 4)
y = torch.rand(10, 4, 5)
z = torch.matmul(x, y)
</code></pre>
<p>If this is correct, we should change the docs since we are assuming not making copies when broadcasting <a href="https://pytorch.org/docs/stable/notes/broadcasting.html#broadcasting-semantics" rel="nofollow noopener">here</a>.</p>

tom: <p>Hi,</p>
<p>that is a great observation and good research to pinpoint the exact line that is problematic!</p>
<p>I think this has two parts:</p>
<ul>
<li>There seems to be indeed something going on with matmul that shouldn’t. The timing below suggests that matmul has 4.5ms while even the naive mul + sum has 0.6ms with the below:</li>
</ul>
<pre><code class="lang-auto">     x = torch.rand(1000, 10, 3, 4)
     y = torch.rand(10, 4, 5)
     y2 = y.unsqueeze(0)
     x_ = x.unsqueeze(-1)
     y_ = y2.unsqueeze(2)
     %timeit z = torch.matmul(x, y)
     %timeit z2 = (x_*y_).sum(-2)
</code></pre>
<ul>
<li>The strategy matmul employs is to reshape the array into something to be fed to <code>at::bmm</code>. Actually, in your example, it is not possible to do this, as you cannot combine the stride 0 of the first axis and the stride of the second axis into one. So a proper fix probably boils down to reworking the <code>matmul</code> rather fundamentally. (Probably a fun project!)</li>
</ul>
<p>If you wanted to, you could file a bug. I do believe there is a bug in the code, not just the docs.</p>
<p>Best regards</p>
<p>Thomas</p>

thuyenvn: <p>The CPU kernel is slow because we use a <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/TH/generic/THTensorMath.cpp#L2308" rel="nofollow noopener">for loop</a> over batch dimension. I think the main reason for making copies and utilizing  <code>bmm</code> is that we don’t have a readily available kernel for the GPU op otherwise. Right now CUDA provides <code>cublasSgemmBatched</code>, <code>cublasDgemmBatched</code> and their strided versions for batched <code>mm</code>.</p>

thuyenvn: <p>Just posted feature request <a href="https://devtalk.nvidia.com/default/topic/1036464/more-flexible-cublas-lt-t-gt-gemmstridedbatched" rel="nofollow noopener">here</a> for <code>cublas&lt;T&gt;gemmStridedBatched</code></p>

ngimel: <p>GemmBatchedEx <a href="https://docs.nvidia.com/cuda/cublas/index.html#cublas-GemmBatchedEx" rel="nofollow noopener">https://docs.nvidia.com/cuda/cublas/index.html#cublas-GemmBatchedEx</a> should be able to do what you want, because it accepts an array of pointers, so B pointers could be repeating in a loop, but cublasGemmBatchedEx is not used in pytorch backend because preparing aforementioned lists of pointers is somewhat painful unless really necessary.</p>

thuyenvn: <p>Actually Pytorch uses something similar prior to CUDA 8.0, and it’s still effective <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/THC/generic/THCTensorMathBlas.cu#L574" rel="nofollow noopener">here</a>. But as you said, <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/THC/generic/THCTensorMathBlas.cu#L450" rel="nofollow noopener">allocating</a> lots of pointers is not that efficient. For the toy example above, it’s allocating <code>30K</code> pointers and actually slower than making copies +  <code>gemmStridedBatched</code>. It would be really nice if we can do <code>gemmStridedBatched</code> without copying.</p>

mikael: <p>Hi <a class="mention" href="/u/tom">@tom</a>,</p>
<p>is naive mul + sum still faster than matmul in PyTorch 1.0 ? I tried your code and I got the opposite.</p>
<p>Thank you,</p>
<p>Mikaël</p>

mark1: <p>Is there still no memory and time-efficient way to do broadcasted matrix multiplication?  It seems like this would be essential for something like a transformer network since otherwise dot-product attention would require you to copy both your keys and your queries?  Is there a special trick to do dot-product attention in particular?</p>

Thomas_Ahle: <p>I was trying to figure out why my benchmarking trace was full of <code>copy_()</code> calls. Turns out it was because of this <code>matmul</code> broadcasting issue. Now I use <code>F.linear</code> instead, which is a bit awkward, since everything has to be transposed, but it is <em>much</em> faster. All the copying is gone.</p>
<p>To be more precise: The point is that you need to structure your code so broadcasted matrix multiplication is done from the right, not the left. Here is a simple example:</p>
<pre><code class="lang-auto">[ins] In [3]: BS, DIM = 100, 100
         ...: A = torch.randn(DIM, DIM)
         ...: B = torch.randn(BS, DIM, DIM)

[ins] In [4]: %timeit A @ B
6.11 ms ± 117 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[ins] In [5]: %timeit B @ A
4.02 ms ± 15.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
</code></pre>