From 3a063ca6b35a491366408ca924464ee764977f8a Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Mon, 6 Jan 2025 04:39:03 -0800
Subject: [PATCH] [XLA:CPU] Remove CompiledFunctionLibrary from JitCompiler.

PiperOrigin-RevId: 712484420
---
 xla/backends/cpu/codegen/BUILD           |  1 +
 xla/backends/cpu/codegen/jit_compiler.cc | 29 +-----------------------
 xla/backends/cpu/codegen/jit_compiler.h  | 24 --------------------
 3 files changed, 2 insertions(+), 52 deletions(-)

diff --git a/xla/backends/cpu/codegen/BUILD b/xla/backends/cpu/codegen/BUILD
index 27f286addb514..43e17eb30b45d 100644
--- a/xla/backends/cpu/codegen/BUILD
+++ b/xla/backends/cpu/codegen/BUILD
@@ -82,6 +82,7 @@ cc_library(
     srcs = ["jit_compiler.cc"],
     hdrs = ["jit_compiler.h"],
     deps = [
+        ":compiled_function_library",
         ":contiguous_section_memory_manager",
         ":cpu_features",
         ":ir_compiler",
diff --git a/xla/backends/cpu/codegen/jit_compiler.cc b/xla/backends/cpu/codegen/jit_compiler.cc
index 0fd205c551313..7f3acba32e57d 100644
--- a/xla/backends/cpu/codegen/jit_compiler.cc
+++ b/xla/backends/cpu/codegen/jit_compiler.cc
@@ -49,6 +49,7 @@ limitations under the License.
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/TargetParser/Host.h"
+#include "xla/backends/cpu/codegen/compiled_function_library.h"
 #include "xla/backends/cpu/codegen/contiguous_section_memory_manager.h"
 #include "xla/backends/cpu/codegen/cpu_features.h"
 #include "xla/backends/cpu/codegen/ir_compiler.h"
@@ -349,32 +350,4 @@ void JitCompiler::TaskDispatcher::shutdown() {
   absl::MutexLock lock(&mu_, absl::Condition(&all_tasks_finished));
 }
 
-JitCompiler::CompiledFunctionLibrary::CompiledFunctionLibrary(
-    std::unique_ptr<llvm::orc::ExecutionSession> execution_session,
-    std::unique_ptr<llvm::orc::RTDyldObjectLinkingLayer> object_layer,
-    absl::flat_hash_map<std::string, ResolvedSymbol> symbols_map)
-    : execution_session_(std::move(execution_session)),
-      object_layer_(std::move(object_layer)),
-      symbols_map_(std::move(symbols_map)) {
-  DCHECK(execution_session_) << "Execution session must not be null";
-}
-
-JitCompiler::CompiledFunctionLibrary::~CompiledFunctionLibrary() {
-  if (auto err = execution_session_->endSession()) {
-    execution_session_->reportError(std::move(err));
-  }
-}
-
-absl::StatusOr<void*> JitCompiler::CompiledFunctionLibrary::ResolveFunction(
-    TypeId type_id, absl::string_view name) {
-  if (auto it = symbols_map_.find(name); it != symbols_map_.end()) {
-    if (it->second.type_id != type_id) {
-      return Internal("Symbol %s has type id %d, expected %d", name,
-                      it->second.type_id.value(), type_id.value());
-    }
-    return it->second.ptr;
-  }
-  return NotFound("Function %s not found (type id: %d)", name, type_id.value());
-}
-
 }  // namespace xla::cpu
diff --git a/xla/backends/cpu/codegen/jit_compiler.h b/xla/backends/cpu/codegen/jit_compiler.h
index 6e9c3b5d5eb5c..e98a999ddeb52 100644
--- a/xla/backends/cpu/codegen/jit_compiler.h
+++ b/xla/backends/cpu/codegen/jit_compiler.h
@@ -159,30 +159,6 @@ class JitCompiler {
     size_t num_dispatched_tasks_ ABSL_GUARDED_BY(mu_) = 0;
   };
 
-  // Function library constructed from the set of jit-compiled symbols.
-  class CompiledFunctionLibrary : public FunctionLibrary {
-   public:
-    struct ResolvedSymbol {
-      TypeId type_id;
-      void* ptr;
-    };
-
-    CompiledFunctionLibrary(
-        std::unique_ptr<llvm::orc::ExecutionSession> execution_session,
-        std::unique_ptr<llvm::orc::RTDyldObjectLinkingLayer> object_layer,
-        absl::flat_hash_map<std::string, ResolvedSymbol> symbols_map);
-
-    ~CompiledFunctionLibrary() final;
-
-    absl::StatusOr<void*> ResolveFunction(TypeId type_id,
-                                          absl::string_view name) final;
-
-   private:
-    std::unique_ptr<llvm::orc::ExecutionSession> execution_session_;
-    std::unique_ptr<llvm::orc::RTDyldObjectLinkingLayer> object_layer_;
-    absl::flat_hash_map<std::string, ResolvedSymbol> symbols_map_;
-  };
-
   JitCompiler(IrCompiler::TargetMachineBuilder target_machine_builder,
               std::shared_ptr<llvm::TargetMachine> target_machine,
               TaskDispatcher* task_dispatcher,