Distribute Literal->Tensor copies across thread pool (pytorch#5825)

* Distribute Literal->Tensor copies across thread pool * Update for pytorch#5799
chunnienc · Dec 14, 2023 · 990fe5a · 990fe5a
1 parent 53cd01b
commit 990fe5a
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 11 deletions.
diff --git a/test/cpp/test_xla_sharding.cpp b/test/cpp/test_xla_sharding.cpp
@@ -27,7 +27,8 @@ namespace {
 bool XlaDataValuesEqual(torch::lazy::BackendDataPtr a,
                         torch::lazy::BackendDataPtr b,
                         at::ScalarType element_type) {
-  std::vector<at::Tensor> tensors = XlaDataToTensors({a, b}, element_type);
+  std::vector<at::Tensor> tensors =
+      XlaDataToTensors({a, b}, {element_type, element_type});
   return TensorCompare(tensors[0], tensors[1]);
 }
 }  // namespace

diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1784,8 +1784,8 @@ void InitXlaModuleBindings(py::module m) {
                shard_handles) {
             shards.push_back(
                 XlaDataToTensors({shard_handle},
-                                 MaybeUpcastToHostTorchType(
-                                     shard_handle->shape().element_type()))
+                                 {MaybeUpcastToHostTorchType(
+                                     shard_handle->shape().element_type())})
                     .front());
             str_devices.push_back(shard_handle->device());
           }

diff --git a/torch_xla/csrc/tensor.cpp b/torch_xla/csrc/tensor.cpp
@@ -467,7 +467,8 @@ at::Tensor XLATensor::ToTensor(bool detached) {
     XLAGraphExecutor::Get()->DeviceBarrier(GetDevice());
     // The GetXlaData() call will trigger an ApplyPendingGraph() if an IR
     // XlaNode is available on the tensor.
-    std::vector<at::Tensor> tensors = XlaDataToTensors({GetXlaData()}, dtype());
+    std::vector<at::Tensor> tensors =
+        XlaDataToTensors({GetXlaData()}, {dtype()});
     tensor = std::move(tensors.front());
     if (!detached) {
       SetTensorData(tensor);

diff --git a/torch_xla/csrc/tensor_util.cpp b/torch_xla/csrc/tensor_util.cpp
@@ -796,13 +796,18 @@ std::vector<xla::Literal> ReleaseGilAndTransferData(
 
 std::vector<at::Tensor> XlaDataToTensors(
     absl::Span<const torch::lazy::BackendDataPtr> xla_data,
-    at::ScalarType dest_element_type) {
+    absl::Span<const at::ScalarType> dest_element_type) {
   std::vector<xla::Literal> literals = ReleaseGilAndTransferData(xla_data);
-  std::vector<at::Tensor> tensors;
-  tensors.reserve(literals.size());
-  for (auto& literal : literals) {
-    tensors.push_back(MakeTensorFromXlaLiteral(literal, dest_element_type));
+  std::vector<at::Tensor> tensors(literals.size());
+  absl::BlockingCounter counter(literals.size());
+  for (size_t i = 0; i < tensors.size(); ++i) {
+    auto copy_fn = [&, i]() {
+      tensors[i] = MakeTensorFromXlaLiteral(literals[i], dest_element_type[i]);
+      counter.DecrementCount();
+    };
+    thread::Schedule(std::move(copy_fn));
   }
+  counter.Wait();
   return tensors;
 }
 

diff --git a/torch_xla/csrc/tensor_util.h b/torch_xla/csrc/tensor_util.h
@@ -34,7 +34,7 @@ std::vector<xla::Literal> ReleaseGilAndTransferData(
 // TODO LTC @wonjoo - Migrate to upstream after Device -> BackendDevice
 std::vector<at::Tensor> XlaDataToTensors(
     absl::Span<const torch::lazy::BackendDataPtr> xla_data,
-    at::ScalarType dest_element_type);
+    absl::Span<const at::ScalarType> dest_element_type);
 
 bool TensorCompare(const at::Tensor& t1, const at::Tensor& t2);
 

diff --git a/torch_xla/csrc/xla_backend_impl.cpp b/torch_xla/csrc/xla_backend_impl.cpp
@@ -93,7 +93,7 @@ class XlaBackendImpl : public torch::lazy::BackendImplInterface {
       const torch::lazy::BackendDataPtr data,
       c10::optional<at::ScalarType> logical_scalar_type) const override {
     // TODO(JackCaoG): handle the logical_scalar_type == nullptr case
-    return XlaDataToTensors({data}, *logical_scalar_type)[0];
+    return XlaDataToTensors({data}, {*logical_scalar_type})[0];
   }
 
   std::unique_ptr<torch::lazy::LoweringContext> CreateLoweringContext(