Skip to content

Commit

Permalink
[XLA:CPU] Remove unused stream_executor host code.
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 713662875
  • Loading branch information
@WillFroom @Google-ML-Automation
WillFroom authored and Google-ML-Automation committed Jan 9, 2025
1 parent defdfa7 commit ed4fd4f
Show file tree
Hide file tree
Showing 12 changed files with 4 additions and 1,622 deletions.
140 changes: 2 additions & 138 deletions xla/stream_executor/host/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,6 @@
load("//xla:xla.bzl", "xla_cc_test")
load("//xla/stream_executor:build_defs.bzl", "stream_executor_friends")
load("//xla/tsl:tsl.bzl", "internal_visibility")
load(
"//xla/tsl/platform:build_config_root.bzl",
"if_llvm_aarch32_available",
"if_llvm_aarch64_available",
"if_llvm_powerpc_available",
"if_llvm_system_z_available",
"if_llvm_x86_available",
)
load("//xla/tsl/platform:rules_cc.bzl", "cc_library")

package(
Expand Down Expand Up @@ -83,14 +75,14 @@ cc_library(
],
deps = [
":host_event",
":host_kernel",
"//xla/stream_executor:device_memory",
"//xla/stream_executor:event",
"//xla/stream_executor:kernel",
"//xla/stream_executor:launch_dim",
"//xla/stream_executor:stream",
"//xla/stream_executor:stream_common",
"//xla/stream_executor:stream_executor_h",
"//xla/tsl/platform:env",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/functional:any_invocable",
"@com_google_absl//absl/log:check",
Expand All @@ -103,66 +95,6 @@ cc_library(
],
)

cc_library(
name = "host_kernel_c_api",
hdrs = ["host_kernel_c_api.h"],
)

cc_library(
name = "host_kernel",
srcs = ["host_kernel.cc"],
hdrs = ["host_kernel.h"],
deps = [
":host_kernel_c_api",
"//xla/stream_executor:device_memory",
"//xla/stream_executor:kernel",
"//xla/stream_executor:kernel_spec",
"//xla/stream_executor:launch_dim",
"//xla/stream_executor:stream",
"//xla/tsl/concurrency:async_value",
"//xla/tsl/platform:env",
"//xla/tsl/platform:logging",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/functional:any_invocable",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"@tsl//tsl/platform:env",
],
)

xla_cc_test(
name = "host_kernel_test",
srcs = ["host_kernel_test.cc"],
deps = [
":host_kernel",
":host_kernel_c_api",
":host_platform",
":jit_host_kernel_function",
":ptr_host_kernel_function",
"//xla/stream_executor:device_memory",
"//xla/stream_executor:kernel",
"//xla/stream_executor:kernel_spec",
"//xla/stream_executor:launch_dim",
"//xla/stream_executor:platform",
"//xla/stream_executor:platform_manager",
"//xla/stream_executor:stream_executor_h",
"//xla/tsl/concurrency:async_value",
"//xla/tsl/lib/core:status_test_util",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@tsl//tsl/platform:env",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:statusor",
"@tsl//tsl/platform:test",
"@tsl//tsl/platform:test_benchmark",
"@tsl//tsl/platform:test_main",
],
)

cc_library(
name = "host_executor",
srcs = [
Expand All @@ -173,7 +105,6 @@ cc_library(
],
deps = [
":host_event",
":host_kernel",
":host_stream",
"//xla/stream_executor:device_description",
"//xla/stream_executor:device_memory",
Expand All @@ -185,15 +116,14 @@ cc_library(
"//xla/stream_executor:platform",
"//xla/stream_executor:stream",
"//xla/stream_executor:stream_executor_common",
"//xla/tsl/platform:env",
"//xla/tsl/platform/profile_utils:profile_utils_cpu_utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:env",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:statusor",
],
alwayslink = True,
)
Expand All @@ -216,69 +146,3 @@ xla_cc_test(
"@tsl//tsl/platform:test_main",
],
)

cc_library(
name = "ptr_host_kernel_function",
srcs = ["ptr_host_kernel_function.cc"],
hdrs = ["ptr_host_kernel_function.h"],
deps = [
":host_executor",
":host_kernel",
":host_kernel_c_api",
"//xla/stream_executor:kernel_spec",
"//xla/stream_executor/platform:initialize",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
],
alwayslink = True, # static kernel function loader registration
)

cc_library(
name = "jit_host_kernel_function",
srcs = ["jit_host_kernel_function.cc"],
hdrs = ["jit_host_kernel_function.h"],
deps = [
":host_executor",
":host_kernel",
":host_kernel_c_api",
"//xla/stream_executor:kernel_spec",
"//xla/stream_executor/platform:initialize",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Analysis",
"@llvm-project//llvm:AsmParser",
"@llvm-project//llvm:Core",
"@llvm-project//llvm:ExecutionEngine",
"@llvm-project//llvm:JITLink",
"@llvm-project//llvm:OrcJIT",
"@llvm-project//llvm:OrcShared",
"@llvm-project//llvm:Passes",
"@llvm-project//llvm:Support",
"@llvm-project//llvm:Target",
"@llvm-project//llvm:TargetParser",
"@llvm-project//llvm:TransformUtils",
"@llvm-project//llvm:ir_headers",
"@tsl//tsl/platform:statusor",
] + if_llvm_aarch32_available([
"@llvm-project//llvm:ARMAsmParser",
"@llvm-project//llvm:ARMCodeGen",
]) + if_llvm_aarch64_available([
"@llvm-project//llvm:AArch64AsmParser",
"@llvm-project//llvm:AArch64CodeGen",
]) + if_llvm_powerpc_available([
"@llvm-project//llvm:PowerPCAsmParser",
"@llvm-project//llvm:PowerPCCodeGen",
]) + if_llvm_system_z_available([
"@llvm-project//llvm:SystemZAsmParser",
"@llvm-project//llvm:SystemZCodeGen",
]) + if_llvm_x86_available([
"@llvm-project//llvm:X86AsmParser",
"@llvm-project//llvm:X86CodeGen",
]),
alwayslink = 1, # static kernel function loader registration
)
29 changes: 1 addition & 28 deletions xla/stream_executor/host/host_executor.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,10 @@ limitations under the License.

#include <cstdint>
#include <memory>
#include <new>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include "absl/log/check.h"
#include "absl/log/log.h"
Expand All @@ -38,18 +36,15 @@ limitations under the License.
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/event.h"
#include "xla/stream_executor/host/host_event.h"
#include "xla/stream_executor/host/host_kernel.h"
#include "xla/stream_executor/host/host_stream.h"
#include "xla/stream_executor/kernel.h"
#include "xla/stream_executor/kernel_spec.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/stream.h"
#include "xla/tsl/platform/profile_utils/cpu_utils.h"
#include "xla/tsl/platform/threadpool.h"
#include "tsl/platform/cpu_info.h"
#include "tsl/platform/env.h"
#include "tsl/platform/mem.h"
#include "tsl/platform/statusor.h"
#include "tsl/platform/threadpool.h"

namespace stream_executor {
namespace host {
Expand All @@ -59,16 +54,6 @@ HostStream* AsHostStream(Stream* stream) {
return dynamic_cast<HostStream*>(stream);
}

static std::vector<HostExecutor::KernelFunctionLoader>&
KernelFunctionLoaderRegistry() {
static auto* registry = new std::vector<HostExecutor::KernelFunctionLoader>();
return *registry;
}

void HostExecutor::RegisterKernelFunctionLoader(KernelFunctionLoader loader) {
KernelFunctionLoaderRegistry().push_back(std::move(loader));
}

absl::Status HostExecutor::Init() {
thread_pool_ = std::make_shared<tsl::thread::ThreadPool>(
tsl::Env::Default(), "host-executor", tsl::port::MaxParallelism());
Expand All @@ -77,18 +62,6 @@ absl::Status HostExecutor::Init() {

absl::StatusOr<std::unique_ptr<Kernel>> HostExecutor::LoadKernel(
const MultiKernelLoaderSpec& spec) {
auto host_kernel = std::make_unique<HostKernel>(thread_pool_);
host_kernel->SetArity(spec.arity());

for (auto& loader : KernelFunctionLoaderRegistry()) {
auto loaded = loader(spec);
if (!loaded.has_value()) continue;

TF_ASSIGN_OR_RETURN(auto kernel_function, *std::move(loaded));
host_kernel->SetKernelFunction(std::move(kernel_function));
return std::move(host_kernel);
}

return absl::InternalError("No method of loading host kernel provided");
}

Expand Down
13 changes: 1 addition & 12 deletions xla/stream_executor/host/host_executor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ limitations under the License.
#define XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTOR_H_

#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <variant>
Expand All @@ -27,15 +26,14 @@ limitations under the License.
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/event.h"
#include "xla/stream_executor/host/host_kernel.h"
#include "xla/stream_executor/host_memory_allocation.h"
#include "xla/stream_executor/kernel.h"
#include "xla/stream_executor/kernel_spec.h"
#include "xla/stream_executor/memory_allocation.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/stream.h"
#include "xla/stream_executor/stream_executor_common.h"
#include "tsl/platform/threadpool.h"
#include "xla/tsl/platform/threadpool.h"

namespace stream_executor {
namespace host {
Expand All @@ -48,15 +46,6 @@ namespace host {
// routines executed under the context of a GPU executor.
class HostExecutor : public StreamExecutorCommon {
public:
// A function that loads a kernel function from a given spec. If spec is not
// supported it returns an empty optional.
using KernelFunctionLoader = std::function<std::optional<
absl::StatusOr<std::unique_ptr<HostKernel::KernelFunction>>>(
const MultiKernelLoaderSpec& spec)>;

// Registers a kernel function loader in a static registry.
static void RegisterKernelFunctionLoader(KernelFunctionLoader loader);

HostExecutor(Platform* platform, int device_ordinal)
: StreamExecutorCommon(platform), device_ordinal_(device_ordinal) {}

Expand Down
Loading

0 comments on commit ed4fd4f

Please sign in to comment.