Skip to content

Python Jiterator supports multiple outputs#78139

Closed
SherlockNoMad wants to merge 6 commits intomasterfrom
bahuang/jiterator_multioutput3
Closed

Python Jiterator supports multiple outputs#78139
SherlockNoMad wants to merge 6 commits intomasterfrom
bahuang/jiterator_multioutput3

Conversation

@SherlockNoMad
Copy link
Contributor

@SherlockNoMad SherlockNoMad commented May 24, 2022

This PR is part3.
Part1: #77902
Part2: #77921

Python Jiterator now supports returning multiple outputs

fn = torch.cuda.jiterator._create_multi_output_jit_fn(
"""
template <typename T>
T binary_2outputs(T i0, T i1, T& out0, T& out1) {
    out0 = i0 + i1;
    out1 = i0 - i1;
}
""",
num_outputs=2)

x = torch.rand(3, device='cuda')
y = torch.rand(3, device='cuda')
out0, out1 = fn(x, y) 

torch.allclose(out0, x+y)
torch.allclose(out1, x-y)

@facebook-github-bot
Copy link
Contributor

facebook-github-bot commented May 24, 2022 •
edited
Loading

🔗 Helpful links

✅ No Failures (0 Pending)

As of commit d4a9002 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

ngimel
ngimel approved these changes May 24, 2022
View reviewed changes
Copy link
Collaborator

@ngimel ngimel left a comment

Choose a reason for hiding this comment

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

This is great!

@SherlockNoMad SherlockNoMad force-pushed the bahuang/jiterator_multioutput3 branch from 66f00b3 to b1e8fad Compare May 24, 2022 04:33
@SherlockNoMad SherlockNoMad force-pushed the bahuang/jiterator_multioutput3 branch from b1e8fad to 1edbfe9 Compare May 24, 2022 18:39
@SherlockNoMad SherlockNoMad force-pushed the bahuang/jiterator_multioutput3 branch from 1edbfe9 to d4a9002 Compare May 24, 2022 18:49
@SherlockNoMad
Copy link
Contributor Author

SherlockNoMad commented May 24, 2022

@pytorchbot merge this.

@github-actions
Copy link
Contributor

github-actions bot commented May 24, 2022

Hey @SherlockNoMad.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

facebook-github-bot pushed a commit that referenced this pull request May 26, 2022
@SherlockNoMad @facebook-github-bot
Python Jiterator supports multiple outputs (#78139) (#78139)
93cf393 
Summary:
This PR is part3.
Part1: #77902
Part2: #77921

Python Jiterator now supports returning multiple outputs

```
fn = torch.cuda.jiterator._create_multi_output_jit_fn(
"""
template <typename T>
T binary_2outputs(T i0, T i1, T& out0, T& out1) {
    out0 = i0 + i1;
    out1 = i0 - i1;
}
""",
num_outputs=2)

x = torch.rand(3, device='cuda')
y = torch.rand(3, device='cuda')
out0, out1 = fn(x, y)

torch.allclose(out0, x+y)
torch.allclose(out1, x-y)
```

Pull Request resolved: #78139
Approved by: https://github.com/ngimel

Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/6db8440f35831a903a1084aec1a3803af0e1dd10

Reviewed By: mehtanirav

Differential Revision: D36668751

Pulled By: SherlockNoMad

fbshipit-source-id: 756ad6da1682030d6db8cad390a3a72ef3b44444
@github-actions github-actions bot deleted the bahuang/jiterator_multioutput3 branch February 16, 2024 01:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ngimel ngimel ngimel approved these changes

@mruberry mruberry Awaiting requested review from mruberry

Assignees

No one assigned

Labels

cla signed Merged module: jiterator

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@SherlockNoMad @facebook-github-bot @ngimel @pytorchmergebot