WeHub Mirror
16d3f15320
ARM CI / build (3.10) (push) Has been cancelled
cffconvert / validate (push) Has been cancelled
Creates a GitHub Issue when a PR Rolled back via Commit to Master / create-issue-on-pr-rollback (push) Has been cancelled
Scorecards supply-chain security / Scorecards analysis (push) Has been cancelled
195 lines
7.5 KiB
Python
195 lines
7.5 KiB
Python
# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
"""Definitions for using tools like saved_model_cli."""
|
|
|
|
load("@rules_cc//cc:cc_library.bzl", "cc_library")
|
|
load("//tensorflow:tensorflow.bzl", "clean_dep", "if_xla_available")
|
|
load("//tensorflow:tensorflow.default.bzl", "tfcompile_target_cpu")
|
|
load("//tensorflow/compiler/aot:tfcompile.bzl", "target_llvm_triple")
|
|
|
|
def _maybe_force_compile(args, force_compile):
|
|
if force_compile:
|
|
return args
|
|
else:
|
|
return if_xla_available(args)
|
|
|
|
def saved_model_compile_aot(
|
|
name,
|
|
directory,
|
|
filegroups,
|
|
cpp_class,
|
|
checkpoint_path = None,
|
|
tag_set = "serve",
|
|
signature_def = "serving_default",
|
|
variables_to_feed = "",
|
|
target_triple = None,
|
|
target_cpu = None,
|
|
multithreading = False,
|
|
force_without_xla_support_flag = True,
|
|
tags = None):
|
|
"""Compile a SavedModel directory accessible from a filegroup.
|
|
|
|
This target rule takes a path to a filegroup directory containing a
|
|
SavedModel and generates a cc_library with an AOT compiled model.
|
|
For extra details, see the help for saved_model_cli's aot_compile_cpu help.
|
|
|
|
**NOTE** Any variables passed to `variables_to_feed` *must be set by the
|
|
user*. These variables will NOT be frozen and their values will be
|
|
uninitialized in the compiled object (this applies to all input
|
|
arguments from the signature as well).
|
|
|
|
Example usage:
|
|
|
|
```
|
|
saved_model_compile_aot(
|
|
name = "aot_compiled_x_plus_y",
|
|
cpp_class = "tensorflow::CompiledModel",
|
|
directory = "//tensorflow/cc/saved_model:testdata/x_plus_y_v2_debuginfo",
|
|
filegroups = [
|
|
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
|
|
]
|
|
)
|
|
|
|
cc_test(
|
|
name = "test",
|
|
srcs = ["test.cc"],
|
|
deps = [
|
|
"//tensorflow/core:test_main",
|
|
":aot_compiled_x_plus_y",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core/platform:logging",
|
|
]),
|
|
)
|
|
|
|
In "test.cc":
|
|
|
|
#include "third_party/tensorflow/python/tools/aot_compiled_x_plus_y.h"
|
|
|
|
TEST(Test, Run) {
|
|
tensorflow::CompiledModel model;
|
|
CHECK(model.Run());
|
|
}
|
|
```
|
|
|
|
Args:
|
|
name: The rule name, and the name prefix of the headers and object file
|
|
emitted by this rule.
|
|
directory: The bazel directory containing saved_model.pb and variables/
|
|
subdirectories.
|
|
filegroups: List of `filegroup` targets; these filegroups contain the
|
|
files pointed to by `directory` and `checkpoint_path`.
|
|
cpp_class: The name of the C++ class that will be generated, including
|
|
namespace; e.g. "my_model::InferenceRunner".
|
|
checkpoint_path: The bazel directory containing `variables.index`. If
|
|
not provided, then `$directory/variables/` is used
|
|
(default for SavedModels).
|
|
tag_set: The tag set to use in the SavedModel.
|
|
signature_def: The name of the signature to use from the SavedModel.
|
|
variables_to_feed: (optional) The names of the variables to feed, a comma
|
|
separated string, or 'all'. If empty, all variables will be frozen and none
|
|
may be fed at runtime.
|
|
|
|
**NOTE** Any variables passed to `variables_to_feed` *must be set by
|
|
the user*. These variables will NOT be frozen and their values will be
|
|
uninitialized in the compiled object (this applies to all input
|
|
arguments from the signature as well).
|
|
target_triple: The LLVM target triple to use (defaults to current build's
|
|
target architecture's triple). Similar to clang's -target flag.
|
|
target_cpu: The LLVM cpu name used for compilation. Similar to clang's
|
|
-mcpu flag.
|
|
multithreading: Whether to compile multithreaded AOT code.
|
|
Note, this increases the set of dependencies for binaries using
|
|
the AOT library at both build and runtime. For example,
|
|
the resulting object files may have external dependencies on
|
|
multithreading libraries like Abseil.
|
|
force_without_xla_support_flag: Whether to compile even when
|
|
`--define=with_xla_support=true` is not set. If `False`, and the
|
|
define is not passed when building, then the created `cc_library`
|
|
will be empty. In this case, downstream targets should
|
|
conditionally build using macro `tfcompile.bzl:if_xla_available`.
|
|
This flag is used by the TensorFlow build to avoid building on
|
|
architectures that do not support XLA.
|
|
tags: List of target tags.
|
|
"""
|
|
saved_model = "{}/saved_model.pb".format(directory)
|
|
target_triple = target_triple or target_llvm_triple()
|
|
target_cpu = target_cpu or tfcompile_target_cpu(name) or ""
|
|
variables_to_feed = variables_to_feed or "''"
|
|
if checkpoint_path:
|
|
checkpoint_cmd_args = (
|
|
"--checkpoint_path \"$$(dirname $(location {}/variables.index))\" "
|
|
.format(checkpoint_path)
|
|
)
|
|
checkpoint_srcs = ["{}/variables.index".format(checkpoint_path)]
|
|
else:
|
|
checkpoint_cmd_args = ""
|
|
checkpoint_srcs = []
|
|
|
|
native.genrule(
|
|
name = "{}_gen".format(name),
|
|
srcs = filegroups + [saved_model] + checkpoint_srcs,
|
|
outs = [
|
|
"{}.h".format(name),
|
|
"{}.o".format(name),
|
|
"{}_metadata.o".format(name),
|
|
"{}_constants.o".format(name),
|
|
"{}_makefile.inc".format(name),
|
|
],
|
|
cmd = (
|
|
"$(location {}) aot_compile_cpu ".format(
|
|
clean_dep("//tensorflow/python/tools:saved_model_cli"),
|
|
) +
|
|
"--dir \"$$(dirname $(location {}))\" ".format(saved_model) +
|
|
checkpoint_cmd_args +
|
|
"--output_prefix $(@D)/{} ".format(name) +
|
|
"--cpp_class {} ".format(cpp_class) +
|
|
"--variables_to_feed {} ".format(variables_to_feed) +
|
|
"--signature_def_key {} ".format(signature_def) +
|
|
"--multithreading {} ".format(multithreading) +
|
|
"--target_triple " + target_triple + " " +
|
|
("--target_cpu " + target_cpu + " " if target_cpu else "") +
|
|
"--tag_set {} ".format(tag_set)
|
|
),
|
|
tags = tags,
|
|
tools = [
|
|
"//tensorflow/python/tools:saved_model_cli",
|
|
],
|
|
)
|
|
cc_library(
|
|
name = name,
|
|
srcs = _maybe_force_compile(
|
|
[
|
|
":{}.o".format(name),
|
|
":{}_metadata.o".format(name),
|
|
":{}_constants.o".format(name),
|
|
],
|
|
force_compile = force_without_xla_support_flag,
|
|
),
|
|
hdrs = _maybe_force_compile(
|
|
[
|
|
":{}.h".format(name),
|
|
],
|
|
force_compile = force_without_xla_support_flag,
|
|
),
|
|
tags = tags,
|
|
deps = _maybe_force_compile(
|
|
[
|
|
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_runtime_standalone",
|
|
],
|
|
force_compile = force_without_xla_support_flag,
|
|
),
|
|
)
|