Cherrypick fix for DLRM real dataset crash (pytorch#2409)

jysohn23 · davidel · web-flow · commit 06d564bac320 · 2020-08-10T08:33:37.000-07:00
Co-authored-by: Davide Libenzi &lt;davide.libenzi@gmail.com&gt;
diff --git a/third_party/xla_client/multi_wait.cc b/third_party/xla_client/multi_wait.cc
@@ -11,7 +11,7 @@ void MultiWait::Done() {
   {
     std::lock_guard<std::mutex> lock(mutex_);
     completed_count_ += 1;
-    notify = completed_count_ >= count_;
+    notify = completed_count_ == count_;
   }
   if (notify) {
     cv_.notify_all();
@@ -45,17 +45,27 @@ void MultiWait::Reset(size_t count) {
 }
 
 std::function<void()> MultiWait::Completer(std::function<void()> func) {
-  auto completer = [this, func = std::move(func)]() {
-    try {
-      func();
-    } catch (...) {
-      std::lock_guard<std::mutex> lock(mutex_);
-      exptr_ = std::current_exception();
-    }
-    Done();
+  auto completer = [this, func = std::move(func)]() { Complete(func); };
+  return completer;
+}
+
+std::function<void()> MultiWait::Completer(std::shared_ptr<MultiWait> mwait,
+                                           std::function<void()> func) {
+  auto completer = [mwait = std::move(mwait), func = std::move(func)]() {
+    mwait->Complete(func);
   };
   return completer;
 }
 
+void MultiWait::Complete(const std::function<void()>& func) {
+  try {
+    func();
+  } catch (...) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    exptr_ = std::current_exception();
+  }
+  Done();
+}
+
 }  // namespace util
 }  // namespace xla
diff --git a/third_party/xla_client/multi_wait.h b/third_party/xla_client/multi_wait.h
@@ -3,6 +3,7 @@
 
 #include <condition_variable>
 #include <functional>
+#include <memory>
 #include <mutex>
 
 #include "tensorflow/compiler/xla/types.h"
@@ -31,10 +32,19 @@ class MultiWait {
 
   // Creates a completer functor which signals the mult wait object once func
   // has completed. Handles exceptions by signaling the multi wait with the
-  // proper status value.
+  // proper status value. This API returns a function which captures a MultiWait
+  // reference, so care must be taken such that the reference remains valid for
+  // the whole lifetime of the returned function.
   std::function<void()> Completer(std::function<void()> func);
 
+  // Similar as the above API, but with explicit capture of the MultiWait shared
+  // pointer.
+  static std::function<void()> Completer(std::shared_ptr<MultiWait> mwait,
+                                         std::function<void()> func);
+
  private:
+  void Complete(const std::function<void()>& func);
+
   std::mutex mutex_;
   std::condition_variable cv_;
   size_t count_ = 0;
diff --git a/third_party/xla_client/xrt_computation_client.cc b/third_party/xla_client/xrt_computation_client.cc
@@ -302,7 +302,7 @@ std::vector<ComputationClient::DataPtr> XrtComputationClient::TransferToServer(
   }
   XLA_COUNTER("XrtPartitionedTransferToServer", 1);
 
-  util::MultiWait mwait(partitions.size());
+  auto mwait = std::make_shared<util::MultiWait>(partitions.size());
   std::vector<DataPtr> results(tensors.size());
   for (size_t i = 0; i < partitions.size(); ++i) {
     auto sender = [&, i]() {
@@ -316,9 +316,10 @@ std::vector<ComputationClient::DataPtr> XrtComputationClient::TransferToServer(
         results[base_index + r] = std::move(partitions_results[r]);
       }
     };
-    env::ScheduleIoClosure(mwait.Completer(std::move(sender)));
+    env::ScheduleIoClosure(
+        util::MultiWait::Completer(mwait, std::move(sender)));
   }
-  mwait.Wait();
+  mwait->Wait();
   return results;
 }
 
@@ -330,7 +331,7 @@ XrtComputationClient::TransferToServerInternal(
   std::mutex lock;
   XrtSessionCache::SessionMap session_map;
   int64 total_size = 0;
-  util::MultiWait mwait(tensors.size());
+  auto mwait = std::make_shared<util::MultiWait>(tensors.size());
   std::map<XrtSession*, SessionWork> session_work_map;
   {
     metrics::TimedSection timed(TransferToServerTransformMetric());
@@ -363,13 +364,14 @@ XrtComputationClient::TransferToServerInternal(
           total_size += tdata.size();
         }
       };
-      env::ScheduleClosure(mwait.Completer(std::move(converter)));
+      env::ScheduleClosure(
+          util::MultiWait::Completer(mwait, std::move(converter)));
     }
-    mwait.Wait();
+    mwait->Wait();
   }
   OutboundDataMetric()->AddSample(total_size);
 
