使用 AOT 编译

什么是 tfcompile？

tfcompile 是一款独立的工具，可预先 (AOT) 将 TensorFlow 图编译为可执行代码。它可以缩减二进制文件的总大小，还可以避免一些运行时开销。tfcompile 的一个典型用例是将推断图编译为适用于移动设备的可执行代码。

TensorFlow 图通常由 TensorFlow 运行时执行。因此，执行图中的每个节点时，会产生一些运行时开销。这也会导致二进制文件的总大小变大，因为除了图本身之外，还需要提供 TensorFlow 运行时的代码。tfcompile 生成的可执行代码不使用 TensorFlow 运行时，并且仅依赖于计算中实际使用的内核。

编译器基于 XLA 框架构建。将 TensorFlow 桥接到 XLA 框架的代码位于 tensorflow/compiler 下。

tfcompile 有什么作用？

tfcompile 接受一个子图（由 TensorFlow 的 Feed 和提取概念标识），并生成实现该子图的函数。feeds 是函数的输入参数，fetches 是函数的输出参数。所有输入都必须由 Feed 完全指定；产生的删减的子图不能包含占位符或变量节点。常见的做法是将所有占位符和变量指定为 Feed，以确保生成的子图不再包含这些节点。生成的函数会被打包为 cc_library，其中包含一个用于导出函数签名的头文件，以及一个包含相应实现的对象文件。用户编写代码以根据需要调用生成的函数。

使用 tfcompile

本部分详细介绍了使用 tfcompile 从 TensorFlow 子图生成可执行二进制文件的概要步骤。具体步骤包括：

• 第 1 步：配置子图以进行编译
• 第 2 步：使用 tf_library build 宏编译子图
• 第 3 步：编写代码以调用子图
• 第 4 步：创建最终二进制文件

第 1 步：配置子图以进行编译

确定与生成的函数的输入和输出参数相对应的 Feed 和提取项。然后，在 tensorflow.tf2xla.Config proto 中配置 feeds 和 fetches。

# Each feed is a positional input argument for the generated function.  The order
# of each entry matches the order of each input argument.  Here “x_hold” and “y_hold”
# refer to the names of placeholder nodes defined in the graph.
feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}

# Each fetch is a positional output argument for the generated function.  The order
# of each entry matches the order of each output argument.  Here “x_y_prod”
# refers to the name of a matmul node defined in the graph.
fetch {
  id { node_name: "x_y_prod" }
}

第 2 步：使用 tf_library build 宏编译子图

此步骤使用 tf_library 构建宏将图表转换为 cc_library。cc_library 包含一个对象文件，其中包含从图表生成的代码，以及一个提供所生成代码的头文件。tf_library 利用 tfcompile 将 TensorFlow 图编译为可执行代码。

load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# Use the tf_library macro to compile your graph into executable code.
tf_library(
    # name is used to generate the following underlying build rules:
    # <name>           : cc_library packaging the generated header and object files
    # <name>_test      : cc_test containing a simple test and benchmark
    # <name>_benchmark : cc_binary containing a stand-alone benchmark with minimal deps;
    #                    can be run on a mobile device
    name = "test_graph_tfmatmul",
    # cpp_class specifies the name of the generated C++ class, with namespaces allowed.
    # The class will be generated in the given namespace(s), or if no namespaces are
    # given, within the global namespace.
    cpp_class = "foo::bar::MatMulComp",
    # graph is the input GraphDef proto, by default expected in binary format.  To
    # use the text format instead, just use the ‘.pbtxt’ suffix.  A subgraph will be
    # created from this input graph, with feeds as inputs and fetches as outputs.
    # No Placeholder or Variable ops may exist in this subgraph.
    graph = "test_graph_tfmatmul.pb",
    # config is the input Config proto, by default expected in binary format.  To
    # use the text format instead, use the ‘.pbtxt’ suffix.  This is where the
    # feeds and fetches were specified above, in the previous step.
    config = "test_graph_tfmatmul.config.pbtxt",
)

如需为此示例生成 GraphDef proto (test_graph_tfmatmul.pb)，请运行 make_test_graphs.py 并使用 --out_dir 标志指定输出位置。

