Reverts 56bc4f8

PiperOrigin-RevId: 820748684

Author: Google-ML-Automation

xla/backends/cpu/buffer_allocation_info.h

@@ -22,8 +22,7 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 
-namespace xla {
-namespace cpu {
+namespace xla::cpu {
 
 // `BufferAllocationInfo` stores information about buffer allocations required
 // by an XLA:CPU executable at run time. It corresponds to a `BufferAllocation`
@@ -194,14 +193,6 @@ class BufferAllocationInfo {
   int32_t result_number_ = -1;
 };
 
-}  // 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
+}  // namespace xla::cpu
 
 #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_deprecated {
+namespace cpu_function_runtime {
 
 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_deprecated
+}  // namespace cpu_function_runtime
 }  // namespace xla
 
 #endif  // XLA_CPU_FUNCTION_RUNTIME_H_

xla/service/cpu/BUILD

@@ -170,12 +170,25 @@ 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",
@@ -414,6 +427,7 @@ 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

@@ -0,0 +1,101 @@
+/* Copyright 2018 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/service/cpu/buffer_info_util.h"
+
+#include <cassert>
+#include <cstdint>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "xla/cpu_function_runtime.h"
+
+namespace xla {
+namespace cpu {
+
+using BufferInfo = cpu_function_runtime::BufferInfo;
+
+std::vector<BufferInfo> CreateBufferInfosFromBufferAssignment(
+    const HloModule& module, const BufferAssignment& buffer_assignment) {
+  std::vector<BufferInfo> buffer_infos;
+  for (const BufferAllocation& allocation : buffer_assignment.Allocations()) {
+    if (allocation.is_thread_local()) {
+      buffer_infos.push_back(BufferInfo::MakeOnStackBuffer(allocation.size()));
+    } else if (allocation.is_constant()) {
+      buffer_infos.push_back(BufferInfo::MakeConstant(allocation.size()));
+    } else if (allocation.is_entry_computation_parameter()) {
+      buffer_infos.push_back(BufferInfo::MakeEntryParameter(
+          /*size=*/allocation.size(),
+          /*param_number=*/allocation.parameter_number()));
+    } else {
+      buffer_infos.push_back(BufferInfo::MakeTempBuffer(allocation.size()));
+    }
+  }
+
+  // Fill in the result parameters' indices, expanding all tuples.
+  auto root_instr = module.entry_computation()->root_instruction();
+  auto output_allocation = buffer_assignment.GetUniqueTopLevelOutputSlice();
+  if (output_allocation->allocation()->is_tuple()) {
+    int out_index = 0;
+    ShapeUtil::ForEachSubshape(
+        root_instr->shape(),
+        [&](const Shape& subshape, const ShapeIndex& index) {
+          if (subshape.IsTuple()) {
+            return;
+          }
+          int64_t result_index =
+              buffer_assignment.GetUniqueSlice(root_instr, index)->index();
+          assert(result_index < buffer_infos.size());
+          buffer_infos[result_index].set_result_parameter_number(out_index++);
+        });
+  }
+
+  return buffer_infos;
+}
+
+std::vector<int32_t> CreateArgIndexTableFromBufferInfos(
+    absl::Span<const BufferInfo> buffer_infos) {
+  std::vector<int32_t> ret;
+  for (int64_t i = 0; i < buffer_infos.size(); i++) {
+    if (!buffer_infos[i].is_entry_parameter()) {
+      continue;
+    }
+    uint64_t param_index = buffer_infos[i].entry_parameter_number();
+    if (param_index >= ret.size()) {
+      ret.resize(param_index + 1);
+    }
+    ret[param_index] = i;
+  }
+  return ret;
+}
+
+std::vector<int32_t> CreateResultIndexTableFromBufferInfos(
+    absl::Span<const BufferInfo> buffer_infos) {
+  std::vector<int32_t> ret;
+  for (int64_t i = 0; i < buffer_infos.size(); i++) {
+    if (!buffer_infos[i].is_result_parameter()) {
+      continue;
+    }
+    uint64_t result_index = buffer_infos[i].result_parameter_number();
+    if (result_index >= ret.size()) {
+      ret.resize(result_index + 1);
+    }
+    ret[result_index] = i;
+  }
+  return ret;
+}
+
+}  // namespace cpu
+}  // namespace xla

xla/service/cpu/buffer_info_util.h

@@ -0,0 +1,48 @@
+/* Copyright 2018 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_CPU_BUFFER_INFO_UTIL_H_
+#define XLA_SERVICE_CPU_BUFFER_INFO_UTIL_H_
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "xla/cpu_function_runtime.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "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<cpu_function_runtime::BufferInfo>
+CreateBufferInfosFromBufferAssignment(
+    const HloModule& module, const BufferAssignment& buffer_assignment);
+
+// Creates and returns a table containing the mapping from entry computation
+// parameters to buffer allocation indices.
+//
+// If this function returns V then entry parameter i has buffer allocation index
+// V[i].
+std::vector<int32_t> CreateArgIndexTableFromBufferInfos(
+    absl::Span<const cpu_function_runtime::BufferInfo> buffer_infos);
+
+std::vector<int32_t> CreateResultIndexTableFromBufferInfos(
+    absl::Span<const cpu_function_runtime::BufferInfo> buffer_infos);
+}  // namespace cpu
+}  // namespace xla
+
+#endif  // XLA_SERVICE_CPU_BUFFER_INFO_UTIL_H_

xla/service/cpu/cpu_aot_compilation_result.cc

@@ -43,6 +43,7 @@ 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"
@@ -56,6 +57,7 @@ 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,
@@ -86,10 +88,14 @@ 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);
 
@@ -108,19 +114,21 @@ 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_allocation_infos), std::move(function_library),
-      std::move(hlo_profile_printer_data)));
+      std::move(buffer_infos), 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/cpu_aot_compilation_result.h

@@ -148,6 +148,10 @@ 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_;
   }
@@ -184,6 +188,7 @@ 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);
@@ -198,6 +203,7 @@ 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_compiler.cc

@@ -172,6 +172,7 @@ 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"