Skip to content

[Helion + torch.compile] Add prologue/epilogue fusion codegen to HelionTemplateBuffer#1725

Closed
yf225 wants to merge 1 commit intoyf225/stack/71from
yf225/stack/72
Closed

[Helion + torch.compile] Add prologue/epilogue fusion codegen to HelionTemplateBuffer#1725
yf225 wants to merge 1 commit intoyf225/stack/71from
yf225/stack/72

Conversation

@yf225
Copy link
Copy Markdown
Contributor

@yf225 yf225 commented Mar 16, 2026 •
edited
Loading

Stacked PRs:

[Helion + torch.compile] Add prologue/epilogue fusion codegen to HelionTemplateBuffer

Add fusion code generation to template_buffer.py:

  • _render_with_hooks: setup fusion hooks, read fusion metadata, build
    extra_params, pass store/load transforms conditionally
  • _build_call_args: add prologue_source_buffers and extra_params support
  • _generate_triton_ast: pass store/load transform and extra_params
  • create(): add on_tensor_leaf/on_non_tensor_leaf callback params
  • _codegen_epilogue_fusion: emit per-epilogue index definitions and
    STORE_OUTPUT placeholders
  • _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
    placeholders
  • _flatten_return_ast: helper for build_multi_outputs traversal
  • lower_helion_kernel: compute epilogue_fusable_outputs from stored
    proxy ids and output fusion metadata

This was referenced Mar 16, 2026
Merged
Merged
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 16, 2026
@yf225 yf225 mentioned this pull request Mar 16, 2026
Closed
@yf225 yf225 marked this pull request as draft March 16, 2026 22:29
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 16, 2026 22:29
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 16, 2026 22:30
@yf225 yf225 mentioned this pull request Mar 16, 2026
Merged
@yf225 yf225 marked this pull request as ready for review March 16, 2026 22:30
@yf225 yf225 marked this pull request as draft March 16, 2026 22:36
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 16, 2026 22:36
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 16, 2026 22:36
@yf225 yf225 marked this pull request as ready for review March 16, 2026 22:36
@yf225 yf225 marked this pull request as draft March 16, 2026 22:44
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 16, 2026 22:44
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 16, 2026 22:44
@yf225 yf225 marked this pull request as ready for review March 16, 2026 22:44
yf225 added a commit that referenced this pull request Mar 16, 2026
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
d6db9b7 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
@yf225 yf225 marked this pull request as draft March 16, 2026 23:46
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 16, 2026 23:46
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 16, 2026 23:46
@yf225 yf225 marked this pull request as ready for review March 16, 2026 23:46
yf225 added a commit that referenced this pull request Mar 16, 2026
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
4e52d91 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
@yf225 yf225 marked this pull request as draft March 17, 2026 00:00
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 17, 2026 00:00
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 17, 2026 04:20
@yf225 yf225 marked this pull request as ready for review March 17, 2026 04:20
@yf225 yf225 marked this pull request as draft March 17, 2026 04:42
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 17, 2026 04:42
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 17, 2026 04:42
@yf225 yf225 marked this pull request as ready for review March 17, 2026 04:42
@yf225 yf225 marked this pull request as draft March 17, 2026 04:56
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 17, 2026 04:56
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 17, 2026 04:56
@yf225 yf225 marked this pull request as ready for review March 17, 2026 04:56
@yf225 yf225 marked this pull request as draft March 17, 2026 04:58
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 17, 2026 04:58
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 17, 2026 04:58
@yf225 yf225 marked this pull request as ready for review March 17, 2026 04:58
@yf225 yf225 requested review from jansel, oulgen and shunting314 March 17, 2026 05:21
yf225 added a commit that referenced this pull request Mar 17, 2026
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
05483ff 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
@yf225 yf225 marked this pull request as draft March 19, 2026 19:03
@yf225 yf225 changed the base branch from yf225/stack/71 to main March 19, 2026 19:03
@yf225 yf225 changed the base branch from main to yf225/stack/71 March 19, 2026 19:03
@yf225 yf225 marked this pull request as ready for review March 19, 2026 19:03
yf225 added a commit that referenced this pull request Mar 19, 2026
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
9a620fb 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
yf225 added a commit that referenced this pull request Mar 19, 2026
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
958b572 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
@yf225
[Helion + torch.compile] Add prologue/epilogue fusion codegen to Heli…
e3ea22a 
…onTemplateBuffer

Add fusion code generation to template_buffer.py:
- _render_with_hooks: setup fusion hooks, read fusion metadata, build
  extra_params, pass store/load transforms conditionally
- _build_call_args: add prologue_source_buffers and extra_params support
- _generate_triton_ast: pass store/load transform and extra_params
- create(): add on_tensor_leaf/on_non_tensor_leaf callback params
- _codegen_epilogue_fusion: emit per-epilogue index definitions and
  STORE_OUTPUT placeholders
- _codegen_prologue_fusion: emit prologue variables and LOAD_INPUT
  placeholders
- _flatten_return_ast: helper for build_multi_outputs traversal
- lower_helion_kernel: compute epilogue_fusable_outputs from stored
  proxy ids and output fusion metadata

stack-info: PR: #1725, branch: yf225/stack/72
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jansel jansel Awaiting requested review from jansel

@oulgen oulgen Awaiting requested review from oulgen

@shunting314 shunting314 Awaiting requested review from shunting314

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@yf225