使用 AOT 編譯

什麼是 tfcompile？

tfcompile 是一項獨立工具，會預先 (AOT) 將 TensorFlow 圖表編譯為可執行的程式碼。這可以降低二進位檔的總大小，也能避免某些執行階段負擔。tfcompile 的常見用途是將推論圖編譯成適合行動裝置的可執行程式碼。

TensorFlow 圖形通常是由 TensorFlow 執行階段執行。因此，在執行圖形中每個節點時，都會產生一些執行階段負擔。這也會導致二進位檔的總大小變大，因為除了圖表本身以外，TensorFlow 執行階段也需要使用的程式碼。tfcompile 產生的可執行程式碼不會使用 TensorFlow 執行階段，而且只有實際用於運算的核心核心依附元件。

編譯器是以 XLA 架構為基礎建構而成。將 TensorFlow 繫結至 XLA 架構的程式碼位於 tensorflow/compiler 下，

tfcompile 有哪些功用？

tfcompile 會採用由 TensorFlow 動態饋給及擷取概念識別的子圖，並產生實作該子圖的函式。feeds 是函式的輸入引數，fetches 則是函式的輸出引數。所有輸入內容都必須由動態饋給完整指定；產生的剪裁子圖表不得包含預留位置或變數節點。將所有預留位置和變數指定為動態饋給是很常見的做法，可確保產生的子圖表不再包含這些節點。產生的函式會封裝為 cc_library (具有匯出函式簽章的標頭檔案) 以及包含實作項目的物件檔案。使用者撰寫程式碼，視情況叫用產生的函式。

使用 tfcompile

本節詳細說明如何使用 tfcompile 從 TensorFlow 子圖表產生可執行二進位檔的高階步驟。步驟如下：

• 步驟 1：設定要編譯的子圖表
• 步驟 2：使用 tf_library 建構巨集編譯子圖表
• 步驟 3：編寫程式碼以叫用子圖表
• 步驟 4：建立最終二進位檔

步驟 1：設定要編譯的子圖表

針對產生的函式，找出與輸入和輸出引數相對應的動態饋給和擷取項目。接著在 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 建構巨集來編譯子圖表

這個步驟會使用 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。

已編譯子圖表中顯示的常數會直接編譯到產生的程式碼中。如要將常數傳遞給產生的函式，而非進行編譯，只需將其做為動態饋給傳遞即可。

如要進一步瞭解 tf_library 建構巨集，請參閱 tfcompile.bzl。

如要進一步瞭解基礎 tfcompile 工具，請參閱 tfcompile_main.cc。

步驟 3：編寫程式碼以叫用子圖表

這個步驟會使用在上一個步驟中 tf_library 建構巨集產生的標頭檔案 (test_graph_tfmatmul.h)，叫用已產生的程式碼。標頭檔案位於與建構套件對應的 bazel-bin 目錄中，且會根據 tf_library 建構巨集上設定的名稱屬性命名。舉例來說，針對 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",
    ]
)