[inductor] make thread order consistent with loop order

[shunting314]

Stack from ghstack (oldest at bottom):

• -> [inductor] make thread order consistent with loop order #106827

I found that for a tiled kernel for tensor with shape [a, b], we map 'a' with XBLOCK and 'b' with YBLOCK. However, 'a' actually should be the outer looper while 'b' corresponding to the inner loop. This order is picked by our loop ordering algorithm. Mapping 'a' with XBLOCK has the semantic like assigning 'a' to the inner loop instead.

For a simple 'A + B.t()' kernel, making the loop order consistent can brings 1.027x speedup ( 1.938ms -> 1.887ms speedup) . Here are the dump of kernels:

• before fix: https://gist.github.com/shunting314/4dacf73cf495cdd7e84dede7c3e0872d
• after fix (this one is done manually): https://gist.github.com/shunting314/441e8839d24e1878c313e539b1ebd551

I tried this on DistillGPT2 and found perf is neutral. But that because DistillGPT2 has a single tiled pointwise kernel in it's backward graph. Will check the dashboard.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106827

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 88a4ea1:

NEW FAILURES - The following jobs have failed:

• linux-bionic-py3.8-clang9 / test (dynamo, 2, 2, linux.2xlarge) (gh)
• linux-focal-py3.8-gcc7 / test (default, 1, 3, linux.2xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[shunting314]

Perf run: link

So looks like the PR helps with some models but also slow down others. My guess is some kernel's loop order makes the behavior without the fix running faster. Also I'm thinking maybe we can update the existing algorithm that picks loop order a bit. The current algorithm will do the following

1. if all strides for dim-a is smaller than all strides for dim-b, we put dim-a later in the order ( this basically describes the order reversing return value of scheduler.pick_loop_order since that's what we usually do for the return value).
2. if for some indices the strides of dim-a is smaller than dim-b, but for other indices, it's the opposite, we put the dimension with smaller index first in the order.

Maybe we can change heuristic 2 to consider how many indices that dim-a has smaller strides. i.e. define X as the number of indices that dim-a has smaller strides than dim-b and Y as the number of indices that dim-b has smaller strides. If X - Y is larger than a threshold, than we decide to put dim-a later in the order.

[shunting314]

Here is the latest perf run .

• TIMM: 1.65x -> 1.69x
• HF: neutral
• TB: 1.75x -> 1.74x

So overall I think this PR is helpful especially on TIMM models.

[shunting314]

@pytorchbot merge -f "keyboard interrupt and test failure on functorch are unrelated"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[shunting314]

I just want to share more why this PR is helpful with the perf. Hopefully the deeper understanding can help me find future optimization opportunities.

Here is a micro-benchmark:

import torch
from torch._inductor import config
from torch._inductor.utils import do_bench

config.benchmark_kernel = True
config.triton.unique_kernel_names = True

def f(*inp):
    x = inp[0]
    for a in inp[1:]:
        x = x + a.permute(0, 2, 1)

    return x.sum(dim=0)


K = 4
inp = [torch.randn(512, 512, 512).cuda() for _ in range(K)]
opt_f = torch.compile(f)
opt_f(*inp)
print(do_bench(lambda: opt_f(*inp)))

The kernel does an outer reduction. Inside the reduction loop, all but one of the memory access prefer dim-1 as the inner most pointwise loop. The PR just make sure we pick dim-1 rather than dim-2 as the innermost pointwise loop in this case. This turns more memory access to coalesced access (with the help of tiling).

Metric and generated wrapper:

• without the PR: 17.46ms, https://gist.github.com/shunting314/a788db61dfe92c0e41c0b1761f39b7d2
• with the PR: 4.28ms, https://gist.github.com/shunting314/3d8a0936fe9c65b86d1f09a25a9cbb25

cc @jansel @Chillee @eellison