Dive into TensorFlow - 解析 TF 核心抽象 op 算子

TF 计算图从逻辑层来讲，由 op 与 tensor 构成。op 是项点代表计算单元，tensor 是边代表 op 之间流动的数据内容，两者配合以数据流图的形式来表达计算图。那么 op 对应的物理层实现是什么？TF 中有哪些 op，以及各自的适用场景是什么？op 到底是如何运行的？接下来让我们一起探索和回答这些问题。

一、初识 op

1.1 op 定义

op 代表计算图中的节点，是 tf.Operation 对象，代表一个计算单元。用户在创建模型和训练代码时，会创建一系列 op 及其依赖关系，并将这些 op 和依赖添加到 tf.Graph 对象中（一般为默认图）。比如：tf.matmul () 就是一个 op，它有两个输入 tensor 和一个输出 tensor。

1.2 op 分类

op 的分类一般有多个视角，比如按是否内置划分、按工作类型划分。

按是否内置划分，一般分为：内置 op 和自定义 op（见 “二、自定义 op” 部分介绍）。

按工作类型划分，一般分为：常见数学 op、数组 op、矩阵 op、有状态 op、神经网络 op、检查点 op、队列与同步 op、控制流 op。TF 白皮书对内置 op 的分类总结如下：





1.3 op 与 kernel

op 一般都有名称且代表一个抽象的计算过程。op 可以设置若干属性，但这些属性必须在编译期提供或推理得到，因为它们用来实例化一个节点对象从而执行真正的计算。属性的经典用法就是拿来支持类型多态，比如两个浮点张量的矩阵乘法与两个整型张量的矩阵乘法。

kernel 是 op 在指定设备类型（CPU/GPU）上的具体实现。TF 二进制库通过注册机制定义了一系列 op 及对应的 kernel 实现，用户可以提供额外的 op 定义与 kernel 实现进行扩充。一般来说，一个 op 对应多个 kernel 实现。

接下来让我们一起用矩阵乘法 MatMul 算子的相关代码来理解 op 与 kernel 的关系（此处不必纠结代码细节，只需体会 op 与 kernel 关系即可）：

// 首先给出op注册的定义。其中输入输出支持泛型，其合法类型在Attr中进行枚举。
// 代码位置 tensorflow1.15.5\tensorflow\core\ops\math_ops.cc
REGISTER_OP("MatMul")
    .Input("a: T")
    .Input("b: T")
    .Output("product: T")
    .Attr("transpose_a: bool = false")
    .Attr("transpose_b: bool = false")
    .Attr(
        "T: {bfloat16, half, float, double, int32, int64, complex64, "
        "complex128}")
    .SetShapeFn(shape_inference::MatMulShape);
    
// MatMul的实现，采用类模板机制
// 代码位置 tensorflow1.15.5\tensorflow\core\kernels\matmul_op.cc
template <typename Device, typename T, bool USE_CUBLAS>
class MatMulOp : public OpKernel {
  public:
    explicit MatMulOp(OpKernelConstruction* ctx)
      : OpKernel(ctx), algorithms_set_already_(false) {
      OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));
      OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
 
      LaunchMatMul<Device, T, USE_CUBLAS>::GetBlasGemmAlgorithm(
        ctx, &algorithms_, &algorithms_set_already_);
      use_autotune_ = MatmulAutotuneEnable();
    }
  // 省略了很多代码...  
  private:
    std::vector<int64> algorithms_;
    bool algorithms_set_already_;
    bool use_autotune_;
    bool transpose_a_;
    bool transpose_b_;
};
 
// MatMul的op定义与kernel实现绑定处理
// 代码位置 tensorflow1.15.5\tensorflow\core\kernels\matmul_op.cc
#define REGISTER_CPU_EIGEN(T)  /*cpu与eigen组合对应实现*/                       \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), \
      MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>);
 
#define REGISTER_CPU(T)      /*cpu对应实现(eigen与非eigen)*/         \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T"),     \
      MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); \
  REGISTER_CPU_EIGEN(T);
 
