PyTorch's matmul triggers different kernels to cublasSgemm

yzh119: <p>I’m trying to write a batched-gemm operator where each input has different shapes, following is a script I used for profiling and debugging:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import myop
import time

m = torch.tensor([100, 128, 20, 284, 89, 10, 82, 92, 10, 20, 30, 100, 128, 20, 284, 89, 10, 82, 92, 10, 20, 30,
                  100, 128, 20, 284, 89, 10, 82, 92, 10, 20, 30, 100, 128, 20, 284, 89, 10, 82, 92, 10, 20, 30])
k = torch.tensor([48, 256, 33, 92, 17, 520,  48, 256, 33, 92, 17, 520, 48, 256, 33, 92, 17, 520, 33, 92, 17, 520,
                  48, 256, 33, 92, 17, 520,  48, 256, 33, 92, 17, 520, 48, 256, 33, 92, 17, 520, 33, 92, 17, 520])
n = torch.tensor([200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30,
                  200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30, 17, 30, 200, 30])

A_arr = []
B_arr = []
C_arr = []

for mi, ni, ki in zip(m.tolist(), n.tolist(), k.tolist()):
    A_arr.append(torch.rand(mi, ki).to(0))
    B_arr.append(torch.rand(ki, ni).to(0))
    C_arr.append(torch.rand(mi, ni).to(0))

A = torch.cat([_.view(-1) for _ in A_arr], 0)
B = torch.cat([_.view(-1) for _ in B_arr], 0)
C = torch.cat([_.view(-1) for _ in C_arr], 0)

torch.cuda.synchronize()
tic = time.time()
for _ in range(100):
    for i in range(len(m)):
        Ai = A_arr[i]
        Bi = B_arr[i]
        Ci = Ai @ Bi

torch.cuda.synchronize()
toc = time.time()
print(toc - tic)

tic = time.time()
for _ in range(100):
    myop.multiple_gemm(A, B, C, m, n, k, False, False)

torch.cuda.synchronize()
toc = time.time()
print(toc - tic)
</code></pre>
<p>where <code>myop</code> is a C++ extension I wrote which basically calls <code>cublasSgemm</code> (non-transpose for both operands) operator for each (a, b, c) in the batch sequentially and only uses a single stream.</p>
<p>However, I found that myop’s execution time is longer than sequentially apply <code>torch.matmul</code> in Python. To investigate the reason I use nsight-system to get more insights:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/c/4c6b2104ae23779777746c81705063bcf7f59d09.png" data-download-href="https://discuss.pytorch.org/uploads/default/4c6b2104ae23779777746c81705063bcf7f59d09" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/c/4c6b2104ae23779777746c81705063bcf7f59d09_2_689x46.png" alt="image" data-base62-sha1="aU1OfG2Y5tlAOFOKbtHjKYozUdr" width="689" height="46" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/c/4c6b2104ae23779777746c81705063bcf7f59d09_2_689x46.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/c/4c6b2104ae23779777746c81705063bcf7f59d09_2_1033x69.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/c/4c6b2104ae23779777746c81705063bcf7f59d09.png 2x" data-dominant-color="D6DDE8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1346×91 22.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I found that PyTorch only triggers <code>sgemm_32x32x32_NN_vec</code> kernel while <code>myop</code> trigger multiple kinds of kernels including <code>maxwell_sgemm_128_64_nn</code> and <code>sgemm_32x32x32_NN_vec</code>, and I can confirm that both ways produce exactly the same result. I was wondering that what’s PyTorch behavior for GEMM on GPU, has it enabled some environment variable that might force cublas to not select kernel in the default way, etc?</p>

ptrblck: <p>No, cublas env variables are not set, so I assume your manual implementation might just call into more cublas operations. Are these operations always called in your implementation?</p>