Skip to content
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

schedulePointwise cleanup: - computeAt + InlinePropagator #1815

Merged
merged 23 commits into from
Jul 14, 2022
Merged

schedulePointwise cleanup: - computeAt + InlinePropagator #1815

merged 23 commits into from
Jul 14, 2022

Conversation

zasdfgbnm
Copy link
Collaborator

@zasdfgbnm zasdfgbnm commented Jul 9, 2022 •
edited
Loading

This PR depends on #1811 and #1813, please merge them first.

@zasdfgbnm zasdfgbnm marked this pull request as ready for review July 9, 2022 22:12
csarofeen
csarofeen reviewed Jul 13, 2022
View reviewed changes
torch/csrc/jit/codegen/cuda/scheduler/pointwise.cpp Outdated

cached_output->computeAt(output, unswitch_pos, ComputeAtMode::BestEffort);
compute_to.push_back(cached_output);
inner_most_tensors.erase(cached_output);
Copy link
Owner

Choose a reason for hiding this comment

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

Can't you leave cached_output in for this next traversal? I think it's just the output we need to protect as it will have the vectorized dimension.

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 was thinking about the same, but if I did so there would be two failures:

[  FAILED  ] 2 tests, listed below:
[  FAILED  ] NVFuserTest.FusionParser_CUDA
[  FAILED  ] NVFuserTest.FusionChannelsLastParser_CUDA

so I reverted back. So seems that the previous behavior was excluding cached_output from getting a new inline position as well.

But these two failures are trivial to fix and should have no real changes in the generated kernels. I pushed a new commit 24397fb for this, please check if this makes sense to you.

Copy link
Owner

Choose a reason for hiding this comment

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

Yeah, small changes in the parser code are common. We just update them to make sure we haven't fundamentally changed the world on most PRs ;-)

csarofeen
csarofeen reviewed Jul 13, 2022
View reviewed changes
torch/csrc/jit/codegen/cuda/scheduler/pointwise.cpp
scheduler_utils::computeAtBetween(
compute_from, compute_to, -1, ComputeAtMode::BestEffort);
// Fix max producer position
MaxProducerPosUpdater updater;
Copy link
Owner

Choose a reason for hiding this comment

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

Should we call MaxProducerPosUpdater at the end of InlinePropagator::traverse?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

InlinePropagator does not have a traverse, MaxRootDomainInfoSpanningTree does, so probably no?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Probably we should move the functionality of MaxProducerPosUpdater into InlinePropagator::setCAPos? I can do that in a separate PR, but doing so will slow down compilation speed as we will be updating the max producer pos of a tensor multiple times.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Or... Is the max producer position only used in expr sort? If so, we can just leave it as unset when scheduling, and add a pass in the lowering for setting it.

Copy link
Owner

@csarofeen csarofeen Jul 14, 2022 •
edited
Loading

Choose a reason for hiding this comment

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

Max producer position is important for multiple transformation passes. To make sure that we aren't modifying tensors that are being relied on by other tensors. It's more of a rule to enforce as we're modifying the schedule it remains legal.

Copy link
Owner

Choose a reason for hiding this comment

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

I think adding it to the tail end of inlining pass is worth the extra cost, as it will help catch user errors in scheduling.

Copy link
Owner

Choose a reason for hiding this comment

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

Agree it can be done in a separate PR.

Base automatically changed from schedulepointwise-cleanup2 to devel July 13, 2022 22:32
csarofeen
csarofeen approved these changes Jul 14, 2022
View reviewed changes
Copy link
Owner

@csarofeen csarofeen left a comment

Choose a reason for hiding this comment

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

LGTM

@zasdfgbnm zasdfgbnm merged commit 4cae122 into devel Jul 14, 2022
@zasdfgbnm zasdfgbnm deleted the pointwise-sched-no-ca branch July 14, 2022 15:36
@zasdfgbnm zasdfgbnm mentioned this pull request Jul 14, 2022
Merged
shmsong pushed a commit to shmsong/pytorch that referenced this pull request Jul 24, 2022
@jjsjann123
[NVFuser] Upstream push 0714
f686d82 
Syncing nvfuser devel branch to upstream master. https://github.com/csarofeen/pytorch/

Code changes includes:

