[xla:runtime] NFC: Extract executable library from jitrt and move it to xla/runtime

PiperOrigin-RevId: 467332094

Author: ezhulenev

backends/jitrt/BUILD

@@ -111,6 +111,7 @@ licenses(["notice"])
 #         "//third_party/tensorflow/compiler/xla/runtime:custom_call",
 #         "//third_party/tensorflow/compiler/xla/runtime:custom_call_registry",
 #         "//third_party/tensorflow/compiler/xla/runtime:diagnostics",
+#         "//third_party/tensorflow/compiler/xla/runtime:executable",
 #         "//third_party/tensorflow/compiler/xla/runtime:execution_engine",
 #         "//third_party/tensorflow/compiler/xla/runtime:memory_mapper",
 #         "//third_party/tensorflow/compiler/xla/runtime:symbolic_shape",

backends/jitrt/include/tfrt/jitrt/jitrt.h

@@ -47,6 +47,7 @@
 #include "third_party/tensorflow/compiler/xla/runtime/constraints.h"
 #include "third_party/tensorflow/compiler/xla/runtime/custom_call.h"
 #include "third_party/tensorflow/compiler/xla/runtime/diagnostics.h"
+#include "third_party/tensorflow/compiler/xla/runtime/executable.h"
 #include "third_party/tensorflow/compiler/xla/runtime/execution_engine.h"
 #include "third_party/tensorflow/compiler/xla/runtime/memory_mapper.h"
 #include "third_party/tensorflow/compiler/xla/runtime/symbolic_shape.h"
