Skip to content

[Helion + torch.compile] Fix MultiOutput write deps and extend fusion score matching#177302

Closed
yf225 wants to merge 5 commits intogh/yf225/138/basefrom
gh/yf225/138/head
Closed

[Helion + torch.compile] Fix MultiOutput write deps and extend fusion score matching#177302
yf225 wants to merge 5 commits intogh/yf225/138/basefrom
gh/yf225/138/head

Conversation

@yf225
Copy link
Copy Markdown
Contributor

@yf225 yf225 commented Mar 12, 2026

Stack from ghstack (oldest at bottom):

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel. This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent. Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

This PR also now contains the changes from #177359, with proper fixes to avoid breaking internal tests.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

…ent fusion matching

Normalize using the same policy as SchedulerNode so that the
index expressions are directly comparable during fusion checks.

[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 12, 2026

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit ef99f9a with merge base e0818d6 (image):

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

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

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 12, 2026

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@yf225 yf225 added the topic: not user facing topic category label Mar 12, 2026
Copy link
Copy Markdown
Contributor

@shunting314 shunting314 left a comment

Choose a reason for hiding this comment

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

Would be good to have a test for the discovered failure

yf225 added 2 commits March 12, 2026 14:00
…for consistent fusion matching"

Normalize using the same policy as SchedulerNode so that the
index expressions are directly comparable during fusion checks.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo

[ghstack-poisoned]
…for consistent fusion matching"

Normalize using the same policy as SchedulerNode so that the
index expressions are directly comparable during fusion checks.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo

[ghstack-poisoned]
@yf225 yf225 changed the title [Helion + torch.compile] Normalize MultiOutput write deps for consistent fusion matching [Helion + torch.compile] Fix MultiOutput write deps and extend fusion score matching Mar 13, 2026
…tend fusion score matching"

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

[ghstack-poisoned]
yf225 added a commit that referenced this pull request Mar 13, 2026
… score matching

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

ghstack-source-id: 7b40701
Pull Request resolved: #177302
…tend fusion score matching"

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

[ghstack-poisoned]
@yf225
Copy link
Copy Markdown
Contributor Author

yf225 commented Mar 13, 2026

@pytorchbot merge -f "unrelated failures"

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

pytorchmergebot pushed a commit that referenced this pull request Mar 13, 2026
…ble base class (#177367)

This is a reland of #177063.

Move common fields and methods from TritonTemplateBuffer up to
TemplateBuffer so that external template backends (e.g. Helion) can
reuse the same mutation-tracking and prologue-fusion infrastructure:

- Add mutated_inputs, allowed_prologue_inps params to TemplateBuffer.__init__
- Build mutation_outputs list in base class (parallel to ExternKernel.mutation_outputs)
- Move get_allowed_prologue_inps() to base class
- Extract _read_deps_from_inputs() helper from extract_read_writes()
- Remove can_fuse_multi_output_epilogue() (always returned False, unused)
- Simplify TritonTemplateBuffer.__init__() to delegate to super()

get_outputs() stays on TritonTemplateBuffer since it is the only
subclass that currently passes mutated_inputs; other subclasses
(CppTemplateBuffer, CuteDSLTemplateBuffer, etc.) manage their own
output lists independently.

Pull Request resolved: #177367
Approved by: https://github.com/shunting314
ghstack dependencies: #177302
yf225 added a commit to yf225/pytorch that referenced this pull request Mar 15, 2026
… score matching

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

ghstack-source-id: 4478c51
Pull Request resolved: pytorch#177302
AaronWang04 pushed a commit to AaronWang04/pytorch that referenced this pull request Mar 24, 2026
… score matching (pytorch#177302)

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

This PR also now contains the changes from pytorch#177359, with proper fixes to avoid breaking internal tests.
Pull Request resolved: pytorch#177302
Approved by: https://github.com/shunting314, https://github.com/jansel
AaronWang04 pushed a commit to AaronWang04/pytorch that referenced this pull request Mar 24, 2026
…ble base class (pytorch#177367)

This is a reland of pytorch#177063.

Move common fields and methods from TritonTemplateBuffer up to
TemplateBuffer so that external template backends (e.g. Helion) can
reuse the same mutation-tracking and prologue-fusion infrastructure:

- Add mutated_inputs, allowed_prologue_inps params to TemplateBuffer.__init__
- Build mutation_outputs list in base class (parallel to ExternKernel.mutation_outputs)
- Move get_allowed_prologue_inps() to base class
- Extract _read_deps_from_inputs() helper from extract_read_writes()
- Remove can_fuse_multi_output_epilogue() (always returned False, unused)
- Simplify TritonTemplateBuffer.__init__() to delegate to super()

get_outputs() stays on TritonTemplateBuffer since it is the only
subclass that currently passes mutated_inputs; other subclasses
(CppTemplateBuffer, CuteDSLTemplateBuffer, etc.) manage their own
output lists independently.

Pull Request resolved: pytorch#177367
Approved by: https://github.com/shunting314
ghstack dependencies: pytorch#177302
EmanueleCoradin pushed a commit to EmanueleCoradin/pytorch that referenced this pull request Mar 30, 2026
… score matching (pytorch#177302)

Give MultiOutput proper MemoryDep writes derived from its own FixedLayout
instead of inheriting StarDep from InputsKernel.  This removes the hack
in FusedSchedulerNode.fuse() that copied index expressions from the
template parent.  Extend score_fusion_memory to use name-based dep
matching for templates so that views/reshapes between template outputs
and epilogues do not block fusion.

This PR also now contains the changes from pytorch#177359, with proper fixes to avoid breaking internal tests.
Pull Request resolved: pytorch#177302
Approved by: https://github.com/shunting314, https://github.com/jansel
EmanueleCoradin pushed a commit to EmanueleCoradin/pytorch that referenced this pull request Mar 30, 2026
…ble base class (pytorch#177367)

This is a reland of pytorch#177063.

Move common fields and methods from TritonTemplateBuffer up to
TemplateBuffer so that external template backends (e.g. Helion) can
reuse the same mutation-tracking and prologue-fusion infrastructure:

- Add mutated_inputs, allowed_prologue_inps params to TemplateBuffer.__init__
- Build mutation_outputs list in base class (parallel to ExternKernel.mutation_outputs)
- Move get_allowed_prologue_inps() to base class
- Extract _read_deps_from_inputs() helper from extract_read_writes()
- Remove can_fuse_multi_output_epilogue() (always returned False, unused)
- Simplify TritonTemplateBuffer.__init__() to delegate to super()

get_outputs() stays on TritonTemplateBuffer since it is the only
subclass that currently passes mutated_inputs; other subclasses
(CppTemplateBuffer, CuteDSLTemplateBuffer, etc.) manage their own
output lists independently.

Pull Request resolved: pytorch#177367
Approved by: https://github.com/shunting314
ghstack dependencies: pytorch#177302
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants