Update on "[Helion + torch.compile] Refactor TemplateBuffer as extensible base class"

yf225 · yf225 · commit ddfde897f2af · 2026-03-10T18:45:46.000-07:00
Move common fields and methods up from TritonTemplateBuffer to
TemplateBuffer so that all template subclasses (Triton, CuteDSL,
external backends) share them:

- Add mutated_inputs, allowed_prologue_inps to TemplateBuffer.__init__
- Move mutation_outputs setup from TritonTemplateBuffer to base class
- Move get_outputs(), get_allowed_prologue_inps() up
- Extract _read_deps_from_inputs() helper from extract_read_writes()
- Remove can_fuse_multi_output_epilogue() (unused)
- Simplify TritonTemplateBuffer to delegate to super().__init__()
- Remove redundant self.outputs from CppTemplateBuffer

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

[ghstack-poisoned]
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -5223,9 +5223,7 @@ def __init__(
         # Annotations dict for storing metadata (e.g., KernelTemplateChoice)
         self.annotations: dict[str, Any] = {}
 
-        # Inputs that the kernel mutates in-place (parallel to
-        # ExternKernel.mutation_outputs — kept separate from self.outputs
-        # so subclasses can freely use self.outputs for other purposes).
+        # Inputs that the kernel mutates in-place
         self.mutated_inputs = mutated_inputs
         self.mutation_outputs: list[MutationOutput] = []
         if mutated_inputs is not None:
@@ -5335,9 +5333,6 @@ def __init__(
         self.subgraph_inps: list[IRNode | sympy.Expr | None] | None = None
         self.subgraph_outs: list[IRNode | None] | None = None
 
-    def get_outputs(self) -> list[Buffer]:
-        return [self, *self.mutation_outputs]
-
     @cache_on_self_and_args("TritonTemplateBuffer")
     def get_free_symbol_uses(
         self, unbacked_only: bool = False
@@ -5362,6 +5357,9 @@ def get_free_symbol_uses(
 
         return res
 
+    def get_outputs(self) -> list[Buffer]:
+        return [self, *self.mutation_outputs]
+
     def __str__(self) -> str:
         out = f"TritonTemplateBuffer(layout={self.layout})"
         return out
