f798728d37
This allows running the pass in the MlirFunctionOptimizationPass pipeline so that it only transforms the MLIR module in case of successful execution. Failure in the pass execution must not affect the MLIR module. Track failures in the execution by the /tensorflow/core/mlir_pass_fallback metric. PiperOrigin-RevId: 356316824 Change-Id: Ida9b761c3dd39202982ad6a5d5804fb0e5433c4d
233 lines
7.0 KiB
Python
233 lines
7.0 KiB
Python
# Description:
|
|
# TensorFlow/TensorFlow Lite/XLA MLIR dialects and tools.
|
|
|
|
load("//tensorflow:tensorflow.bzl", "filegroup")
|
|
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
|
|
load(
|
|
"//tensorflow:tensorflow.bzl",
|
|
"tf_cc_binary",
|
|
"tf_cc_test",
|
|
)
|
|
|
|
package(
|
|
default_visibility = [
|
|
"//tensorflow/compiler/tf2xla:__subpackages__",
|
|
"@llvm-project//mlir:friends",
|
|
],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
package_group(
|
|
name = "subpackages",
|
|
packages = ["//tensorflow/compiler/mlir/..."],
|
|
)
|
|
|
|
exports_files(glob(["g3doc/*.md"]))
|
|
|
|
# To reference all tablegen files here when checking for updates to them.
|
|
filegroup(
|
|
name = "td_files",
|
|
srcs = glob(["**/*.td"]),
|
|
)
|
|
|
|
cc_library(
|
|
name = "string_container_utils",
|
|
hdrs = ["utils/string_container_utils.h"],
|
|
deps = [
|
|
"@com_google_absl//absl/strings",
|
|
"@llvm-project//llvm:Support",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "array_container_utils",
|
|
hdrs = ["utils/array_container_utils.h"],
|
|
deps = [
|
|
"@com_google_absl//absl/types:span",
|
|
"@llvm-project//llvm:Support",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "name_utils",
|
|
srcs = ["utils/name_utils.cc"],
|
|
hdrs = ["utils/name_utils.h"],
|
|
deps = [
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:IR",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "op_or_arg_name_mapper",
|
|
srcs = ["op_or_arg_name_mapper.cc"],
|
|
hdrs = ["op_or_arg_name_mapper.h"],
|
|
deps = [
|
|
":name_utils",
|
|
"@com_google_absl//absl/strings",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:IR",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "tf_mlir_opt_main",
|
|
srcs = ["tf_mlir_opt_main.cc"],
|
|
deps = [
|
|
":init_mlir",
|
|
"//tensorflow/compiler/mlir/hlo:all_passes",
|
|
"//tensorflow/compiler/mlir/hlo:hlo_dialect_registration",
|
|
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:tensorflow_passes",
|
|
"//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
|
|
"//tensorflow/core:lib",
|
|
"@llvm-project//mlir:AllPassesAndDialectsNoRegistration",
|
|
"@llvm-project//mlir:MlirOptLib",
|
|
"@llvm-project//mlir:Shape",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "passes",
|
|
visibility = [
|
|
":__subpackages__",
|
|
"//tensorflow/python:__subpackages__",
|
|
],
|
|
deps = [
|
|
"@llvm-project//mlir:Affine",
|
|
"@llvm-project//mlir:QuantOps",
|
|
# Link jit lib to link JIT devices required to run
|
|
# xla-legalize-tf-with-tf2xla pass.
|
|
"//tensorflow/compiler/jit",
|
|
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
|
|
"//tensorflow/compiler/mlir/lite:tensorflow_lite_legalize_tf",
|
|
"//tensorflow/compiler/mlir/lite:tensorflow_lite_optimize",
|
|
"//tensorflow/compiler/mlir/lite:tensorflow_lite_quantize",
|
|
"//tensorflow/compiler/mlir/lite/quantization:quantization_passes",
|
|
"//tensorflow/compiler/mlir/lite/quantization/tensorflow:tf_to_quant",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_pass",
|
|
"//tensorflow/compiler/mlir/tensorflow:tensorflow_passes",
|
|
"//tensorflow/compiler/mlir/tensorflow:tensorflow_test_passes",
|
|
"//tensorflow/compiler/mlir/tensorflow:tf_dialect_passes",
|
|
"//tensorflow/compiler/mlir/tensorflow:tf_legalize_hlo",
|
|
"//tensorflow/compiler/mlir/tfjs:tensorflow_js_passes",
|
|
"//tensorflow/compiler/mlir/tosa:tf_passes",
|
|
"//tensorflow/compiler/mlir/tosa:tfl_passes",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "init_mlir",
|
|
srcs = ["init_mlir.cc"],
|
|
hdrs = ["init_mlir.h"],
|
|
deps = [
|
|
"//tensorflow/core:lib",
|
|
"@llvm-project//llvm:Support",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "mlir_graph_optimization_pass",
|
|
srcs = ["mlir_graph_optimization_pass.cc"],
|
|
hdrs = ["mlir_graph_optimization_pass.h"],
|
|
deps = [
|
|
"//tensorflow/compiler/mlir:mlir_bridge_rollout_policy",
|
|
"//tensorflow/compiler/mlir/tensorflow",
|
|
"//tensorflow/compiler/mlir/tensorflow:convert_graphdef",
|
|
"//tensorflow/compiler/mlir/tensorflow:device_util",
|
|
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
|
|
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
|
|
"//tensorflow/core:core_cpu",
|
|
"//tensorflow/core:lib",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Shape",
|
|
"@llvm-project//mlir:StandardOps",
|
|
],
|
|
alwayslink = 1,
|
|
)
|
|
|
|
cc_library(
|
|
name = "mlir_graph_optimization_pass_registration",
|
|
srcs = [
|
|
"mlir_graph_optimization_pass_registration.cc",
|
|
],
|
|
deps = [
|
|
":mlir_graph_optimization_pass",
|
|
"//tensorflow/core:core_cpu",
|
|
],
|
|
alwayslink = 1,
|
|
)
|
|
|
|
tf_cc_binary(
|
|
name = "tf-opt",
|
|
deps = [
|
|
":passes",
|
|
":tf_mlir_opt_main",
|
|
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration",
|
|
"//tensorflow/compiler/mlir/tensorflow:tf_graph_optimization_pass",
|
|
"//tensorflow/compiler/mlir/xla:all_xla_passes_for_testing",
|
|
],
|
|
)
|
|
|
|
tf_cc_binary(
|
|
name = "tf-mlir-translate",
|
|
srcs = ["tf_mlir_translate_main.cc"],
|
|
deps = [
|
|
":init_mlir",
|
|
"//tensorflow/compiler/mlir/tensorflow:convert_graphdef",
|
|
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
|
|
"//tensorflow/compiler/mlir/tensorflow:tf_xla_mlir_translate",
|
|
"//tensorflow/compiler/mlir/tensorflow:translate_cl_options",
|
|
"//tensorflow/compiler/mlir/tensorflow:translate_lib",
|
|
"//tensorflow/compiler/mlir/tensorflow:translate_registration",
|
|
"//tensorflow/compiler/mlir/tensorflow:translate_tf_dialect_op",
|
|
"//tensorflow/compiler/mlir/xla:xla_mlir_translate",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:tensorflow",
|
|
"//tensorflow/stream_executor/lib",
|
|
"@com_google_absl//absl/strings",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:AllPassesAndDialectsNoRegistration",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Support",
|
|
"@llvm-project//mlir:Translation",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "mlir_bridge_rollout_policy",
|
|
srcs = ["mlir_bridge_rollout_policy.cc"],
|
|
hdrs = ["mlir_bridge_rollout_policy.h"],
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
"//tensorflow/compiler/jit:flags",
|
|
"//tensorflow/core:graph",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"@com_google_absl//absl/types:optional",
|
|
],
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "mlir_graph_optimization_pass_test",
|
|
srcs = ["mlir_graph_optimization_pass_test.cc"],
|
|
deps = [
|
|
":mlir_graph_optimization_pass",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core:test_main",
|
|
"@llvm-project//mlir:IR",
|
|
],
|
|
)
|
|
|
|
filegroup(
|
|
name = "litfiles",
|
|
srcs = glob(["runlit*py"]),
|
|
)
|
|
|
|
exports_files(["run_lit.sh"])
|