Avoid using constant folding fallback hook concurrently

Ran into case where failed constant folds resulted in non-deterministic test
behavior (e.g., failed to fold unrelated op).

Need to dig into this a bit more.

PiperOrigin-RevId: 286299132
Change-Id: I290dade5fba0fd5b08dae4d696878b9196fc516d

tensorflow/compiler/mlir/tensorflow/BUILD

@@ -625,6 +625,7 @@ cc_library(
         "//tensorflow/c:tf_status",
         "//tensorflow/c/eager:c_api",
         "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
         "//tensorflow/stream_executor",
         "//tensorflow/stream_executor/lib",
         "@llvm//:support",

tensorflow/compiler/mlir/tensorflow/transforms/constant_fold.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/c/eager/c_api.h"
 #include "tensorflow/c/tf_status.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/eval_util.h"
+#include "tensorflow/core/platform/mutex.h"
 
 namespace mlir {
 namespace TF {
@@ -59,6 +60,10 @@ LogicalResult ConstantFoldFallbackHook(
     inputs.push_back(input.cast<ElementsAttr>());
   }
 
+  // Avoid overlapping folds with the same context.
+  // TODO(jpienaar): Avoid using global context & mutex here.
+  static auto* mu = new tensorflow::mutex();
+  tensorflow::mutex_lock l(*mu);
   return tensorflow::EvaluateOperation(inst, inputs, ctx, &results);
 }