Default explicit donation for step barriers (#8982)

Author: rpsilva-aws

test/dynamo/test_dynamo_aliasing.py

@@ -185,41 +185,6 @@ def test_buffer_donation_on_non_data_tensor(self):
     self.assertNotIn('XlaSetBufferDonation', met.counter_names())
 
 
-class TestNonDynamoBufferDonationAliasing(unittest.TestCase):
-
-  def dummy_fn(self, input):
-    return torch.cos(torch.sin(input))
-
-  # Currently let's skip buffer donation api for the non-dynamo use case
-  def test_buffer_donation_skip_for_non_dynamo(self):
-    device = xm.xla_device()
-    input = torch.randn(5, 5).to(device)
-    xm.mark_step()
-    met.clear_all()
-
-    # We should be able to set buffer donation for input tensor, but when mark_step
-    # triggered, the buffer donation should be ignored.
-    self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, True))
-    res = self.dummy_fn(input)
-    xm.mark_step()
-    # Make sure that input buffer is not aliased and can be used for other compuations.
-    # Also make sure that buffer_donation will not trigger recompilation in non-dynamo.
-    self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, False))
-    res2 = self.dummy_fn(input)
-    xm.mark_step()
-    torch.allclose(res.cpu(), res2.cpu())
-    self.assertEqual(met.metric_data('CompileTime')[0], 1)
-
-  def test_no_op_mark_step_keep_buffer_donation(self):
-    device = xm.xla_device()
-    input = torch.randn(5, 5).to(device)
-    self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, True))
-    xm.mark_step()
-    self.assertTrue(torch_xla._XLAC._get_buffer_donation(input))
-    xm.mark_step()
-    self.assertTrue(torch_xla._XLAC._get_buffer_donation(input))
-
-
 if __name__ == '__main__':
   test = unittest.main()
   sys.exit(0 if test.result.wasSuccessful() else 1)

test/test_input_output_aliases.py

@@ -1,15 +1,16 @@
+import contextlib
+import copy
 import os
 import sys
+import unittest
+from absl.testing import parameterized
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch_xla
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met
-import unittest
-import contextlib
-import copy
 
 
 def create_xla_config_context(set_func, get_func):
@@ -34,7 +35,7 @@ def config_context(value):
 
 
 # TODO(alanwaketan): add test for views.
-class InputOutputAliasesTest(unittest.TestCase):
+class InputOutputAliasesTest(parameterized.TestCase):
 
   def test_non_view(self):
     xla_device = xm.xla_device()
@@ -233,34 +234,59 @@ def test_device_data_cache_no_aliasing(self):
     self.assertEqual(t1.item(), 43)
 
   def test_user_config_donation_with_ltc_donation(self):
-    with alias_with_buffer_donor_config_context(True):
+    met.clear_all()
+    xla_device = xm.xla_device()
+    t0 = torch.randn(4, 2, 2).to(xla_device)
+    t1 = torch.randn(4, 2, 2).to(xla_device)
+    self.assertTrue(torch_xla._XLAC._set_buffer_donation(t0, True))
+    self.assertTrue(torch_xla._XLAC._get_buffer_donation(t0))
+    self.assertFalse(torch_xla._XLAC._get_buffer_donation(t1))
+    t2 = t0 + t1
+    t1 += 2
+    xm.mark_step(wait=True)
+
+    # We surface the C++ runtime error by checking that the backend data is
+    # no longer present for the IR node.
+    self.assertTrue(torch_xla._XLAC._is_placecholder(t0))
+    self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 2.0)
+
+  @parameterized.parameters(True, False)
+  def test_user_config_donation_with_ltc_donation_graph_sync(
+      self, enable_buffer_donor_config):
+    with alias_with_buffer_donor_config_context(enable_buffer_donor_config):
       met.clear_all()
       xla_device = xm.xla_device()
       t0 = torch.randn(4, 2, 2).to(xla_device)
       t1 = torch.randn(4, 2, 2).to(xla_device)
       self.assertTrue(torch_xla._XLAC._set_buffer_donation(t0, True))
       self.assertTrue(torch_xla._XLAC._get_buffer_donation(t0))
       self.assertFalse(torch_xla._XLAC._get_buffer_donation(t1))
