In-place operations on an DLPack aliased XLA tensor does not propagate. about xla HOT 8 OPEN

This behavior should be the results of our functionalization pass. @alanwaketan to confirm the expected behavior. Either way, let's have a dlpack documentation/tutorial that goes through example use cases and fully explains correct behavior @ysiraichi.

from xla.

Thanks for the issue. I checked buffer pointer at more places:

>>> t0 = torch.arange(10, device=xm.xla_device())
>>> xm.mark_step(wait=True)
>>> 
>>> capsule = xdlpack.to_dlpack(t0)
>>> t1 = xdlpack.from_dlpack(capsule)
>>> print(torch_xla._XLAC._unsafe_buffer_pointer(t0)== torch_xla._XLAC._unsafe_buffer_pointer(t1))
True
>>> 
>>> t0[0] = 100
>>> xm.mark_step()
>>> 
>>> print(torch_xla._XLAC._unsafe_buffer_pointer(t0)== torch_xla._XLAC._unsafe_buffer_pointer(t1))
True
>>> print(t0.eq(t1).all().item())
False
>>> 
>>> print(torch_xla._XLAC._unsafe_buffer_pointer(t0)== torch_xla._XLAC._unsafe_buffer_pointer(t1))
False

Could you elaborate on That's because even though functionalization emulates views and mutation, PyTorch/XLA doesn't really have the concept of views and can't mutate a given tensor.? Do you mean when we do t0[0]=100, the underlying pjrt buffer is not mutated hence t1 is not updated, even though t0 and t1 share the same storage? Let me also look into what torch_xla does when we do t0[0]=100

from xla.

Yes, exactly. In summary, functionalized lazy tensors is composed of:

Tensor(
    impl=FunctionalTensorWrapper(
        value=Tensor(
            impl=XLATensorImpl(
                tensor=XLATensor(handle or tensor_data or ir_value)
            )
        )
    )
)

Suppose t0 and t1 share the same storage using the DLPack API. Whenever an in-place operation is called, e.g. t0.add_(1), the functionalization layer actually calls the functional variant (XLANativeFunctions::add), which generates a new XLATensor. Later, that is wrapped by a new FunctionalTensorWrapper (let's call it temp). In the end, the functionalization layer replaces the FunctionalTensorWrapper::value of t0 by the one inside temp. Thus, t0 ends up with the updated value, while t1 remains with the old one.

from xla.

Try this: https://github.com/pytorch/xla/blob/master/torch_xla/csrc/aten_xla_type.cpp#L2703

from xla.

Hmm. Not sure I get it. Could you explain a bit more?

from xla.

That's a helper where we can bridge information through intermediate tensors created by functionalization for in-place ops.

from xla.

When we do the in-place op t0[0] = 100, I see XLANativeFunctions::_propagate_xla_data invoked twice by:

• at::functionalization::fill__Scalar
• at::functionalization::copy_

in sequence. So it seems the help is already being used?

from xla.