wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
tensorflow--tensorflow/tensorflow/compiler/mlir/tfr/BUILD
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

382 lines
12 KiB
Python
Raw Permalink Blame History

load(
"@llvm-project//mlir:tblgen.bzl",
"gentbl_cc_library",
"td_library",
)
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_binary",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_py_strict_test", "tf_python_pybind_extension")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":friends",
],
licenses = ["notice"],
)
package_group(
name = "friends",
packages = [
"//tensorflow/c/...",
"//tensorflow/compiler/...",
# Allow visibility from the mlir language server.
"//learning/brain/mlir/mlir_lsp_server/...",
],
)
td_library(
name = "tfr_ops_td_files",
srcs = [
"ir/tfr_ops.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:CallInterfacesTdFiles",
"@llvm-project//mlir:ControlFlowInterfacesTdFiles",
"@llvm-project//mlir:FunctionInterfacesTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
"@llvm-project//mlir:QuantizationOpsTdFiles",
"@llvm-project//mlir:ShapeOpsTdFiles",
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
gentbl_cc_library(
name = "tfr_ops_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"ir/tfr_ops.h.inc": ["-gen-op-decls"],
"ir/tfr_ops.cc.inc": ["-gen-op-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "ir/tfr_ops.td",
deps = [
":tfr_ops_td_files",
],
)
gentbl_cc_library(
name = "tfr_decompose_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/generated_decompose.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/decompose_patterns.td",
deps = [
":tfr_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
cc_library(
name = "tfr",
srcs = [
"ir/tfr_ops.cc",
"ir/tfr_ops.cc.inc",
],
hdrs = [
"ir/tfr_ops.h",
"ir/tfr_ops.h.inc",
"ir/tfr_types.h",
],
deps = [
":tfr_ops_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "utils",
srcs = [
"utils/utils.cc",
],
hdrs = [
"utils/utils.h",
],
deps = [
":tfr",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "passes",
srcs = [
"passes/canonicalize.cc",
"passes/decompose.cc",
"passes/generated_decompose.inc",
"passes/raise_to_tf.cc",
"passes/rewrite_quantized_io.cc",
],
hdrs = [
"passes/passes.h",
],
deps = [
":tfr",
":utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AffineUtils",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:SCFToControlFlow",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
tf_cc_binary(
name = "tfr-opt",
srcs = ["passes/tfr_opt.cc"],
deps = [
":passes",
":tfr",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir:passes",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:MlirOptLib",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:RegisterAllPasses",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "//tensorflow/compiler/mlir:run_lit.sh",
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
":tfr-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
],
)
cc_library(
name = "tfr_decompose_ctx",
srcs = ["integration/tfr_decompose_ctx.cc"],
hdrs = ["integration/tfr_decompose_ctx.h"],
deps = [
":passes",
":tfr",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_attr",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/core:framework",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
],
)
tf_cc_test(
name = "tfr_decompose_ctx_test",
srcs = ["integration/tfr_decompose_ctx_test.cc"],
deps = [
":tfr_decompose_ctx",
"//tensorflow/core:framework",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/types:span",
"@llvm-project//mlir:IR",
"@xla//xla/hlo/testlib:test",
],
)
tf_python_pybind_extension(
name = "tfr_wrapper",
srcs = ["python/tfr_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"tfr_wrapper.pyi",
],
starlark_only = True,
visibility = [
"//tensorflow/python:__pkg__",
],
deps = [
":tfr",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/python/lib/core:pybind11_lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
"@pybind11",
],
)
py_library(
name = "composite",
srcs = ["python/composite.py"],
strict_deps = True,
)
py_library(
name = "tfr_gen",
srcs = ["python/tfr_gen.py"],
strict_deps = True,
deps = [
":tfr_wrapper",
"//tensorflow:tensorflow_py", # buildcleaner: keep
"//tensorflow/core:protos_all_py",
"//tensorflow/python/autograph/converters:control_flow",
"//tensorflow/python/autograph/converters:return_statements",
"//tensorflow/python/autograph/impl:api",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/autograph/pyct:transpiler",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_fndefs",
"//tensorflow/python/autograph/pyct/static_analysis:type_inference",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:load_library",
"//tensorflow/python/framework:op_def_registry",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
tf_py_strict_test(
name = "tfr_gen_test",
size = "medium",
srcs = ["python/tfr_gen_test.py"],
deps = [
":composite",
":tfr_gen",
"//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper",
"//tensorflow/compiler/mlir/tfr/resources:test_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "op_reg_gen",
srcs = ["python/op_reg_gen.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/autograph/pyct:transpiler",
"//tensorflow/python/framework:op_def_registry",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
tf_py_strict_test(
name = "op_reg_gen_test",
size = "small",
srcs = ["python/op_reg_gen_test.py"],
deps = [
":composite",
":op_reg_gen",
"//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "test_utils",
srcs = ["python/test_utils.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
gen_op_libraries(
name = "one_op",
src = "define_op_template.py",
deps = [
"//tensorflow/python/platform:flags",
"@absl_py//absl:app",
],
)
Reference in New Issue View Git Blame Copy Permalink