[XLA] Move tensorflow/compiler/tf2xla/cpu_function_runtime.h to tensorflow/compiler/xla/cpu_function_runtime.h

PiperOrigin-RevId: 241337813

tensorflow/compiler/aot/BUILD

@@ -33,13 +33,13 @@ cc_library(
         ":aot_only_var_handle_op",
         ":embedded_protocol_buffers",
         "//tensorflow/compiler/tf2xla",
-        "//tensorflow/compiler/tf2xla:cpu_function_runtime",
         "//tensorflow/compiler/tf2xla:tf2xla_proto",
         "//tensorflow/compiler/tf2xla:tf2xla_util",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla/kernels:xla_cpu_only_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
@@ -70,6 +70,7 @@ tf_cc_test(
     ],
     deps = [
         ":tfcompile_lib",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",

tensorflow/compiler/aot/codegen.cc

@@ -25,8 +25,8 @@ limitations under the License.
 #include "absl/strings/str_replace.h"
 #include "absl/types/span.h"
 #include "tensorflow/compiler/aot/embedded_protocol_buffers.h"
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/tf2xla/tf2xla_util.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/service/compiler.h"
 #include "tensorflow/compiler/xla/service/cpu/buffer_info_util.h"
 #include "tensorflow/compiler/xla/shape_util.h"
@@ -38,7 +38,7 @@ namespace tfcompile {
 
 namespace {
 
-using BufferInfo = cpu_function_runtime::BufferInfo;
+using BufferInfo = xla::cpu_function_runtime::BufferInfo;
 
 bool IsAlpha(char c) {
   return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z');
@@ -372,7 +372,7 @@ std::vector<string> BufferInfosToCppExpression(
                            ? "~0ULL"
                            : absl::StrCat(encoded.second, "ULL");
                    return absl::StrCat(
-                       "::tensorflow::cpu_function_runtime::BufferInfo({",
+                       "::xla::cpu_function_runtime::BufferInfo({",
                        encoded.first, "ULL, ", encoded_second_as_str, "})");
                  });
   return buffer_infos_as_strings;
@@ -407,13 +407,15 @@ Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config,
   TF_RETURN_IF_ERROR(GenArgMethods(config, ps, compile_result, &methods_arg));
   TF_RETURN_IF_ERROR(GenResultMethods(config, ps, &methods_result));
   TF_RETURN_IF_ERROR(GenVariableMethods(config, ps, &methods_variable));
-  const size_t arg_bytes_aligned = cpu_function_runtime::AlignedBufferBytes(
+  const size_t arg_bytes_aligned =
-      buffer_infos_for_args.data(), buffer_infos_for_args.size(),
+      xla::cpu_function_runtime::AlignedBufferBytes(
-      /*allocate_entry_params=*/true);
+          buffer_infos_for_args.data(), buffer_infos_for_args.size(),
+          /*allocate_entry_params=*/true);
   const size_t arg_bytes_total = TotalBufferBytes(buffer_infos_for_args);
-  const size_t temp_bytes_aligned = cpu_function_runtime::AlignedBufferBytes(
+  const size_t temp_bytes_aligned =
-      buffer_infos_for_temps.data(), buffer_infos_for_temps.size(),
+      xla::cpu_function_runtime::AlignedBufferBytes(
-      /*allocate_entry_params=*/true);
+          buffer_infos_for_temps.data(), buffer_infos_for_temps.size(),
+          /*allocate_entry_params=*/true);
   const size_t temp_bytes_total = TotalBufferBytes(buffer_infos_for_temps);
 
   // Create rewrite strings for namespace start and end.
@@ -607,8 +609,8 @@ class {{CLASS}} final : public tensorflow::XlaCompiledCpuFunction {
   // Number of buffers for the compiled computation.
   static constexpr size_t kNumBuffers = {{NUM_BUFFERS}};
 
-  static const ::tensorflow::cpu_function_runtime::BufferInfo* BufferInfos() {
+  static const ::xla::cpu_function_runtime::BufferInfo* BufferInfos() {
-    static const ::tensorflow::cpu_function_runtime::BufferInfo
+    static const ::xla::cpu_function_runtime::BufferInfo
       kBufferInfos[kNumBuffers] = {
 {{BUFFER_INFOS_AS_STRING}}
       };

tensorflow/compiler/aot/codegen_test.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 #include "llvm/Support/TargetSelect.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/types.pb.h"
@@ -34,7 +35,7 @@ namespace tensorflow {
 namespace tfcompile {
 namespace {
 
-using ::tensorflow::cpu_function_runtime::BufferInfo;
+using ::xla::cpu_function_runtime::BufferInfo;
 
 void ExpectErrorContains(const Status& status, absl::string_view str) {
   EXPECT_NE(Status::OK(), status);

tensorflow/compiler/aot/codegen_test_h.golden

@@ -274,15 +274,15 @@ class MyClass final : public tensorflow::XlaCompiledCpuFunction {
   // Number of buffers for the compiled computation.
   static constexpr size_t kNumBuffers = 6;
 
-  static const ::tensorflow::cpu_function_runtime::BufferInfo* BufferInfos() {
+  static const ::xla::cpu_function_runtime::BufferInfo* BufferInfos() {
-    static const ::tensorflow::cpu_function_runtime::BufferInfo
+    static const ::xla::cpu_function_runtime::BufferInfo
       kBufferInfos[kNumBuffers] = {
-::tensorflow::cpu_function_runtime::BufferInfo({5ULL, ~0ULL}),
+::xla::cpu_function_runtime::BufferInfo({5ULL, ~0ULL}),
-::tensorflow::cpu_function_runtime::BufferInfo({34ULL, 0ULL}),
+::xla::cpu_function_runtime::BufferInfo({34ULL, 0ULL}),
-::tensorflow::cpu_function_runtime::BufferInfo({9ULL, ~0ULL}),
+::xla::cpu_function_runtime::BufferInfo({9ULL, ~0ULL}),
-::tensorflow::cpu_function_runtime::BufferInfo({386ULL, 1ULL}),
+::xla::cpu_function_runtime::BufferInfo({386ULL, 1ULL}),
-::tensorflow::cpu_function_runtime::BufferInfo({13ULL, ~0ULL}),
+::xla::cpu_function_runtime::BufferInfo({13ULL, ~0ULL}),
-::tensorflow::cpu_function_runtime::BufferInfo({481ULL, ~0ULL})
+::xla::cpu_function_runtime::BufferInfo({481ULL, ~0ULL})
       };
     return kBufferInfos;
   }

tensorflow/compiler/tf2xla/BUILD

@@ -103,22 +103,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "cpu_function_runtime",
-    srcs = ["cpu_function_runtime.cc"],
-    hdrs = ["cpu_function_runtime.h"],
-    visibility = [
-        "//tensorflow/compiler/aot:__pkg__",
-        "//tensorflow/compiler/xla/service/cpu:__pkg__",
-    ],
-    deps = [
-        # Keep dependencies to a minimum here; this library is used in every AOT
-        # binary produced by tfcompile.
-        "//tensorflow/compiler/xla:executable_run_options",
-        "//tensorflow/core:framework_lite",
-    ],
-)
-
 cc_library(
     name = "xla_compiled_cpu_function",
     srcs = ["xla_compiled_cpu_function.cc"],
@@ -127,7 +111,7 @@ cc_library(
     deps = [
         # Keep dependencies to a minimum here; this library is used in every AOT
         # binary produced by tfcompile.
-        ":cpu_function_runtime",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla:executable_run_options",
         "//tensorflow/core:framework_lite",
     ],
@@ -137,7 +121,7 @@ tf_cc_test(
     name = "cpu_function_runtime_test",
     srcs = ["cpu_function_runtime_test.cc"],
     deps = [
-        ":cpu_function_runtime",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/core:framework",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
@@ -153,6 +137,7 @@ cc_library(
         ":tf2xla",
         ":tf2xla_proto",
         ":xla_compiled_cpu_function",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:xla_data_proto",

tensorflow/compiler/tf2xla/cpu_function_runtime_test.cc

@@ -13,15 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
-
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
 namespace {
 
-using cpu_function_runtime::BufferInfo;
+using ::xla::cpu_function_runtime::BufferInfo;
 
 TEST(XlaCompiledCpuFunctionTest, AlignmentValue) {
   // We've chosen 64 byte alignment for the tfcompile runtime to mimic the
@@ -29,7 +28,7 @@ TEST(XlaCompiledCpuFunctionTest, AlignmentValue) {
   // The tfcompile runtime also has a requirement that comes from the xla
   // generated code, on the relation: buffer_size >= 16 ? 2 * sizeof(void*) : 8
   // So any value that we choose must abide by that constraint as well.
-  EXPECT_EQ(cpu_function_runtime::kAlign, Allocator::kAllocatorAlignment);
+  EXPECT_EQ(xla::cpu_function_runtime::kAlign, Allocator::kAllocatorAlignment);
 }
 
 std::vector<BufferInfo> SizesToBufferInfos(const intptr_t* sizes, size_t n) {
@@ -91,7 +90,7 @@ TEST(XlaCompiledCpuFunctionTest, MallocFreeContiguousBuffers) {
   // Test empty sizes.
   void* base = MallocContiguousBuffersFromSizes(nullptr, 0, nullptr, false);
   EXPECT_EQ(base, nullptr);
-  cpu_function_runtime::FreeContiguous(base);
+  xla::cpu_function_runtime::FreeContiguous(base);
 
   // Test non-empty sizes with 0 sum.
   static constexpr intptr_t sizesA[1] = {-1};
@@ -99,7 +98,7 @@ TEST(XlaCompiledCpuFunctionTest, MallocFreeContiguousBuffers) {
   base = MallocContiguousBuffersFromSizes(sizesA, 1, bufA, false);
   EXPECT_EQ(base, nullptr);
   EXPECT_EQ(bufA[0], nullptr);
-  cpu_function_runtime::FreeContiguous(base);
+  xla::cpu_function_runtime::FreeContiguous(base);
 
   // Test non-empty sizes with non-0 sum.
   static constexpr intptr_t sizesB[1] = {3};
@@ -111,7 +110,7 @@ TEST(XlaCompiledCpuFunctionTest, MallocFreeContiguousBuffers) {
   bufB0_bytes[0] = 'A';
   bufB0_bytes[1] = 'B';
   bufB0_bytes[2] = 'C';
-  cpu_function_runtime::FreeContiguous(base);
+  xla::cpu_function_runtime::FreeContiguous(base);
 
   // Test non-empty sizes with non-0 sum, and annotate_initialized.
   static constexpr intptr_t sizesC[1] = {3};
@@ -123,7 +122,7 @@ TEST(XlaCompiledCpuFunctionTest, MallocFreeContiguousBuffers) {
   bufC0_bytes[0] = 'A';
   bufC0_bytes[1] = 'B';
   bufC0_bytes[2] = 'C';
-  cpu_function_runtime::FreeContiguous(base);
+  xla::cpu_function_runtime::FreeContiguous(base);
 
   // Test mixed sizes.
   static constexpr intptr_t sizesD[7] = {1, -1, 32, -1, 64, 2, 3};
@@ -146,7 +145,7 @@ TEST(XlaCompiledCpuFunctionTest, MallocFreeContiguousBuffers) {
       }
     }
   }
-  cpu_function_runtime::FreeContiguous(base);
+  xla::cpu_function_runtime::FreeContiguous(base);
 }
 
 void CheckRoundTripIsOk(const BufferInfo& buffer_info) {

tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
 
 #include <cassert>
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 
 namespace tensorflow {
 
@@ -34,7 +35,7 @@ XlaCompiledCpuFunction::XlaCompiledCpuFunction(const StaticData& static_data,
   bool allocate_entry_params =
       alloc_mode == AllocMode::ARGS_RESULTS_PROFILES_AND_TEMPS;
   // Allocate arg and temp buffers.
-  alloc_buffer_table_ = cpu_function_runtime::MallocContiguousBuffers(
+  alloc_buffer_table_ = xla::cpu_function_runtime::MallocContiguousBuffers(
       static_data.buffer_infos_, static_data.num_buffers_,
       /*allocate_entry_params=*/allocate_entry_params, buffer_table_,
       /*annotate_initialized=*/true);
@@ -55,7 +56,7 @@ bool XlaCompiledCpuFunction::Run() {
 }
 
 XlaCompiledCpuFunction::~XlaCompiledCpuFunction() {
-  cpu_function_runtime::FreeContiguous(alloc_buffer_table_);
+  xla::cpu_function_runtime::FreeContiguous(alloc_buffer_table_);
   delete[] buffer_table_;
   delete[] profile_counters_;
 }

tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h

@@ -19,7 +19,7 @@ limitations under the License.
 #include <cassert>
 #include <string>
 
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/executable_run_options.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -66,7 +66,7 @@ class XlaCompiledCpuFunction {
     RawFunction raw_function_;
 
     // Contains information about the buffers used by the XLA computation.
-    const cpu_function_runtime::BufferInfo* buffer_infos_ = nullptr;
+    const xla::cpu_function_runtime::BufferInfo* buffer_infos_ = nullptr;
     size_t num_buffers_ = 0;
 
     // Entry parameter i is described by
@@ -260,7 +260,7 @@ class XlaCompiledCpuFunction {
 
   static void set_static_data_buffer_infos(
       StaticData* static_data,
-      const cpu_function_runtime::BufferInfo* buffer_infos) {
+      const xla::cpu_function_runtime::BufferInfo* buffer_infos) {
     static_data->buffer_infos_ = buffer_infos;
   }
 
@@ -324,7 +324,7 @@ class XlaCompiledCpuFunction {
   void** const buffer_table_;
 
   // Describes the buffers used by the XLA computation.
-  const cpu_function_runtime::BufferInfo* const buffer_infos_;
+  const xla::cpu_function_runtime::BufferInfo* const buffer_infos_;
 
   // Argument i needs to be placed in buffer_table_[arg_index_to_temp_index_[i]]
   // for XLA generated code to be able to find it.

tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc

@@ -24,6 +24,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/client/client_library.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
 #include "tensorflow/compiler/xla/client/xla_computation.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/service/cpu/buffer_info_util.h"
 #include "tensorflow/compiler/xla/service/cpu/cpu_executable.h"
 #include "tensorflow/compiler/xla/shape_util.h"
@@ -120,7 +121,7 @@ XlaJitCompiledCpuFunction::Compile(
       cpu_executable->buffer_assignment();
 
   // Compute buffer infos and the result index, needed to run the raw function.
-  std::vector<cpu_function_runtime::BufferInfo> buffer_infos =
+  std::vector<xla::cpu_function_runtime::BufferInfo> buffer_infos =
       xla::cpu::CreateBufferInfosFromBufferAssignment(buffer_assignment);
   std::vector<int32> arg_index_table =
       xla::cpu::CreateArgIndexTableFromBufferInfos(buffer_infos);

tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h

@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
 #include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/platform/types.h"
@@ -67,7 +68,7 @@ class XlaJitCompiledCpuFunction {
   XlaCompiledCpuFunction::StaticData static_data_;
 
   // The backing array for buffer infos.
-  std::vector<cpu_function_runtime::BufferInfo> buffer_infos_;
+  std::vector<xla::cpu_function_runtime::BufferInfo> buffer_infos_;
 
   // The backing array for the arg index table.
   std::vector<int32> arg_index_table_;

tensorflow/compiler/xla/BUILD

@@ -854,6 +854,14 @@ cc_library(
         ],
 )
 
+cc_library(
+    name = "cpu_function_runtime",
+    srcs = ["cpu_function_runtime.cc"],
+    hdrs = ["cpu_function_runtime.h"],
+    visibility = [":friends"],
+    deps = ["//tensorflow/core:framework_lite"],
+)
+
 tf_cc_test(
     name = "debug_options_parsers_test",
     size = "small",

tensorflow/compiler/xla/cpu_function_runtime.cc

@@ -13,11 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 
 #include "tensorflow/core/platform/dynamic_annotations.h"
 
-namespace tensorflow {
+namespace xla {
 namespace {
 // Inline memory allocation routines here, because depending on '//base' brings
 // in libraries which use c++ streams, which adds considerable code size on
@@ -105,4 +105,4 @@ void FreeContiguous(void* contiguous) {
   }
 }
 }  // namespace cpu_function_runtime
-}  // namespace tensorflow
+}  // namespace xla

tensorflow/compiler/xla/cpu_function_runtime.h

@@ -13,14 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_COMPILER_TF2XLA_CPU_FUNCTION_RUNTIME_H_
+#ifndef TENSORFLOW_COMPILER_XLA_CPU_FUNCTION_RUNTIME_H_
-#define TENSORFLOW_COMPILER_TF2XLA_CPU_FUNCTION_RUNTIME_H_
+#define TENSORFLOW_COMPILER_XLA_CPU_FUNCTION_RUNTIME_H_
 
 #include "tensorflow/core/platform/types.h"
 
 #include <cassert>
 
-namespace tensorflow {
+namespace xla {
 namespace cpu_function_runtime {
 // Stores information about one buffer used by an XLA:CPU compiled function.
 // These buffers are used for holding inputs to the computation, outputs from
@@ -28,10 +28,11 @@ namespace cpu_function_runtime {
 class BufferInfo {
  public:
   // Creates a BufferInfo from a serialized encoding generated by `Encode`.
-  explicit BufferInfo(std::pair<uint64, uint64> encoding)
+  explicit BufferInfo(
+      std::pair<tensorflow::uint64, tensorflow::uint64> encoding)
       : entry_param_number_(encoding.second) {
     Kind kind;
-    uint64 size;
+    tensorflow::uint64 size;
     Unpack(encoding.first, &kind, &size);
     kind_ = kind;
     size_ = size;
@@ -47,7 +48,7 @@ class BufferInfo {
   bool is_entry_parameter() const { return kind() == Kind::kEntryParameter; }
 
   // Returns the entry parameter number of this buffer.
-  uint64 entry_parameter_number() const {
+  tensorflow::uint64 entry_parameter_number() const {
     assert(is_entry_parameter());
     return entry_param_number_;
   }
@@ -61,16 +62,16 @@ class BufferInfo {
   bool is_on_stack_buffer() const { return kind() == Kind::kOnStackBuffer; }
 
   // Returns the size for this buffer.
-  uint64 size() const { return size_; }
+  tensorflow::uint64 size() const { return size_; }
 
   // Encodes this BufferInfo into two 64 bit integers that can be used to
   // reconstruct the BufferInfo later using the constructor.  We need this
   // because we use BufferInfo in places where using protocol buffers would
   // negatively impact binary size.
-  std::pair<uint64, uint64> Encode() const {
+  std::pair<tensorflow::uint64, tensorflow::uint64> Encode() const {
     static_assert(sizeof(*this) == 16, "");
-    uint64 upper = Pack(kind(), size_);
+    tensorflow::uint64 upper = Pack(kind(), size_);
-    uint64 lower = entry_param_number_;
+    tensorflow::uint64 lower = entry_param_number_;
     return {upper, lower};
   }
 
@@ -84,19 +85,20 @@ class BufferInfo {
 
   // Factory methods:
 
-  static BufferInfo MakeTempBuffer(uint64 size) {
+  static BufferInfo MakeTempBuffer(tensorflow::uint64 size) {
     return BufferInfo(Kind::kTempBuffer, /*size=*/size,
                       /*entry_param_number=*/-1);
   }
-  static BufferInfo MakeConstant(uint64 size) {
+  static BufferInfo MakeConstant(tensorflow::uint64 size) {
     return BufferInfo(Kind::kConstant, /*size=*/size,
                       /*entry_param_number=*/-1);
   }
-  static BufferInfo MakeEntryParameter(uint64 size, uint64 param_number) {
+  static BufferInfo MakeEntryParameter(tensorflow::uint64 size,
+                                       tensorflow::uint64 param_number) {
     return BufferInfo(Kind::kEntryParameter, /*size=*/size,
                       /*entry_param_number=*/param_number);
   }
-  static BufferInfo MakeOnStackBuffer(uint64 size) {
+  static BufferInfo MakeOnStackBuffer(tensorflow::uint64 size) {
     return BufferInfo(Kind::kOnStackBuffer, /*size=*/size,
                       /*entry_param_number=*/-1);
   }
@@ -104,7 +106,7 @@ class BufferInfo {
  private:
   BufferInfo() = default;
 
-  enum class Kind : uint64 {
+  enum class Kind : tensorflow::uint64 {
     kConstant,
     kTempBuffer,
     kEntryParameter,
@@ -113,21 +115,24 @@ class BufferInfo {
 
   Kind kind() const { return static_cast<Kind>(kind_); }
 
-  explicit BufferInfo(Kind kind, uint64 size, uint64 entry_param_number)
+  explicit BufferInfo(Kind kind, tensorflow::uint64 size,
+                      tensorflow::uint64 entry_param_number)
       : kind_(kind), size_(size), entry_param_number_(entry_param_number) {}
 
-  static uint64 Pack(Kind kind, uint64 size) {
+  static tensorflow::uint64 Pack(Kind kind, tensorflow::uint64 size) {
-    return (static_cast<uint64>(size) << 2) | static_cast<uint64>(kind);
+    return (static_cast<tensorflow::uint64>(size) << 2) |
+           static_cast<tensorflow::uint64>(kind);
   }
 
-  static void Unpack(uint64 packed, Kind* kind, uint64* size) {
+  static void Unpack(tensorflow::uint64 packed, Kind* kind,
+                     tensorflow::uint64* size) {
     *size = packed >> 2;
     *kind = static_cast<Kind>((packed << 62) >> 62);
   }
 
   Kind kind_ : 2;
-  uint64 size_ : 62;
+  tensorflow::uint64 size_ : 62;
-  int64 entry_param_number_;
+  tensorflow::int64 entry_param_number_;
 };
 
 // Align to 64-bytes, to mimic tensorflow::Allocator::kAllocatorAlignment.
@@ -160,6 +165,6 @@ void* MallocContiguousBuffers(const BufferInfo* buffer_infos, size_t n,
 // MallocContiguousBuffers.
 void FreeContiguous(void* contiguous);
 }  // namespace cpu_function_runtime
-}  // namespace tensorflow
+}  // namespace xla
 
-#endif  // TENSORFLOW_COMPILER_TF2XLA_CPU_FUNCTION_RUNTIME_H_
+#endif  // TENSORFLOW_COMPILER_XLA_CPU_FUNCTION_RUNTIME_H_

tensorflow/compiler/xla/service/cpu/BUILD

@@ -62,7 +62,7 @@ cc_library(
     srcs = ["buffer_info_util.cc"],
     hdrs = ["buffer_info_util.h"],
     deps = [
-        "//tensorflow/compiler/tf2xla:cpu_function_runtime",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla/service:buffer_assignment",
         "//tensorflow/core:lib",
         "@com_google_absl//absl/types:span",
@@ -93,13 +93,13 @@ cc_library(
         "@com_google_absl//absl/strings",
         ":target_machine_features",
         "@com_google_absl//absl/types:span",
-        "//tensorflow/compiler/tf2xla:cpu_function_runtime",
         "//tensorflow/compiler/xla/service:hlo_casting_utils",
         "//tensorflow/compiler/xla/service:dump",
         "//tensorflow/compiler/xla/service:map_inliner",
         "//tensorflow/compiler/xla/service:hlo_get_dimension_size_rewriter",
         "//tensorflow/compiler/xla/service:conditional_to_select",
         "//tensorflow/compiler/xla/service:scatter_expander",
+        "//tensorflow/compiler/xla:cpu_function_runtime",
         "//tensorflow/compiler/xla:literal",
         "//tensorflow/compiler/xla:protobuf_util",
         "//tensorflow/compiler/xla:status_macros",

tensorflow/compiler/xla/service/cpu/buffer_info_util.cc

@@ -14,11 +14,12 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/compiler/xla/service/cpu/buffer_info_util.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 
 namespace xla {
 namespace cpu {
 
-using BufferInfo = ::tensorflow::cpu_function_runtime::BufferInfo;
+using BufferInfo = cpu_function_runtime::BufferInfo;
 
 std::vector<BufferInfo> CreateBufferInfosFromBufferAssignment(
     const BufferAssignment& buffer_assignment) {

tensorflow/compiler/xla/service/cpu/buffer_info_util.h

@@ -17,14 +17,14 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_BUFFER_INFO_UTIL_H_
 
 #include "absl/types/span.h"
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/service/buffer_assignment.h"
 
 namespace xla {
 namespace cpu {
 // Creates and returns a list of BufferInfo instances containing relevant
 // information from `buffer_assignment`.
-std::vector<::tensorflow::cpu_function_runtime::BufferInfo>
+std::vector<cpu_function_runtime::BufferInfo>
 CreateBufferInfosFromBufferAssignment(
     const BufferAssignment& buffer_assignment);
 
@@ -34,8 +34,7 @@ CreateBufferInfosFromBufferAssignment(
 // If this function returns V then entry parameter i has buffer allocation index
 // V[i].
 std::vector<int32> CreateArgIndexTableFromBufferInfos(
-    absl::Span<const ::tensorflow::cpu_function_runtime::BufferInfo>
+    absl::Span<const cpu_function_runtime::BufferInfo> buffer_infos);
-        buffer_infos);
 }  // namespace cpu
 }  // namespace xla
 

tensorflow/compiler/xla/service/cpu/cpu_compiler.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/literal.h"
 #include "tensorflow/compiler/xla/map_util.h"
 #include "tensorflow/compiler/xla/protobuf_util.h"
@@ -113,7 +114,7 @@ limitations under the License.
 
 namespace xla {
 namespace cpu {
-using BufferInfo = ::tensorflow::cpu_function_runtime::BufferInfo;
+using BufferInfo = cpu_function_runtime::BufferInfo;
 
 CpuAotCompilationOptions::CpuAotCompilationOptions(
     string triple, string cpu_name, string features, string entry_point_name,

tensorflow/compiler/xla/service/cpu/cpu_compiler.h

@@ -20,7 +20,7 @@ limitations under the License.
 
 #include "absl/types/span.h"
 #include "llvm/Target/TargetMachine.h"
-#include "tensorflow/compiler/tf2xla/cpu_function_runtime.h"
+#include "tensorflow/compiler/xla/cpu_function_runtime.h"
 #include "tensorflow/compiler/xla/service/cpu/target_machine_features.h"
 #include "tensorflow/compiler/xla/service/executable.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
@@ -81,7 +81,7 @@ class CpuAotCompilationResult : public AotCompilationResult {
  public:
   CpuAotCompilationResult(
       ObjectFileData object_file_data,
-      std::vector<::tensorflow::cpu_function_runtime::BufferInfo> buffer_infos,
+      std::vector<cpu_function_runtime::BufferInfo> buffer_infos,
       int64 result_buffer_index,
       std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data);
   ~CpuAotCompilationResult();
@@ -91,8 +91,7 @@ class CpuAotCompilationResult : public AotCompilationResult {
   }
 
   const ObjectFileData& object_file_data() const { return object_file_data_; }
-  const std::vector<::tensorflow::cpu_function_runtime::BufferInfo>&
+  const std::vector<cpu_function_runtime::BufferInfo>& buffer_infos() const {
-  buffer_infos() const {
     return buffer_infos_;
   }
   int64 result_buffer_index() const { return result_buffer_index_; }
@@ -103,8 +102,7 @@ class CpuAotCompilationResult : public AotCompilationResult {
 
   // A list of BufferInfo objects describing the buffers used by the XLA
   // computation.
-  const std::vector<::tensorflow::cpu_function_runtime::BufferInfo>
+  const std::vector<cpu_function_runtime::BufferInfo> buffer_infos_;
-      buffer_infos_;
 
   // Contains which buffer index into |buffer_sizes| was designated to the
   // result of the computation.  This buffer should be passed into the output