Repository: pytorch/pytorch

Issue #131192: custom ops don't reinplace when mutated arg is a view of a graph input
Created: 2024-07-19T15:16:35Z
import torch
from torch import Tensor # E: invalid syntax  [syntax]

@torch.library.custom_op("mylib::foo", mutates_args={"x"})
def foo(x: Tensor) -> None:
    x.sin_()

@torch.compile(fullgraph=True)
def f(x):
    x0 = x[0]
    foo(x0)

x = torch.randn(2)
f(x)
"""
# produces the following
def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (2, ), (1, ))
    buf0 = empty_strided_cpu((1, ), (1, ), torch.float32)
    cpp_fused_0(arg0_1, buf0)
    # Source Nodes: [], Original ATen: []
    torch.ops.mylib.foo.default(reinterpret_tensor(buf0, (), (), 0))
    cpp_fused_1(buf0, arg0_1)
    del arg0_1
    return ()
"""
changing the "foo" to x0.sin_() does lead to inductor reinplacing the code:
def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (2, ), (1, ))
    cpp_fused_sin_0(arg0_1, arg0_1)
    del arg0_1
    return ()

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

Duplicate issue #132995: [Inductor] Reinplace pass failed to support input from view with customized inplace op
Created: 2024-08-08T13:16:23Z
🐛 Describe the bug
Re-inplace pass in post-grad should support the reinplace of customized kernel.
Take this as example: https://gist.github.com/leslie-fang-intel/d62da8d11192e0565ad3e3739f1bdd19#file-test_pass-py
The final generate code should only has 1 inplace operator.
def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1 = args
    args.clear()
    assert_size_stride(arg0_1, (1024, ), (1, ))
    assert_size_stride(arg1_1, (128, 8, 16, 64), (8192, 1024, 64, 1))
    assert_size_stride(arg2_1, (128, 8, 16, 64), (8192, 1024, 64, 1))
    assert_size_stride(arg3_1, (1024, 8, 64), (1536, 64, 1))
    assert_size_stride(arg4_1, (1024, 8, 64), (1536, 64, 1))
    # Source Nodes: [], Original ATen: []
    torch.ops.mylib.reshape_and_cache.default(arg4_1, arg3_1, arg2_1, arg1_1, arg0_1)
    del arg0_1
    del arg1_1
    del arg2_1
    del arg3_1
    del arg4_1
    return ()

However, if we slightly change the input of this customized operator as a view of another tensor as in https://gist.github.com/leslie-fang-intel/d62da8d11192e0565ad3e3739f1bdd19#file-test_failed-py
The reinplace pass failed to work. Here is pattern of test case 2 before post grad passes:
class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "bf16[2, 1, 128, 8192]", arg1_1: "i32[1024]", arg2_1: "bf16[1024, 8, 64]", arg3_1: "bf16[1024, 8, 64]"):
         # File: /localdisk/leslie/torch_inductor_community/torch_script/inductor/vllm/reinplace_op/test_inplace_v2.py:16 in test, code: key_cache = kv_cache[0]
        select: "bf16[1, 128, 8192]" = torch.ops.aten.select.int(arg0_1, 0, 0)
        
         # File: /localdisk/leslie/torch_inductor_community/torch_script/inductor/vllm/reinplace_op/test_inplace_v2.py:17 in test, code: key_cache = key_cache.view(num_blocks, num_head, -1, head_size)
        view: "bf16[128, 8, 16, 64]" = torch.ops.aten.view.default(select, [128, 8, -1, 64]);  select = None
        
         # File: /localdisk/leslie/torch_inductor_community/torch_script/inductor/vllm/reinplace_op/test_inplace_v2.py:18 in test, code: value_cache = kv_cache[1]
        select_1: "bf16[1, 128, 8192]" = torch.ops.aten.select.int(arg0_1, 0, 1)
        
         # File: /localdisk/leslie/torch_inductor_community/torch_script/inductor/vllm/reinplace_op/test_inplace_v2.py:19 in test, code: value_cache = value_cache.view(num_blocks, num_head, -1, head_size)
        view_1: "bf16[128, 8, 16, 64]" = torch.ops.aten.view.default(select_1, [128, 8, -1, 64]);  select_1 = None
        
         # File: /localdisk/leslie/torch_inductor_lz/pytorch/torch/_library/custom_ops.py:663 in __call__, code: return self._opoverload(*args, **kwargs)
        auto_functionalized = torch._higher_order_ops.auto_functionalize.auto_functionalized(torch.ops.mylib.reshape_and_cache.default, key = arg3_1, value = arg2_1, key_cache = view, value_cache = view_1, slot_mapping = arg1_1);  arg3_1 = arg2_1 = view = view_1 = arg1_1 = None
        getitem_1: "bf16[128, 8, 16, 64]" = auto_functionalized[1]
        getitem_2: "bf16[128, 8, 16, 64]" = auto_functionalized[2];  auto_functionalized = None
        view_2: "bf16[1, 128, 8192]" = torch.ops.aten.view.default(getitem_1, [1, 128, 8192]);  getitem_1 = None
        select_scatter: "bf16[2, 1, 128, 8192]" = torch.ops.aten.select_scatter.default(arg0_1, view_2, 0, 0);  view_2 = None
        view_4: "bf16[1, 128, 8192]" = torch.ops.aten.view.default(getitem_2, [1, 128, 8192]);  getitem_2 = None
        select_scatter_1: "bf16[2, 1, 128, 8192]" = torch.ops.aten.select_scatter.default(select_scatter, view_4, 0, 1);  select_scatter = view_4 = None
        copy_: "bf16[2, 1, 128, 8192]" = torch.ops.aten.copy_.default(arg0_1, select_scatter_1);  arg0_1 = select_scatter_1 = None
        return ()

