Skip to content

Commit

Permalink
Split CompileOnlyIfRtClient into its own directory
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 726476996
  • Loading branch information
junwhanahn authored and Google-ML-Automation committed Feb 13, 2025
1 parent b195247 commit 96a9745
Show file tree
Hide file tree
Showing 5 changed files with 414 additions and 325 deletions.
1 change: 1 addition & 0 deletions xla/python/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -392,6 +392,7 @@ cc_library(
"//xla/pjrt:transpose",
"//xla/pjrt/distributed",
"//xla/pjrt/distributed:client",
"//xla/python/compile_only_ifrt:client",
"//xla/python/ifrt",
"//xla/python/ifrt:attribute_map",
"//xla/python/ifrt:custom_call_program",
Expand Down
34 changes: 34 additions & 0 deletions xla/python/compile_only_ifrt/BUILD
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
load("//xla/tsl:tsl.bzl", "internal_visibility")

package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = internal_visibility([
"//xla/python:__subpackages__",
]),
)

cc_library(
name = "client",
srcs = ["client.cc"],
hdrs = ["client.h"],
deps = [
"//xla:shape_util",
"//xla/pjrt:host_memory_spaces",
"//xla/pjrt:pjrt_device_description",
"//xla/pjrt:pjrt_layout",
"//xla/python/ifrt",
"//xla/python/ifrt:attribute_map",
"//xla/python/pjrt_ifrt",
"//xla/python/pjrt_ifrt:pjrt_attribute_map_util",
"//xla/python/pjrt_ifrt:pjrt_dtype",
"//xla/service:computation_placer_hdr",
"//xla/tsl/concurrency:ref_count",
"//xla/tsl/platform:errors",
"//xla/tsl/platform:statusor",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
],
)
25 changes: 25 additions & 0 deletions xla/python/compile_only_ifrt/client.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
/* Copyright 2025 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/python/compile_only_ifrt/client.h"

namespace xla {

[[maybe_unused]] char CompileOnlyMemory::ID = 0;
[[maybe_unused]] char CompileOnlyDevice::ID = 0;
[[maybe_unused]] char CompileOnlyIfrtCompiler::ID = 0;
[[maybe_unused]] char CompileOnlyIfRtClient::ID = 0;

} // namespace xla
Loading

0 comments on commit 96a9745

Please sign in to comment.