-  mwait.Reset(session_work_map.size());
+  mwait->Reset(session_work_map.size());
   std::vector<DataPtr> results(tensors.size());
   for (auto& session_session_work : session_work_map) {
     XrtSession* session = session_session_work.first;
@@ -388,9 +390,10 @@ XrtComputationClient::TransferToServerInternal(
       }
       CreateDataHandlesCounter()->AddValue(outputs.size());
     };
-    env::ScheduleIoClosure(mwait.Completer(std::move(runner)));
+    env::ScheduleIoClosure(
+        util::MultiWait::Completer(mwait, std::move(runner)));
   }
-  mwait.Wait();
+  mwait->Wait();
   return results;
 }
 
@@ -426,7 +429,7 @@ std::vector<Literal> XrtComputationClient::TransferFromServer(
     session_work->index_mapping.push_back(i);
   }
 
-  util::MultiWait mwait(session_work_map.size());
+  auto mwait = std::make_shared<util::MultiWait>(session_work_map.size());
   std::atomic<int64> total_size(0);
   std::vector<Literal> results(handles.size());
   for (auto& session_session_work : session_work_map) {
@@ -446,9 +449,10 @@ std::vector<Literal> XrtComputationClient::TransferFromServer(
         total_size += results[li].size_bytes();
       }
     };
-    env::ScheduleIoClosure(mwait.Completer(std::move(runner)));
+    env::ScheduleIoClosure(
+        util::MultiWait::Completer(mwait, std::move(runner)));
   }
-  mwait.Wait();
+  mwait->Wait();
   InboundDataMetric()->AddSample(total_size.load());
   return results;
 }
@@ -458,7 +462,7 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(
   metrics::TimedSection timed(CompileMetric());
 
   std::mutex lock;
-  util::MultiWait mwait(instances.size());
+  auto mwait = std::make_shared<util::MultiWait>(instances.size());
   std::vector<ProgramShape> program_shapes(instances.size());
   std::vector<ComputationPtr> results(instances.size());
   std::vector<CompilationCacheKey> cache_keys(instances.size());
@@ -499,10 +503,10 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(
         results[i] = computation_ptr;
       }
     };
-    env::ScheduleClosure(mwait.Completer(std::move(builder)));
+    env::ScheduleClosure(util::MultiWait::Completer(mwait, std::move(builder)));
   }
-  mwait.Wait();
-  mwait.Reset(session_work_map.size());
+  mwait->Wait();
+  mwait->Reset(session_work_map.size());
 
   for (auto& session_and_work : session_work_map) {
     XrtSession* session = session_and_work.first;
@@ -532,9 +536,10 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(
         CreateCompileHandlesCounter()->AddValue(1);
       }
     };
-    env::ScheduleIoClosure(mwait.Completer(std::move(session_runner)));
+    env::ScheduleIoClosure(
+        util::MultiWait::Completer(mwait, std::move(session_runner)));
   }
-  mwait.Wait();
+  mwait->Wait();
   return results;
 }
 
@@ -626,7 +631,7 @@ XrtComputationClient::RunComputations(
   }
   XLA_CHECK_EQ(computations.size(), devices.size());
 
-  util::MultiWait mwait(session_replicas.size());
+  auto mwait = std::make_shared<util::MultiWait>(session_replicas.size());
   std::vector<std::vector<DataPtr>> results(devices.size());
   for (auto& sess_replica : session_replicas) {
     XrtSession* session = sess_replica.first;
@@ -655,9 +660,10 @@ XrtComputationClient::RunComputations(
             GetEffectiveDevice(devices[replica]));
       }
     };