- codegen improvements:
  1. Indexing refactor -> Remove reference tensor in predicate indexing logic
  2. MMA Rfactor support for cross-warp and cross-CTA split on K dimension
  3. Grouping grid allreduces across iterations
  4. Swizzle op formulation for non-affine swizzles
  5. Use scheduler_utils to cache inputs and outputs in schedulePointwise
- scheduler refactor
  1. New compute at interface
- transformation propagation refactor on MaxInfoSpanningTree
  1. Added sibling path that is required to generate consistent replay for some cases where `MaxInfoSpanningTree` is used with a selector.
  2. Optimization to skip Transform propagator
  3. SpanningTreePrinter for debugging
- parser update
  1. Fixes `div`
  2. Added `_to_copy`
  3. Broadcast in dim with expand to support expanding to concrete size
  4. Dropout prob extremal patch
- executor patch on caching strides for output allocation

Squashed commits to WAR github API
Commits that's actually in this PR from the devel branch:

```
3b87896 Fix allocation of work buffers and `fused_reduction::ParallelReduce` with unswitch (csarofeen#1818)
4cae122 schedulePointwise cleanup: - computeAt + InlinePropagator (csarofeen#1815)
3df9742 Use scheduler_utils to cache inputs and outputs in schedulePointwise (csarofeen#1811)
03180aa improve broadcast resolution (csarofeen#1792)
bee6c69 bug fix (csarofeen#1819)
4413c8f Support PYTORCH_NVFUSER_DUMP=transform_propagator (csarofeen#1812)
de6b7ca Fix negative position in InlinePropagator (csarofeen#1813)
10a996c Remove redundant check in schedulePointwise (csarofeen#1810)
acd5ed4 Swizzle op formulation for non-affine swizzles (csarofeen#1441)
3ed8330 Kernel args patch to show zero_init buffer (csarofeen#1809)
037a75a Dropout prob extremal patch (csarofeen#1804)
282c429 spam nvrtc options (csarofeen#1783)
3ba6a5f Broadcast in dim with expand (csarofeen#1794)
fd4be12 remove dead indexing code (csarofeen#1806)
fa4e6a4 Check siblings in getMaxPosAll (csarofeen#1805)
025c840 Grouping grid allreduces across iterations (csarofeen#1755)
37c579e Temporarily disable test requring large shared memory. (csarofeen#1802)
5f375d0 More cleanup on InlinePropagator (csarofeen#1800)
8d384da Indexing refactor stage 2 : Remove reference tensor in predicate indexing logic (csarofeen#1784)
f008140 MMA Rfactor support for cross-warp and cross-CTA split on K dimension (csarofeen#1554)
76b3cca Add parsing support for `_to_copy` to handle AMP casts. (csarofeen#1756)
ef04f6c Coding style cleanups (csarofeen#1798)
38c7f3c InlinePropagator please don't replay (csarofeen#1797)
3f2c263 validateDomain in TransformPropagator (csarofeen#1796)
c077085 Use TransformPropagatorWithCheck in many tests (csarofeen#1795)
d0d0908 Some further cleanup for the new computeAt interface (csarofeen#1793)
45f5203 Fix TransformReplay::getMatchedLeafPosWithoutReplay* (csarofeen#1791)
28cbaf9 New compute at interface (csarofeen#1743)
635ebfc Add SpanningTreePrinter (csarofeen#1786)
59f3c32 Output allocate patch (csarofeen#1790)
fe93bf5 Transform propagator skip replay when possible (csarofeen#1782)
ebf23a5 Fix isIntegralType error msg (csarofeen#1789)
0c82ecf Disable register reuse across serial broadcast ops (csarofeen#1787)
33a824d Adding sibling path for MaxInfoSpanningTree (csarofeen#1776)
86f46aa Fix div(Val, TensorView) (csarofeen#1778)
d3de227 Fix FusionMaxRootDomainInfoSpanningTreePrintTwice_CUDA (csarofeen#1781)
ecc7a87 Extend mma dimension and layout checking to support strided batched matmul and tensor contractions (csarofeen#1761)
```

[ghstack-poisoned]
shmsong pushed a commit to shmsong/pytorch that referenced this pull request Jul 24, 2022
@jjsjann123
Update on "[NVFuser] Upstream push 0714"
d720403 
Syncing nvfuser devel branch to upstream master. https://github.com/csarofeen/pytorch/

