From 7b477b8752d51f08ab599652cd02c8dd938d2e3e Mon Sep 17 00:00:00 2001
From: Ben Barsdell <bbarsdell@nvidia.com>
Date: Tue, 20 Oct 2020 12:39:47 +1100
Subject: [PATCH] Move blasLt SetBiasPointer call inside mutex lock

- This was not thread-safe before.
---
 tensorflow/stream_executor/cuda/cuda_blas.cc | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index 8b9d9174f22..afea6790187 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -3493,13 +3493,6 @@ bool CUDABlas::DoBlasLtMatmulInternal(
                "epilogue for the given bias pointer.";
     return false;
   }
-  if (bias != nullptr) {
-    if (!cuda_plan.SetBiasPointer(bias.opaque())) {
-      VLOG(2) << "DoBlasLtMatmul returning false because setting the bias "
-                 "pointer failed.";
-      return false;
-    }
-  }
   const void *alpha_ptr = alpha.is_pointer() ? alpha.opaque_pointer().opaque()
                                              : alpha.opaque_value();
   const void *beta_ptr =
@@ -3524,6 +3517,14 @@ bool CUDABlas::DoBlasLtMatmulInternal(
 
   absl::MutexLock lock(&mu_);
 
+  if (bias != nullptr) {
+    if (!cuda_plan.SetBiasPointer(bias.opaque())) {
+      VLOG(2) << "DoBlasLtMatmul returning false because setting the bias "
+                 "pointer failed.";
+      return false;
+    }
+  }
+
   CHECK(blasLt_ != nullptr);
 
   gpu::ScopedActivateExecutorContext sac{parent_};