Fix input output alias for custom inplace ops

test/test_input_output_aliases.py

@@ -41,6 +41,18 @@ def test_aliasing_with_cloned(self):
  torch.allclose(t1 - 1, t1_cloned)
  self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 1.0)
 
+ def test_aliasing_across_custom_inplace(self):
+ xla_device = xm.xla_device()
+ met.clear_all()
+ t1 = torch.randn(4, 5).to(xla_device)
+ t1 *= t1
+ xm.mark_step()
+ self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 1.0)
+ xm.optimization_barrier_([t1])
+ t1 *= 100
+ xm.mark_step()
+ self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 2.0)
+
  def test_aliasing_across_mark_step(self):
  xla_device = xm.xla_device()
  met.clear_all()

torch_xla/csrc/tensor.cpp

@@ -387,6 +387,10 @@ torch::lazy::Value XLATensor::GetIrValue() const {
  // which wants the XLA data will still find it, w/out having to fetch it
  // via a computation client from-server call.
  AssignIrValue(CreateTensorNode(handle, /*read_only=*/false));
+ // CreateTensorNode will set the data info of the tensor to the current
+ // unique_id. Here the alias id needs to be updated so that input output
+ // alias can correctly work on the xla's custom inplace operation.
+ data()->alias_id = GetUniqueId();
  return data()->ir_value;
  }
  std::optional<at::Tensor> tensor_data = CurrentTensorData();