[xla:cpu] Migrate tf2xla to BufferAllocationInfo

PiperOrigin-RevId: 820707093

Author: ezhulenev

xla/backends/cpu/buffer_allocation_info.h

@@ -22,7 +22,8 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 
-namespace xla::cpu {
+namespace xla {
+namespace cpu {
 
 // `BufferAllocationInfo` stores information about buffer allocations required
 // by an XLA:CPU executable at run time. It corresponds to a `BufferAllocation`
@@ -193,6 +194,14 @@ class BufferAllocationInfo {
   int32_t result_number_ = -1;
 };
 
-}  // namespace xla::cpu
+}  // namespace cpu
+
+// TODO(ezhulenev): This is a temporary hack to keep `tfcompile` code working.
+namespace cpu_function_runtime {
+using BufferInfo = ::xla::cpu::BufferAllocationInfo;
+using EncodedBufferInfo = ::xla::cpu::BufferAllocationInfo::Encoded;
+}  // namespace cpu_function_runtime
+
+}  // namespace xla
 
 #endif  // XLA_BACKENDS_CPU_BUFFER_ALLOCATION_INFO_H_

xla/cpu_function_runtime.h

@@ -22,7 +22,7 @@ limitations under the License.
 #include <cstdlib>
 
 namespace xla {
-namespace cpu_function_runtime {
+namespace cpu_function_runtime_deprecated {
 
 struct EncodedBufferInfo {
   uint64_t packed_kind_and_size = 0;
@@ -174,7 +174,7 @@ class BufferInfo {
   int32_t result_param_number_ = -1;
 };
 
-}  // namespace cpu_function_runtime
+}  // namespace cpu_function_runtime_deprecated
 }  // namespace xla
 
 #endif  // XLA_CPU_FUNCTION_RUNTIME_H_

xla/service/cpu/BUILD

@@ -170,25 +170,12 @@ cc_library(
     alwayslink = True,  # Contains per-platform transfer manager registration
 )
 
-cc_library(
-    name = "buffer_info_util",
-    srcs = ["buffer_info_util.cc"],
-    hdrs = ["buffer_info_util.h"],
-    deps = [
-        "//xla:cpu_function_runtime",
-        "//xla/hlo/ir:hlo",
-        "//xla/service:buffer_assignment",
-        "@com_google_absl//absl/types:span",
-    ],
-)
-
 cc_library(
     name = "cpu_compiler_pure",
     srcs = ["cpu_compiler.cc"],
     hdrs = ["cpu_compiler.h"],
     copts = tsl_copts(),
     deps = [
-        ":buffer_info_util",
         ":conv_canonicalization",
         ":cpu_aot_compilation_result",
         ":cpu_aot_loader",
@@ -427,7 +414,6 @@ cc_library(
     srcs = ["cpu_aot_compilation_result.cc"],
     hdrs = ["cpu_aot_compilation_result.h"],
     deps = [
-        ":buffer_info_util",
         ":cpu_executable",
         ":executable_proto_cc",
         "//xla:cpu_function_runtime",

xla/service/cpu/buffer_info_util.cc

@@ -43,7 +43,6 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/compiler.h"
-#include "xla/service/cpu/buffer_info_util.h"
 #include "xla/service/cpu/cpu_executable.h"
 #include "xla/service/cpu/executable.pb.h"
 #include "xla/service/executable.h"
@@ -57,7 +56,6 @@ limitations under the License.
 #include "xla/util.h"
 
 namespace xla::cpu {
-using BufferInfo = cpu_function_runtime::BufferInfo;
 
 CpuAotCompilationOptions::CpuAotCompilationOptions(
     std::string triple, std::string cpu_name, std::string features,
@@ -88,14 +86,10 @@ CpuAotCompilationResult::Create(
   TF_ASSIGN_OR_RETURN(ThunkSequenceProto thunk_proto,
                       thunk_sequence_serdes.ToProto(thunks));
 
-  std::vector<cpu_function_runtime::BufferInfo> buffer_infos;
   std::vector<cpu::BufferAllocationInfo> buffer_allocation_infos;
   std::optional<size_t> temp_allocation_index;
 
   if (buffer_assignment) {
-    buffer_infos =
-        CreateBufferInfosFromBufferAssignment(*hlo_module, *buffer_assignment);
-
     buffer_allocation_infos =
         CreateBufferAllocationInfos(*hlo_module, *buffer_assignment);
 
@@ -114,21 +108,19 @@ CpuAotCompilationResult::Create(
   return absl::WrapUnique(new CpuAotCompilationResult(
       hlo_module, buffer_assignment, function_name, std::move(obj_files),
       std::move(symbols), thunk_proto, std::move(temp_allocation_index),
-      std::move(buffer_infos), std::move(buffer_allocation_infos),
-      std::move(function_library), std::move(hlo_profile_printer_data)));
+      std::move(buffer_allocation_infos), std::move(function_library),
+      std::move(hlo_profile_printer_data)));
 }
 
 CpuAotCompilationResult::CpuAotCompilationResult(
     const HloModule* hlo_module, const BufferAssignment* buffer_assignment,
     absl::string_view function_name, std::vector<ObjFileProto> obj_files,
     std::vector<SymbolProto> symbols, const ThunkSequenceProto& thunks,
     std::optional<size_t> temp_allocation_index,
-    std::vector<cpu_function_runtime::BufferInfo> buffer_infos,
     std::vector<BufferAllocationInfo> buffer_allocation_infos,
     std::unique_ptr<FunctionLibrary> function_library,
     std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data)
     : temp_allocation_index_(temp_allocation_index),
-      buffer_infos_(std::move(buffer_infos)),
       buffer_allocation_infos_(std::move(buffer_allocation_infos)),
       function_library_(std::move(function_library)),
       hlo_profile_printer_data_(std::move(hlo_profile_printer_data)) {

xla/service/cpu/buffer_info_util.h

@@ -148,10 +148,6 @@ class CpuAotCompilationResult : public AotCompilationResult {
     return temp_allocation_index_;
   }
 
-  const std::vector<cpu_function_runtime::BufferInfo>& buffer_infos() const {
-    return buffer_infos_;
-  }
-
   absl::Span<const BufferAllocationInfo> buffer_allocation_infos() const {
     return buffer_allocation_infos_;
   }
@@ -188,7 +184,6 @@ class CpuAotCompilationResult : public AotCompilationResult {
       absl::string_view function_name, std::vector<ObjFileProto> obj_files,
       std::vector<SymbolProto> symbols, const ThunkSequenceProto& thunks,
       std::optional<size_t> temp_allocation_index,
-      std::vector<cpu_function_runtime::BufferInfo> buffer_infos,
       std::vector<BufferAllocationInfo> buffer_allocation_infos,
       std::unique_ptr<FunctionLibrary> function_library,
       std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data);
@@ -203,7 +198,6 @@ class CpuAotCompilationResult : public AotCompilationResult {
   CompilationResultProto proto_;
   std::unique_ptr<HloModule> module_;
   std::optional<size_t> temp_allocation_index_;
-  std::vector<cpu_function_runtime::BufferInfo> buffer_infos_;
   std::vector<BufferAllocationInfo> buffer_allocation_infos_;
 
   std::unique_ptr<FunctionLibrary> function_library_;

xla/service/cpu/cpu_aot_compilation_result.cc

@@ -172,7 +172,6 @@ limitations under the License.
 #include "xla/service/conditional_simplifier.h"
 #include "xla/service/conditional_to_select.h"
 #include "xla/service/copy_insertion.h"
-#include "xla/service/cpu/buffer_info_util.h"
 #include "xla/service/cpu/conv_canonicalization.h"
 #include "xla/service/cpu/cpu_aot_compilation_result.h"
 #include "xla/service/cpu/cpu_aot_loader.h"