[ROCm] Fix largeIndexBlockSize

[pragupta]

On ROCm, hipification converts std::min to ::min, but ::min is not returning the right result. This impacts index_add_ operation on a large tensor, we end up picking the large values instead of max supported block size (128). This leads to GPU accessing memory out of bounds.

While we wait for ::min to be fixed, we can use < operator to compare instead of relying on ::min.

Example Code w/ failure:

D=6144
hidden_states = torch.zeros([16384, 6144],           device="cuda:0", dtype=torch.bfloat16)
index         = torch.randint(0, 16384, (1, 32, 16384), device="cuda:0", dtype=torch.int64)
output        = torch.empty([1, 32, 16384, 6144],    device="cuda:0", dtype=torch.bfloat16)
hidden_states.index_add_(0, index.view(-1), output.view(-1, D))

Traceback (most recent call last):
RuntimeError: HIP error: invalid configuration argument

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd

[pytorch-bot]

🔗 Helpful Links

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

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

• [DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit bd6b78d with merge base 6bdbc86 ():
💚 Looks good so far! There are no failures yet. 💚

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

[pragupta]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased pg-msft-fix onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout pg-msft-fix && git pull --rebase)

[eqy]

If there is a reproducer snippet, add a test case please

[pragupta]

If there is a reproducer snippet, add a test case please

@eqy -- I've added the reproducer code and the error generated with it in the PR description.

Example Code w/ failure:

D=6144
hidden_states = torch.zeros([16384, 6144],           device="cuda:0", dtype=torch.bfloat16)
index         = torch.randint(0, 16384, (1, 32, 16384), device="cuda:0", dtype=torch.int64)
output        = torch.empty([1, 32, 16384, 6144],    device="cuda:0", dtype=torch.bfloat16)
hidden_states.index_add_(0, index.view(-1), output.view(-1, D))

Traceback (most recent call last):
RuntimeError: HIP error: invalid configuration argument

[pragupta]

If there is a reproducer snippet, add a test case please

@eqy -- I misunderstood earlier. Added UT for this edge case. Thanks

[pragupta]

@eq -- if the UT looks good now, can you please provide your review on this one?

[pruthvistony]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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