Does torch.compile try to reorder matmuls for efficiency

PeaBrane: <p>I’m just curious if <code>torch.compile</code> is able to perform <a href="https://optimized-einsum.readthedocs.io/en/stable/" rel="noopener nofollow ugc"><code>opt_einsum</code></a> style optimizations, where the order of matrix multiplications is optimized to reduce compute.</p>
<p>The minimal example here is</p>
<pre><code class="lang-auto">@torch.compile
def matmul(A, B, C):
    return A @ B @ C
</code></pre>
<p>In the case where, say, <code>A</code> is 1000 x 100, <code>B</code> is 100 x 10, and <code>C</code> is 10 x 1, it is clearly more efficient to perform the matmul as <code>A @ (B @ C)</code>, where the last two matrices can be multiplied first due to <a href="https://proofwiki.org/wiki/Matrix_Multiplication_is_Associative" rel="noopener nofollow ugc">associativity</a>. I’m wondering if this is something <code>torch.compile</code> can optimize for.</p>

KFrank: <p>Hi Pea!</p>
<aside class="quote no-group quote-modified" data-username="PeaBrane" data-post="1" data-topic="192112" data-full="true">
<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/peabrane/48/40023_2.png" class="avatar"> PeaBrane:</div>
<blockquote>
<pre><code class="lang-auto">@torch.compile
def matmul(A, B, C):
    return A @ B @ C
</code></pre>
<p>…<br>
it is clearly more efficient to perform the matmul as <code>A @ (B @ C)</code>, where the last two matrices can be multiplied first<br>
…<br>
is something <code>torch.compile</code> can optimize for.</p>
</blockquote>
</aside>
<p>This does not appear to be something that <code>torch.compile</code> optimizes.</p>
<p>It’s speculation on my part, but I don’t think that <code>torch.compile</code> knows<br>
that it is doing matrix multiplication (nor realizes that matrix multiplication<br>
is associative).</p>
<p>Note that <a href="https://pytorch.org/docs/2.1/generated/torch.linalg.multi_dot.html" rel="noopener nofollow ugc">torch.linalg.multi_dot()</a> has this matrix-associativity optimization<br>
as its specific purpose.</p>
<p>Here is a script that shows no <code>torch.compile</code> speed-up for your minimal<br>
example, while showing speed-up when explicitly performing <code>B @ C</code> first<br>
and when letting <code>multi_dot()</code> perform the optimization:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())
print (torch.cuda.get_device_capability())

import time

def matmul_no_compile (A, B, C):
    return A @ B @ C                 # equivalent to (A @ B) @ c

@torch.compile (backend = 'eager')   # 'eager' for older gpu
def matmul_compile (A, B, C):
    return A @ B @ C                 # equivalent to (A @ B) @ c

def matmul_rl (A, B, C):
    return A @ (B @ C)               # explicit efficient ordering

Aa = []
Bb = []
Cc = []

print ('generating data ...')
for  i in range (10):
    Aa.append (torch.randn (10000, 1000, device = 'cuda'))
    Bb.append (torch.randn (1000, 100, device = 'cuda'))
    Cc.append (torch.randn (100, 10, device = 'cuda'))

print ('run timing loops ...')

# matmul_no_compile

# run warmup
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_no_compile (A, B, C)

# timing loop
torch.cuda.synchronize()
t0 = time.time()
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_no_compile (A, B, C)

torch.cuda.synchronize()
t1 = time.time()
print ('matmul_no_compile:  elapsed time: ', t1 - t0)

# matmul_compile

# run warmup
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_compile (A, B, C)

# timing loop
torch.cuda.synchronize()
t0 = time.time()
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_compile (A, B, C)

torch.cuda.synchronize()
t1 = time.time()
print ('matmul_compile:     elapsed time: ', t1 - t0)

# matmul_rl

# run warmup
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_rl (A, B, C)

# timing loop
torch.cuda.synchronize()
t0 = time.time()
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = matmul_rl (A, B, C)

torch.cuda.synchronize()
t1 = time.time()
print ('matmul_rl:          elapsed time: ', t1 - t0)

# multi_dot -- knows how to reorder matrix multiplications

# run warmup
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = torch.linalg.multi_dot ((A, B, C))

# timing loop
torch.cuda.synchronize()
t0 = time.time()
for  i in range (10):
    for  A, B, C in zip (Aa, Bb, Cc):
        D = torch.linalg.multi_dot ((A, B, C))

torch.cuda.synchronize()
t1 = time.time()
print ('multi_dot:          elapsed time: ', t1 - t0)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.1.0
11.8
GeForce GTX 1050 Ti
(6, 1)
generating data ...
run timing loops ...
matmul_no_compile:  elapsed time:  0.18052339553833008
matmul_compile:     elapsed time:  0.1767871379852295
matmul_rl:          elapsed time:  0.056723833084106445
multi_dot:          elapsed time:  0.05632662773132324
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

