[Helion + torch.compile] Add store/load transform hooks and prologue/epilogue fusion codegen by yf225 · Pull Request #1724 · pytorch/helion

yf225 · 2026-03-16T21:55:43Z

Stacked PRs:

[Helion + torch.compile] Enable torch.compile fusion tests #1727
->[Helion + torch.compile] Add store/load transform hooks and prologue/epilogue fusion codegen #1724

[Helion + torch.compile] Add store/load transform hooks and prologue/epilogue fusion codegen

…rameter plumbing Add infrastructure for prologue/epilogue fusion in generated Triton code: - store_transform/load_transform hooks in memory_ops for hl.store/hl.load - Extra params, removed args, and protected arg names in codegen pipeline - dim_index_exprs on SubscriptIndexing for epilogue fusion offsets - _ensure_inductor_fusion_config() to enable fusion via config_patches stack-info: PR: #1724, branch: yf225/stack/71

…rameter plumbing Add infrastructure for prologue/epilogue fusion in generated Triton code: - store_transform/load_transform hooks in memory_ops for hl.store/hl.load - Extra params, removed args, and protected arg names in codegen pipeline - dim_index_exprs on SubscriptIndexing for epilogue fusion offsets - _ensure_inductor_fusion_config() to enable fusion via config_patches stack-info: PR: #1724, branch: yf225/stack/71 # Conflicts: # helion/_compiler/generate_ast.py

shunting314

maybe you did in a later PR but in high level how to you pass in the load/store transform and extra params?

helion/language/memory_ops.py

…rameter plumbing Add infrastructure for prologue/epilogue fusion in generated Triton code: - store_transform/load_transform hooks in memory_ops for hl.store/hl.load - Extra params, removed args, and protected arg names in codegen pipeline - dim_index_exprs on SubscriptIndexing for epilogue fusion offsets - _ensure_inductor_fusion_config() to enable fusion via config_patches stack-info: PR: #1724, branch: yf225/stack/71 # Conflicts: # helion/_compiler/generate_ast.py

yf225 · 2026-03-19T23:13:40Z

maybe you did in a later PR but in high level how to you pass in the load/store transform and extra params?

Yes I now merged them into the same PR (this one), with this high-level flow:

HelionTemplateBuffer._render_with_hooks() (template_buffer.py) calls kernel._setup_fusion_hooks() to populate fusion metadata (epilogue indices, prologue variables, source buffers, etc.).
_generate_triton_ast() (also in template_buffer.py) checks whether fusion metadata was populated (e.g., self._epilogue_idx_by_param, self._prologue_fused_params) and conditionally passes self._codegen_epilogue_fusion as store_transform and self._codegen_prologue_fusion as load_transform to generate_ast().
generate_ast() forwards them into GenerateAST.__init__() where they become self.store_transform / self.load_transform.
During codegen, memory_ops.py checks state.codegen.store_transform / state.codegen.load_transform (the hooks added in this PR) and invokes them at the hl.store / hl.load call sites.

…rameter plumbing Add infrastructure for prologue/epilogue fusion in generated Triton code: - store_transform/load_transform hooks in memory_ops for hl.store/hl.load - Extra params, removed args, and protected arg names in codegen pipeline - dim_index_exprs on SubscriptIndexing for epilogue fusion offsets - _ensure_inductor_fusion_config() to enable fusion via config_patches stack-info: PR: #1724, branch: yf225/stack/71 # Conflicts: # helion/_compiler/generate_ast.py

helion/_compiler/_inductor/template_buffer.py

shunting314 · 2026-03-20T21:01:46Z

@claude review this PR. Try to focus on making corner cases covered and make the integration between different components clean

yf225

[Claude Review] Overall: well-structured PR with clean callback pattern for store/load transforms. See inline comments for corner cases and integration concerns.

helion/_compiler/_inductor/template_buffer.py

helion/language/memory_ops.py

helion/_compiler/_inductor/template_buffer.py

helion/_compiler/generate_ast.py

yf225

[deleted]

…epilogue fusion codegen stack-info: PR: #1724, branch: yf225/stack/71

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 16, 2026

