About PyTorch C++ operator backend selection

Ther-nullptr: <p>I am curious about the C++ backend used by PyTorch. In version 2.0, I have noticed at least 3 kinds of C++ backends used by PyTorch——CUBLAS(The oldest backend but not open-sourced), CUTLASS(an open-sourced C++ operator library developed by NVIDIA) and Triton(a python-like language and compiler for writing highly efficient custom Deep-Learning primitives, developed by OpenAI). I want to ask some questions about this(for 2.0+ version PyTorch):</p>
<ol>
<li>What are the principles of PyTorch’s back-end selection? e.g. all of the 3 backends implement the GEMM operator, PyTorch will choose which one? Where is the source code?</li>
<li>Since the operators have some hyper-parameters (e.g. tile size), the best hyper-parameters maybe very different according to the size of operator. So how do PyTorch determines these hyper-parameters runtimely? Where is the source code to determine these params?<br>
Take my profiling examples as below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/f/6f22f5aeb891da3ed1405ebbcaf0c4823312915e.png" data-download-href="https://discuss.pytorch.org/uploads/default/6f22f5aeb891da3ed1405ebbcaf0c4823312915e" title="02fbb2d8c302db4a434cbce2cc8c654"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f22f5aeb891da3ed1405ebbcaf0c4823312915e_2_690x272.png" alt="02fbb2d8c302db4a434cbce2cc8c654" data-base62-sha1="fR9TNHm46bxObUUXlMTQD2jrnhc" width="690" height="272" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f22f5aeb891da3ed1405ebbcaf0c4823312915e_2_690x272.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f22f5aeb891da3ed1405ebbcaf0c4823312915e_2_1035x408.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f22f5aeb891da3ed1405ebbcaf0c4823312915e_2_1380x544.png 2x" data-dominant-color="17355D"><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">02fbb2d8c302db4a434cbce2cc8c654</span><span class="informations">2674×1057 257 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><br>
I have profiled GEMM of (1024,1024,1024), (2048, 2048, 2048), (4096, 4096, 4096), (8192, 8192, 8192), and notice they all use CUTLASS, with different configs(e.g. different tiling sizes, for shape of 2048 is 128x128, while for shape of 4096 is 128x256).</li>
</ol>
<p>source code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from triton.testing import do_bench

def get_flops(N, get_kernels=True):
    A = torch.randn(N, N, device='cuda', dtype=torch.float16)
    B = torch.randn(N, N, device='cuda', dtype=torch.float16)

    def f():
        return torch.mm(A, B)

    if get_kernels:
        with torch.profiler.profile() as prof:
            f()

        for e in prof.events():
            # if True:
            if "gemm" in e.name or "triton" in e.name or "gemv" in e.name:
                print(f"{N}: {e}")
                timer = e.cuda_time/1e3
    timer = do_bench(f)
    iters_per_second = 1e3/timer
    flops = A.shape[0] * A.shape[1] * B.shape[1] * 2
    flops_achieved = iters_per_second * flops/1e12
    print(f"{N}: {flops_achieved:.2f}TF/s")

for N in [1024, 2048, 4096, 8192]:
    get_flops(N)
</code></pre>

nvbkdw: <p>did you figure it out?</p>