Implement multi-output autotuning in Torchinductor

[edderstack]

Add support for multi-output autotuning in Torchinductor, allowing templates and extern kernels to share the same operator contract without additional copy operations. This change includes updates to choice construction, benchmarking, and IR generation to handle multiple outputs seamlessly.

[meta-cla]

Hi @edderstack!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[eellison]

@edderstack have you looked into the Subgraph Decomposition autotuning ? this can handle arbitrary, multiple operations and outputs.

[edderstack]

Hi @eellison , thanks for the suggestion. I just took a look through the Subgraph Decomposition autotuning code, but I don't think it solves the specific problem of autotuning Triton templates and extern choices under a shared multi-output contract.
It does allow candidates to be arbitrary FX subgraphs with multiple operations, but the autotuned custom-op/subgraph path is still single-output-oriented:

• Layout inference (in SubgraphTemplate._infer_custom_op_layout) still asserts a single-tensor output
• SubgraphChoiceCaller.output_node() returns one SubgraphBuffer
• SubgraphBuffer itself is constructed with a single layout
• The custom op API explicitly says multi-output custom ops are not yet supported (labelling it as a TODO in custom_op.py), then assumes that the fake output is one tensor

[eellison]

I think the single output is an unnecessary constraint. I think is also true for TritonTemplates. E.g with flex attention I believe we work around by treating extra outputs as mutating writes.

[edderstack]

Yeah, I think the current single-output behaviour is a an API constraint rather than a fundamental limitation. My proposal would be to make the additional outputs, now treated as mutating writes, first-class outputs so Triton templates, extern choices and potentially subgraph decomposition choices (with some additional work in the future) could all be conpared under the same output spec.

[eellison]

This makes sense but id be curious to see specific use cases and speedups first.

[edderstack]

Yeah, fair enough; I had included some examples of use cases (under "Example Multi-Output Operations"), but I have changed the heading to make it clearer what these are, added a few extra examples and added a section to each that explains what the effect of the proposed changes would be on them. It is more difficult to give an idea of speedups: multi-output autotuning would not automatically make these kernels faster if the selected kernel does not change with the new options. The argument for a speedup is that it would allow inductor to fairly autotune template candidates against external/library candidates without including additional copy operations in the autotuning.