Add API to donate input buffer for dynamo execution

[JackCaoG]

This pr adds an api torch_xla._XLAC._set_buffer_donation which can be used to mark the buffer associated with current tensor to be donated in the next execution. This api is currently only enabled for dynamo(torch.compile) use case, ti is a no-op in LazyTensor world(since LTC already has the auto aliasing based on inplace op and dynamo's functionization pass remove all inplace ops).

Example usage should be

  def dummy_inplace_add(self, input):
    input += 1
    return

  def test_manual_buffer_donation(self):
    device = xm.xla_device()
    input = torch.randn(5, 5).to(device)
    dummy_inplace_add_compiled = torch.compile(
        self.dummy_inplace_add, backend='openxla')

    met.clear_all()
    # input is a device_data, we should be able to set the buffer donation field.
    self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, True))
    # make sure buffer donation setting is correctly updated
    self.assertTrue(torch_xla._XLAC._get_buffer_donation(input))
	
	for _ in range(100):
      # You don't need to keep calling this function if you function does not cause dynamo recompilation.
      # check below for the reason.
      self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, True))
      dummy_inplace_add_compiled(input)

Please note a couple things

1. _set_buffer_donation is called on a tensor, but being applied to the buffer associated with this tensor. If the tensor you passed in is not a buffer(for example an intermediate tensor that has not been evulated) this api will be no-op and return false.
2. Buffer aliasing is being set up during compilation time. Torch Dynamo also does not track this field, so even if you change _set_buffer_donation after first execution(torch.compile compilation triggered at first execution of the compiled function), aliasing will not change.

[JackCaoG]

I chatted with @bdhirsh , ideally we want to figure out which input tensor can be aliased without user calling this api. There should be a way for us to retrive the aliasing information from functionization pass in aot-autograd.

The context is that If we need to alias an input buffer to output buffer, we need to make sure input buffer can not be accessed by original tensor. For example

a += 1

is fine because we know once the in place ops happens, a will point to a different buffer, so it is safe to reuse that buffer to something else.

if it is

b = a + 1

I can't alias because a's buffer is still needed. once aliasing happens, original buffer's value will be invalidated. if we alias in the

b = a + 1

case. a's value will be incorrect after this computation. Dynamo comes with functionization which will remove all inplace ops from the fx graph passed down, which is a problem in this case.

[JackCaoG]

this is not ideal.. I want to add a new config to coll.config but that struct now lives on upstream...

[lsy323]

LGTM, left a few questions

[JackCaoG]

I will merge this change to unblock the user, fix comments in a follow up pr.

[alanwaketan]

Sorry for being late, but looks good to me.