yf225 force-pushed the yf225/stack/71 branch from e55096b to f2b3843 Compare March 16, 2026 21:55

yf225 force-pushed the yf225/stack/70 branch from 78aa014 to 48878cf Compare March 16, 2026 21:55

yf225 marked this pull request as draft March 16, 2026 22:29

yf225 changed the base branch from yf225/stack/70 to main March 16, 2026 22:29

yf225 changed the base branch from main to yf225/stack/70 March 16, 2026 22:30

yf225 mentioned this pull request Mar 16, 2026

[Helion + torch.compile] Enable torch.compile fusion tests #1727

Merged

yf225 marked this pull request as ready for review March 16, 2026 22:30

yf225 marked this pull request as draft March 16, 2026 22:36

yf225 changed the base branch from yf225/stack/70 to main March 16, 2026 22:36

yf225 changed the base branch from main to yf225/stack/70 March 16, 2026 22:36

yf225 marked this pull request as ready for review March 16, 2026 22:36

yf225 marked this pull request as draft March 16, 2026 22:43

yf225 changed the base branch from yf225/stack/70 to main March 16, 2026 22:43

yf225 force-pushed the yf225/stack/71 branch from f2b3843 to 76bb051 Compare March 16, 2026 22:44

yf225 changed the base branch from main to yf225/stack/70 March 16, 2026 22:44

yf225 marked this pull request as ready for review March 16, 2026 22:44

yf225 marked this pull request as draft March 16, 2026 23:46

yf225 changed the base branch from yf225/stack/70 to main March 16, 2026 23:46

yf225 force-pushed the yf225/stack/71 branch from 76bb051 to 5e09efc Compare March 16, 2026 23:46

yf225 changed the base branch from main to yf225/stack/70 March 16, 2026 23:46

yf225 marked this pull request as ready for review March 16, 2026 23:46

yf225 marked this pull request as draft March 17, 2026 00:00

yf225 changed the base branch from yf225/stack/70 to main March 17, 2026 00:00

yf225 force-pushed the yf225/stack/71 branch from 5e09efc to 019bfed Compare March 17, 2026 00:01

yf225 force-pushed the yf225/stack/71 branch from 4c9b725 to 8c1c6b2 Compare March 17, 2026 04:56

yf225 changed the base branch from main to yf225/stack/70 March 17, 2026 04:56

yf225 marked this pull request as ready for review March 17, 2026 04:56

yf225 marked this pull request as draft March 17, 2026 04:58

yf225 changed the base branch from yf225/stack/70 to main March 17, 2026 04:58

yf225 force-pushed the yf225/stack/71 branch from 8c1c6b2 to 7b4f580 Compare March 17, 2026 04:58

yf225 changed the base branch from main to yf225/stack/70 March 17, 2026 04:58

yf225 marked this pull request as ready for review March 17, 2026 04:58

yf225 requested review from jansel, oulgen and shunting314 March 17, 2026 05:20

shunting314 reviewed Mar 19, 2026

View reviewed changes

helion/language/memory_ops.py Show resolved Hide resolved

shunting314 reviewed Mar 19, 2026

View reviewed changes

helion/language/memory_ops.py Outdated Show resolved Hide resolved

yf225 marked this pull request as draft March 19, 2026 19:03

yf225 changed the base branch from yf225/stack/70 to main March 19, 2026 19:03

yf225 force-pushed the yf225/stack/71 branch from 7b4f580 to a64c890 Compare March 19, 2026 19:03

yf225 marked this pull request as ready for review March 19, 2026 19:03

yf225 marked this pull request as draft March 19, 2026 23:35

shunting314 reviewed Mar 20, 2026

View reviewed changes

yf225 commented Mar 20, 2026

View reviewed changes

shunting314 approved these changes Mar 20, 2026

View reviewed changes

yf225 mentioned this pull request Mar 21, 2026

[Helion + torch.compile] Enable torch.compile fusion tests #1778

Closed

jansel approved these changes Mar 22, 2026

View reviewed changes

[Helion + torch.compile] Add store/load transform hooks and prologue/…

ce24879

…epilogue fusion codegen stack-info: PR: #1724, branch: yf225/stack/71

Conversation

yf225 commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!