🔗 Helpful Links

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

• 📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ghstack-mergeability-check and Check labels failing with 'Resource not accessible by integration'

✅ No Failures

As of commit 70792b5 with merge base 1465757 ():
💚 Looks good so far! There are no failures yet. 💚

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

Reproducer:

import time
import torch

shapes = [(1, 2, 3, 420, 648, 128),
    (1, 2, 3, 420, 648, 128),
    (5079670, 128)
]

dims = [(3,4),
    (-3, -2), 
    (1)
]

for i, shape in enumerate(shapes):
    x = torch.randn(shape, device='cuda', dtype=torch.float)
    for _ in range(20):
        _ = torch.sum(x, dims[i])
    torch.cuda.synchronize()

    start_time = time.time()
    for _ in range(100):
        _ = torch.sum(x, dims[i])
    torch.cuda.synchronize()
    end_time = time.time()
    mean_time = (end_time - start_time)/100
    print(f"Avg time for shape {shape}: {mean_time * 1e6:.2f} us")

Before
Avg time for shape (1, 2, 3, 420, 648, 128): 4408.10 us
Avg time for shape (1, 2, 3, 420, 648, 128): 4428.89 us
Avg time for shape (5079670, 128): 1458.86 us

After:
Avg time for shape (1, 2, 3, 420, 648, 128): 223.73 us
Avg time for shape (1, 2, 3, 420, 648, 128): 218.85 us
Avg time for shape (5079670, 128): 1461.55 us

To add the ciflow label ciflow/periodic-rocm-mi300 please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

Question : Was the choice of 2048 as the threshold for "values per thread" purely heuristic? It would be helpful to add a comment or reference explaining why this value was chosen and whether it is empirically optimal.

It was indeed empirically determined. I'll add a comment.

Another question : Is there an upper bound for config.ctas_per_output *= 2;

From the previous semantics there doesn't seem to be the case.