#define REGISTER_GPU(T)     /*gpu对应实现(cublas与非cublas)*/       \
  REGISTER_KERNEL_BUILDER(                                         \
      Name("MatMul").Device(DEVICE_GPU).TypeConstraint<T>("T"),    \
      MatMulOp<GPUDevice, T, true /* cublas, true by default */>); \
  REGISTER_KERNEL_BUILDER(Name("MatMul")                           \
                              .Device(DEVICE_GPU)                  \
                              .TypeConstraint<T>("T")              \
                              .Label("cublas"),                    \
                          MatMulOp<GPUDevice, T, true /* cublas */>)

二、自定义 op

用户编写的模型训练代码一般由 TF 原生的 op 算子及其依赖关系组成，但有时候我们定义的计算逻辑在 TF 中没有相应的 op 实现。根据 TensorFlow 官网的建议，我们应当先组合 python op 算子或 python 函数进行尝试。完成尝试之后再决定要不要自定义 op。

2.1 自定义 op 场景

一般来说，需要自定义 op 的场景有如下 3 个：

• 用 TF 原生 op 组合来表达新计算逻辑的过程比较复杂或不可能

• 用 TF 原生 op 组合来表达新计算逻辑，其计算性能较低

• 在新版编译器中也较难实现 op 融合的计算逻辑需要我们手动实现融合

在此举个例子方便大家理解。假如我们要实现一个新计算实逻：中位数池化（median pooling），过程中要在滑动窗口不断求得中位数。检索 TF 文档没有发现对应 op，因此我们先考虑用 TF python op 组合来实现它，果然通过 ExtractImagePatches and TopK 就可以实现这个功能。经测试前述组合方案并不是计算和存储高效的，因此我们就有必要将 median pooling 在一个 op 中进行高效实现。

2.2 自定义 op 流程

自定义 op 一般遵循 5 个基本步骤：

1. 注册 op，具体包括：指定名称、输入 / 输出声明、形状函数。

2. 定义 kernel（即 op 的实现）并与 op 绑定。一个 op 有多个 kernel 实现，具体由输入输出类型、硬件（CPU、GPU）决定。

3. 创建 python 包装器，一般由 op 注册机制自动完成。

4. 编写 op 的梯度计算函数（可选项）。

5. 测试 op，通过 python 测试较为方便，当然也可通过 C++ 进行测试。

接下来我们就以官网最简单的 ZeroOut 同步式自定义 op（继承 OpKernel）为例，结合代码来讲述上述 5 个步骤。下面先给出步骤 1 和步骤 2 用 C++ 实现的代码（官方推荐用 bazel 编译 so 文件）：

// 步骤1：注册op
REGISTER_OP("ZeroOut")
.Input("to_zero: int32")
.Output("zeroed: int32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->input(0));    //c's input and output type is std::vector
    return Status::OK();
    });
 
// 步骤2：定义kernel（常规CPU设备）,并把kernel与op绑定
class ZeroOutOp : public OpKernel {
public:
    explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
 
    void Compute(OpKernelContext* context) override {
        // Grab the input tensor from OpKernelContext instance
        const Tensor& input_tensor = context->input(0); 
        auto input = input_tensor.flat();
 
        // Create an output tensor
        Tensor* output_tensor = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
            &output_tensor));    // OP_REQUIRES_OK第二个参数一般为方法调用,此处为输出张量分配内存空间
        auto output_flat = output_tensor->flat();
 
        // Set all but the first element of the output tensor to 0.
        const int N = input.size();
        for (int i = 1; i < N; i++) {
            output_flat(i) = 0;
        }
 
        // Preserve the first input value if possible.
        if (N > 0) output_flat(0) = input(0);
    }
};
 
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

步骤 3 加载上述 so 文件（自动完成前后端 op 映射）；步骤 4 是可选项，此处不需要；步骤 5 基于 python api 测试 op 功能。相应代码如下：

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')    # 加载so文件生成python module
with tf.Session(''):
  zero_out_module.zero_out([[1, 2], [3, 4]]).eval()
 
# Prints
array([[1, 0], [0, 0]], dtype=int32)

2.3 高级话题

关于 op 的技术话题还有很多，我们在此简述一些要点：

1. 如果实现了一个多线程 CPU kernel，则可以利用 work_sharder.h 中的 Shard 函数。

2. 大多数 op 以同步方式工作，只需继承 OpKernel 改写 Compute () 方法，且此方法必须线程安全。

3. 如果一个 op 因为其它 op 的运行而阻塞，则这个 op 可以采用异步方式工作，继承 AsyncOpKernel 改写 ComputeAsync () 方法，且此方法必须线程安全。异步 op 最经典的例子就是跨设备通信 send/recv pair 中的 RecvOp。