典型图包含 Variables（表示通过训练学到的权重），但 tfcompile 无法编译包含 Variables 的子图。freeze_graph.py 工具使用存储在检查点文件中的值将变量转换为常量。为方便起见，tf_library 宏支持运行该工具的 freeze_checkpoint 参数。如需查看更多示例，请参阅 tensorflow/compiler/aot/tests/BUILD。

在已编译的子图中显示的常量会直接编译到生成的代码中。如需将常量传递到生成的函数中，而不是编译它们，只需将它们作为 Feed 传入即可。

如需详细了解 tf_library build 宏，请参阅 tfcompile.bzl。

如需详细了解底层 tfcompile 工具，请参阅 tfcompile_main.cc。

第 3 步：编写代码以调用子图

此步骤使用上一步中 tf_library build 宏生成的头文件 (test_graph_tfmatmul.h) 来调用生成的代码。头文件位于 build 软件包对应的 bazel-bin 目录中，并根据在 tf_library build 宏上设置的名称属性进行命名。例如，为 test_graph_tfmatmul 生成的标头将为 test_graph_tfmatmul.h。以下是所生成内容的简化版本。生成的文件位于 bazel-bin 中，其中包含其他有用注释。

namespace foo {
namespace bar {

// MatMulComp represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code.
class MatMulComp {
 public:
  // AllocMode controls the buffer allocation mode.
  enum class AllocMode {
    ARGS_RESULTS_AND_TEMPS,  // Allocate arg, result and temp buffers
    RESULTS_AND_TEMPS_ONLY,  // Only allocate result and temp buffers
  };

  MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS);
  ~MatMulComp();

  // Runs the computation, with inputs read from arg buffers, and outputs
  // written to result buffers. Returns true on success and false on failure.
  bool Run();

  // Arg methods for managing input buffers. Buffers are in row-major order.
  // There is a set of methods for each positional argument.
  void** args();

  void set_arg0_data(float* data);
  float* arg0_data();
  float& arg0(size_t dim0, size_t dim1);

  void set_arg1_data(float* data);
  float* arg1_data();
  float& arg1(size_t dim0, size_t dim1);

  // Result methods for managing output buffers. Buffers are in row-major order.
  // Must only be called after a successful Run call. There is a set of methods
  // for each positional result.
  void** results();


  float* result0_data();
  float& result0(size_t dim0, size_t dim1);
};

}  // end namespace bar
}  // end namespace foo

生成的 C++ 类在 foo::bar 命名空间中称为 MatMulComp，因为这是 tf_library 宏中指定的 cpp_class。所有生成的类都具有类似的 API，唯一的区别在于处理参数和结果缓冲区的方法。这些方法因缓冲区的数量和类型而异，缓冲区由 tf_library 宏的 feed 和 fetch 参数指定。

在生成的类中管理三种类型的缓冲区：args 表示输入，results 表示输出，temps 表示在内部用于执行计算的临时缓冲区。默认情况下，生成的类的每个实例会为您分配和管理所有这些缓冲区。AllocMode 构造函数参数可用于更改此行为。所有缓冲区均与 64 字节边界对齐。

生成的 C++ 类只是由 XLA 生成的低级别代码的封装容器。

基于 tfcompile_test.cc 调用生成的函数的示例：

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "third_party/tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());


  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

第 4 步：创建最终二进制文件

此步骤会将第 2 步中 tf_library 生成的库与第 3 步中编写的代码相结合，以创建最终二进制文件。以下是 bazel BUILD 文件示例。

# Example of linking your binary
# Also see //tensorflow/compiler/aot/tests/BUILD
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# The same tf_library call from step 2 above.
tf_library(
    name = "test_graph_tfmatmul",
    ...
)

# The executable code generated by tf_library can then be linked into your code.
cc_binary(
    name = "my_binary",
    srcs = [
        "my_code.cc",  # include test_graph_tfmatmul.h to access the generated header
    ],
    deps = [
        ":test_graph_tfmatmul",  # link in the generated object file
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)