Skip to content

sampled_addmm: BSR support #101163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 22 commits into from
Closed

sampled_addmm: BSR support #101163

wants to merge 22 commits into from

Conversation

nikitaved
Copy link
Collaborator

@nikitaved nikitaved commented May 11, 2023 •
edited
Loading

This PR implements a sampled_addmm kernel that works with a BSR mask.

Stack from ghstack (oldest at bottom):

@nikitaved
sampled_addmm: BSR support
c8e47bb 
[ghstack-poisoned]
@nikitaved nikitaved mentioned this pull request May 11, 2023
Closed
@pytorch-bot
Copy link

pytorch-bot bot commented May 11, 2023 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 8a25f86:
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added the release notes: sparse release notes category label May 11, 2023
@nikitaved nikitaved marked this pull request as draft May 11, 2023 07:42
@nikitaved nikitaved mentioned this pull request May 11, 2023
Closed
nikitaved added a commit that referenced this pull request May 11, 2023
@nikitaved
sampled_addmm: BSR support
01318de 
ghstack-source-id: ffa39d1
Pull Request resolved: #101163
nikitaved added a commit that referenced this pull request May 12, 2023
@nikitaved
sampled_addmm: BSR support
f6108d5 
ghstack-source-id: 3cdcc6f
Pull Request resolved: #101163
nikitaved added a commit that referenced this pull request May 15, 2023
@nikitaved
sampled_addmm: BSR support
cbf7427 
ghstack-source-id: de96103
Pull Request resolved: #101163
@nikitaved
Update on "sampled_addmm: BSR support"
f611578 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
This was referenced May 22, 2023
Closed
Closed
@nikitaved
Update on "sampled_addmm: BSR support"
e25a2cc 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
cpuhrsch
cpuhrsch reviewed May 23, 2023
View reviewed changes
test/test_sparse_csr.py Outdated
batches = [(), (2,), (2, 2)]
size = [128, 256, 0]

def sampled_addmm_ref(input, mat1, mat2, alpha, beta):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could another way to implement this be torch.addmm(input.to_dense(), mat1, mat2, alpha, beta).sparse_mask(input.to_sparse()).to_sparse_bsr(input.values().shape[-2:])?

Sure it's less efficient, but a simpler reference implementation could give us more confidence?

Copy link
Collaborator Author

@nikitaved nikitaved May 23, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We test against CSR from sparse.sampled_addmm for more confidence. The proposed solution also only works for non-batched inputs, so, for simplicity, we still need to loop over batches... Alternatively, we can remove the reference and just test against CSR with half promoted to float. Would that be more preffered?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, testing against CSR with float would work too. That's even simpler.

cpuhrsch
cpuhrsch requested changes May 23, 2023
View reviewed changes
Copy link
Contributor

@cpuhrsch cpuhrsch left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit worried about the complexity of the reference addmm implementation just so we can have a simple comparison point.

cpuhrsch
cpuhrsch reviewed May 23, 2023
View reviewed changes
test/test_sparse_csr.py Outdated
res_csr = torch.sparse.sampled_addmm(csr, mat1csr, mat2csr, alpha=alpha, beta=beta)
self.assertEqual(res_tri.to_dense(), res_csr.to_dense())

# Check grid consistency
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this mean?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will update the comment to clarify on that.

nikitaved added 2 commits May 24, 2023 08:29
@nikitaved
Update on "sampled_addmm: BSR support"
8fee341 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
@nikitaved
Update on "sampled_addmm: BSR support"
4bb5aee 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
@nikitaved nikitaved requested a review from cpuhrsch May 24, 2023 14:21
@nikitaved nikitaved mentioned this pull request May 24, 2023
Closed
@nikitaved
Update on "sampled_addmm: BSR support"
4f08ba3 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
@nikitaved nikitaved mentioned this pull request May 24, 2023
Closed
@cpuhrsch
Copy link
Contributor

Looks like there's one lint error left otherwise this is good to go. Thanks for writing this!

nikitaved added 2 commits May 25, 2023 08:31
@nikitaved
Update on "sampled_addmm: BSR support"
64307a6 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
@nikitaved
Update on "sampled_addmm: BSR support"
8a25f86 
This PR implements a `sampled_addmm` kernel that works with a BSR mask.




[ghstack-poisoned]
@nikitaved
Copy link
Collaborator Author

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

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

xwang233 referenced this pull request May 31, 2023
@nikitaved @pytorchmergebot
nn.Linear: dispatch to bsr_dense_mm for half and bfloat16 (#94825)
1adb6fa 
Pull Request resolved: #94825
Approved by: https://github.com/albanD, https://github.com/cpuhrsch
@xwang233 xwang233 mentioned this pull request May 31, 2023
Closed
pytorchmergebot pushed a commit that referenced this pull request Jun 1, 2023
@xwang233 @pytorchmergebot
Skip test test_triton_bsr_dense_bmm if not TEST_WITH_TORCHINDUCTOR [v…
6340aa5 
…2] (#102660)

Test was originally skipped in #98462

Not sure why it was removed in #94825

Now the test hits CUDA illegal memory access on H100 again after #101163

Pull Request resolved: #102660
Approved by: https://github.com/zou3519
@facebook-github-bot facebook-github-bot deleted the gh/nikitaved/45/head branch June 8, 2023 18:07
@crcrpar crcrpar mentioned this pull request Jun 28, 2023
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@amjames amjames amjames left review comments

@cpuhrsch cpuhrsch cpuhrsch approved these changes

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged open source release notes: sparse release notes category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@nikitaved @cpuhrsch @pytorchmergebot @amjames @pytorchbot