Fix race in tensorflow's use of gemmlowp contexts.

Change: 129782045

1 files changed, 7 insertions, 0 deletions

diff --git a/tensorflow/contrib/quantization/kernels/quantization_utils.h b/tensorflow/contrib/quantization/kernels/quantization_utils.h
index 45fda79ce5..3b6a4901ba 100644
--- a/tensorflow/contrib/quantization/kernels/quantization_utils.h
+++ b/tensorflow/contrib/quantization/kernels/quantization_utils.h
@@ -496,6 +496,13 @@ class TensorflowGemmlowpWorkersPool {
   TensorflowGemmlowpWorkersPool(thread::ThreadPool* workers)
       : workers_(workers) {}
 
+  ~TensorflowGemmlowpWorkersPool() {
+    // This workaround ensures that all worker tasks have exited methods in the
+    // BlockingCounter. Without this, there is a race where the context is torn
+    // down while the counter is in use.
+    counter_to_decrement_when_ready_.Reset(0);
+  }
+
   void Prepare(int workers_count) {
     counter_to_decrement_when_ready_.Reset(workers_count);
   }