PeaBrane: <p>Thank you Frank!</p>
<p>I’ve adapted your code to run on larger matrices, with tensorfloat32, and including <code>torch.einsum</code> in the mix. And it turns out that you are completely right that <code>torch.compile</code> seems to make no effort in reordering matmuls for efficiency.</p>
<pre><code class="lang-auto">import torch

print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())
print (torch.cuda.get_device_capability())

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

import time


def matmul_no_compile(A, B, C):
    return A @ B @ C                 # equivalent to (A @ B) @ c

@torch.compile
def matmul_compile(A, B, C):
    return A @ B @ C                 # equivalent to (A @ B) @ c

def matmul_rl(A, B, C):
    return A @ (B @ C)               # explicit efficient ordering

def multi_dot(A, B, C):
    return torch.linalg.multi_dot((A, B, C))

def einsum(A, B, C):
    return torch.einsum('ij,jk,kl-&gt;il', A, B, C)

Aa, Bb, Cc = [], [], []

print ('generating data ...')
for  i in range (10):
    Aa.append (torch.randn (2 ** 16, 2 ** 12, device = 'cuda'))
    Bb.append (torch.randn (2 ** 12, 2 ** 12, device = 'cuda'))
    Cc.append (torch.randn (2 ** 12, 2 ** 8, device = 'cuda'))
    
print(f"matrices shapes: {Aa[0].shape} {Bb[0].shape} {Cc[0].shape}")

print ('run timing loops ...')

# matmul_no_compile

def run_benchmark(fn, Aa, Bb, Cc):
    # warmup
    for i in range(10):
        for (A, B, C) in zip(Aa, Bb, Cc):
            D = fn(A, B, C)
    
    # timeing
    torch.cuda.synchronize()
    t0 = time.time()
    for  i in range (10):
        for  A, B, C in zip (Aa, Bb, Cc):
            D = fn(A, B, C)
            
    torch.cuda.synchronize()
    t1 = time.time()
    print(f'{fn.__name__}: {t1 - t0}')


run_benchmark(matmul_no_compile, Aa, Bb, Cc)
run_benchmark(matmul_compile, Aa, Bb, Cc)
run_benchmark(matmul_rl, Aa, Bb, Cc)
run_benchmark(multi_dot, Aa, Bb, Cc)
run_benchmark(einsum, Aa, Bb, Cc)
</code></pre>
<p>output</p>
<pre><code class="lang-auto">2.1.0
11.8
NVIDIA A30
(8, 0)
generating data ...
matrices shapes: torch.Size([65536, 4096]) torch.Size([4096, 4096]) torch.Size([4096, 256])
run timing loops ...
matmul_no_compile: 4.105190992355347
matmul_compile: 4.124130725860596
matmul_rl: 0.31899595260620117
multi_dot: 0.31679391860961914
einsum: 0.32097911834716797
</code></pre>
<p>I guess the question is now: should this be something <code>torch.compile</code> be trying to support? I’m not familiar enough with compilation to know how difficult this would be, but I’m guessing the pattern of a chain of matmuls should be preserved down to the IR (?) and it should be possible to reorder the matmuls there (at least for simple patterns such as a linear chain).</p>

KFrank: <p>Hi Pea!</p>
<aside class="quote no-group" data-username="PeaBrane" data-post="3" data-topic="192112" data-full="true">
<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/peabrane/48/40023_2.png" class="avatar"> PeaBrane:</div>
<blockquote>
<p><code>torch.compile</code> seems to make no effort in reordering matmuls for efficiency.</p>
</blockquote>
</aside>
<p>I’m glad that you confirmed this on a more modem gpu (and hence a<br>
different compilation backend).</p>
<aside class="quote no-group">
<blockquote>
<p>I guess the question is now: should this be something <code>torch.compile</code> be trying to support?</p>
</blockquote>
</aside>
<p>That wouldn’t seem unreasonable to me.  I don’t think that proposing<br>
this as a feature request on github would be poorly received.</p>
<aside class="quote no-group">
<blockquote>
<p>a chain of matmuls should be preserved down to the IR (?) and it should be possible to reorder the matmuls there</p>
</blockquote>
</aside>
<p>I don’t know anything about the internals of <code>torch.compile</code>, but this all<br>
sounds plausible to me.  The logic is straightforward enough (witness<br>
<code>multi_dot()</code>).  I just don’t know whether <code>torch.compile</code> has hooks in<br>
the right places to make it easy.</p>
<p>I don’t think, however, that chains of matrix multiplications (without<br>
intervening non-linearities or additive biases) are that common of<br>
a use case in pytorch, so implementing something like this – even<br>
if relatively easy – might not be a priority.</p>
<p>(If you do file a feature request, I would suggest including a couple of<br>
real-world use cases to motivate the feature.)</p>
<p>Best.</p>
<p>K. Frank</p>