4. 如果要为 op 配置一些静态属性，可使用 Attr，它有一套特有的支持类型。典型应用是支持泛型。

5. 实现 GPU kernel 有两部分内容：OpKernel 和 CUDA kernel，相应的加载代码。

6. 编译自定义 op，首先要配置头文件搜索路径与库文件搜索路径，接着指定编译和链接选项，最后还要确保 ABI 兼容性。

7. Resource（资源）代表相同设备上 op 共享的内容，比如：张量值、kv 存储表、队列、读取器、网络连接等。代表资源的类必须继承 ResourceBase，然后注册 ResourceHandleOp 生成资源句柄，普通 op 以 resouce 类型的 Input 进行引入。

三、op 工作原理

3.1 op 运行框架

整体来看，op 与 kernel 都有其结构描述与统一的注册管理中心。而 OpDefBuilder 有两个包装类 OpDefBuilderWrapper 和 OpDefBuilderReceiver，前者支持 op 构建的链式语法，后者接受 op 构建结果并进行注册。众所周知，op 是编译期概念，而 kernel 是运行期概念，在 AI 编译器的后端处理流程中会进行 op 的算子选择，此过程会基于一系列策略为 op 匹配最合适的 kernel 实现。





3.2 若干技术细节

首先，我们来看一下大家在使用 TensorFlow 过程中经常碰到的 libtensorflow_framework.so。按照 tf1.15.5/tensorflow/BUILD 中的描述，libtensorflow_framework.so 定义了 op 和 kernel 的注册机制而不涉及具体实现。

// rootdir=tensorflow1.15.5
// ${rootdir}/tensorflow/BUILD
/*
# A shared object which includes registration mechanisms for ops and
# kernels. Does not include the implementations of any ops or kernels. Instead,
# the library which loads libtensorflow_framework.so
# (e.g. _pywrap_tensorflow_internal.so for Python, libtensorflow.so for the C
# API) is responsible for registering ops with libtensorflow_framework.so. In
# addition to this core set of ops, user libraries which are loaded (via
# TF_LoadLibrary/tf.load_op_library) register their ops and kernels with this
# shared object directly.
*/
tf_cc_shared_object(
    name = "tensorflow_framework",
    framework_so = [],
    linkopts = select({
        "//tensorflow:macos": [],
        "//tensorflow:windows": [],
        "//tensorflow:freebsd": [
            "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)",
            "-lexecinfo",
        ],
        "//conditions:default": [
            "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)",
        ],
    }),
    linkstatic = 1,
    per_os_targets = True,
    soversion = VERSION,
    visibility = ["//visibility:public"],
    deps = [
        "//tensorflow/cc/saved_model:loader_lite_impl",
        "//tensorflow/core:core_cpu_impl",
        "//tensorflow/core:framework_internal_impl",    /* 展开此target进行查看 */
        "//tensorflow/core:gpu_runtime_impl",
        "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry_impl",
        "//tensorflow/core:lib_internal_impl",
        "//tensorflow/stream_executor:stream_executor_impl",
        "//tensorflow:tf_framework_version_script.lds",
    ] + tf_additional_binary_deps(),
)
 
