343 lines
10 KiB
Python
343 lines
10 KiB
Python
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
|
|
load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
|
|
load("//tensorflow/core/platform:build_config.bzl", "if_llvm_aarch64_available")
|
|
|
|
package(
|
|
default_visibility = ["//visibility:private"],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
# Don't depend on this directly; this is only used for the benchmark test
|
|
# generated by tf_library.
|
|
cc_library(
|
|
name = "tf_library_test_main",
|
|
testonly = 1,
|
|
visibility = ["//visibility:public"],
|
|
deps = ["//tensorflow/core:test_main"],
|
|
)
|
|
|
|
filegroup(
|
|
name = "quantize_header",
|
|
srcs = ["quantize.h"],
|
|
visibility = ["//visibility:public"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "tfcompile_lib",
|
|
srcs = [
|
|
"codegen.cc",
|
|
"compile.cc",
|
|
"flags.cc",
|
|
],
|
|
hdrs = [
|
|
"codegen.h",
|
|
"compile.h",
|
|
"flags.h",
|
|
"quantize.h",
|
|
],
|
|
defines = if_llvm_aarch64_available(["TF_LLVM_AARCH64_AVAILABLE=1"]),
|
|
visibility = ["//tensorflow/python:__pkg__"],
|
|
deps = [
|
|
":aot_only_var_handle_op",
|
|
":embedded_protocol_buffers",
|
|
"@com_google_absl//absl/base",
|
|
"@com_google_absl//absl/memory",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/types:span",
|
|
"//tensorflow/compiler/tf2xla",
|
|
"//tensorflow/compiler/tf2xla:mlir_tf2xla",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_proto_cc",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_util",
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
|
|
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
|
|
"//tensorflow/compiler/xla:cpu_function_runtime",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto_cc",
|
|
"//tensorflow/compiler/xla/client:client_library",
|
|
"//tensorflow/compiler/xla/client:compile_only_client",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"//tensorflow/compiler/xla/service:compiler",
|
|
"//tensorflow/compiler/xla/service/cpu:buffer_info_util",
|
|
"//tensorflow/compiler/xla/service/cpu:cpu_compiler",
|
|
"//tensorflow/core:core_cpu_internal",
|
|
"//tensorflow/core:framework_internal",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@llvm-project//llvm:ARMCodeGen", # fixdeps: keep
|
|
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
|
|
"@llvm-project//llvm:Target",
|
|
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
|
|
"//tensorflow/core:regexp_internal",
|
|
] + if_llvm_aarch64_available([
|
|
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
|
|
]),
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "codegen_test",
|
|
srcs = ["codegen_test.cc"],
|
|
data = [
|
|
"codegen_test_h.golden",
|
|
"codegen_test_o.golden",
|
|
],
|
|
deps = [
|
|
":tfcompile_lib",
|
|
"//tensorflow/compiler/xla:cpu_function_runtime",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core:test_main",
|
|
"//tensorflow/core/platform:resource_loader",
|
|
"@com_google_absl//absl/strings",
|
|
"@llvm-project//llvm:Support", # fixdeps: keep
|
|
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
|
|
],
|
|
)
|
|
|
|
tf_cc_binary(
|
|
name = "tfcompile",
|
|
visibility = ["//visibility:public"],
|
|
deps = [":tfcompile_main"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "llvm_targets",
|
|
visibility = ["//tensorflow/python:__pkg__"],
|
|
deps = [
|
|
"@llvm-project//llvm:ARMCodeGen", # fixdeps: keep
|
|
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
|
|
"@llvm-project//llvm:Target",
|
|
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
|
|
] + if_llvm_aarch64_available([
|
|
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
|
|
]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "tfcompile_main",
|
|
srcs = ["tfcompile_main.cc"],
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
":tfcompile_lib",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_proto_cc",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_util",
|
|
"//tensorflow/compiler/xla:debug_options_flags",
|
|
"//tensorflow/compiler/xla/service:compiler",
|
|
"//tensorflow/core:core_cpu",
|
|
"//tensorflow/core:core_cpu_internal",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_internal",
|
|
"//tensorflow/core:graph",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@com_google_absl//absl/strings",
|
|
],
|
|
)
|
|
|
|
# NOTE: Most end-to-end tests are in the "tests" subdirectory, to ensure that
|
|
# tfcompile.bzl correctly handles usage from outside of the package that it is
|
|
# defined in.
|
|
|
|
# A simple test of tf_library from a text protobuf, to enable benchmark_test.
|
|
# This test uses an incompleted graph with a node that is not defined. The
|
|
# compilation works because the undefined node is a feed node.
|
|
tf_library(
|
|
name = "test_graph_tfadd",
|
|
testonly = 1,
|
|
config = "test_graph_tfadd.config.pbtxt",
|
|
cpp_class = "AddComp",
|
|
graph = "test_graph_tfadd.pbtxt",
|
|
mlir_components = "None",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
tf_library(
|
|
name = "test_graph_tfadd_mlir_bridge",
|
|
testonly = 1,
|
|
config = "test_graph_tfadd.config.pbtxt",
|
|
cpp_class = "AddComp",
|
|
graph = "test_graph_tfadd.pbtxt",
|
|
mlir_components = "Bridge",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
# A test of tf_library that includes a graph with an unknown op, but where
|
|
# the compilation works because the node with the unknown op is not needed
|
|
# for the fetches.
|
|
tf_library(
|
|
name = "test_graph_tfunknownop",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "None",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
tf_library(
|
|
name = "test_graph_tfunknownop_mlir_bridge",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "Bridge",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
# A test of tf_library that includes a graph with an unknown op, but where
|
|
# the compilation works because the node with the unknown op is only used as
|
|
# an input of a feed node.
|
|
tf_library(
|
|
name = "test_graph_tfunknownop2",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop2.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "None",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
tf_library(
|
|
name = "test_graph_tfunknownop2_mlir_bridge",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop2.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "Bridge",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
# A test of tf_library that includes a graph with an unknown op, but where
|
|
# the compilation works because the node with the unknown op is a feed node.
|
|
tf_library(
|
|
name = "test_graph_tfunknownop3",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop3.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "None",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
tf_library(
|
|
name = "test_graph_tfunknownop3_mlir_bridge",
|
|
testonly = 1,
|
|
config = "test_graph_tfunknownop3.config.pbtxt",
|
|
cpp_class = "UnknownOpAddComp",
|
|
graph = "test_graph_tfunknownop.pbtxt",
|
|
mlir_components = "Bridge",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
)
|
|
|
|
# Utility library for benchmark binaries, used by the *_benchmark rules that are
|
|
# added by the tfcompile bazel macro.
|
|
cc_library(
|
|
name = "benchmark",
|
|
srcs = ["benchmark.cc"],
|
|
hdrs = ["benchmark.h"],
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
# The purpose of the benchmark library is to support building an aot
|
|
# binary with minimal dependencies, to demonstrate small binary sizes.
|
|
#
|
|
# KEEP THE DEPENDENCIES MINIMAL.
|
|
"//tensorflow/core:framework_lite",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "benchmark_extra_android",
|
|
tags = [
|
|
"manual",
|
|
],
|
|
visibility = ["//visibility:public"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "embedded_protocol_buffers",
|
|
srcs = ["embedded_protocol_buffers.cc"],
|
|
hdrs = ["embedded_protocol_buffers.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla/service/llvm_ir:llvm_util",
|
|
"//tensorflow/core:lib",
|
|
"@com_google_absl//absl/memory",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/types:span",
|
|
"@llvm-project//llvm:Core",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//llvm:Target",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "aot_only_var_handle_op",
|
|
srcs = ["aot_only_var_handle_op.cc"],
|
|
hdrs = ["aot_only_var_handle_op.h"],
|
|
visibility = [
|
|
"//tensorflow/compiler/tf2xla:__pkg__",
|
|
],
|
|
deps = [
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//tensorflow/core:framework",
|
|
],
|
|
alwayslink = 1,
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "benchmark_test",
|
|
srcs = ["benchmark_test.cc"],
|
|
tags = ["manual"],
|
|
deps = [
|
|
":benchmark",
|
|
":test_graph_tfadd",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core:test_main",
|
|
],
|
|
)
|
|
|
|
test_suite(
|
|
name = "all_tests",
|
|
tags = ["manual"],
|
|
tests = [
|
|
":benchmark_test",
|
|
":codegen_test",
|
|
":test_graph_tfadd_mlir_bridge_test",
|
|
":test_graph_tfadd_test",
|
|
":test_graph_tfunknownop2_mlir_bridge_test",
|
|
":test_graph_tfunknownop2_test",
|
|
":test_graph_tfunknownop3_mlir_bridge_test",
|
|
":test_graph_tfunknownop3_test",
|
|
":test_graph_tfunknownop_mlir_bridge_test",
|
|
":test_graph_tfunknownop_test",
|
|
"//tensorflow/compiler/aot/tests:all_tests",
|
|
],
|
|
)
|
|
|
|
exports_files([
|
|
"benchmark_main.template", # used by tf_library(...,gen_benchmark=True)
|
|
"test.cc", # used by tf_library(...,gen_test=True)
|
|
])
|