and the final generated code of case2, there are 2 copies before and after this customized kernel.
def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1 = args
    args.clear()
    assert_size_stride(arg0_1, (2, 1, 128, 8192), (1048576, 1048576, 8192, 1))
    assert_size_stride(arg1_1, (1024, ), (1, ))
    assert_size_stride(arg2_1, (1024, 8, 64), (1536, 64, 1))
    assert_size_stride(arg3_1, (1024, 8, 64), (1536, 64, 1))
    buf0 = empty_strided_cpu((1048576, ), (1, ), torch.bfloat16)
    buf1 = empty_strided_cpu((2097152, ), (1, ), torch.bfloat16)
    cpp_fused_0(arg0_1, buf0, buf1)
    # Source Nodes: [], Original ATen: []
    torch.ops.mylib.reshape_and_cache.default(arg3_1, arg2_1, reinterpret_tensor(buf0, (128, 8, 16, 64), (8192, 1024, 64, 1), 0), reinterpret_tensor(buf1, (128, 8, 16, 64), (8192, 1024, 64, 1), 1048576), arg1_1)
    del arg1_1
    del arg2_1
    del arg3_1
    cpp_fused_1(buf1, buf0, arg0_1, arg0_1)
    del arg0_1
    return ()

Versions
[conda] blas                      1.0                         mkl
[conda] intel-extension-for-pytorch 2.5.0+gitd8dda2f           dev_0    <develop>
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-include               2024.0.0                 pypi_0    pypi
[conda] mkl-service               2.4.0           py310h5eee18b_1
[conda] mkl-static                2024.0.0                 pypi_0    pypi
[conda] mkl_fft                   1.3.8           py310h5eee18b_0
[conda] mkl_random                1.2.4           py310hdb19cb5_0
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] optree                    0.12.1                   pypi_0    pypi
[conda] torch                     2.5.0a0+gitae708e9           dev_0    <develop>
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torchao                   0.3.1                     dev_0    <develop>
[conda] torchaudio                2.2.0a0+17a7081          pypi_0    pypi
[conda] torchfix                  0.4.0                    pypi_0    pypi
[conda] torchmetrics              1.3.0.post0              pypi_0    pypi
[conda] torchtune                 0.0.1                    pypi_0    pypi
[conda] torchvision               0.19.0a0+5181a85          pypi_0    pypi
[conda] triton-nightly            3.0.0.post20240716052845          pypi_0    pypi

Closing comment (2024-08-08T13:37:55Z): This is a duplicate of #131192, @laithsakka is working on it. Let's consolidate discussion there.