// ${rootdir}/tensorflow/core/BUILD
tf_cuda_library(
    name = "framework_internal_impl",
    srcs = FRAMEWORK_INTERNAL_PRIVATE_HEADERS + glob(   // 可以查看FRAMEWORK_INTERNAL_PRIVATE_HEADERS内容
        [
            "example/**/*.cc",
            "framework/**/*.cc",
            "util/**/*.cc",
            "graph/edgeset.cc",
            "graph/graph.cc",
            "graph/graph_def_builder.cc",
            "graph/node_builder.cc",
            "graph/tensor_id.cc",
            "graph/while_context.h",
            "graph/while_context.cc",
        ],
    // 省略了诸多代码
)
 
// FRAMEWORK_INTERNAL_PRIVATE_HEADERS的内容
FRAMEWORK_INTERNAL_PRIVATE_HEADERS =  [
    "graph/edgeset.h",
    "graph/graph.h",
    "graph/graph_def_builder.h",
    "graph/node_builder.h",
    "graph/tensor_id.h",
] + glob(
    [
        "example/**/*.h",
        "framework/**/*.h",   // 这里就是重点，查看${rootdir}/tensorflow/core/framework/op.h和opkernel.h
        "util/**/*.h",
    ]
) 
 
// 先来看op.h
#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)
#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)
#define REGISTER_OP_UNIQ(ctr, name)                                          \
  static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr    \
      TF_ATTRIBUTE_UNUSED =                                                  \
          ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( \
              name)>(name)
              
// 再来看看opkernel.h
#define REGISTER_KERNEL_BUILDER(kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__)
 
#define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)
 
#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...)        \
  constexpr bool should_register_##ctr##__flag =                      \
      SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__);                        \
  static ::tensorflow::kernel_factory::OpKernelRegistrar              \
      registrar__body__##ctr##__object(                               \
          should_register_##ctr##__flag                               \
              ? ::tensorflow::register_kernel::kernel_builder.Build() \
              : nullptr,                                              \
          #__VA_ARGS__,                                               \
          [](::tensorflow::OpKernelConstruction* context)             \
              -> ::tensorflow::OpKernel* {                            \
            return new __VA_ARGS__(context);                          \
          });

参照上述同样的流程，我们可以发现 libtensorflow.so 中涉及 op 与 kernel 的具体实现，同时也包括 Session 的具体实现。

最后，我们再来讲讲 REGISTER_OP 宏背后的具体原理。我们在上面已经给出了此宏的定义，此处针对它的实现展开谈谈：

// 先来看op.h
#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)
#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)
#define REGISTER_OP_UNIQ(ctr, name)                                          \
  static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr    \
      TF_ATTRIBUTE_UNUSED =                                                  \
          ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( \
              name)>(name)
 
// REGISTER_OP的一般用法如下
REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });
 
// op定义的链式规则是通过OpDefBuilderWrapper类实现的
class OpDefBuilderWrapper<true> {
 public:
  explicit OpDefBuilderWrapper(const char name[]) : builder_(name) {}
 
  OpDefBuilderWrapper<true>& Input(string spec) {
    builder_.Input(std::move(spec));
    return *this;                        // 显而易见，调用Input仍然返回OpDefBuilderWrapper<true>本身
  }
  OpDefBuilderWrapper<true>& Output(string spec) {
    builder_.Output(std::move(spec));
    return *this;
  }
 
  OpDefBuilderWrapper<true>& SetShapeFn(
      Status (*fn)(shape_inference::InferenceContext*)) {
    builder_.SetShapeFn(fn);
    return *this;
  }
  const ::tensorflow::OpDefBuilder& builder() const { return builder_; }
 
 private:
  mutable ::tensorflow::OpDefBuilder builder_;
};
 
// 当通过链式规划构建好op后，再通过OpDefBuilderReceiver完成op的注册
// op.h
struct OpDefBuilderReceiver {
  // To call OpRegistry::Global()->Register(...), used by the
  // REGISTER_OP macro below.
  // Note: These are implicitly converting constructors.
  OpDefBuilderReceiver(
      const OpDefBuilderWrapper<true>& wrapper);  // NOLINT(runtime/explicit)
  constexpr OpDefBuilderReceiver(const OpDefBuilderWrapper<false>&) {
  }  // NOLINT(runtime/explicit)
};
 
// op.cc，然后在OpDefBuilderReceiver构造函数内部完成OpDefBuilderWrapper的全局注册
OpDefBuilderReceiver::OpDefBuilderReceiver(
    const OpDefBuilderWrapper<true>& wrapper) {
  OpRegistry::Global()->Register(
      [wrapper](OpRegistrationData* op_reg_data) -> Status {
        return wrapper.builder().Finalize(op_reg_data);
      });
}

四、总结

本文为大家系统讲解了 TensorFlow 的核心抽象 op 及其 kernel 实现。需要自定义 op 的具体场景，以及 op 的运行框架及若干技术细节。读罢此文，读者应该有如下几点收获：

• TensorFlow 中 op 是编译期概念，kernel 是运行期概念，两者各自的定义与注册方式，以及相应的映射逻辑。

• 掌握 TensorFlow 的高阶玩法：自定义 op。这将使你之前工作的不可能变为可能，由低效转化为高效。

• 掌握 op 与 kernel 注册的宏定义来自何方，以及宏定义背后具体的运行框架。