Make optimization_barrier take vector and add test for correctness

Author: JackCaoG

test/test_operations.py

@@ -748,6 +748,19 @@ def test_masked_select_shape(self):
     self.assertEqual(x_dim0_shape.item(), 3)
 
 
+class TestOptimizationBarrier(XlaTestCase):
+
+  def test_optimization_barrier_correctness(self):
+    device = xm.xla_device()
+    # only test optimization_barrier on TPU
+    if xm.xla_device_hw(device) != 'TPU':
+      return
+    x = torch.randn(5, 5, device=device)
+    y = torch.randn(5, 5, device=device)
+    (x1, y1) = xm.optimization_barrier([x, y])
+    self.assertEqual(x + y, x1 + y1)
+
+
 class TestDataType(XlaTestCase):
 
   def test_mixed_dtype_tuple(self):

torch_xla/core/xla_model.py

@@ -1032,6 +1032,6 @@ def optimization_barrier(tensors):
   the gradient checkpointing.
 
   Args:
-    tensors (torch.Tensor): `torch.Tensor`s to add barrier to.
+    tensors (List[torch.Tensor]): List of `torch.Tensor` to add barrier to.
   """
   return torch_xla._XLAC._xla_optimization_barrier(tensors)

torch_xla/csrc/init_python_bindings.cpp

@@ -276,11 +276,17 @@ std::pair<at::Tensor, std::shared_ptr<ir::Value>> CollectivePermute(
       std::make_shared<ir::Value>(new_token));
 }
 
-at::Tensor OptimizationBarrier(const at::Tensor& input) {
-  at::Tensor result = bridge::AtenFromXlaTensor(
-      XLATensor::optimization_barrier(bridge::GetXlaTensor(input)));
-  return torch::autograd::make_variable(
-      result, /*requires_grad=*/input.requires_grad());
+std::vector<at::Tensor> OptimizationBarrier(
+    const std::vector<at::Tensor>& tensors) {
+  std::vector<at::Tensor> result;
+  result.reserve(tensors.size());
+  for (auto& tensor : tensors) {
+    result.push_back(torch::autograd::make_variable(
+        bridge::AtenFromXlaTensor(
+            XLATensor::optimization_barrier(bridge::GetXlaTensor(tensor))),
+        /*requires_grad=*/tensor.requires_grad()));
+  }
+  return result;
 }
 
 void SyncTensors(const std::vector<at::Tensor>& tensors,
@@ -1035,8 +1041,9 @@ void InitXlaModuleBindings(py::module m) {
           }
           return new_token;
         });
-  m.def("_xla_optimization_barrier",
-        [](const at::Tensor& input) { return OptimizationBarrier(input); });
+  m.def("_xla_optimization_barrier", [](const std::vector<at::Tensor>& inputs) {
+    return OptimizationBarrier(inputs);
+  });
   m.def("_xla_set_default_device", [](const std::string& device) {
     return SetCurrentThreadDevice(device);
   });