430 lines
14 KiB
Python
430 lines
14 KiB
Python
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
|
|
load(
|
|
"//tensorflow/core/platform:build_config.bzl",
|
|
"tf_proto_library",
|
|
)
|
|
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
|
|
|
|
package(
|
|
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
|
|
default_visibility = [
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
|
|
"//tensorflow/compiler/mlir/tf2xla/api/v1:__subpackages__",
|
|
"//tensorflow/compiler/mlir/tf2xla/api/v2:__subpackages__",
|
|
"//tensorflow/core/tfrt/saved_model:__pkg__",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "compilation_timer",
|
|
hdrs = ["compilation_timer.h"],
|
|
deps = [
|
|
"//tensorflow/core/platform:profile_utils_cpu_utils",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "compilation_timer_test",
|
|
srcs = ["compilation_timer_test.cc"],
|
|
deps = [
|
|
":compilation_timer",
|
|
"@com_google_absl//absl/time",
|
|
"@com_google_googletest//:gtest_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "test_matchers",
|
|
testonly = True,
|
|
hdrs = ["test_matchers.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_googletest//:gtest",
|
|
"@xla//xla/tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "test_matchers_test",
|
|
srcs = ["test_matchers_test.cc"],
|
|
deps = [
|
|
":test_matchers",
|
|
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
|
|
"//tensorflow/compiler/tf2xla:xla_helpers",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core/lib/monitoring:cell_reader",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@xla//xla/hlo/builder:xla_computation",
|
|
"@xla//xla/service:hlo_proto_cc",
|
|
"@xla//xla/tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "mlir_pass_instrumentation",
|
|
srcs = ["mlir_pass_instrumentation.cc"],
|
|
hdrs = ["mlir_pass_instrumentation.h"],
|
|
deps = [
|
|
"//tensorflow/core/platform:logging",
|
|
"@com_google_absl//absl/log",
|
|
"@llvm-project//mlir:Pass",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "mlir_pass_instrumentation_test",
|
|
srcs = ["mlir_pass_instrumentation_test.cc"],
|
|
deps = [
|
|
":mlir_pass_instrumentation",
|
|
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
|
|
"//tensorflow/core:test",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_absl//absl/strings:string_view",
|
|
"@com_google_googletest//:gtest_main",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "legalize_tf_mlir",
|
|
srcs = ["legalize_tf_mlir.cc"],
|
|
hdrs = ["legalize_tf_mlir.h"],
|
|
deps = [
|
|
":compilation_timer",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
|
|
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
|
|
"//tensorflow/compiler/tf2xla:layout_util",
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//tensorflow/compiler/tf2xla:xla_helpers",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
|
|
"//tensorflow/core/tpu:tpu_compile",
|
|
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Pass",
|
|
"@stablehlo//:register",
|
|
"@tsl//tsl/platform:error_logging",
|
|
"@xla//xla:shape_util",
|
|
"@xla//xla/mlir_hlo:hlo_dialect_registration",
|
|
"@xla//xla/tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "legalize_tf_mlir_test",
|
|
srcs = ["legalize_tf_mlir_test.cc"],
|
|
deps = [
|
|
":legalize_tf_mlir",
|
|
":test_matchers",
|
|
"//tensorflow/compiler/jit",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//tensorflow/compiler/tf2xla:xla_helpers",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:test_main",
|
|
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
|
|
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_googletest//:gtest",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Pass",
|
|
"@stablehlo//:register",
|
|
"@xla//xla:shape_util",
|
|
"@xla//xla/mlir_hlo:hlo_dialect_registration",
|
|
"@xla//xla/tsl/platform:errors",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "legalize_tf_to_hlo",
|
|
srcs = ["legalize_tf_to_hlo.cc"],
|
|
hdrs = ["legalize_tf_to_hlo.h"],
|
|
deps = [
|
|
":legalize_tf_mlir",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
|
|
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
|
|
"//tensorflow/compiler/tf2xla:layout_util",
|
|
"//tensorflow/compiler/tf2xla:xla_helpers",
|
|
"//tensorflow/compiler/tf2xla:xla_tpu_backend_registration",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core/platform:status",
|
|
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
|
|
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:DataLayoutInterfaces",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Pass",
|
|
"@stablehlo//:register",
|
|
"@xla//xla:shape_util",
|
|
"@xla//xla/client:compile_only_client",
|
|
"@xla//xla/mlir_hlo:hlo_dialect_registration",
|
|
"@xla//xla/tsl/platform:errors",
|
|
"@xla//xla/tsl/platform:statusor",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "legalize_tf_to_hlo_test",
|
|
srcs = ["legalize_tf_to_hlo_test.cc"],
|
|
deps = [
|
|
":legalize_tf_to_hlo",
|
|
":test_matchers",
|
|
"//tensorflow/compiler/jit",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//tensorflow/compiler/tf2xla:xla_helpers",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:test_main",
|
|
"//tensorflow/core/lib/monitoring:cell_reader",
|
|
"//tensorflow/core/protobuf:for_core_protos_cc",
|
|
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
|
|
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_googletest//:gtest",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Pass",
|
|
"@stablehlo//:register",
|
|
"@xla//xla:shape_util",
|
|
"@xla//xla/client:client_library",
|
|
"@xla//xla/mlir_hlo:hlo_dialect_registration",
|
|
"@xla//xla/stream_executor:platform",
|
|
"@xla//xla/stream_executor:platform_manager",
|
|
"@xla//xla/tsl/platform:errors",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "clustering_bridge_passes",
|
|
srcs = ["clustering_bridge_passes.cc"],
|
|
hdrs = ["clustering_bridge_passes.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/jit:flags_headers",
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
|
|
"//tensorflow/compiler/mlir/tensorflow/transforms/sparsecore:sparsecore_passes",
|
|
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
|
|
"@com_google_absl//absl/base:core_headers",
|
|
"@com_google_absl//absl/log",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:FuncDialect",
|
|
"@llvm-project//mlir:Pass",
|
|
"@llvm-project//mlir:Transforms",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "clustering_bridge_passes_test",
|
|
srcs = ["clustering_bridge_passes_test.cc"],
|
|
deps = [
|
|
":clustering_bridge_passes",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@llvm-project//mlir:Pass",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "logging_hooks",
|
|
srcs = ["logging_hooks.cc"],
|
|
hdrs = ["logging_hooks.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
|
|
"//tensorflow/core:framework",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:Pass",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "logging_hooks_test",
|
|
srcs = ["logging_hooks_test.cc"],
|
|
data = [
|
|
"testdata/dead_const.mlir",
|
|
],
|
|
deps = [
|
|
":logging_hooks",
|
|
"//tensorflow/compiler/mlir:register_common_dialects",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core/platform:resource_loader",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@llvm-project//mlir:FuncDialect",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Parser",
|
|
"@llvm-project//mlir:Pass",
|
|
"@llvm-project//mlir:Transforms",
|
|
"@xla//xla/tsl/lib/core:status_test_util",
|
|
"@xla//xla/tsl/platform:status",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "mlir_bridge_pass_util",
|
|
srcs = ["mlir_bridge_pass_util.cc"],
|
|
hdrs = ["mlir_bridge_pass_util.h"],
|
|
visibility = ["//tensorflow/compiler/tf2xla:__pkg__"],
|
|
deps = [
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_defs",
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/common_runtime:function_body",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/strings:string_view",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:FuncDialect",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Support",
|
|
"@xla//xla/tsl/platform:status",
|
|
],
|
|
alwayslink = 1,
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "mlir_bridge_pass_util_test",
|
|
srcs = ["mlir_bridge_pass_util_test.cc"],
|
|
deps = [
|
|
":mlir_bridge_pass_util",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@llvm-project//mlir:FuncDialect",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Parser",
|
|
"//tensorflow/cc:array_ops",
|
|
"//tensorflow/cc:function_ops",
|
|
"//tensorflow/cc:functional_ops",
|
|
"//tensorflow/cc:ops",
|
|
"//tensorflow/cc:scope",
|
|
"//tensorflow/cc:tpu_ops",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/tf2xla:tf2xla_defs",
|
|
"//tensorflow/core/protobuf:for_core_protos_cc",
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core/platform:types",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:test_main",
|
|
"//tensorflow/core:testlib",
|
|
"//tensorflow/core/platform:enable_tf2_utils",
|
|
# "//tensorflow/core/platform:resource_loader",
|
|
"@xla//xla/tsl/lib/core:status_test_util",
|
|
],
|
|
)
|
|
|
|
tf_proto_library(
|
|
name = "reproducer_proto",
|
|
srcs = ["reproducer.proto"],
|
|
protodeps = [
|
|
"//tensorflow/core:protos_all",
|
|
"//tensorflow/core/protobuf/tpu:compile_metadata_proto",
|
|
],
|
|
visibility = [
|
|
"//learning/brain/mlir/bridge:__pkg__",
|
|
"//tensorflow/compiler/mlir/tf2xla/api/v2:__pkg__",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "node_order",
|
|
srcs = ["node_order.cc"],
|
|
hdrs = ["node_order.h"],
|
|
deps = [
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:lib",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "node_order_test",
|
|
size = "small",
|
|
srcs = [
|
|
"node_order_test.cc",
|
|
],
|
|
deps = [
|
|
":node_order",
|
|
"//tensorflow/cc:cc_ops",
|
|
"//tensorflow/cc:cc_ops_internal",
|
|
"//tensorflow/cc:function_ops",
|
|
"//tensorflow/cc:sendrecv_ops",
|
|
"//tensorflow/core",
|
|
"//tensorflow/core:core_cpu",
|
|
"//tensorflow/core:core_cpu_internal",
|
|
"//tensorflow/core:direct_session_internal",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_internal",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_internal",
|
|
"//tensorflow/core:ops",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core:test_main",
|
|
"//tensorflow/core:testlib",
|
|
"@com_google_absl//absl/log:check",
|
|
"@com_google_googletest//:gtest",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "graph_to_tf_executor_util",
|
|
srcs = ["graph_to_tf_executor_util.cc"],
|
|
hdrs = ["graph_to_tf_executor_util.h"],
|
|
deps = [
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_types_hdr",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/common_runtime:function_body",
|
|
"//tensorflow/core/platform:enable_tf2_utils",
|
|
"//tensorflow/core/protobuf:for_core_protos_cc",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/strings:string_view",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "graph_to_tf_executor_util_test",
|
|
srcs = ["graph_to_tf_executor_util_test.cc"],
|
|
deps = [
|
|
":graph_to_tf_executor_util",
|
|
"//tensorflow/cc:array_ops",
|
|
"//tensorflow/cc:cc_ops",
|
|
"//tensorflow/cc:functional_ops",
|
|
"//tensorflow/cc:ops",
|
|
"//tensorflow/cc:scope",
|
|
"//tensorflow/cc:tpu_ops",
|
|
"//tensorflow/compiler/tf2xla/ops:xla_ops",
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/framework:tensor_testutil",
|
|
"//tensorflow/core/platform:enable_tf2_utils",
|
|
"//tensorflow/core/platform:types",
|
|
"@com_google_absl//absl/log:check",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/strings",
|
|
"@com_google_googletest//:gtest_main",
|
|
"@xla//xla/tsl/lib/core:status_test_util",
|
|
],
|
|
)
|