451 lines
14 KiB
Python
451 lines
14 KiB
Python
# Common computation builders for XLA.
|
|
|
|
load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "xla_test")
|
|
|
|
licenses(["notice"]) # Apache 2.0
|
|
|
|
package(default_visibility = ["//tensorflow/compiler/xla/client:friends"])
|
|
|
|
# Filegroup used to collect source files for dependency checking.
|
|
filegroup(
|
|
name = "c_srcs",
|
|
data = glob([
|
|
"**/*.cc",
|
|
"**/*.h",
|
|
]),
|
|
)
|
|
|
|
# Generate test_suites for all backends, named "${backend}_tests".
|
|
generate_backend_suites()
|
|
|
|
cc_library(
|
|
name = "arithmetic",
|
|
srcs = ["arithmetic.cc"],
|
|
hdrs = ["arithmetic.h"],
|
|
deps = [
|
|
":constants",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"@com_google_absl//absl/strings",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "arithmetic_test",
|
|
srcs = ["arithmetic_test.cc"],
|
|
deps = [
|
|
":arithmetic",
|
|
"//tensorflow/compiler/xla:literal_util",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "cholesky",
|
|
srcs = ["cholesky.cc"],
|
|
hdrs = ["cholesky.h"],
|
|
deps = [
|
|
":math",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client/lib:constants",
|
|
"//tensorflow/compiler/xla/client/lib:loops",
|
|
"//tensorflow/compiler/xla/client/lib:matrix",
|
|
"//tensorflow/compiler/xla/client/lib:slicing",
|
|
"//tensorflow/compiler/xla/client/lib:triangular_solve",
|
|
"//tensorflow/core:lib",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "cholesky_test",
|
|
srcs = ["cholesky_test.cc"],
|
|
tags = ["optonly"],
|
|
deps = [
|
|
":arithmetic",
|
|
":cholesky",
|
|
":matrix",
|
|
"//tensorflow/compiler/xla:array2d",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:literal_test_util",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
"//tensorflow/core:test",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "constants",
|
|
srcs = ["constants.cc"],
|
|
hdrs = ["constants.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:literal_util",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "constants_test",
|
|
srcs = ["constants_test.cc"],
|
|
deps = [
|
|
":constants",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "conv_grad_size_util",
|
|
srcs = ["conv_grad_size_util.cc"],
|
|
hdrs = ["conv_grad_size_util.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla/client:padding",
|
|
"//tensorflow/core:lib",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "loops",
|
|
srcs = ["loops.cc"],
|
|
hdrs = ["loops.h"],
|
|
deps = [
|
|
":constants",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/types:span",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "math",
|
|
srcs = ["math.cc"],
|
|
hdrs = ["math.h"],
|
|
deps = [
|
|
":constants",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "math_test",
|
|
srcs = ["math_test.cc"],
|
|
deps = [
|
|
":math",
|
|
"//tensorflow/compiler/xla:literal_util",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "matrix",
|
|
srcs = ["matrix.cc"],
|
|
hdrs = ["matrix.h"],
|
|
deps = [
|
|
":arithmetic",
|
|
":constants",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"@com_google_absl//absl/types:span",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "matrix_test",
|
|
srcs = ["matrix_test.cc"],
|
|
deps = [
|
|
":matrix",
|
|
":slicing",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "pooling",
|
|
srcs = ["pooling.cc"],
|
|
hdrs = ["pooling.h"],
|
|
deps = [
|
|
":arithmetic",
|
|
":constants",
|
|
":conv_grad_size_util",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"@com_google_absl//absl/container:inlined_vector",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "pooling_test",
|
|
srcs = ["pooling_test.cc"],
|
|
deps = [
|
|
":pooling",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
"@com_google_absl//absl/container:inlined_vector",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "prng",
|
|
srcs = ["prng.cc"],
|
|
hdrs = ["prng.h"],
|
|
deps = [
|
|
":constants",
|
|
":math",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/core:lib",
|
|
"@com_google_absl//absl/base",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "qr",
|
|
srcs = ["qr.cc"],
|
|
hdrs = ["qr.h"],
|
|
deps = [
|
|
":arithmetic",
|
|
":constants",
|
|
":loops",
|
|
":math",
|
|
":matrix",
|
|
":slicing",
|
|
"//tensorflow/compiler/xla:literal_util",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/core:lib",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "qr_test",
|
|
srcs = ["qr_test.cc"],
|
|
tags = ["optonly"],
|
|
deps = [
|
|
":matrix",
|
|
":qr",
|
|
"//tensorflow/compiler/xla:array2d",
|
|
"//tensorflow/compiler/xla:array3d",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:literal_test_util",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
"//tensorflow/core:test",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "slicing",
|
|
srcs = ["slicing.cc"],
|
|
hdrs = ["slicing.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"@com_google_absl//absl/types:span",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "slicing_test",
|
|
srcs = ["slicing_test.cc"],
|
|
deps = [
|
|
":slicing",
|
|
"//tensorflow/compiler/xla:literal_util",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "sorting",
|
|
srcs = ["sorting.cc"],
|
|
hdrs = ["sorting.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "sorting_test",
|
|
srcs = ["sorting_test.cc"],
|
|
deps = [
|
|
":sorting",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "quantize",
|
|
hdrs = ["quantize.h"],
|
|
deps = [
|
|
":constants",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/core:lib",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "quantize_test",
|
|
srcs = ["quantize_test.cc"],
|
|
# TODO(b/122119490): re-enable TAP after fixing.
|
|
tags = [
|
|
"notap",
|
|
],
|
|
deps = [
|
|
":quantize",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "testing",
|
|
srcs = ["testing.cc"],
|
|
hdrs = ["testing.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:execution_options_util",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client",
|
|
"//tensorflow/compiler/xla/client:global_data",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"//tensorflow/compiler/xla/tests:test_utils",
|
|
"//tensorflow/core:lib",
|
|
"@com_google_absl//absl/strings",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "triangular_solve",
|
|
srcs = ["triangular_solve.cc"],
|
|
hdrs = ["triangular_solve.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:status_macros",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"//tensorflow/compiler/xla/client/lib:constants",
|
|
"//tensorflow/compiler/xla/client/lib:math",
|
|
"//tensorflow/compiler/xla/client/lib:matrix",
|
|
"//tensorflow/compiler/xla/client/lib:slicing",
|
|
"//tensorflow/core:lib",
|
|
],
|
|
)
|
|
|
|
xla_test(
|
|
name = "triangular_solve_test",
|
|
srcs = ["triangular_solve_test.cc"],
|
|
tags = [
|
|
"enable_for_xla_interpreter",
|
|
"noasan", # sometimes times out, http://b/78650012
|
|
],
|
|
deps = [
|
|
":math",
|
|
":matrix",
|
|
":triangular_solve",
|
|
"//tensorflow/compiler/xla:array2d",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla:shape_util",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:test",
|
|
"//tensorflow/compiler/xla:types",
|
|
"//tensorflow/compiler/xla:xla_data_proto",
|
|
"//tensorflow/compiler/xla/client:global_data",
|
|
"//tensorflow/compiler/xla/client:local_client",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/tests:client_library_test_base",
|
|
"//tensorflow/compiler/xla/tests:literal_test_util",
|
|
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:test",
|
|
],
|
|
)
|