-    env::ScheduleIoClosure(mwait.Completer(std::move(session_runner)));
+    env::ScheduleIoClosure(
+        util::MultiWait::Completer(mwait, std::move(session_runner)));
   }
-  mwait.Wait();
+  mwait->Wait();
   return results;
 }
 
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -480,7 +480,7 @@ py::bytes ReadTfFile(tensorflow::RandomAccessFile* file, uint64_t offset,
         std::min<size_t>(num_threads, std::thread::hardware_concurrency());
     size_t block_size = size / num_threads;
 
-    xla::util::MultiWait mwait(num_threads);
+    auto mwait = std::make_shared<xla::util::MultiWait>(num_threads);
     for (size_t i = 0; i < num_threads; ++i) {
       auto reader = [&, i]() {
         uint64_t base = static_cast<uint64_t>(i) * block_size;
@@ -491,9 +491,10 @@ py::bytes ReadTfFile(tensorflow::RandomAccessFile* file, uint64_t offset,
         XLA_CHECK_OK(
             file->Read(offset + base, tsize, &result, buffer.get() + base));
       };
-      xla::env::ScheduleIoClosure(mwait.Completer(std::move(reader)));
+      xla::env::ScheduleIoClosure(
+          xla::util::MultiWait::Completer(mwait, std::move(reader)));
     }
-    mwait.Wait();
+    mwait->Wait();
   }
   return py::bytes(buffer.get(), size);
 }