-      t3 = t0 + t1
+      t2 = t0 + t1
       t1 += 2
-      xm.mark_step(wait=True)
+      # We use _xla_sync_multi to explicitly disable sync_xla_data, which will
+      # in turn avoid using LTC aliasings. This ensures that the resulting
+      # aliasings are due to the buffer donation.
+      torch_xla._XLAC._xla_sync_multi([t0, t1, t2], [str(xla_device)], True,
+                                      False)
 
       # We surface the C++ runtime error by checking that the backend data is
       # no longer present for the IR node.
-      self.assertTrue(torch_xla._XLAC._is_placecholder(t0))
-      self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 2.0)
+      self.assertEqual(
+          torch_xla._XLAC._is_placecholder(t0), enable_buffer_donor_config)
+      self.assertEqual(
+          met.metric_data("InputOutputAliasCount")[1],
+          enable_buffer_donor_config)
 
   def test_user_config_donation_with_ltc_donation_overlap(self):
-    with alias_with_buffer_donor_config_context(True):
-      met.clear_all()
-      xla_device = xm.xla_device()
-      t0 = torch.randn(4, 2, 2).to(xla_device)
-      self.assertTrue(torch_xla._XLAC._set_buffer_donation(t0, True))
-      self.assertTrue(torch_xla._XLAC._get_buffer_donation(t0))
-      t0 += 2
-      xm.mark_step()
+    met.clear_all()
+    xla_device = xm.xla_device()
+    t0 = torch.randn(4, 2, 2).to(xla_device)
+    self.assertTrue(torch_xla._XLAC._set_buffer_donation(t0, True))
+    self.assertTrue(torch_xla._XLAC._get_buffer_donation(t0))
+    t0 += 2
+    xm.mark_step()
 
-      self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 1.0)
+    self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 1.0)
 
   def test_user_config_donation(self):
     with alias_with_buffer_donor_config_context(True):
@@ -304,6 +330,15 @@ def test_user_config_donation_no_op_mark_step(self):
       xm.mark_step()
       self.assertTrue(torch_xla._XLAC._get_buffer_donation(t0))
 
+  def test_no_op_mark_step_keep_buffer_donation(self):
+    xla_device = xm.xla_device()
+    input = torch.randn(5, 5).to(xla_device)
+    self.assertTrue(torch_xla._XLAC._set_buffer_donation(input, True))
+    xm.mark_step()
+    self.assertTrue(torch_xla._XLAC._get_buffer_donation(input))
+    xm.mark_step()
+    self.assertTrue(torch_xla._XLAC._get_buffer_donation(input))
+
 
 if __name__ == '__main__':
   test = unittest.main()

torch_xla/csrc/xla_graph_executor.cpp

@@ -1305,9 +1305,11 @@ std::vector<size_t> XLAGraphExecutor::GetBufferDonors(
     return {};
   }
 
+  bool donate_ltc_data =
+      coll.config.sync_ltc_data && coll.config.force_ltc_data;
   std::vector<size_t> ltc_buffer_donor_indices;
-  if (coll.config.sync_ltc_data && coll.config.force_ltc_data) {
-    // We can only alias at the step barrier, when force_ltc_data is true.
+  if (donate_ltc_data) {
+    // We can only alias at the step barrier, when donate_ltc_data is true.
     // Consider the case:
     //   1. Tensor A(DEVICE_DATA)
     //   2. Tensor B = A + 0.9
@@ -1336,7 +1338,10 @@ std::vector<size_t> XLAGraphExecutor::GetBufferDonors(
   }
 
   std::vector<size_t> user_config_buffer_donor_indices;
-  if (GetAliasWithBufferDonorConfig()) {
+  if (donate_ltc_data || GetAliasWithBufferDonorConfig()) {
+    // In case any tensor is explicitly marked for donation, we ensure that it
+    // is donated during step barrier, or if explicitly forced to donate via
+    // GetAliasWithBufferDonorConfig().
     user_config_buffer_donor_indices =
         GetBufferDonorIndexFromUserConfig(parameters_data);
   }