Files
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

592 lines
23 KiB
Python

"""Build macro that compiles a TensorFlow graph into a cc_library.
To use from your BUILD file, add the following line to load the macro:
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
Then call the macro like this:
tf_library(
name = "test_graph_tfmatmul",
config = "test_graph_tfmatmul.config.pbtxt",
cpp_class = "MatMulComp",
graph = ":test_graph_tfmatmul.pb",
)
"""
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"if_android",
"if_google",
"if_oss",
"tf_cc_test",
"tf_copts",
)
load("//tensorflow:tensorflow.default.bzl", "tfcompile_dfsan_abilists", "tfcompile_dfsan_enabled", "tfcompile_friends", "tfcompile_target_cpu")
visibility(tfcompile_friends())
def _tfcompile_model_library_rule_impl(ctx):
header_file = ctx.outputs.header_out
metadata_object_file = ctx.actions.declare_file("%s_tfcompile_metadata.o" % ctx.attr.model_name)
function_object_file = ctx.actions.declare_file("%s_tfcompile_function.o" % ctx.attr.model_name)
constant_buffers_object_file = ctx.actions.declare_file("%s_tfcompile_constant_buffers.o" % ctx.attr.model_name)
session_module_pb = ctx.actions.declare_file("%s_session_module.pb" % ctx.attr.model_name)
out_files = [header_file, metadata_object_file, function_object_file, constant_buffers_object_file, session_module_pb]
compiler_log_file = None
if ctx.attr.gen_compiler_log:
compiler_log_file = ctx.actions.declare_file("%s_compiler.log" % ctx.attr.model_name)
out_files.append(compiler_log_file)
output_dict = {}
output_dict["header_files"] = [header_file]
output_dict["object_files"] = [metadata_object_file, function_object_file, constant_buffers_object_file]
if compiler_log_file:
output_dict["log_files"] = [compiler_log_file]
output_flags = [
"--out_header=" + header_file.path,
"--out_metadata_object=" + metadata_object_file.path,
"--out_function_object=" + function_object_file.path,
"--out_constant_buffers_object=" + constant_buffers_object_file.path,
"--out_session_module=" + session_module_pb.path,
]
additional_xla_flags = ctx.attr.xla_flags
tfcompile_env = {
"XLA_FLAGS": ("--xla_cpu_enable_fast_math=true " +
"--xla_cpu_fast_math_honor_nans=false " +
"--xla_cpu_fast_math_honor_infs=false " +
"--xla_cpu_fast_math_honor_functions=false " +
"--xla_cpu_fast_math_honor_division=false " +
"--xla_cpu_enable_fast_min_max=true " +
"--xla_cpu_experimental_ynn_fusion_type= " +
additional_xla_flags + " " +
"$${XLA_FLAGS:-}' "),
"CUDA_VISIBLE_DEVICES": "",
}
dfsan_flags = []
dfsan_deps = []
# DFSan is only supported on linux.
if ctx.attr.is_linux and ctx.attr.dfsan:
dfsan_flags = [
"--sanitize_dataflow",
"--sanitize_abilists_dataflow=" + ",".join([f.path for f in ctx.files.dfsan_abilists]),
]
dfsan_deps = ctx.files.dfsan_abilists
cpu_flags = ["--target_cpu=" + ctx.attr.target_cpu] if ctx.attr.target_cpu else []
flags = [
"--graph=" + ctx.file.tfcompile_graph.path,
"--config=" + ctx.file.tfcompile_config.path,
"--entry_point=" + ctx.attr.entry_point,
"--cpp_class=" + ctx.attr.cpp_class,
"--target_triple=" + ctx.attr.target_triple,
] + cpu_flags + output_flags + ctx.attr.extra_flags + dfsan_flags
post_command = ""
if ctx.attr.gen_compiler_log:
post_command += " --vmodule=cpu_compiler=5 2> >(tee -a " + compiler_log_file.path + " >&2) "
full_cmd = (
ctx.executable.tfcompile_tool.path + " " + " ".join(flags) + " " + ctx.attr.flags + post_command
)
ctx.actions.run_shell(
inputs = ctx.files.srcs,
outputs = out_files,
tools = [ctx.executable.tfcompile_tool] + dfsan_deps,
env = tfcompile_env,
command = full_cmd,
progress_message = "tfcompile for model %s (%s)" % (ctx.attr.model_name, ctx.file.tfcompile_graph.path),
mnemonic = "TensorflowCompile",
)
return [
DefaultInfo(
files = depset(out_files),
),
OutputGroupInfo(**output_dict),
]
# Use tf_library macro instead of using this rule directly.
_tfcompile_model_library = rule(
implementation = _tfcompile_model_library_rule_impl,
attrs = {
"model_name": attr.string(),
"srcs": attr.label_list(mandatory = True, allow_files = True),
"header_out": attr.output(),
"cmd": attr.string(),
"tfcompile_tool": attr.label(cfg = "exec", executable = True, allow_files = True),
"tfcompile_graph": attr.label(allow_single_file = True),
"tfcompile_config": attr.label(allow_single_file = True),
"entry_point": attr.string(),
"cpp_class": attr.string(),
"target_triple": attr.string(),
"target_cpu": attr.string(),
# The tfcompile_flags passed into tf_library macro may be a string
# containing multiple flags (and there are cases that do this).
"flags": attr.string(),
# Extra flags are built in the tf_library macro as a list.
"extra_flags": attr.string_list(),
"dfsan": attr.bool(default = False),
"dfsan_abilists": attr.label_list(default = [], allow_files = True),
"is_linux": attr.bool(),
"gen_compiler_log": attr.bool(),
"xla_flags": attr.string(),
},
)
def _tf_library(
name,
graph,
config,
debug_info = None,
freeze_checkpoint = None,
freeze_saver = None,
cpp_class = None,
gen_test = True,
gen_benchmark = True,
gen_compiler_log = False,
visibility = None,
testonly = None,
tfcompile_flags = None,
tfcompile_tool = "//tensorflow/compiler/aot:tfcompile",
include_standard_runtime_deps = True,
enable_xla_hlo_profiling = False,
enable_tracemes = False,
mlir_components = "None",
deps = None,
tags = [],
copts = [],
xla_flags = None):
if not cpp_class:
fail("cpp_class must be specified")
tfcompile_graph = graph
if freeze_checkpoint or freeze_saver:
if not freeze_checkpoint:
fail("freeze_checkpoint must be specified when freeze_saver is " +
"specified")
freeze_name = "freeze_" + name
freeze_file = freeze_name + ".pb"
# First run tfcompile to generate the list of out_nodes.
#
# Here and below, we set CUDA_VISIBLE_DEVICES='' to prevent the code we
# launch from using any GPUs which might be present. This is important
# because builds may run concurrently with tests, and tests need to be
# able to assume that they have control of the full GPU.
out_nodes_file = "out_nodes_" + freeze_name
native.genrule(
name = ("gen_" + out_nodes_file),
srcs = [config],
outs = [out_nodes_file],
cmd = ("CUDA_VISIBLE_DEVICES='' " +
"$(location " + tfcompile_tool + ")" +
" --config=$(location " + config + ")" +
" --dump_fetch_nodes > $@"),
tools = [tfcompile_tool],
# Run tfcompile on the build host, rather than forge, since it's
# typically way faster on the local machine.
local = 1,
tags = tags,
)
# Now run freeze_graph to convert variables into constants.
freeze_args = (
" --input_graph=$(location " + graph + ")" +
" --checkpoint_version=1" +
" --input_binary=" + str(not graph.endswith(".pbtxt")) +
" --input_checkpoint=$(location " + freeze_checkpoint + ")" +
" --output_graph=$(location " + freeze_file + ")" +
" --output_node_names=$$(<$(location " + out_nodes_file +
"))"
)
freeze_saver_srcs = []
if freeze_saver:
freeze_args += " --input_saver=$(location " + freeze_saver + ")"
freeze_saver_srcs.append(freeze_saver)
native.genrule(
name = freeze_name,
srcs = [
graph,
freeze_checkpoint,
out_nodes_file,
] + freeze_saver_srcs,
outs = [freeze_file],
cmd = (
"CUDA_VISIBLE_DEVICES='' " +
"$(location " +
"//tensorflow/python/tools:freeze_graph)" +
freeze_args
),
tools = ["//tensorflow/python/tools:freeze_graph"],
tags = tags,
)
tfcompile_graph = freeze_file
# Rule that runs tfcompile to produce the header and object file.
header_file = name + ".h"
# The XLA backends morph kernel name prefix __ that is not in the form of
# __xla_.
ep = ("__xla_" + native.package_name() + "__" + name).replace("/", "_")
if type(tfcompile_flags) == type(""):
flags = tfcompile_flags
else:
flags = " ".join([
"'" + arg.replace("'", "'\\''") + "'"
for arg in (tfcompile_flags or [])
])
# Do this before we append the `select` into `flags`, because doing so
# transforms `flags` into a variable of type `select`, and we can't call
# `find` on such an object.
need_xla_data_proto = flags and flags.find("--gen_program_shape") != -1
if enable_xla_hlo_profiling:
profiling_flags = ["--xla_hlo_profile"]
else:
profiling_flags = []
if enable_tracemes:
traceme_flags = ["--xla_cpu_enable_xprof_traceme=true"]
else:
traceme_flags = ["--xla_cpu_enable_xprof_traceme=false"]
mlir_flags = ["--mlir_components=" + mlir_components]
srcs = [tfcompile_graph, config]
debug_info_flags = []
if debug_info:
srcs.append(debug_info)
debug_info_flags = ["--debug_info=$(location " + debug_info + ")"]
tfcompile_gen = "gen_" + name
_tfcompile_model_library(
name = tfcompile_gen,
model_name = name,
srcs = srcs,
gen_compiler_log = gen_compiler_log,
header_out = header_file,
tfcompile_tool = tfcompile_tool,
tfcompile_graph = tfcompile_graph,
tfcompile_config = config,
entry_point = ep,
cpp_class = cpp_class,
target_cpu = tfcompile_target_cpu(name),
target_triple = target_llvm_triple(),
flags = flags,
extra_flags = debug_info_flags + profiling_flags + mlir_flags + traceme_flags,
dfsan = tfcompile_dfsan_enabled(),
dfsan_abilists = tfcompile_dfsan_abilists(),
is_linux = select({
"//tensorflow:linux_x86_64": True,
"//conditions:default": False,
}),
visibility = visibility,
testonly = testonly,
tags = tags,
xla_flags = xla_flags,
)
tfcompile_gen_object_files = tfcompile_gen + "_object_files"
native.filegroup(
name = tfcompile_gen_object_files,
srcs = [tfcompile_gen],
output_group = "object_files",
visibility = visibility,
testonly = testonly,
)
use_xla_nanort_runtime = False
if tfcompile_flags and "--use_xla_nanort_runtime" in tfcompile_flags:
use_xla_nanort_runtime = True
# The cc_library rule packaging up the header and object file, and needed
# kernel implementations.
cc_library(
name = name,
srcs = [tfcompile_gen_object_files],
hdrs = [header_file],
visibility = visibility,
testonly = testonly,
deps = [
# These deps are required by all tf_library targets even if
# include_standard_runtime_deps is False. Without them, the
# generated code will fail to compile.
"//third_party/absl/log:check",
"//third_party/absl/synchronization",
"//tensorflow/core:framework_lite",
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
"@xla//xla:types",
"@xla//xla/backends/cpu/runtime:kernel_c_api",
"@xla//xla/backends/cpu/runtime:rng_state_lib",
] + (need_xla_data_proto and [
# If we're generating the program shape, we must depend on the
# proto.
"@xla//xla:xla_data_proto_cc",
] or []) + (enable_xla_hlo_profiling and [
"@xla//xla/service:hlo_profile_printer_data_cc",
] or []) + (include_standard_runtime_deps and [
# TODO(cwhipkey): only depend on kernel code that the model actually
# needed.
"@xla//xla/backends/cpu/runtime:dot_lib",
"@xla//xla/backends/cpu/runtime:sort_lib",
"@xla//xla/backends/cpu/runtime:topk_lib",
"@xla//xla/backends/cpu/runtime:convolution_lib",
"@xla//xla/service/cpu:runtime_matmul",
"@xla//xla/service/cpu:runtime_single_threaded_matmul",
"@eigen_archive//:eigen3",
] or []) + (use_xla_nanort_runtime and [
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function_thunks",
] or []) + (deps or []),
tags = tags,
copts = copts,
)
# Variables used for gen_test and gen_benchmark.
cpp_class_split = cpp_class.rsplit("::", 2)
if len(cpp_class_split) == 1:
no_ns_name = cpp_class_split[0]
else:
no_ns_name = cpp_class_split[1]
sed_replace = (
"-e \"s|{{TFCOMPILE_HEADER}}|$(location " + header_file + ")|g\" " +
"-e \"s|{{TFCOMPILE_CPP_CLASS}}|" + cpp_class + "|g\" " +
"-e \"s|{{TFCOMPILE_NAME}}|" + no_ns_name + "|g\" "
)
if gen_test:
test_name = name + "_test"
test_file = test_name + ".cc"
template_file = "//tensorflow/compiler/aot:test"
template_file += if_oss("", "_google") + ".cc"
# Rule to rewrite the template_file to produce the test_file.
native.genrule(
name = ("gen_" + test_name),
testonly = 1,
srcs = [
template_file,
header_file,
],
outs = [test_file],
cmd = (
"sed " + sed_replace +
" $(location " + template_file + ") " +
"> $(OUTS)"
),
tags = tags,
)
# The cc_test rule for the generated code. To ensure that this works
# reliably across build configurations, we must use tf_cc_test instead
# of native.cc_test. This is related to how we build
# //tensorflow/core:lib -- see the note in
# tensorflow/core/BUILD for more details.
tf_cc_test(
name = test_name,
srcs = [test_file],
deps = [
":" + name,
"//tensorflow/compiler/aot:tf_library_test_main",
"@xla//xla:executable_run_options",
"@eigen_archive//:eigen3",
] + if_oss([
"//tensorflow/core:lib",
"//tensorflow/core:test",
]) + if_google([
"@com_google_googletest//:gtest",
"//tensorflow/core/platform:byte_order",
"//tensorflow/core/platform:platform_port",
]),
tags = tags,
extra_copts = copts,
visibility = visibility,
)
if gen_benchmark:
benchmark_name = name + "_benchmark"
benchmark_file = benchmark_name + ".cc"
benchmark_main = ("//tensorflow/compiler/aot:" +
"benchmark_main.template")
# Rule to rewrite benchmark.cc to produce the benchmark_file.
native.genrule(
name = ("gen_" + benchmark_name),
srcs = [
benchmark_main,
header_file,
],
testonly = testonly,
outs = [benchmark_file],
cmd = ("sed " + sed_replace +
" $(location " + benchmark_main + ") " +
"> $(OUTS)"),
tags = tags,
)
# The cc_benchmark rule for the generated code. This does not need the
# tf_cc_binary since we (by deliberate design) do not depend on
# //tensorflow/core:lib.
#
# Note: to get smaller size on android for comparison, compile with:
# --copt=-fvisibility=hidden
# --copt=-D_LIBCPP_TYPE_VIS=_LIBCPP_HIDDEN
# --copt=-D_LIBCPP_EXCEPTION_ABI=_LIBCPP_HIDDEN
cc_binary(
name = benchmark_name,
srcs = [benchmark_file],
testonly = testonly,
copts = copts + tf_copts(),
linkopts = if_android(["-pie", "-s"]),
deps = [
":" + name,
"//tensorflow/compiler/aot:benchmark",
"@xla//xla:executable_run_options",
"@eigen_archive//:eigen3",
] + if_android([
"//tensorflow/compiler/aot:benchmark_extra_android",
]),
tags = tags,
visibility = visibility,
)
def tf_library(
name,
graph,
config,
debug_info = None,
freeze_checkpoint = None,
freeze_saver = None,
cpp_class = None,
gen_test = True,
gen_benchmark = True,
gen_compiler_log = False,
visibility = None,
testonly = None,
tfcompile_flags = None,
tfcompile_tool = "//tensorflow/compiler/aot:tfcompile",
include_standard_runtime_deps = True,
enable_xla_hlo_profiling = False,
enable_tracemes = False,
mlir_components = "None",
deps = None,
tags = [],
copts = [],
xla_flags = None):
"""Compiles a TensorFlow graph into an executable with fast math enabled.
Given an invocation of tf_library(name="foo", ...), generates the following
build targets:
foo: A cc_library containing the generated header and
computation.
foo_test: A cc_test with simple tests and benchmarks. Only created if
gen_test=True.
foo_benchmark: A cc_binary that runs a minimal-dependency benchmark,
useful for mobile devices or other platforms that can't
compile the full test libraries. Only created if
gen_benchmark=True.
The output header is called <name>.h.
Args:
name: The name of the build rule.
graph: The TensorFlow GraphDef to compile. If the file ends in '.pbtxt'
it is expected to be in the human-readable proto text format, otherwise
it is expected to be in the proto binary format.
config: File containing tensorflow.tf2xla.Config proto. If the file ends
in '.pbtxt' it is expected to be in the human-readable proto text
format, otherwise it is expected to be in the proto binary format.
debug_info: Debug info to include in the output.
freeze_checkpoint: If provided, run freeze_graph with this checkpoint to
convert variables into constants.
freeze_saver: If provided, run freeze_graph with this saver, in SaverDef
binary form, to convert variables into constants.
cpp_class: The name of the generated C++ class, wrapping the generated
function. The syntax of this flag is
[[<optional_namespace>::],...]<class_name>. This mirrors the C++ syntax
for referring to a class, where multiple namespaces may precede the
class name, separated by double-colons. The class will be generated in
the given namespace(s), or if no namespaces are given, within the global
namespace.
gen_test: If True, also generate a cc_test rule that builds a simple
test and benchmark.
gen_benchmark: If True, also generate a binary with a simple benchmark.
Unlike the output of gen_test, this benchmark can be run on android.
gen_compiler_log: If True, dumps XLA:CPU debug output to a log file.
visibility: Bazel build visibility.
testonly: Bazel testonly attribute.
tfcompile_flags: Extra flags to pass to tfcompile to control compilation.
tfcompile_tool: The tfcompile binary. A non-default can be passed to
use a tfcompile built with extra dependencies.
include_standard_runtime_deps: If True, the standard list of
kernel/runtime deps is added to deps. If False, deps must contain the
full set of deps needed by the generated library.
enable_xla_hlo_profiling: Enable XLA HLO profiling in the generated
program, and emit metadata that lets us pretty-print the gathered
profile counters.
enable_tracemes: Tell tfcompile to generate calls to
TraceMe::Activity{Start|End} around HLO instructions that can be used by
Xprof to construct profiler timelines.
mlir_components: When the value is "None", no components use MLIR. When
the value is "Bridge", use MLIR to translate GraphDef to HLO.
deps: a list of deps to include on the build rules for the generated
library, added to the standard deps if standard_runtime_deps is True.
tags: tags to apply to subsidiary build rules.
copts: list of copts to pass to cc rules.
"""
_tf_library(
name,
graph,
config,
debug_info,
freeze_checkpoint,
freeze_saver,
cpp_class,
gen_test,
gen_benchmark,
gen_compiler_log,
visibility,
testonly,
tfcompile_flags,
tfcompile_tool,
include_standard_runtime_deps,
enable_xla_hlo_profiling,
enable_tracemes,
mlir_components,
deps,
tags,
copts,
xla_flags,
)
def target_llvm_triple():
"""Returns the target LLVM triple to be used for compiling the target."""
# TODO(toddw): Add target_triple for other targets. For details see:
# http://llvm.org/docs/doxygen/html/Triple_8h_source.html
return select({
"//tensorflow:android_armeabi": "armv5-none-android",
"//tensorflow:android_arm": "armv7-none-android",
"//tensorflow:android_arm64": "aarch64-none-android",
"//tensorflow:android_x86": "i686-none-android",
"//tensorflow:ios": "arm64-none-ios",
"//tensorflow:ios_x86_64": "x86_64-apple-ios",
"//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu",
"//tensorflow:linux_aarch64": "aarch64-none-linux-gnu",
"//tensorflow:macos_x86_64": "x86_64-none-darwin",
"//tensorflow:macos_arm64": "aarch64-none-darwin",
"//tensorflow:windows": "x86_64-none-windows",
"//tensorflow:linux_s390x": "systemz-none-linux-gnu",
# internal placeholder,
"//conditions:default": "x86_64-pc-linux",
})