diff --git a/torch_xla/csrc/tensor_util.cpp b/torch_xla/csrc/tensor_util.cpp
@@ -423,15 +423,16 @@ void CopyTensors(const void* src_buffer, const xla::Shape& src_shape,
     std::vector<xla::int64> iter_dims = GetIterationDimensions(dest_shape);
     std::vector<CopyPartition> parts =
         CreateCopyPartitions(dest_shape.dimensions(), iter_dims.front());
-    xla::util::MultiWait mwait(parts.size());
+    auto mwait = std::make_shared<xla::util::MultiWait>(parts.size());
     for (size_t i = 0; i < parts.size(); ++i) {
       auto copy_fn = [&, i]() {
         SlicedCopy<SType, DType>(dest_shape.dimensions(), src_data, src_strides,
                                  dest_data, dest_strides, iter_dims, parts[i]);
       };
-      xla::env::ScheduleClosure(mwait.Completer(std::move(copy_fn)));
+      xla::env::ScheduleClosure(
+          xla::util::MultiWait::Completer(mwait, std::move(copy_fn)));
     }
-    mwait.Wait();
+    mwait->Wait();
   }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -302,7 +302,7 @@ std::vector<ComputationClient::DataPtr> XrtComputationClient::TransferToServer(`
`302`	`302`	`}`
`303`	`303`	`XLA_COUNTER("XrtPartitionedTransferToServer", 1);`
`304`	`304`
`305`		`- util::MultiWait mwait(partitions.size());`
	`305`	`+ auto mwait = std::make_shared<util::MultiWait>(partitions.size());`
`306`	`306`	`std::vector<DataPtr> results(tensors.size());`
`307`	`307`	`for (size_t i = 0; i < partitions.size(); ++i) {`
`308`	`308`	`auto sender = [&, i]() {`
`@@ -316,9 +316,10 @@ std::vector<ComputationClient::DataPtr> XrtComputationClient::TransferToServer(`
`316`	`316`	`results[base_index + r] = std::move(partitions_results[r]);`
`317`	`317`	`}`
`318`	`318`	`};`
`319`		`- env::ScheduleIoClosure(mwait.Completer(std::move(sender)));`
	`319`	`+ env::ScheduleIoClosure(`
	`320`	`+ util::MultiWait::Completer(mwait, std::move(sender)));`
`320`	`321`	`}`
`321`		`- mwait.Wait();`
	`322`	`+ mwait->Wait();`
`322`	`323`	`return results;`
`323`	`324`	`}`
`324`	`325`
`@@ -330,7 +331,7 @@ XrtComputationClient::TransferToServerInternal(`
`330`	`331`	`std::mutex lock;`
`331`	`332`	`XrtSessionCache::SessionMap session_map;`
`332`	`333`	`int64 total_size = 0;`
`333`		`- util::MultiWait mwait(tensors.size());`
	`334`	`+ auto mwait = std::make_shared<util::MultiWait>(tensors.size());`
`334`	`335`	`std::map<XrtSession*, SessionWork> session_work_map;`
`335`	`336`	`{`
`336`	`337`	`metrics::TimedSection timed(TransferToServerTransformMetric());`
`@@ -363,13 +364,14 @@ XrtComputationClient::TransferToServerInternal(`
`363`	`364`	`total_size += tdata.size();`
`364`	`365`	`}`
`365`	`366`	`};`
`366`		`- env::ScheduleClosure(mwait.Completer(std::move(converter)));`
	`367`	`+ env::ScheduleClosure(`
	`368`	`+ util::MultiWait::Completer(mwait, std::move(converter)));`
`367`	`369`	`}`
`368`		`- mwait.Wait();`
	`370`	`+ mwait->Wait();`
`369`	`371`	`}`
`370`	`372`	`OutboundDataMetric()->AddSample(total_size);`
`371`	`373`
`372`		`- mwait.Reset(session_work_map.size());`
	`374`	`+ mwait->Reset(session_work_map.size());`
`373`	`375`	`std::vector<DataPtr> results(tensors.size());`
`374`	`376`	`for (auto& session_session_work : session_work_map) {`
`375`	`377`	`XrtSession* session = session_session_work.first;`
`@@ -388,9 +390,10 @@ XrtComputationClient::TransferToServerInternal(`
`388`	`390`	`}`
`389`	`391`	`CreateDataHandlesCounter()->AddValue(outputs.size());`
`390`	`392`	`};`
`391`		`- env::ScheduleIoClosure(mwait.Completer(std::move(runner)));`
	`393`	`+ env::ScheduleIoClosure(`
	`394`	`+ util::MultiWait::Completer(mwait, std::move(runner)));`
`392`	`395`	`}`
`393`		`- mwait.Wait();`
	`396`	`+ mwait->Wait();`
`394`	`397`	`return results;`
`395`	`398`	`}`
`396`	`399`
`@@ -426,7 +429,7 @@ std::vector<Literal> XrtComputationClient::TransferFromServer(`
`426`	`429`	`session_work->index_mapping.push_back(i);`
`427`	`430`	`}`
`428`	`431`
`429`		`- util::MultiWait mwait(session_work_map.size());`
	`432`	`+ auto mwait = std::make_shared<util::MultiWait>(session_work_map.size());`
`430`	`433`	`std::atomic<int64> total_size(0);`
`431`	`434`	`std::vector<Literal> results(handles.size());`
`432`	`435`	`for (auto& session_session_work : session_work_map) {`
`@@ -446,9 +449,10 @@ std::vector<Literal> XrtComputationClient::TransferFromServer(`
`446`	`449`	`total_size += results[li].size_bytes();`
`447`	`450`	`}`
`448`	`451`	`};`
`449`		`- env::ScheduleIoClosure(mwait.Completer(std::move(runner)));`
	`452`	`+ env::ScheduleIoClosure(`
	`453`	`+ util::MultiWait::Completer(mwait, std::move(runner)));`
`450`	`454`	`}`
`451`		`- mwait.Wait();`
	`455`	`+ mwait->Wait();`
`452`	`456`	`InboundDataMetric()->AddSample(total_size.load());`
`453`	`457`	`return results;`
`454`	`458`	`}`
`@@ -458,7 +462,7 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(`
`458`	`462`	`metrics::TimedSection timed(CompileMetric());`
`459`	`463`
`460`	`464`	`std::mutex lock;`
`461`		`- util::MultiWait mwait(instances.size());`
	`465`	`+ auto mwait = std::make_shared<util::MultiWait>(instances.size());`
`462`	`466`	`std::vector<ProgramShape> program_shapes(instances.size());`
`463`	`467`	`std::vector<ComputationPtr> results(instances.size());`
`464`	`468`	`std::vector<CompilationCacheKey> cache_keys(instances.size());`
`@@ -499,10 +503,10 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(`
`499`	`503`	`results[i] = computation_ptr;`
`500`	`504`	`}`
`501`	`505`	`};`
`502`		`- env::ScheduleClosure(mwait.Completer(std::move(builder)));`
	`506`	`+ env::ScheduleClosure(util::MultiWait::Completer(mwait, std::move(builder)));`
`503`	`507`	`}`
`504`		`- mwait.Wait();`
`505`		`- mwait.Reset(session_work_map.size());`
	`508`	`+ mwait->Wait();`
	`509`	`+ mwait->Reset(session_work_map.size());`
`506`	`510`
`507`	`511`	`for (auto& session_and_work : session_work_map) {`
`508`	`512`	`XrtSession* session = session_and_work.first;`
`@@ -532,9 +536,10 @@ std::vector<ComputationClient::ComputationPtr> XrtComputationClient::Compile(`
`532`	`536`	`CreateCompileHandlesCounter()->AddValue(1);`
`533`	`537`	`}`
`534`	`538`	`};`
`535`		`- env::ScheduleIoClosure(mwait.Completer(std::move(session_runner)));`
	`539`	`+ env::ScheduleIoClosure(`
	`540`	`+ util::MultiWait::Completer(mwait, std::move(session_runner)));`
`536`	`541`	`}`
`537`		`- mwait.Wait();`
	`542`	`+ mwait->Wait();`
`538`	`543`	`return results;`
`539`	`544`	`}`
`540`	`545`
`@@ -626,7 +631,7 @@ XrtComputationClient::RunComputations(`
`626`	`631`	`}`
`627`	`632`	`XLA_CHECK_EQ(computations.size(), devices.size());`
`628`	`633`
`629`		`- util::MultiWait mwait(session_replicas.size());`
	`634`	`+ auto mwait = std::make_shared<util::MultiWait>(session_replicas.size());`
`630`	`635`	`std::vector<std::vector<DataPtr>> results(devices.size());`
`631`	`636`	`for (auto& sess_replica : session_replicas) {`
`632`	`637`	`XrtSession* session = sess_replica.first;`
`@@ -655,9 +660,10 @@ XrtComputationClient::RunComputations(`
`655`	`660`	`GetEffectiveDevice(devices[replica]));`
`656`	`661`	`}`
`657`	`662`	`};`
`658`		`- env::ScheduleIoClosure(mwait.Completer(std::move(session_runner)));`
	`663`	`+ env::ScheduleIoClosure(`
	`664`	`+ util::MultiWait::Completer(mwait, std::move(session_runner)));`
`659`	`665`	`}`
`660`		`- mwait.Wait();`
	`666`	`+ mwait->Wait();`
`661`	`667`	`return results;`
`662`	`668`	`}`
`663`	`669`
Original file line number	Diff line number	Diff line change
`@@ -423,15 +423,16 @@ void CopyTensors(const void* src_buffer, const xla::Shape& src_shape,`
`423`	`423`	`std::vector<xla::int64> iter_dims = GetIterationDimensions(dest_shape);`
`424`	`424`	`std::vector<CopyPartition> parts =`
`425`	`425`	`CreateCopyPartitions(dest_shape.dimensions(), iter_dims.front());`
`426`		`- xla::util::MultiWait mwait(parts.size());`
	`426`	`+ auto mwait = std::make_shared<xla::util::MultiWait>(parts.size());`
`427`	`427`	`for (size_t i = 0; i < parts.size(); ++i) {`
`428`	`428`	`auto copy_fn = [&, i]() {`
`429`	`429`	`SlicedCopy<SType, DType>(dest_shape.dimensions(), src_data, src_strides,`
`430`	`430`	`dest_data, dest_strides, iter_dims, parts[i]);`
`431`	`431`	`};`
`432`		`- xla::env::ScheduleClosure(mwait.Completer(std::move(copy_fn)));`
	`432`	`+ xla::env::ScheduleClosure(`
	`433`	`+ xla::util::MultiWait::Completer(mwait, std::move(copy_fn)));`
`433`	`434`	`}`
`434`		`- mwait.Wait();`
	`435`	`+ mwait->Wait();`
`435`	`436`	`}`
`436`	`437`	`}`
`437`	`438`