PyTorch example

[msaroufim]

Description

This PR provides an end to end example of integrating PyTorch with cuda-python. In particular this PR

1. Ensures cuda-python reuses the same CUDA stream as PyTorch similarly to what Triton does. This logic is likely useful enough to encapsulate into its own helper on the core Stream object in cuda-python
2. Shows how to pass in pointers to torch tensors to a kernel

Limitation

1. scalars needs to be encapsulated in a torch.Tensor otherwise I was running into correctness issues

Example of running this locally

(cuda-p) [marksaroufim@devvm3083]~/cuda-python/cuda_core/examples% python pytorch_example.py PyTorch stream: <torch.cuda.Stream device=cuda:0 cuda_stream=0x0>
Using PyTorch stream ID: 0
Successfully created CUDA Python stream from PyTorch stream
Single precision test passed!
Double precision test passed!
All tests passed successfully!

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[leofang]

Many thanks to your contributions, @msaroufim! 🙏 Please give us a bit of time to unblock you -- This repository, including cuda.core, originally was not allowed to accept external contributions. However, we've received legal approval to make changes for cuda.core, and #583 is the last step to open up the gate. Once we finalize the license change we can review/merge your PRs!

[msaroufim]

No worries @leofang I understand, great repo! So looking forward to contributing more soon!

[msaroufim]

@leofang is this ready for a review? FWIW this is the PR I'm most interested in of the 3 I sent

[kkraus14]

@msaroufim apologies for the delay in reviewing this.

The examples are run as part of the test suite so I think we need to add pytorch to the test dependencies: the requirements-cuXX.txt here: https://github.com/NVIDIA/cuda-python/tree/main/cuda_core/tests

And then also modify the test utilities to properly skip if torch isn't installed in the environment:

cuda-python/cuda_core/tests/example_tests/utils.py

Line 40 in 9e868ba

formin ("cupy",):

[msaroufim]

Thanks @kkraus14 mind kicking off a test run for me please?

[msaroufim]

/ok to test 1378e6a

[kkraus14]

/ok to test 1378e6a

The testing runs on infrastructure that requires an NVIDIA maintainer to trigger.

[kkraus14]

/ok to test 1378e6a

[kkraus14]

/ok to test 01ad808

[leofang]

Great addition, thanks @msaroufim! I'll add another StridedMemoryView based example later (after #150 is done) and it will also be able to handle PyTorch tensors too, without users fetching the pointer as kernel arg. Will ping you for review.

[msaroufim]

Regarding the testing changes @leofang please feel free to revert depending on what makes sense to you, I just tested your changes now locally and everything worked well

[leofang]

/ok to test 7d3582f

[github-actions]

Doc Preview CI
Preview removed because the pull request was closed or merged.