@@ -173,14 +174,6 @@ namespace jitrt {
 // concrete shapes or values if needed.
 class JitExecutable;
 
-// Forward declare the Executable class that represents a fully compiled module,
-// which in practice means that it has a function pointer to the compiled
-// function, and knows how to execute it, and return results to the caller.
-class Executable;
-
-// Converts a custom call library into the execution engine symbols binding.
-ExecutionEngine::SymbolsBinding GetSymbolsBinding(DirectCustomCallLibrary lib);
-
 struct CompilationOptions {
   // Calling convention defines an ABI for JitRt to call a compiled kernel. See
   // documentation and example below.
@@ -316,261 +309,6 @@ struct CompilationOptions {
 // otherwise.
 Expected<MemrefDesc> ConvertTensorToMemrefDesc(const Tensor& tensor);
 
-//----------------------------------------------------------------------------//
-// Result of compiling MLIR module to executable kernel function.
-//----------------------------------------------------------------------------//
-
-namespace internal {
-class JitCompilationContext;
-}  // namespace internal
-
-class Executable {
- public:
-  // Forward declare types defined below.
-  struct ArgumentsMemoryLayout;
-  struct ResultsMemoryLayout;
-  struct CallFrame;
-  struct ExecuteOpts;
-  struct KernelContext;
-
-  // Initializes call frame by adding all arguments according to the compiled
-  // kernel ABI. Also allocates storage for the returned values according to the
-  // results memory layout.
-  //
-  // If `verify_arguments` is true (in debug mode it's always on, independent of
-  // the argument value) this function also verifies that operands passed at run
-  // time matches the executable entrypoint signature (e.g. all statically known
-  // dimensions of the memrefs matches the operands). Returns an error if finds
-  // a mismatch.
-  //
-  // This function leaves the kernel context argument (the first argument of a
-  // kernel function) uninitialized. It will be initialized in the `Execute`
-  // function right before the actual execution.
-  Error InitializeCallFrame(ArgumentsRef arguments, CallFrame* call_frame,
-                            bool verify_arguments = true) const;
-
-  // Converts returned values owned by the call frame using provided result
-  // converter. If compiled function execution finished with an error (error
-  // flag is `true` in the call frame) returns error for all results.
-  Error ReturnResults(const ResultConverter& results,
-                      CallFrame* call_frame) const;
-
-  // Executes compiled function with given arguments.
-  //
-  // If `verify_arguments` is true (in debug mode it's always on, independent of
-  // the argument value) this function also verifies that arguments passed at
-  // run time matches the executable entrypoint signature. If some of the
-  // arguments do not match the expected type, this function allocates error
-  // async values for all results and returns an error.
-  //
-  // Returns compiled function results via the user-provided results converter.
-  // If execution completed in the error state, returns error for all results.
-  Error Execute(ArgumentsRef arguments, const ResultConverter& results,
-                const ExecuteOpts& opts, bool verify_arguments = true) const;
-
-  // Executes compiled function using user provided call frame.
-  //
-  // It is the caller responsibility to handle the compiled function results
-  // stored in the call frame.
-  void Execute(CallFrame& call_frame, const ExecuteOpts& opts) const;
-
-  bool IsAsync() const { return results_memory_layout_.has_async_results; }
-
-  llvm::StringRef name() const { return name_; }
-
-  Optional<size_t> specialization() const { return specialization_; }
-
-  // Returns the number of results in the runtime signature.
-  unsigned num_results() const;
-
-  // Signature of the compiled module entrypoint function before lowering to
-  // the runtime dialects. See JitExecutable's `signature_` for more details.
-  const FunctionType& signature() const;
-
-  // Signature of the compiled module entrypoint function after lowering it from
-  // high level dialects to the dialects supported by the jitrt runtime.
-  // See JitExecutable's `signature_` for more details.
-  const FunctionType& runtime_signature() const;
-
-  std::chrono::milliseconds time_to_compile() const;
-
-  // Get the object file behind this executable (on linux for example, it will
-  // be https://en.wikipedia.org/wiki/Executable_and_Linkable_Format
-  // executable). Can be null.
-  std::unique_ptr<llvm::MemoryBuffer> obj_file() const;
-
-  // CallFrame provides a pointer-stable storage for packed function arguments
-  // and storage for returned values.
-  struct CallFrame {
-    // Pointers to compiled kernel arguments.
-    llvm::SmallVector<void*, 32> args;
-
-    // We use single block of memory to store compiled kernel results. We need
-    // to be able to store pointers to async values and tokens, and strided
-    // memrefs which at runtime are represented as StridedMemrefType<T, rank>.
-    //
-    // Currently we only need to provide result storage for pointers and memref
-    // sizes and strides (int64_t type). If we'll need to support more complex
-    // return types we'll have to be more careful about alignment requirements.
-    static_assert(sizeof(uintptr_t) == sizeof(int64_t),
-                  "uintptr_t size must be the same as int64_t");
-
-    // Memory where the compiled kernel will write its results.
-    llvm::SmallVector<uint8_t, 128> results;
-
-    // Tracks whether any of the outputs were set.
-    bool has_set_outputs = false;
-
-    // Indicates whether the kernel function execution finished with an error.
-    bool is_error = false;
-
-    // The error message which is available only if `is_error` is true. The
-    // assumption is that the error message string is owned by the compiled
-    // binary and the call frame can safely keep a non-owning pointer.
-    llvm::StringRef error;
-  };
-
-  // Requirements for passing arguments to the compiled function.
-  struct ArgumentsMemoryLayout {
-    // Currently we always pass arguments as an array of pointers.
-    size_t num_args_ptrs = 0;
-  };
-
-  // Requirements for the contiguous block of memory to store compiled function
-  // results. When we invoke a compiled fuction we allocate a block of memory,
-  // and pass pointers to pre-computed offsets as output arguments to the
-  // function.
-  struct ResultsMemoryLayout {
-    bool has_async_results = false;     // true iff returns async results
-    size_t size = 0;                    // number of bytes required
-    llvm::SmallVector<size_t> offsets;  // offsets in the block of memory
-  };
-
-  // Options for configuring compiled kernel execution.
-  struct ExecuteOpts {
-    // Async task runner for executing async runtime tasks. Typically it
-    // schedules async tasks into the underlying thread pool. It's the caller's
-    // responsibility to guarantee that it will outlive the execution of all
-    // async tasks started by the executable.
-    AsyncTaskRunner* async_task_runner = nullptr;
-
-    // A container for passing arbitrary user-provided data to the custom call
-    // handlers. Must outlive all async tasks launched by this executable.
-    CustomCall::UserData* custom_call_data = nullptr;
-
-    // Diagnostic engine is responsible for passing runtime diagnostics back
-    // to the caller through the diagnostic handler.
-    DiagnosticEngine* diagnostic_engine = nullptr;
-  };
-
-  // Loads executable from an object file. It is the caller responsibility to
-  // guarantee that signatures do match the compiled function in the object
-  // file, otherwise it will surely lead to crash.
-  static Expected<Executable> LoadFromObjFile(
-      llvm::StringRef name, std::unique_ptr<llvm::MemoryBuffer> obj_file,
-      llvm::StringRef entrypoint, FunctionType signature,
-      FunctionType runtime_signature,
-      ExecutionEngine::SymbolsBinding runtime_symbol_map = {},
-      llvm::StringRef memory_region_name = "");
-
-  // Verifies that all operands types in the entrypoint function signature are
-  // supported at run time . Returns a pre-computed layout for the function
-  // arguments. If some arguments are not supported returns an error.
-  static Expected<ArgumentsMemoryLayout> GetArgumentsMemoryLayout(
-      const FunctionType& signature);
-
-  // Verifies that all results types in the entrypoint function signature are
-  // supported at run time . Returns a pre-computed layout for the function
-  // results. If some results are not supported returns an error.
-  static Expected<ResultsMemoryLayout> GetResultsMemoryLayout(
-      const FunctionType& signature);
-
-  // TODO(ezhulenev): The following three functions should be decoupled from
-  // the jitrt header file (maybe move them to runtime.h?) so that custom call
-  // implementations do not have to depend on the `jitrt` target.
-
-  // Returns the user data passed via the ExecuteOpts to the compiled kernel.
-  static CustomCall::UserData* GetUserData(xla::runtime::KernelContext* ctx);
-
-  // Returns the diagnostic engine passed via the ExecuteOpts to the compiled
-  // kernel.
-  static DiagnosticEngine* GetDiagnosticEngine(
-      xla::runtime::KernelContext* ctx);
-
-  // Calls the custom call handler with the given runtime context, arguments and
-  // attributes.
-  static mlir::LogicalResult Call(xla::runtime::KernelContext* ctx,
-                                  CustomCall& call, void** args, void** attrs);
-
- private:
-  friend class internal::JitCompilationContext;
-
-  Executable(llvm::StringRef name,
-             std::unique_ptr<XlaRuntimeMemoryMapper> memory_mapper,
-             std::unique_ptr<ExecutionEngine> engine, FunctionType signature,
-             FunctionType runtime_signature,
-             ArgumentsMemoryLayout arguments_memory_layout,
-             ResultsMemoryLayout results_memory_layout,
-             Optional<size_t> specialization,
-             std::chrono::milliseconds time_to_compile)
-      : name_(name.str()),
-        memory_mapper_(std::move(memory_mapper)),
-        engine_(std::move(engine)),
-        fptr_(engine_->entrypoint()),
-        signature_(std::move(signature)),
-        runtime_signature_(std::move(runtime_signature)),
-        arguments_memory_layout_(std::move(arguments_memory_layout)),
-        results_memory_layout_(std::move(results_memory_layout)),
-        specialization_(specialization),
-        time_to_compile_(time_to_compile) {
-    assert(fptr_ != nullptr && "kernel function must be not null");
-  }
-
-  std::string name_;  // name of the compiled kernel module
-
-  // Called by `engine_`'s destructor; must appear before it.
-  std::unique_ptr<XlaRuntimeMemoryMapper> memory_mapper_;  // optional
-
-  // JitRt execution engine owns the LLVM ORC jit compilation stack.
-  std::unique_ptr<ExecutionEngine> engine_;
-
-  // Compiled function owned by the `engine_`.
-  ExecutionEngine::EntrypointFunctionPtr fptr_;
-
-  // Signature of the compiled module entrypoint function before lowering to
-  // the runtime dialects (see JitExecutable `signature_` for more details).
-  FunctionType signature_;
-
-  // Signature of the compiled module entrypoint function after lowering it from
-  // high level dialects to the dialects supported by the jitrt runtime.
-  //
-  // - Operands and results types converted to the types with well-defined ABI
-  //   (e.g. tensors converted to memrefs).
-  //
-  // - First argument is always a kernel context added to the function by the
-  //   lowering pipeline.
-  //
-  // From this signature executable infers how to pack runtime operands
-  // according to the expected memory layout, and how to convert results
-  // returned from the JIT-compiled function into high level types (e.g. how to
-  // convert StridedMemrefType into Tensorflow Tensor).
-  //
-  // To infer the type of the returned value, executable looks at the type
-  // defined by the `runtime_signature_` to get the memory layout of the
-  // returned value, and at the type defined by the `signature_` to get the type
-  // expected by the runtime.
-  FunctionType runtime_signature_;
-
-  ArgumentsMemoryLayout arguments_memory_layout_;
-  ResultsMemoryLayout results_memory_layout_;
-
-  // Specialization id if this executable is a specialization, or an empty
-  // optional if this executable is a default one.
-  Optional<size_t> specialization_;
-  // The time it took to compile this binary.
-  std::chrono::milliseconds time_to_compile_;
-};
-
 //----------------------------------------------------------------------------//
 // JitExecutable to manage multiple compiled executables.
 //----------------------------------------------------------------------------//