Code changes includes:

- codegen improvements:
  1. Indexing refactor -> Remove reference tensor in predicate indexing logic
  2. MMA Rfactor support for cross-warp and cross-CTA split on K dimension
  3. Grouping grid allreduces across iterations
  4. Swizzle op formulation for non-affine swizzles
  5. Use scheduler_utils to cache inputs and outputs in schedulePointwise
- scheduler refactor
  1. New compute at interface
- transformation propagation refactor on MaxInfoSpanningTree
  1. Added sibling path that is required to generate consistent replay for some cases where `MaxInfoSpanningTree` is used with a selector.
  2. Optimization to skip Transform propagator
  3. SpanningTreePrinter for debugging
- parser update
  1. Fixes `div`
  2. Added `_to_copy`
  3. Broadcast in dim with expand to support expanding to concrete size
  4. Dropout prob extremal patch
- executor patch on caching strides for output allocation

Squashed commits to WAR github API
Commits that's actually in this PR from the devel branch:

```
3b87896 Fix allocation of work buffers and `fused_reduction::ParallelReduce` with unswitch (csarofeen#1818)
4cae122 schedulePointwise cleanup: - computeAt + InlinePropagator (csarofeen#1815)
3df9742 Use scheduler_utils to cache inputs and outputs in schedulePointwise (csarofeen#1811)
03180aa improve broadcast resolution (csarofeen#1792)
bee6c69 bug fix (csarofeen#1819)
4413c8f Support PYTORCH_NVFUSER_DUMP=transform_propagator (csarofeen#1812)
de6b7ca Fix negative position in InlinePropagator (csarofeen#1813)
10a996c Remove redundant check in schedulePointwise (csarofeen#1810)
acd5ed4 Swizzle op formulation for non-affine swizzles (csarofeen#1441)
3ed8330 Kernel args patch to show zero_init buffer (csarofeen#1809)
037a75a Dropout prob extremal patch (csarofeen#1804)
282c429 spam nvrtc options (csarofeen#1783)
3ba6a5f Broadcast in dim with expand (csarofeen#1794)
fd4be12 remove dead indexing code (csarofeen#1806)
fa4e6a4 Check siblings in getMaxPosAll (csarofeen#1805)
025c840 Grouping grid allreduces across iterations (csarofeen#1755)
37c579e Temporarily disable test requring large shared memory. (csarofeen#1802)
5f375d0 More cleanup on InlinePropagator (csarofeen#1800)
8d384da Indexing refactor stage 2 : Remove reference tensor in predicate indexing logic (csarofeen#1784)
f008140 MMA Rfactor support for cross-warp and cross-CTA split on K dimension (csarofeen#1554)
76b3cca Add parsing support for `_to_copy` to handle AMP casts. (csarofeen#1756)
ef04f6c Coding style cleanups (csarofeen#1798)
38c7f3c InlinePropagator please don't replay (csarofeen#1797)
3f2c263 validateDomain in TransformPropagator (csarofeen#1796)
c077085 Use TransformPropagatorWithCheck in many tests (csarofeen#1795)
d0d0908 Some further cleanup for the new computeAt interface (csarofeen#1793)
45f5203 Fix TransformReplay::getMatchedLeafPosWithoutReplay* (csarofeen#1791)
28cbaf9 New compute at interface (csarofeen#1743)
635ebfc Add SpanningTreePrinter (csarofeen#1786)
59f3c32 Output allocate patch (csarofeen#1790)
fe93bf5 Transform propagator skip replay when possible (csarofeen#1782)
ebf23a5 Fix isIntegralType error msg (csarofeen#1789)
0c82ecf Disable register reuse across serial broadcast ops (csarofeen#1787)
33a824d Adding sibling path for MaxInfoSpanningTree (csarofeen#1776)
86f46aa Fix div(Val, TensorView) (csarofeen#1778)
d3de227 Fix FusionMaxRootDomainInfoSpanningTreePrintTwice_CUDA (csarofeen#1781)
ecc7a87 Extend mma dimension and layout checking to support strided batched matmul and tensor contractions (csarofeen#1761)
```

RUN_TORCHBENCH: nvfuser

