[Functionalization] The inplace decompose op is meaningless

[hanwen-sun]

Question:

When I use torch_xla to accelerate model training, I find the inlace op will go through PyTorch's functionalization. The inlace op will be dispatched to the python level decompose implementation registered here:

pytorch/torch/_refs/__init__.py

Line 504 in 8f70bf7

def _make_inplace(fn):

. However, after the python level implementation, it will then go to the c++ level implementation:

pytorch/aten/src/ATen/FunctionalTensorWrapper.cpp

Line 219 in 8f70bf7

void FunctionalTensorWrapper::replace_(const Tensor& other, bool from_lazy_regenerate) {

and the output is determined by the c++ level implementation. It seems that the output of decompose op is meaningless.
I find a pr related to the python level implementation:

• Functionalization integration xla#4158
• Prefer python meta function over c++ meta function #87426

and a pr related to the c++ level implementation:

• Add a Functionalization pass #63048

And I guess when you need to use functionalization inlace op will meet this problem.

System info:

Pytorch version: 2.2.2
torch_xla version: 2.3.0
python version: 3.10.13

reproducible code:

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()

tensor = torch.tensor([1.0, 2.0, 3.0])

xla_tensor_a = tensor.to(dev)
xla_tensor_b = tensor.to(dev)

xla_tensor_a.mul_(xla_tensor_b)

you can check the timeline or add some log in the python level and c++ level implementation
here's the timeline of mul_ op:

I want to know is there any problem in this implementation? Or can someone tell me why the functionalization of Pytorch do this, thank you!

cc @bdhirsh @ezyang