Internal change

PiperOrigin-RevId: 544222320

.bazelrc

@@ -22,6 +22,13 @@ build --@rules_cuda//cuda:compiler="Please specify a compiler (--config=clang/gc
 build --incompatible_enforce_config_setting_visibility
 build --incompatible_config_setting_private_default_visibility
 
+# Disable clang extention that rejects type definitions within offsetof. 
+# This was added in clang-16 by https://reviews.llvm.org/D133574.
+# Can be removed once upb is updated, since a type definition is used within
+# offset of in the current version of ubp.
+# See https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183.
+build --copt=-Wno-gnu-offsetof-extensions
+
 # LLVM's LIT test runner needs %PATHEXT% on Windows
 test:windows --test_env=PATHEXT
 

BUILD

@@ -17,6 +17,9 @@ package_group(
     name = "friends",
     packages = [
         "//...",
+        "//third_party/tensorflow/compiler/mlir/tfrt/...",
+        "//third_party/tensorflow/core/runtime_fallback/...",
+        "//third_party/tensorflow/core/tfrt/...",
     ],
 )
 
@@ -128,7 +131,6 @@ bool_flag(
 config_setting(
     name = "gpu_enabled_oss",
     flag_values = {":enable_gpu": "True"},
-    visibility = ["//visibility:private"],
 )
 
 # Config setting to conditionally link GPU targets.
@@ -149,7 +151,6 @@ alias(
 #         "//devtools/blueprint/bluze/public:bluze_ncl",
 #         "//devtools/blueprint/ncl:sanitizer",
 #     ],
-#     deps = ["//devtools/bazel/subteams/configurability/android_platforms_migration:flags_ncl"],
 # )
 #
 # filegroup(
@@ -726,6 +727,7 @@ tfrt_cc_library(
         ":bef_location_emitter",
         ":core_runtime_opdefs",
         ":dtype",
+        ":metrics",
         ":stream_analysis",
         ":support",
         "@llvm-project//llvm:Support",
@@ -1525,6 +1527,7 @@ tfrt_cc_library(
     deps = [
         ":basic_kernels_opdefs",
         ":compiler_tfrt_op_interfaces",
+        "@com_google_absl//absl/log",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

WORKSPACE

@@ -26,15 +26,13 @@ load("@bazel_skylib//lib:versions.bzl", "versions")
 versions.check(minimum_bazel_version = "5.1.0")
 
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
-load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure", "llvm_disable_optional_support_deps")
+load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 
 maybe(
     llvm_configure,
     name = "llvm-project",
 )
 
-llvm_disable_optional_support_deps()
-
 load("@rules_cuda//cuda:dependencies.bzl", "rules_cuda_dependencies")
 
 rules_cuda_dependencies()

backends/common/BUILD

@@ -158,7 +158,7 @@ tfrt_cc_library(
         ],
         "//conditions:default": [],
     }),
-    visibility = ["@tf_runtime//:friends"],
+    visibility = ["//visibility:public"],
     deps = [
         "@eigen_archive//:eigen3",
         "@llvm-project//llvm:Support",

backends/common/include/tfrt/common/compat/eigen/eigen_kernel.h

@@ -85,8 +85,7 @@ template <typename Tin, typename Tout, typename Fn>
 llvm::Error UnaryEigenKernel(
     const DenseHostTensor& input,
     // `output` supplies the buffer in which to write the output.
-    DenseHostTensor* output, Fn fn, const ExecutionContext& exec_ctx) {
-  HostContext* host = exec_ctx.host();
+    DenseHostTensor* output, Fn fn, HostContext& host_ctx) {
   auto input_view = DHTArrayView<Tin>(&input);
   auto output_view = MutableDHTArrayView<Tout>(output);
   const TensorShape& shape_input = input.shape();
@@ -98,7 +97,7 @@ llvm::Error UnaryEigenKernel(
   auto in = AsEigenConstTensor(input_view);
   auto out = AsEigenTensor(output_view);
   auto expr = fn(in, out);
-  out.device(host->GetOrCreateSharedContext<EigenHostContext>().Device()) =
+  out.device(host_ctx.GetOrCreateSharedContext<EigenHostContext>().Device()) =
       expr;
   return llvm::Error::success();
 }
@@ -141,8 +140,7 @@ template <typename Tin, typename Tout, typename Fn>
 llvm::Error BinaryEigenKernel(
     const DenseHostTensor& left, const DenseHostTensor& right,
     // `output` supplies the buffer in which to write the output.
-    DenseHostTensor* output, Fn fn, const ExecutionContext& exec_ctx) {
-  HostContext* host = exec_ctx.host();
+    DenseHostTensor* output, Fn fn, HostContext& host_ctx) {
   auto left_view = DHTArrayView<Tin>(&left);
   auto right_view = DHTArrayView<Tin>(&right);
   auto output_view = MutableDHTArrayView<Tout>(output);
@@ -157,23 +155,22 @@ llvm::Error BinaryEigenKernel(
   auto out_tensor = AsEigenTensor(output_view);
   auto expr = fn(left_tensor, right_tensor, out_tensor);
   out_tensor.device(
-      host->GetOrCreateSharedContext<EigenHostContext>().Device()) = expr;
+      host_ctx.GetOrCreateSharedContext<EigenHostContext>().Device()) = expr;
   return llvm::Error::success();
 }
 
 template <typename Tin, typename Tout, typename Fn>
 llvm::Error BinaryEigenKernelBroadcast(
     const DenseHostTensor& left, Tin s,
     // `output` supplies the buffer in which to write the output.
-    DenseHostTensor* output, Fn fn, const ExecutionContext& exec_ctx) {
-  HostContext* host = exec_ctx.host();
+    DenseHostTensor* output, Fn fn, HostContext& host_ctx) {
   auto left_view = DHTArrayView<Tin>(&left);
   auto output_view = MutableDHTArrayView<Tout>(output);
   auto left_tensor = AsEigenConstTensor(left_view);
   auto out_tensor = AsEigenTensor(output_view);
   auto expr = fn(left_tensor, s, out_tensor);
   out_tensor.device(
-      host->GetOrCreateSharedContext<EigenHostContext>().Device()) = expr;
+      host_ctx.GetOrCreateSharedContext<EigenHostContext>().Device()) = expr;
   return llvm::Error::success();
 }
 

backends/common/lib/ops/tf/dnn_ops_util.cc

@@ -47,7 +47,7 @@ llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
 }
 
 ChannelOrder GetTfChannelOrder(std::optional<string_view> data_format) {
-  if (!data_format.has_value() || !data_format->startswith_insensitive("nc"))
+  if (!data_format.has_value() || !data_format->starts_with_insensitive("nc"))
     return ChannelOrder::ChannelLast;
   return ChannelOrder::ChannelFirst;
 }

backends/cpu/BUILD

@@ -24,7 +24,7 @@ tfrt_cc_library(
     ],
     alwayslink_static_registration_src = "lib/core_runtime/static_registration.cc",
     # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"],
-    visibility = ["@tf_runtime//:friends"],
+    visibility = ["//visibility:public"],
     deps = [
         "@llvm-project//llvm:Support",
         "@tf_runtime//:core_runtime",

backends/cpu/cpp_tests/kernels/cwise_binary_kernels_test.cc

@@ -63,7 +63,8 @@ void BinaryKernel(benchmark::State& state, int num_threads,
     tfrt::latch done(1);
 
     ::tfrt::cpu::BinaryKernel<Functor, compat::AsyncEigenEvaluator>(
-        *lhs, *rhs, &*res, exec_ctx, [&](Error err) { done.count_down(); });
+        *lhs, *rhs, &*res, *exec_ctx.host(),
+        [&](Error err) { done.count_down(); });
 
     done.wait();
   }

backends/cpu/lib/kernels/cpu_kernels.h

@@ -206,9 +206,9 @@ AsyncValueRef<Chain> Relu(const DenseHostTensor& A, DenseHostTensor* B,
 // Computes B = Relu(A) in sync style.
 template <typename T>
 llvm::Error SyncRelu(const DenseHostTensor& A, DenseHostTensor* B,
-                     const ExecutionContext& exec_ctx) {
+                     HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.cwiseMax(static_cast<T>(0)); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 //===----------------------------------------------------------------------===//
@@ -326,9 +326,9 @@ AsyncValueRef<Chain> BiasAdd(const DenseHostTensor& input,
 // Computes B = Sigmoid(A) in sync style.
 template <typename T>
 llvm::Error SyncSigmoid(const DenseHostTensor& A, DenseHostTensor* B,
-                        const ExecutionContext& exec_ctx) {
+                        HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.sigmoid(); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 //===----------------------------------------------------------------------===//
@@ -337,8 +337,7 @@ llvm::Error SyncSigmoid(const DenseHostTensor& A, DenseHostTensor* B,
 
 // Computes B = Sum(A) in sync style for 2D tensors (only).
 template <typename T>
-llvm::Error SyncSum2D(const DenseHostTensor& A, DenseHostTensor* B, int axis,
-                      const ExecutionContext& exec_ctx) {
+llvm::Error SyncSum2D(const DenseHostTensor& A, DenseHostTensor* B, int axis) {
   if (A.shape().GetRank() != 2 || B->shape().GetRank() != 2) {
     return MakeStringError("Inputs must be rank-2 tensors.");
   }
@@ -364,9 +363,9 @@ llvm::Error SyncSum2D(const DenseHostTensor& A, DenseHostTensor* B, int axis,
 // Computes B = Exp(A) in sync style.
 template <typename T>
 llvm::Error SyncExp(const DenseHostTensor& A, DenseHostTensor* B,
-                    const ExecutionContext& exec_ctx) {
+                    HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.exp(); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 //===----------------------------------------------------------------------===//
@@ -376,50 +375,50 @@ llvm::Error SyncExp(const DenseHostTensor& A, DenseHostTensor* B,
 // Computes c = Greater(A, B) in sync style.
 template <typename T>
 llvm::Error SyncGreater(const DenseHostTensor& A, const DenseHostTensor& B,
-                        DenseHostTensor* C, const ExecutionContext& exec_ctx) {
+                        DenseHostTensor* C, HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b, auto& c) { return a > b; };
   if (B.shape().GetRank() == 0) {
     return ::tfrt::compat::BinaryEigenKernelBroadcast<T, bool>(
-        A, B.data<T>()[0], C, std::move(fn), exec_ctx);
+        A, B.data<T>()[0], C, std::move(fn), host_ctx);
   } else {
     return ::tfrt::compat::BinaryEigenKernel<T, bool>(A, B, C, std::move(fn),
-                                                      exec_ctx);
+                                                      host_ctx);
   }
 }
 
 // Computes c = Min(A, B) in sync style.
 template <typename T>
 llvm::Error SyncMinimum(const DenseHostTensor& A, const DenseHostTensor& B,
-                        DenseHostTensor* C, const ExecutionContext& exec_ctx) {
+                        DenseHostTensor* C, HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b, auto& c) { return a.cwiseMin(b); };
 
   if (B.shape().GetRank() == 0) {
     return ::tfrt::compat::BinaryEigenKernelBroadcast<T, T>(
-        A, B.data<T>()[0], C, std::move(fn), exec_ctx);
+        A, B.data<T>()[0], C, std::move(fn), host_ctx);
   } else {
     return ::tfrt::compat::BinaryEigenKernel<T, T>(A, B, C, std::move(fn),
-                                                   exec_ctx);
+                                                   host_ctx);
   }
 }
 
 // Computes c = Max(A, B) in sync style.
 template <typename T>
 llvm::Error SyncMaximum(const DenseHostTensor& A, const DenseHostTensor& B,
-                        DenseHostTensor* C, const ExecutionContext& exec_ctx) {
+                        DenseHostTensor* C, HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b, auto& c) { return a.cwiseMax(b); };
 
   if (B.shape().GetRank() == 0) {
     return ::tfrt::compat::BinaryEigenKernelBroadcast<T, T>(
-        A, B.data<T>()[0], C, std::move(fn), exec_ctx);
+        A, B.data<T>()[0], C, std::move(fn), host_ctx);
   } else {
     return ::tfrt::compat::BinaryEigenKernel<T, T>(A, B, C, std::move(fn),
-                                                   exec_ctx);
+                                                   host_ctx);
   }
 }
 
 template <typename T>
 llvm::Error SyncSoftplus(const DenseHostTensor& input, DenseHostTensor* output,
-                         const ExecutionContext& exec_ctx) {
+                         HostContext& host_ctx) {
   auto fn = [](auto& a, auto& c) {
     static const T threshold =
         Eigen::numext::log(Eigen::NumTraits<T>::epsilon()) + T(2);
@@ -436,24 +435,23 @@ llvm::Error SyncSoftplus(const DenseHostTensor& input, DenseHostTensor* output,
                          features_exp.log1p()));
   };
   return ::tfrt::compat::UnaryEigenKernel<T, T>(input, output, std::move(fn),
-                                                exec_ctx);
+                                                host_ctx);
 }
 
 // Computes B = Sqrt(A) in sync style.
 template <typename T>
 llvm::Error SyncSqrt(const DenseHostTensor& A, DenseHostTensor* B,
-                     const ExecutionContext& exec_ctx) {
+                     HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.sqrt(); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 // Computes B = Mean(A) in sync style for 2D tensors (only). Note this supports
 // reduction only along 1 axis for now
 template <typename T, int NDIMS>
 llvm::Error SyncMean2D(const DenseHostTensor& input,
                        Eigen::array<int, 1> reduction_axes,
-                       DenseHostTensor* result,
-                       const ExecutionContext& exec_ctx) {
+                       DenseHostTensor* result, HostContext& host_ctx) {
   auto input_view = DHTIndexableView<T, NDIMS>(
       input.data<T>(), FixedRankShape<NDIMS>(input.shape()));
 
@@ -463,7 +461,7 @@ llvm::Error SyncMean2D(const DenseHostTensor& input,
   auto result_tensor_eigen = AsEigenTensor(result_view);
   Eigen::internal::SumReducer<T> sum_reducer;
   result_tensor_eigen.device(
-      exec_ctx.host()->GetOrCreateSharedContext<EigenHostContext>().Device()) =
+      host_ctx.GetOrCreateSharedContext<EigenHostContext>().Device()) =
       a_tensor_eigen.reduce(reduction_axes, sum_reducer) /
       static_cast<T>(a_tensor_eigen.size() / result_tensor_eigen.size());
 
@@ -473,26 +471,26 @@ llvm::Error SyncMean2D(const DenseHostTensor& input,
 // Computes B = Log1p(A) in sync style.
 template <typename T>
 llvm::Error SyncLog1p(const DenseHostTensor& A, DenseHostTensor* B,
-                      const ExecutionContext& exec_ctx) {
+                      HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.log1p(); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 // Computes B = Tanh(A) in sync style.
 template <typename T>
 llvm::Error SyncTanh(const DenseHostTensor& A, DenseHostTensor* B,
-                     const ExecutionContext& exec_ctx) {
+                     HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.tanh(); };
-  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), exec_ctx);
+  return ::tfrt::compat::UnaryEigenKernel<T, T>(A, B, std::move(fn), host_ctx);
 }
 
 // Computes B = Cast(A) in sync style.
 template <typename SrcT, typename DstT>
 llvm::Error SyncCast(const DenseHostTensor& A, DenseHostTensor* B,
-                     const ExecutionContext& exec_ctx) {
+                     HostContext& host_ctx) {
   auto fn = [](auto& a, auto& b) { return a.template cast<DstT>(); };
   return ::tfrt::compat::UnaryEigenKernel<SrcT, DstT>(A, B, std::move(fn),
-                                                      exec_ctx);
+                                                      host_ctx);
 }
 
 }  // namespace cpu

backends/cpu/lib/kernels/cwise_binary_kernels.h

@@ -343,12 +343,11 @@ struct BinaryKernelImpl {
 
 template <typename BinaryFunctor, typename EigenEvaluator, typename OnDone>
 void BinaryKernel(const HostTensor& lhs, const HostTensor& rhs,
-                  HostTensor* output, const ExecutionContext& exec_ctx,
-                  OnDone on_done) {
+                  HostTensor* output, HostContext& host_ctx, OnDone on_done) {
   using T = typename BinaryFunctor::Input;
 
   internal::BinaryKernelImpl<BinaryFunctor, EigenEvaluator> impl(
-      EigenEvaluator{exec_ctx.host()});
+      EigenEvaluator{&host_ctx});
 
   if (isa<ScalarHostTensor<T>>(lhs) && isa<ScalarHostTensor<T>>(rhs)) {
     impl.ScalarScalar(lhs, rhs, output, std::move(on_done));
@@ -379,7 +378,7 @@ AsyncValueRef<Chain> BinaryKernel(const HostTensor& lhs, const HostTensor& rhs,
   };
 
   BinaryKernel<BinaryFunctor, compat::AsyncEigenEvaluator>(
-      lhs, rhs, output, exec_ctx, std::move(on_done));
+      lhs, rhs, output, *exec_ctx.host(), std::move(on_done));
 
   return chain;
 }
@@ -389,8 +388,8 @@ Error SyncBinaryKernel(const HostTensor& lhs, const HostTensor& rhs,
                        HostTensor* output, const ExecutionContext& exec_ctx) {
   Error error = Error::success();
   auto on_done = [&](Error err) { error = std::move(err); };
-  BinaryKernel<BinaryFunctor, compat::SyncEigenEvaluator>(lhs, rhs, output,
-                                                          exec_ctx, on_done);
+  BinaryKernel<BinaryFunctor, compat::SyncEigenEvaluator>(
+      lhs, rhs, output, *exec_ctx.host(), on_done);
   return error;
 }
 

backends/cpu/lib/kernels/image/image_kernels.cc

@@ -39,7 +39,7 @@ static AsyncValueRef<DenseHostTensor> DecodeJpeg(
   EnqueueWork(exec_ctx, [data = data.ValueRef(), output = output.CopyRef(),
                          exec_ctx] {
     TFRT_TRACE_SCOPE(Default, "DecodeJpeg");
-    if (!llvm::StringRef(data.get()).startswith("\xff\xd8\xff")) {
+    if (!llvm::StringRef(data.get()).starts_with("\xff\xd8\xff")) {
       auto diag = EmitError(exec_ctx, "image does not have jpeg format");
       output.SetError(diag.status);
       return;

backends/cpu/lib/kernels/image/jpeg/jpeg_mem.h

@@ -23,6 +23,7 @@ limitations under the License.
 #ifndef TFRT_BACKENDS_CPU_LIB_KERNELS_IMAGE_JPEG_JPEG_MEM_H_
 #define TFRT_BACKENDS_CPU_LIB_KERNELS_IMAGE_JPEG_JPEG_MEM_H_
 
+#include <cstdio>
 #include <functional>
 
 extern "C" {