wehub-resource-sync
336 lines
10 KiB
Python
336 lines
10 KiB
Python
load(
|
|
"//tensorflow:tensorflow.bzl",
|
|
"if_oss",
|
|
"tf_cc_test",
|
|
)
|
|
|
|
# Description:
|
|
# Utilities for splitting and joining large protos > 2GB.
|
|
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
|
|
|
|
package(
|
|
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
|
|
default_visibility = [
|
|
"//learning/brain/tfrt/aot:__subpackages__",
|
|
"//learning/brain/tfrt/tf_tpu:__subpackages__",
|
|
"//tensorflow/cc/saved_model/image_format:__subpackages__",
|
|
"//tensorflow/tools/proto_splitter:__subpackages__",
|
|
],
|
|
licenses = ["notice"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "split",
|
|
hdrs = ["split.h"],
|
|
deps = [
|
|
":util",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"//tensorflow/tools/proto_splitter:versions_proto_cc",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@tsl//tsl/platform:protobuf",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "composable_splitter_base",
|
|
srcs = ["composable_splitter_base.cc"],
|
|
hdrs = ["composable_splitter_base.h"],
|
|
deps = [
|
|
":max_size",
|
|
":split",
|
|
":util",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/strings:cord",
|
|
"@riegeli//riegeli/base:initializer",
|
|
"@riegeli//riegeli/bytes:cord_writer",
|
|
"@riegeli//riegeli/bytes:fd_writer",
|
|
"@riegeli//riegeli/bytes:string_writer",
|
|
"@riegeli//riegeli/records:record_writer",
|
|
"@tsl//tsl/platform:errors",
|
|
"@tsl//tsl/platform:protobuf",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "composable_splitter",
|
|
hdrs = ["composable_splitter.h"],
|
|
deps = [
|
|
":composable_splitter_base",
|
|
":util",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"@tsl//tsl/platform:protobuf",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "composable_splitter_test",
|
|
srcs = ["composable_splitter_test.cc"],
|
|
data = [],
|
|
tags = [
|
|
"no_mac", # b/291933687
|
|
"no_windows", # b/291001524
|
|
],
|
|
deps = [
|
|
":composable_splitter",
|
|
":test_util",
|
|
":util",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/platform:path",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/strings:cord",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@riegeli//riegeli/base:initializer",
|
|
"@riegeli//riegeli/bytes:cord_reader",
|
|
"@riegeli//riegeli/bytes:fd_reader",
|
|
"@riegeli//riegeli/bytes:string_reader",
|
|
"@riegeli//riegeli/records:record_reader",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@tsl//tsl/platform:status_matchers",
|
|
"@xla//xla/tsl/lib/core:status_test_util",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "util",
|
|
srcs = ["util.cc"],
|
|
hdrs = ["util.h"],
|
|
visibility = [
|
|
"//learning/brain/tfrt/aot:__subpackages__",
|
|
"//tensorflow/cc/saved_model:__subpackages__",
|
|
"//tensorflow/cc/saved_model/image_format:__subpackages__",
|
|
"//tensorflow/tools/proto_splitter:__subpackages__",
|
|
],
|
|
deps = [
|
|
"//tensorflow/core/platform",
|
|
"//tensorflow/core/platform:env",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"@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",
|
|
"@com_google_absl//absl/strings:cord",
|
|
"@com_google_absl//absl/strings:str_format",
|
|
"@riegeli//riegeli/base:initializer",
|
|
"@riegeli//riegeli/base:types",
|
|
"@riegeli//riegeli/bytes:fd_reader",
|
|
"@riegeli//riegeli/bytes:string_reader",
|
|
"@riegeli//riegeli/records:record_reader",
|
|
"@tsl//tsl/platform:errors",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@tsl//tsl/platform:statusor",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "util_test",
|
|
srcs = ["util_test.cc"],
|
|
data = [
|
|
"//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
|
|
],
|
|
tags = [
|
|
"no_mac", # b/291933687
|
|
"no_windows", # b/291001524
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
":util",
|
|
"//tensorflow/core/platform:macros",
|
|
"//tensorflow/core/platform:path",
|
|
"//tensorflow/core/platform:test",
|
|
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
|
|
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@tsl//tsl/platform:status",
|
|
"@tsl//tsl/platform:status_matchers",
|
|
"@tsl//tsl/platform:statusor",
|
|
"@xla//xla/tsl/lib/core:status_test_util",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "graph_def_splitter",
|
|
srcs = ["graph_def_splitter.cc"],
|
|
hdrs = ["graph_def_splitter.h"],
|
|
deps = [
|
|
":composable_splitter",
|
|
":composable_splitter_base",
|
|
":large_node_splitter",
|
|
":max_size",
|
|
":repeated_field_splitter",
|
|
":size_splitter",
|
|
":util",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@com_google_absl//absl/memory",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_absl//absl/strings",
|
|
"@tsl//tsl/platform:errors",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@tsl//tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "size_splitter",
|
|
hdrs = ["size_splitter.h"],
|
|
deps = [
|
|
":composable_splitter",
|
|
":util",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_absl//absl/types:optional",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "repeated_field_splitter",
|
|
hdrs = ["repeated_field_splitter.h"],
|
|
deps = [
|
|
":composable_splitter",
|
|
":max_size",
|
|
":size_splitter",
|
|
":util",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@xla//xla/tsl/platform:errors",
|
|
"@xla//xla/tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "graph_def_splitter_test",
|
|
srcs = ["graph_def_splitter_test.cc"],
|
|
data = [
|
|
"//tensorflow/tools/proto_splitter/testdata:function-large-nodes.pb",
|
|
"//tensorflow/tools/proto_splitter/testdata:function-lots-of-nodes.pb",
|
|
"//tensorflow/tools/proto_splitter/testdata:graph-def-and-function.pb",
|
|
"//tensorflow/tools/proto_splitter/testdata:split-large-constant.pb",
|
|
"//tensorflow/tools/proto_splitter/testdata:split-large-nodes.pb",
|
|
"//tensorflow/tools/proto_splitter/testdata:split-lots-nodes.pb",
|
|
],
|
|
tags = [
|
|
"no_mac", # b/291933687
|
|
"no_windows", # b/291001524
|
|
],
|
|
deps = [
|
|
":graph_def_splitter",
|
|
":max_size",
|
|
":test_util",
|
|
":util",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core/platform:path",
|
|
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
|
|
"@com_google_absl//absl/strings:cord",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@tsl//tsl/platform:protobuf",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "max_size",
|
|
srcs = ["max_size.cc"],
|
|
hdrs = ["max_size.h"],
|
|
deps = ["@com_google_absl//absl/synchronization"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "saved_model_splitter",
|
|
srcs = ["saved_model_splitter.cc"],
|
|
hdrs = ["saved_model_splitter.h"],
|
|
deps = [
|
|
":composable_splitter",
|
|
":graph_def_splitter",
|
|
":large_node_splitter",
|
|
":max_size",
|
|
":util",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@com_google_absl//absl/status",
|
|
"@tsl//tsl/platform:errors",
|
|
"@tsl//tsl/platform:protobuf",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "saved_model_splitter_test",
|
|
srcs = ["saved_model_splitter_test.cc"],
|
|
data = [
|
|
"//tensorflow/cc/saved_model:saved_model_test_files",
|
|
],
|
|
tags = ["no_windows"], # b/291001524
|
|
deps = [
|
|
":max_size",
|
|
":saved_model_splitter",
|
|
":util",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core/platform:path",
|
|
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@tsl//tsl/platform:protobuf",
|
|
] + if_oss([
|
|
"//tensorflow/tools/proto_splitter:protos_impl",
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "large_node_splitter",
|
|
hdrs = ["large_node_splitter.h"],
|
|
deps = [
|
|
":composable_splitter",
|
|
":composable_splitter_base",
|
|
":max_size",
|
|
":size_splitter",
|
|
":util",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@tsl//tsl/platform:errors",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "test_util",
|
|
testonly = True,
|
|
hdrs = ["test_util.h"],
|
|
deps = [
|
|
"//tensorflow/core:test",
|
|
"@com_google_absl//absl/log:check",
|
|
"@com_google_absl//absl/status",
|
|
"@tsl//tsl/platform:errors",
|
|
"@tsl//tsl/platform:protobuf",
|
|
"@tsl//tsl/platform:statusor",
|
|
],
|
|
)
|