Differential Revision: [D38043938](https://our.internmc.facebook.com/intern/diff/D38043938)

[ghstack-poisoned]
csarofeen pushed a commit that referenced this pull request Aug 4, 2022
@jjsjann123 @pytorchmergebot
[NVFuser] Upstream push 0714 (pytorch#81861)
4a000ff 
Syncing nvfuser devel branch to upstream master. https://github.com/csarofeen/pytorch/

Code changes includes:

- codegen improvements:
  1. Indexing refactor -> Remove reference tensor in predicate indexing logic
  2. MMA Rfactor support for cross-warp and cross-CTA split on K dimension
  3. Grouping grid allreduces across iterations
  4. Swizzle op formulation for non-affine swizzles
  5. Use scheduler_utils to cache inputs and outputs in schedulePointwise
- scheduler refactor
  1. New compute at interface
- transformation propagation refactor on MaxInfoSpanningTree
  1. Added sibling path that is required to generate consistent replay for some cases where `MaxInfoSpanningTree` is used with a selector.
  2. Optimization to skip Transform propagator
  3. SpanningTreePrinter for debugging
- parser update
  1. Fixes `div`
  2. Added `_to_copy`
  3. Broadcast in dim with expand to support expanding to concrete size
  4. Dropout prob extremal patch
- executor patch on caching strides for output allocation

Squashed commits to WAR github API
Commits that's actually in this PR from the devel branch:

```
3b87896 Fix allocation of work buffers and `fused_reduction::ParallelReduce` with unswitch (#1818)
4cae122 schedulePointwise cleanup: - computeAt + InlinePropagator (#1815)
3df9742 Use scheduler_utils to cache inputs and outputs in schedulePointwise (#1811)
03180aa improve broadcast resolution (#1792)
bee6c69 bug fix (#1819)
4413c8f Support PYTORCH_NVFUSER_DUMP=transform_propagator (#1812)
de6b7ca Fix negative position in InlinePropagator (#1813)
10a996c Remove redundant check in schedulePointwise (#1810)
acd5ed4 Swizzle op formulation for non-affine swizzles (#1441)
3ed8330 Kernel args patch to show zero_init buffer (#1809)
037a75a Dropout prob extremal patch (#1804)
282c429 spam nvrtc options (#1783)
3ba6a5f Broadcast in dim with expand (#1794)
fd4be12 remove dead indexing code (#1806)
fa4e6a4 Check siblings in getMaxPosAll (#1805)
025c840 Grouping grid allreduces across iterations (#1755)
37c579e Temporarily disable test requring large shared memory. (#1802)
5f375d0 More cleanup on InlinePropagator (#1800)
8d384da Indexing refactor stage 2 : Remove reference tensor in predicate indexing logic (#1784)
f008140 MMA Rfactor support for cross-warp and cross-CTA split on K dimension (#1554)
76b3cca Add parsing support for `_to_copy` to handle AMP casts. (#1756)
ef04f6c Coding style cleanups (#1798)
38c7f3c InlinePropagator please don't replay (#1797)
3f2c263 validateDomain in TransformPropagator (#1796)
c077085 Use TransformPropagatorWithCheck in many tests (#1795)
d0d0908 Some further cleanup for the new computeAt interface (#1793)
45f5203 Fix TransformReplay::getMatchedLeafPosWithoutReplay* (#1791)
28cbaf9 New compute at interface (#1743)
635ebfc Add SpanningTreePrinter (#1786)
59f3c32 Output allocate patch (#1790)
fe93bf5 Transform propagator skip replay when possible (#1782)
ebf23a5 Fix isIntegralType error msg (#1789)
0c82ecf Disable register reuse across serial broadcast ops (#1787)
33a824d Adding sibling path for MaxInfoSpanningTree (#1776)
86f46aa Fix div(Val, TensorView) (#1778)
d3de227 Fix FusionMaxRootDomainInfoSpanningTreePrintTwice_CUDA (#1781)
ecc7a87 Extend mma dimension and layout checking to support strided batched matmul and tensor contractions (#1761)
```

RUN_TORCHBENCH: nvfuser

Differential Revision: [D38043938](https://our.internmc.facebook.com/intern/diff/D38043938)
Pull Request resolved: pytorch#81861
Approved by: https://github.com/davidberard98
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@csarofeen csarofeen csarofeen approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@zasdfgbnm @csarofeen