[pull] master from tensorflow:master

tensorflow/compiler/jit/BUILD

@@ -588,10 +588,9 @@ cc_library(
 tf_cuda_cc_test(
     name = "pjrt_compile_util_test",
     srcs = ["pjrt_compile_util_test.cc"],
-    tags = [
+    tags = tf_cuda_tests_tags() + [
         "config-cuda-only",
         "no_oss",  # This test only runs with GPU.
-        "requires-gpu-nvidia",
         "xla",
     ],
     deps = [
@@ -1976,10 +1975,9 @@ tf_cc_test(
 tf_cuda_cc_test(
     name = "device_compiler_test",
     srcs = ["device_compiler_test.cc"],
-    tags = [
+    tags = tf_cuda_tests_tags() + [
         "config-cuda-only",
         "no_oss",  # This test only runs with GPU.
-        "requires-gpu-nvidia",
         "xla",
     ],
     deps = [

tensorflow/compiler/tests/BUILD

@@ -2935,7 +2935,6 @@ tf_cuda_cc_test(
         "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
         "no_rocm",  # ROCmSoftwarePlatform #958
         "noasan",  # TODO(b/201651800)
-        "requires-gpu-nvidia",
     ] + tf_cuda_tests_tags(),
     deps = [":randomized_tests_library"],
 )
@@ -2956,7 +2955,6 @@ tf_cuda_cc_test(
         "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
         "no_rocm",  # ROCmSoftwarePlatform #958
         "noasan",  # TODO(b/201651800)
-        "requires-gpu-nvidia",
     ] + tf_cuda_tests_tags(),
     deps = [":randomized_tests_library"],
 )

tensorflow/core/common_runtime/eager/BUILD

@@ -326,11 +326,15 @@ tf_cc_test(
         "//tensorflow/core/kernels:cast_op",
         "//tensorflow/core/kernels:logging_ops",
         "//tensorflow/core/kernels:math",
+        "//tensorflow/core/platform:env",
         "//tensorflow/core/platform:refcount",
         "//tensorflow/core/platform:status",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:refcount",
+        "@xla//xla/tsl/platform:env",
         "@xla//xla/tsl/platform:status",
     ],
 )

tensorflow/core/common_runtime/eager/context.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <functional>
 #include <memory>
+#include <string>
 #include <thread>
 #include <unordered_map>
 #include <unordered_set>
@@ -27,6 +28,7 @@ limitations under the License.
 // clang-format off
 // Required for IS_MOBILE_PLATFORM
 #include "absl/status/status.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/core/common_runtime/function.h"
@@ -618,6 +620,28 @@ void EagerContext::WaitForAndCloseRemoteContexts() {
 #endif  // !IS_MOBILE_PLATFORM
 }
 
+namespace internal {
+bool IsTensorFlowExecutorThread() {
+  std::string name;
+  return tsl::Env::Default()->GetCurrentThreadName(&name) &&
+         absl::StartsWith(name, "tf_");
+}
+}  // namespace internal
+
+void EagerContext::Release() {
+  // If we are on a TF executor thread, we must not execute the destructor
+  // synchronously to avoid a deadlock when it waits on the same thread pool.
+  // We unconditionally detach a thread rather than checking RefCountIsOne()
+  // to prevent TOCTOU races where another thread drops the last reference
+  // concurrently.
+  if (internal::IsTensorFlowExecutorThread()) {
+    env_->StartDetachedThread(ThreadOptions(), "eager_context_cleanup",
+                              [this]() { this->Unref(); });
+  } else {
+    Unref();
+  }
+}
+
 EagerContext::~EagerContext() {
   // TODO(iga): Add a separate API method to shutdown EagerContext so that we
   // don't send RPCs and block in destructor.

tensorflow/core/common_runtime/eager/context.h

@@ -79,6 +79,11 @@ namespace eager {
 class RemoteMgr;
 }  // namespace eager
 
+namespace internal {
+// Returns true if the calling thread's name begins with "tf_".
+bool IsTensorFlowExecutorThread();
+}  // namespace internal
+
 // Check the value of the environment variable,
 // `TF_REMOTE_HANDLE_SKIP_WAIT_FOR_READY` from its cached copy in memory and if
 // not cached, reads from the environment variable.
@@ -105,7 +110,7 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
       CollectiveExecutorMgrInterface* collective_executor_mgr = nullptr,
       bool run_eager_op_as_function = false, bool jit_compile_rewrite = false);
 
-  void Release() override { Unref(); }
+  void Release() override;
 
   AbstractTensorInterface* CreateInt64Scalar(int64_t value) override;
   AbstractTensorInterface* CreateUint64Scalar(uint64_t value) override;

tensorflow/core/common_runtime/eager/context_test.cc

@@ -15,18 +15,22 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/eager/context.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include "absl/status/status.h"
+#include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
 #include "tensorflow/c/eager/abstract_tensor_handle.h"
 #include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/c/eager/immediate_execution_operation.h"
 #include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/status.h"
 #include "tensorflow/core/common_runtime/composite_device.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
@@ -36,11 +40,13 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session_options.h"
+#include "tsl/platform/refcount.h"
 
 namespace tensorflow {
 namespace {
@@ -69,12 +75,15 @@ class EagerContextTest : public ::testing::Test {
   EagerContext* context() { return context_.get(); }
 
   void InitContext(const SessionOptions& opts,
-                   ContextDevicePlacementPolicy policy, bool async = false) {
+                   ContextDevicePlacementPolicy policy, bool async = false,
+                   tsl::core::RefCountPtr<Rendezvous> rendezvous = nullptr) {
     ASSERT_EQ(context_, nullptr);
-    InitDeviceManager();
+    if (!device_manager_) {
+      InitDeviceManager();
+    }
     context_ = core::RefCountPtr<EagerContext>(new EagerContext(
         opts, policy, async, device_manager_.get(),
-        /*device_mgr_owned=*/false, /*rendezvous=*/nullptr,
+        /*device_mgr_owned=*/false, std::move(rendezvous),
         /*cluster_flr=*/nullptr, /*collective_executor_mgr=*/nullptr,
         /*run_eager_op_as_function=*/true));
   }
@@ -414,5 +423,70 @@ TEST_F(EagerContextTest, ReuseGlobalRendezvous) {
   TestGlobalRendezvous(context(), true);
 }
 
+TEST_F(EagerContextTest, IsTensorFlowExecutorThreadTest) {
+  {
+    std::unique_ptr<Thread> t(tsl::Env::Default()->StartThread(
+        ThreadOptions(), "tf_test",
+        []() { EXPECT_TRUE(internal::IsTensorFlowExecutorThread()); }));
+  }
+  {
+    std::unique_ptr<Thread> t(tsl::Env::Default()->StartThread(
+        ThreadOptions(), "other_name",
+        []() { EXPECT_FALSE(internal::IsTensorFlowExecutorThread()); }));
+  }
+}
+
+class NotificationRendezvous : public Rendezvous {
+ public:
+  explicit NotificationRendezvous(absl::Notification* n, int64_t* thread_id)
+      : n_(n), destructor_thread_id_(thread_id) {}
+  ~NotificationRendezvous() override {
+    if (destructor_thread_id_) {
+      *destructor_thread_id_ = tsl::Env::Default()->GetCurrentThreadId();
+    }
+    n_->Notify();
+  }
+
+  absl::Status Send(const ParsedKey& key, const Args& args, const Tensor& val,
+                    const bool is_dead) override {
+    return absl::OkStatus();
+  }
+  void RecvAsync(const ParsedKey& key, const Args& args,
+                 DoneCallback done) override {}
+  void StartAbort(const absl::Status& status) override {}
+
+ private:
+  absl::Notification* n_;
+  int64_t* destructor_thread_id_ = nullptr;
+};
+
+TEST_F(EagerContextTest, ReleaseOnTfThread) {
+  absl::Notification n;
+  int64_t rendezvous_destructor_thread_id = 0;
+  int64_t tf_thread_id = 0;
+  auto rendezvous = core::RefCountPtr<Rendezvous>(
+      new NotificationRendezvous(&n, &rendezvous_destructor_thread_id));
+  InitDeviceManager();
+  InitContext(SessionOptions(), DEVICE_PLACEMENT_EXPLICIT, false,
+              std::move(rendezvous));
+
+  {
+    std::unique_ptr<Thread> t(
+        tsl::Env::Default()->StartThread(ThreadOptions(), "tf_test", [&]() {
+          tf_thread_id = tsl::Env::Default()->GetCurrentThreadId();
+          // Drop the last reference to context_. Because the thread is named
+          // "tf_test", EagerContext::Release() will detach a thread to execute
+          // Unref().
+          context_.release()->Release();
+        }));
+  }
+
+  // Wait deterministically for the detached thread to finish the deletion
+  // by waiting for the EagerContext's rendezvous member to be destroyed.
+  n.WaitForNotification();
+
+  EXPECT_NE(rendezvous_destructor_thread_id, tf_thread_id);
+}
+
 }  // namespace
 }  // namespace tensorflow

tensorflow/dtensor/python/tests/test_util.py

@@ -54,7 +54,9 @@
 
 DEFAULT_TOL = 1e-5
 
-_DEFAULT_GPU_MEMORY_LIMIT = 1024  # 1G
+# Limit each virtual GPU to 512MB so 8 virtual devices plus CUDA/cuSolver
+# runtime overhead fit within smaller GPU memory partitions.
+_DEFAULT_GPU_MEMORY_LIMIT = 512
 
 
 def get_use_xla_spmd(device_type):

tensorflow/lite/c/CMakeLists.txt

@@ -18,17 +18,10 @@ project(tensorflow-lite-c C CXX)
 
 option(TFLITE_C_BUILD_SHARED_LIBS "Build shared libraries" ON)
 
-set(TF_SOURCE_DIR "" CACHE PATH
-  "Directory that contains the TensorFlow project"
+get_filename_component(TFLITE_SOURCE_DIR
+  "${CMAKE_CURRENT_LIST_DIR}/.."
+  ABSOLUTE
 )
-if (NOT TF_SOURCE_DIR)
-  get_filename_component(TF_SOURCE_DIR
-    "${CMAKE_CURRENT_LIST_DIR}/../../../"
-    ABSOLUTE
-  )
-endif()
-
-set(TFLITE_SOURCE_DIR "${TF_SOURCE_DIR}/tensorflow/lite")
 
 add_subdirectory(
   "${TFLITE_SOURCE_DIR}"
@@ -67,7 +60,7 @@ add_library(tensorflowlite_c ${TFLITE_C_LIBTYPE}
   ${TFLITE_SOURCE_DIR}/core/c/c_api_experimental.cc
   ${TFLITE_SOURCE_DIR}/core/c/common.cc
   ${TFLITE_SOURCE_DIR}/core/c/operator.cc
-  ${TF_SOURCE_DIR}/tensorflow/compiler/mlir/lite/core/c/tflite_types.h
+  ${TFLITE_SOURCE_DIR}/converter/core/c/tflite_types.h
   builtin_op_data.h
   c_api.h
   c_api_experimental.h

tensorflow/lite/core/c/common.cc

@@ -265,6 +265,20 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
       dst_params->quantized_dimension = src_params->quantized_dimension;
       break;
     }
+    case kTfLiteBlockwiseQuantizationV2: {
+      dst.params = calloc(1, sizeof(TfLiteBlockwiseQuantizationV2));
+      if (!dst.params) return TfLiteQuantization();
+      const TfLiteBlockwiseQuantizationV2* const src_params =
+          reinterpret_cast<TfLiteBlockwiseQuantizationV2*>(src.params);
+      TfLiteBlockwiseQuantizationV2* const dst_params =
+          reinterpret_cast<TfLiteBlockwiseQuantizationV2*>(dst.params);
+      dst_params->blocksize = src_params->blocksize;
+      dst_params->scale = src_params->scale;
+      dst_params->zero_point = src_params->zero_point;
+      dst_params->quantized_dimension = src_params->quantized_dimension;
+      dst_params->block_shape = TfLiteIntArrayCopy(src_params->block_shape);
+      break;
+    }
     case kTfLiteMultiAxisQuantization: {
       dst.params = calloc(1, sizeof(TfLiteMultiAxisQuantization));
       if (!dst.params) return TfLiteQuantization();
@@ -386,6 +400,17 @@ void TfLiteQuantizationFree(TfLiteQuantization* quantization) {
         reinterpret_cast<TfLiteBlockwiseQuantization*>(quantization->params);
     free(q_params);
   }
+  if (quantization->type == kTfLiteBlockwiseQuantizationV2) {
+    TfLiteBlockwiseQuantizationV2* q_params =
+        reinterpret_cast<TfLiteBlockwiseQuantizationV2*>(quantization->params);
+    if (q_params) {
+      if (q_params->block_shape) {
+        TfLiteIntArrayFree(q_params->block_shape);
+        q_params->block_shape = nullptr;
+      }
+      free(q_params);
+    }
+  }
   if (quantization->type == kTfLiteMultiAxisQuantization) {
     TfLiteMultiAxisQuantization* q_params =
         reinterpret_cast<TfLiteMultiAxisQuantization*>(quantization->params);

tensorflow/lite/core/c/common.h

@@ -358,6 +358,9 @@ typedef enum TfLiteQuantizationType {
   kTfLiteBlockwiseQuantization = 2,
   /// Multi-axis quantization.
   kTfLiteMultiAxisQuantization = 3,
+  /// N-D blockwise quantization.
+  /// Corresponds to TfLiteBlockwiseQuantizationV2.
+  kTfLiteBlockwiseQuantizationV2 = 4,
 } TfLiteQuantizationType;
 
 /// Structure specifying the quantization used by the tensor, if-any.
@@ -397,6 +400,24 @@ typedef struct TfLiteBlockwiseQuantization {
   int32_t quantized_dimension;
 } TfLiteBlockwiseQuantization;
 
+/// Parameters for N-D blockwise quantization.
+/// `block_shape->data[d]` is the extent of a single block along dimension `d`,
+/// so the scale and zero point tensors have shape
+/// `ceil(tensor_shape / block_shape)`. `block_shape` is owned by this struct.
+typedef struct TfLiteBlockwiseQuantizationV2 {
+  // Index of the tensor containing the scales.
+  int32_t scale;
+  // Index of the tensor containing the zero points.
+  int32_t zero_point;
+  // Quantization blocksize. Only expresses blocking along the last dimension.
+  // Ignored when `block_shape` is non-null.
+  int32_t blocksize;
+  int32_t quantized_dimension;
+  // N-D block shape, with one entry per dimension of the quantized tensor.
+  // Owned by this struct.
+  TfLiteIntArray* block_shape;
+} TfLiteBlockwiseQuantizationV2;
+
 /// Parameters for multi-axis quantization. The scales and zero_points fields
 /// refer to tensor indices in the same subgraph. If zero_points is -1, zero
 /// point is assumed to be 0. For per-channel quantization, blocksize is 0.