Skip to content

Commit

Permalink
Internal visibility changes
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 726175970
  • Loading branch information
@junwhanahn @Google-ML-Automation
junwhanahn authored and Google-ML-Automation committed Feb 12, 2025
1 parent 1c9c163 commit 5ef6822
Show file tree
Hide file tree
Showing 7 changed files with 153 additions and 22 deletions.
140 changes: 135 additions & 5 deletions xla/python/ifrt/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,7 @@ load("//xla/tsl/platform:build_config.bzl", "tf_proto_library")

package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = internal_visibility([
":users",
":friends",
":internal",
]),
default_visibility = internal_visibility([":internal"]),
)

exports_files([
Expand Down Expand Up @@ -70,6 +66,11 @@ cc_library(
"value.h",
],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":array_spec_proto_cc",
":attribute_map",
Expand Down Expand Up @@ -137,6 +138,11 @@ cc_library(
srcs = ["attribute_map.cc"],
hdrs = ["attribute_map.h"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":attribute_map_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
Expand All @@ -161,6 +167,11 @@ tf_proto_library(
name = "execute_options_proto",
srcs = ["execute_options.proto"],
protodeps = [":attribute_map_proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

xla_cc_test(
Expand Down Expand Up @@ -261,6 +272,11 @@ cc_library(
testonly = True,
srcs = ["test_util.cc"],
hdrs = ["test_util.h"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":ifrt",
"//xla/tsl/concurrency:ref_count",
Expand All @@ -283,6 +299,10 @@ cc_library(
testonly = True,
srcs = ["device_test_util.cc"],
hdrs = ["device_test_util.h"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
":ifrt",
":mock",
Expand All @@ -301,6 +321,10 @@ cc_library(
name = "no_impl_test_main",
testonly = True,
srcs = ["no_impl_test_main.cc"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
"@com_google_googletest//:gtest",
],
Expand All @@ -310,6 +334,10 @@ cc_library(
name = "array_impl_test_lib",
testonly = True,
srcs = ["array_impl_test_lib.cc"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
":ifrt",
":test_util",
Expand Down Expand Up @@ -341,12 +369,21 @@ xla_cc_test(
tf_proto_library(
name = "attribute_map_proto",
srcs = ["attribute_map.proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

cc_library(
name = "client_impl_test_lib",
testonly = True,
srcs = ["client_impl_test_lib.cc"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
":ifrt",
":test_util",
Expand All @@ -370,6 +407,10 @@ cc_library(
name = "tuple_impl_test_lib",
testonly = True,
srcs = ["tuple_impl_test_lib.cc"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
":ifrt",
":test_util",
Expand Down Expand Up @@ -397,6 +438,11 @@ cc_library(
testonly = True,
srcs = ["mock.cc"],
hdrs = ["mock.h"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":attribute_map",
":ifrt",
Expand All @@ -422,6 +468,11 @@ cc_library(
srcs = ["serdes.cc"],
hdrs = ["serdes.h"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":serdes_proto_cc",
"//xla/tsl/platform:statusor",
Expand Down Expand Up @@ -456,13 +507,23 @@ xla_cc_test(
tf_proto_library(
name = "serdes_proto",
srcs = ["serdes.proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

cc_library(
name = "program_serdes",
srcs = ["program_serdes.cc"],
hdrs = ["program_serdes.h"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":ifrt",
":serdes",
Expand All @@ -475,6 +536,11 @@ cc_library(
name = "sharding_serdes",
srcs = ["sharding_serdes.cc"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":ifrt",
":serdes",
Expand Down Expand Up @@ -512,6 +578,11 @@ tf_proto_library(
":shape_proto",
":sharding_proto",
],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

xla_cc_test(
Expand All @@ -535,6 +606,11 @@ xla_cc_test(
tf_proto_library(
name = "device_proto",
srcs = ["device.proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

xla_cc_test(
Expand All @@ -556,6 +632,11 @@ xla_cc_test(
tf_proto_library(
name = "dtype_proto",
srcs = ["dtype.proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

xla_cc_test(
Expand All @@ -575,6 +656,11 @@ tf_proto_library(
name = "remap_plan_proto",
srcs = ["remap_plan.proto"],
protodeps = [":array_spec_proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

xla_cc_test(
Expand Down Expand Up @@ -605,6 +691,10 @@ cc_library(
name = "remap_impl_test_lib",
testonly = True,
srcs = ["remap_impl_test_lib.cc"],
visibility = internal_visibility([
":friends",
":internal",
]),
deps = [
":ifrt",
":test_util",
Expand All @@ -626,12 +716,22 @@ cc_library(
tf_proto_library(
name = "shape_proto",
srcs = ["shape.proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

tf_proto_library(
name = "sharding_proto",
srcs = ["sharding.proto"],
protodeps = [":serdes_proto"],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

tf_proto_library(
Expand All @@ -642,13 +742,23 @@ tf_proto_library(
":shape_proto",
"//xla/python/ifrt/ir:sharding_param_proto",
],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

cc_library(
name = "plugin_program",
srcs = ["plugin_program.cc"],
hdrs = ["plugin_program.h"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":ifrt",
"@llvm-project//llvm:Support",
Expand All @@ -659,6 +769,11 @@ cc_library(
name = "plugin_program_serdes",
srcs = ["plugin_program_serdes.cc"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":plugin_program",
":serdes",
Expand Down Expand Up @@ -692,6 +807,11 @@ cc_library(
srcs = ["custom_call_program.cc"],
hdrs = ["custom_call_program.h"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":ifrt",
"//xla/tsl/concurrency:ref_count",
Expand All @@ -707,12 +827,22 @@ tf_proto_library(
":array_spec_proto",
":device_proto",
],
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
)

cc_library(
name = "custom_call_program_serdes",
srcs = ["custom_call_program_serdes.cc"],
compatible_with = get_compatible_with_portable(),
visibility = internal_visibility([
":friends",
":internal",
":users",
]),
deps = [
":array_spec_proto_cc",
":custom_call_program",
Expand Down
1 change: 1 addition & 0 deletions xla/python/ifrt/hlo/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = internal_visibility([
"//xla/python/ifrt:users",
"//xla/python/ifrt:friends",
"//xla/python/ifrt:internal",
]),
Expand Down
Loading

0 comments on commit 5ef6822

Please sign in to comment.