[Graph Partition] fix graph partition input signature for fallback kernels

[BoyuanFeng]

Scheduler relies on node.last_usage to free buffers. last_usage may contain a buffer that is allocated in previous graph partition AND not directly accessed in the current graph partition.

Example

def f(x):
    y = x + 1
    z = torch.ops.aten.view.dtype(y, torch.float8_e4m3fn)
    z_cpu = z.cpu()
    u_cuda = z_cpu.cuda()
    return u_cuda

In the generated code, we have

def partition_0(args):
    ...
    # Topologically Sorted Source Nodes: [y, z], Original ATen: [aten.add, aten.view]
    buf1 = torch.ops.aten.view.dtype(buf0, torch.float8_e4m3fn) # < ------ buf1 is a view of buf0
    buf2 = buf1 # <------- buf2 is buf1 
    assert_size_stride(buf2, (8, ), (1, ), 'torch.ops.aten.view.dtype')
    assert_alignment(buf2, 16, 'torch.ops.aten.view.dtype')
    return (buf2, )

def call(self, args):
    ...
    (buf2,) = self.partitions[0](partition0_args)
    ...
    buf3.copy_(buf2, False)
    del buf0
    del buf1
    del buf2  # <---- `del buf2` leads to `del buf0`. BUT `buf0` is not returned from partition_0.
    ...

Note: view is treated as a fallback kernel due to its special dtype.

pytorch/torch/_inductor/lowering.py

Lines 841 to 843 in de09bab

	if src_bits != dst_bits:
	# fallback to aten eager implementation for differing bitwidths
	return fallback_handler(aten.view.dtype)(x, dtype)

Fix

This PR fixes the issue by also returning these buffers to be freed later.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• AWS was down, GHA infrastructure effected / recovering

❌ 1 New Failure, 1 Unrelated Failure

As of commit 976a2c7 with merge base de09bab ():

NEW FAILURE - The following job has failed:

• pull / linux-jammy-py3.10-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh)
  test/distributed/test_dist2.py::ProcessGroupNCCLTest::test_allgather

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• trunk / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge) (gh) (trunk failure)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_w2l_model

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

[BoyuanFeng]

@pytorchbot merge -f "skip unrelated distributed test failure"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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