TensorFlow
Op Registration
Kernel Linking
Machine Learning
Deep Learning

Understand Op Registration and Kernel Linking in TensorFlow

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

TensorFlow operations are split conceptually into two parts: the op definition and the device-specific kernel implementation. Op registration tells TensorFlow what the operation looks like at the graph level, while kernel linking makes the actual executable implementation available for CPU, GPU, or other devices.

Op Definition Versus Kernel Implementation

An op definition describes the public contract of the operation. It tells TensorFlow things such as:

  • the op name
  • input and output tensors
  • attributes such as types or shapes
  • optional shape-inference rules

A kernel implementation is the code that actually runs when the op is executed. You can think of the op definition as the interface and the kernel as one concrete backend implementation.

In custom TensorFlow ops written in C++, the registration of the op often looks like this:

cpp
1#include "tensorflow/core/framework/op.h"
2#include "tensorflow/core/framework/shape_inference.h"
3
4using namespace tensorflow;
5
6REGISTER_OP("AddOne")
7    .Input("x: float")
8    .Output("y: float")
9    .SetShapeFn([](shape_inference::InferenceContext* c) {
10      c->set_output(0, c->input(0));
11      return OkStatus();
12    });

This does not yet tell TensorFlow how to compute the result. It only registers the op contract.

Register the Kernel Separately

The kernel is registered with a device type such as CPU or GPU. That is how TensorFlow knows which implementation to dispatch at runtime.

cpp
1#include "tensorflow/core/framework/op_kernel.h"
2
3using namespace tensorflow;
4
5class AddOneOp : public OpKernel {
6 public:
7  explicit AddOneOp(OpKernelConstruction* context) : OpKernel(context) {}
8
9  void Compute(OpKernelContext* context) override {
10    const Tensor& input = context->input(0);
11    Tensor* output = nullptr;
12    OP_REQUIRES_OK(context, context->allocate_output(0, input.shape(), &output));
13
14    auto in = input.flat<float>();
15    auto out = output->flat<float>();
16    for (int i = 0; i < in.size(); ++i) {
17      out(i) = in(i) + 1.0f;
18    }
19  }
20};
21
22REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_CPU), AddOneOp);

Now TensorFlow knows both the graph-level meaning of AddOne and the CPU implementation it can execute.

What Linking Has To Do With It

Registration macros run when the shared library is loaded. If the compiled object file containing the kernel is never linked into the final TensorFlow extension or binary, the registration code never runs and TensorFlow behaves as if that kernel does not exist.

That is why custom ops can fail in two different ways:

  • the op is unknown because the op definition was not registered
  • the op exists, but no kernel is available for the current device and dtype

In build systems such as Bazel, this means the relevant source files and libraries must actually be part of the produced shared object.

Device Dispatch Is the Real Reason for the Separation

The split between op definition and kernel registration exists because one op may have multiple kernels. The same op name can map to a CPU kernel, a GPU kernel, or multiple dtype-specialized kernels.

For example, one op might have:

  • a CPU kernel for float
  • a CPU kernel for int32
  • a GPU kernel for float

TensorFlow matches the op invocation with the correct registered kernel based on device placement and attribute constraints.

That design is what lets the high-level graph stay stable while execution details vary by backend.

Why This Matters When Debugging Custom Ops

If a custom op fails to load, it helps to ask the questions in the right order:

  • was the op definition library loaded at all
  • was the kernel object linked into that library
  • was a kernel registered for the requested device and dtype
  • does the op name in Python match the registered C++ name exactly

Those questions usually separate build issues from API issues quickly.

A Python load step often looks like this:

python
1import tensorflow as tf
2
3module = tf.load_op_library("./add_one.so")
4result = module.add_one(tf.constant([1.0, 2.0, 3.0]))
5print(result.numpy())

If the shared object loads but execution fails with a missing-kernel message, the registration path is only partially correct.

Common Pitfalls

A common mistake is defining the op with REGISTER_OP and assuming that is enough. It is not. TensorFlow still needs a registered kernel implementation.

Another mistake is registering a kernel only for CPU and then trying to force execution on GPU. The op name exists, but the required kernel does not.

Build configuration is another frequent problem. If the object file containing the registration macro is not linked into the final shared library, the registration code never runs.

Finally, shape functions are often neglected. The op may execute fine, but poor shape inference can make graph building and debugging much harder than necessary.

Summary

  • 'REGISTER_OP defines the graph-level interface of a TensorFlow operation.'
  • 'REGISTER_KERNEL_BUILDER connects that op to an executable implementation on a device.'
  • Linking matters because registration code runs only when the compiled library is actually loaded.
  • One op can have several kernels for different devices and dtypes.
  • Most custom-op debugging starts by separating op-definition problems from missing-kernel problems.

Course illustration
Course illustration

All Rights Reserved.