[inductor] Fix edge case in JIT vs. AOT fusion after finalizing MultiTemplateBuffer

[ColinPeppler]

Context

Here's a peripheral scenario causing the JIT-pass and AOT-pass to pick different fusions.

# JIT -- buf3 is a MultiTemplateBuffer
V.graph.buffers = [buf0, buf1, buf2, buf3, buf4]
                                ^          ^
# JIT pass calls finalize_multi_template_buffers()
V.graph.buffers = [buf0, buf1, buf2, buf4, *buf3*]

# AOT, note proximity_score(buf2, buf4) is "better" for fusion than JIT
V.graph.buffers = [buf0, buf1, buf2, buf4, *buf3*]
                                ^    ^

It happens like this:

• JIT starts with the original set nodes using V.graph.buffers
• In JIT, finalize_multi_template_buffers() is called which can change the order of the buffers.
• This makes the order of buffers/scheduler nodes different.
• Now, each node's min/max-order is different than before.
• As a result, the proximity between two nodes is different.

  pytorch/torch/_inductor/scheduler.py

  Lines 2316 to 2335 in ad67553

  	def score_fusion(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode):
  	"""
  	Assign a score (higher comes first) to the fusion of node1
  	and node2. When different fusions conflict with each other,
  	this is the way we decide what order to run them in.
  	Our current score is based on:
  	- Estimate of the saved memory operations
  	- Fusions closer together in original order
  	"""
  	memory_score = self.score_fusion_memory(node1, node2)
  	proximity_score = -max(
  	abs(node1.min_order - node2.max_order),
  	abs(node2.min_order - node1.max_order),
  	)
  	return (
  	node1.is_template() == config.epilogue_fusion_first and memory_score > 0,
  	node1.is_reduction() == node2.is_reduction() and memory_score > 0,
  	memory_score,
  	proximity_score,

Error

$ TORCH_LOGS="+fusion" python test/inductor/test_max_autotune.py -k test_jit_fusion_matches_aot_fusion
======================================================================
FAIL: test_jit_fusion_matches_aot_fusion (__main__.TestMaxAutotune)
----------------------------------------------------------------------
Traceback (most recent call last):
  ...
  File "/data/users/colinpeppler/pytorch/torch/_inductor/graph.py", line 1718, in compile_to_fn
    code, linemap = self.codegen_with_cpp_wrapper()
  File "/data/users/colinpeppler/pytorch/torch/_inductor/graph.py", line 1618, in codegen_with_cpp_wrapper
    return self.codegen()
  File "/data/users/colinpeppler/pytorch/torch/_inductor/graph.py", line 1636, in codegen
    self.scheduler.codegen()
  File "/data/users/colinpeppler/pytorch/torch/_dynamo/utils.py", line 210, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/colinpeppler/pytorch/torch/_inductor/scheduler.py", line 2602, in codegen
    self.get_backend(device).codegen_node(node)  # type: ignore[possibly-undefined]
  File "/data/users/colinpeppler/pytorch/torch/_inductor/codegen/cuda_combined_scheduling.py", line 66, in codegen_node
    return self._triton_scheduling.codegen_node(node)
  File "/data/users/colinpeppler/pytorch/torch/_inductor/codegen/triton.py", line 3377, in codegen_node
    return self.codegen_node_schedule(node_schedule, buf_accesses, numel, rnumel)
  File "/data/users/colinpeppler/pytorch/torch/_inductor/codegen/triton.py", line 3602, in codegen_node_schedule
    final_kernel.call_kernel(final_kernel.kernel_name)
  File "/data/users/colinpeppler/pytorch/torch/_inductor/codegen/triton.py", line 3055, in call_kernel
    grid = wrapper.generate_default_grid(name, grid)
  File "/data/users/colinpeppler/pytorch/torch/_inductor/codegen/cpp_wrapper_cuda.py", line 174, in generate_default_grid
    params is not None
AssertionError: cuda kernel parameters for triton_poi_fused_add_0 should already exist at this moment, only found dict_keys(['Placeholder.DESCRIPTIVE_NAME', 'triton_poi_fused_add_mul_0', 'triton_poi_fused_pow_1'])

Stack from ghstack (oldest at bottom):

• -> [inductor] Fix edge case in JIT vs. AOT fusion after finalizing MultiTemplateBuffer #126622

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ You can merge normally! (7 Unrelated Failures)

As of commit 0d381a5 with merge base a0429c0 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  hf_BigBird

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#120841)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#120841)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#120841)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#120841)
  sam_fast

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

[ColinPeppler]

other option is to swap the old buffer with new buffer.

but I think sorting buffers by name at the start will be more deterministic

[ColinPeppler]

nvm going with this option

[ColinPeppler]

tried sorting nodes by name frist but ran into test errors.

inductor/test_torchinductor.py::CpuTests::test_buffer_batch_norm_cpu

[chenyang78]

Please fix the lint issue. Otherwise, LGTM. Thanks!

[ColinPeppler]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: 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.

Details for Dev Infra team

Raised by workflow job

[ColinPeppler]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

[eellison]

Looks good - where did this come up ?

[houseroad]

The issue failed some internal models

[ColinPeppler]

Looks good - where did this come up ?

An internal model in production. Interesting enough the error came up after #125772.

What I saw in the logs.

#JIT
[__fusion] fusing buf1109_buf1113_buf1123_buf1127 with buf1152_buf1156
[__fusion] cannot fuse buf1109_buf1113_buf1123_buf1127_buf1152_buf1156 with buf1182_buf1186: will increase peak memory
Generating code for node buf1109_buf1113_buf1123_buf1127_buf1152_buf1156 with estimated runtime 32.787698

# AOT
[__fusion] fusing buf1152_buf1156 with buf1182_buf1186
[__fusion] cannot fuse buf1109_buf1113_buf1123_buf1127 with buf1152_buf1156_buf1182_buf1186: will increase peak memory