Dive into TensorFlow系列（2）- 解析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<ShapeHandle>
    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<int32>();

        // 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<int32>();

        // 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注册的宏定义来自何方，以及宏定义背后具体的运行框架。

参考资料

1.《TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems》： https://arxiv.org/abs/1603.04467

2.Graphs and Sessions： https://github.com/tensorflow...

3.Adding a New Op： https://github.com/tensorflow...

4.跨设备通信send/recv： https://github.com/tensorflow...

5.OpKernel definition： https://github.com/tensorflow...

6.tensorflow源码解析之framework-resource： https://www.cnblogs.com/jican...

7.tensorflow源码解析之framework-op： https://www.cnblogs.com/jican...

Dive into TensorFlow系列（2）- 解析TF核心抽象op算子

一、初识op

1.1 op定义

1.2 op分类

1.3 op与kernel

二、自定义op

2.1 自定义op场景

2.2 自定义op流程

2.3 高级话题

三、op工作原理

3.1 op运行框架

3.2 若干技术细节

四、总结

参考资料

京东云开发者

引用和评论

大数据从业者必知必会的Hive SQL调优技巧

ChatGPT

高质量代码的原则

使用 AI 在医疗影像分析中的应用探索

Scrum 四个会议及正确召开方式

详解Diffusion扩散模型：理论、架构与实现

估算工作时间的 3 条规则（在软件开发项目中）