Skip to content

Prepare python jiterator for multiple outputs#77921

Closed
SherlockNoMad wants to merge 1 commit intomasterfrom
bahuang/jiterator_mulioutput2
Closed

Prepare python jiterator for multiple outputs#77921
SherlockNoMad wants to merge 1 commit intomasterfrom
bahuang/jiterator_mulioutput2

Conversation

@SherlockNoMad
Copy link
Contributor

@SherlockNoMad SherlockNoMad commented May 20, 2022

Part 2 for #77902

Further prepare python jiterator for supporting multiple outputs

@SherlockNoMad SherlockNoMad requested review from mruberry and ngimel May 20, 2022 00:39
@facebook-github-bot
Copy link
Contributor

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

🔗 Helpful links

✅ No Failures (0 Pending)

As of commit f4f09ac (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.

@SherlockNoMad SherlockNoMad force-pushed the bahuang/jiterator_mulioutput2 branch from 4e0607a to f4f09ac Compare May 20, 2022 05:56
ngimel
ngimel approved these changes May 21, 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.

Looks good, if the tests pass

aten/src/ATen/cuda/jiterator_impl.h

template<int N>
static std::unique_ptr<OffsetCalculator<N>> make_unique_input_offset_calculator(const TensorIteratorBase& iter) {
template<bool IS_INPUT, int N>
Copy link
Collaborator

@ngimel ngimel May 21, 2022

Choose a reason for hiding this comment

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

doesn't look like IS_INPUT needs to be template arg?

Copy link
Contributor Author

@SherlockNoMad SherlockNoMad May 21, 2022 •
edited
Loading

Choose a reason for hiding this comment

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

IIUC, compiler will try to optimize away the dead code path for such template arg.

so the switch is done at compilation time, rather than at run time.

@SherlockNoMad
Copy link
Contributor Author

SherlockNoMad commented May 21, 2022

@pytorchbot merge this

@github-actions
Copy link
Contributor

github-actions bot commented May 21, 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.

@SherlockNoMad SherlockNoMad mentioned this pull request May 24, 2022
Closed
facebook-github-bot pushed a commit that referenced this pull request May 24, 2022
@SherlockNoMad @facebook-github-bot
Prepare python jiterator for multiple outputs (#77921) (#77921)
f336cf9 
Summary:
Part 2 for #77902

Further prepare python jiterator for supporting multiple outputs

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

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

Reviewed By: seemethere

Differential Revision: D36603024

Pulled By: seemethere

fbshipit-source-id: 190e409f0cb68d675da3b676a6fa658c0e32c107
pytorchmergebot pushed a commit that referenced this pull request May 24, 2022
@SherlockNoMad @pytorchmergebot
Python Jiterator supports multiple outputs (#78139)
6db8440 
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
swang392 pushed a commit that referenced this pull request May 25, 2022
@SherlockNoMad @swang392
Prepare python jiterator for multiple outputs (#77921)
5cafbba 
Part 2 for #77902

Further prepare python jiterator for supporting multiple outputs
Pull Request resolved: #77921
Approved by: https://github.com/ngimel
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_mulioutput2 branch February 16, 2024 01:56
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