Kernels and ops

tensorflow/core/kernels/ is the largest single directory in the repo by file count — roughly 10 000 C++ source files implementing TensorFlow ops for CPU, GPU, and assorted accelerators. tensorflow/core/ops/ holds the schemas (the OpDefs).

Purpose

• Define the schema of every TF op (OpDef registration with REGISTER_OP("...")).
• Implement kernels that compute those ops on a specific device + dtype combination.
• Provide shape inference functions used at graph-build time.
• Provide gradient registrations for ops that participate in autodiff.

Directory layout

tensorflow/core/ops/
├── math_ops.cc                # Add, Sub, Mul, MatMul, ...
├── array_ops.cc               # Concat, Slice, Pad, ...
├── nn_ops.cc                  # Conv2D, MaxPool, BatchNorm, ...
├── data_flow_ops.cc           # Queues, conditional, merge, ...
├── string_ops.cc
├── parsing_ops.cc             # Example/SequenceExample parsers
├── ragged_ops.cc
├── sparse_ops.cc
├── resource_variable_ops.cc
├── compat/                    # Op-compat tracking files
└── ... (~80 op-family files)

tensorflow/core/kernels/
├── BUILD                      # 222 KB — every kernel target is here
├── matmul_op.cc/.h
├── matmul_op_real.cc
├── matmul_op_test.cc
├── matmul_op_impl.h
├── conv_ops.cc / conv_grad_ops.cc / conv_ops_3d.cc / conv_ops_gpu.h
├── cwise_op_*.cc              # Componentwise ops (Add, Mul, Sub, Sigmoid...)
├── cwise_op_gpu_*.cu.cc       # GPU CUDA implementations
├── reduction_ops.{cc,h,cu.cc} # Reduce sum/mean/max/...
├── data/                      # tf.data dataset/iterator kernels
├── mkl/                       # Intel MKL-DNN paths
├── sparse/                    # Sparse ops
├── linalg/                    # LinAlg / Eigen-backed ops
├── ragged/                    # Ragged tensor ops
└── ~10 000 more files

The kernel files are roughly grouped by op family. Files ending *_gpu.cu.cc are CUDA implementations compiled by nvcc (or clang -x cuda). Files in mkl/ use Intel MKL-DNN/oneDNN.

How a kernel registers

// 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, complex64, complex128}")
    .SetShapeFn(shape_inference::MatMulShape);

// tensorflow/core/kernels/matmul_op_real.cc
template <typename Device, typename T>
class MatMulOp : public OpKernel {
 public:
  explicit MatMulOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
  }
  void Compute(OpKernelContext* ctx) override {
    const Tensor& a = ctx->input(0);
    const Tensor& b = ctx->input(1);
    // ... shape checks, allocate output, dispatch to Eigen/cuBLAS ...
  }
 private:
  bool transpose_a_, transpose_b_;
};

REGISTER_KERNEL_BUILDER(
    Name("MatMul").Device(DEVICE_CPU).TypeConstraint<float>("T"),
    MatMulOp<CPUDevice, float>);

Key macros (tensorflow/core/framework/op_kernel.h):

Macro	Use
REGISTER_OP("Foo")	Declare an op (schema only; in tensorflow/core/ops/)
REGISTER_KERNEL_BUILDER	Bind a kernel to (op, device, type)
REGISTER_SYSTEM_OP	Internal/no-stable-API op
OP_REQUIRES(ctx, cond, status)	Validate input; set status and return on failure
OP_REQUIRES_OK(ctx, expr)	Run expr; if it returns an error, set & return
TF_RETURN_IF_ERROR(expr)	Same idea outside OpKernel::Compute

Shape inference

Shape inference functions live alongside the REGISTER_OP and run at graph-build time:

.SetShapeFn(shape_inference::MatMulShape);

The shape function takes a shape_inference::InferenceContext* and outputs ShapeHandles. Common helpers live in tensorflow/core/framework/common_shape_fns.h. They allow the Python frontend to surface the symbolic shape (tensor.shape returns a TensorShape) before the op actually runs.

Gradients

Gradient ops are registered separately:

• For Python-level autodiff (most common), gradients live in tensorflow/python/ops/<family>_grad.py and are decorated with @RegisterGradient("Op").
• For C++-level, gradients are in tensorflow/cc/gradients/.
• For MLIR/XLA paths, gradients are reconstructed from primitive ops at compile time.

Type dispatch

Kernels are generally templated on a Device (CPUDevice, GPUDevice) and a numeric type (float, double, int32, bfloat16, Eigen::half, etc.). Type combinations explode quickly — the project uses macros (e.g., TF_CALL_REAL_NUMBER_TYPES, TF_CALL_GPU_ALL_TYPES) defined in tensorflow/core/framework/register_types.h to instantiate the templates without writing each by hand.

CUDA / GPU kernels

• GPU implementations live in *_gpu.cu.cc files compiled with NVCC.
• They typically launch CUDA kernels through Eigen's GpuDevice or directly via cudaLaunchKernel.
• BLAS-like ops (MatMul, Conv2D) call into cuBLAS / cuDNN through tensorflow/core/util/cuda_solvers* and tensorflow/core/kernels/conv_ops_gpu.h.
• tensorflow/stream_executor/ (now part of xla/stream_executor) is the device-abstraction layer kernels use to launch on streams.

Op vocabulary at a glance

A handful of frequently-seen op names and their kernel files:

Op	Kernel file
Add, Mul	tensorflow/core/kernels/cwise_op_add.cc, cwise_op_mul.cc
MatMul	tensorflow/core/kernels/matmul_op*.{cc,h}
Conv2D	tensorflow/core/kernels/conv_ops*.{cc,h,cu.cc}
BiasAdd	tensorflow/core/kernels/bias_op.cc
Relu/Sigmoid	tensorflow/core/kernels/relu_op.cc, cwise_op_sigmoid.cc
Softmax	tensorflow/core/kernels/softmax_op*.{cc,cu.cc}
BatchMatMul	tensorflow/core/kernels/batch_matmul_op*.{cc,h}
Slice	tensorflow/core/kernels/slice_op*.{cc,cu.cc}
MaxPool	tensorflow/core/kernels/maxpooling_op*.{cc,cu.cc} (the largest single kernel file at ~78 KB)
Dataset ops	tensorflow/core/kernels/data/...

Integration points

• Auto-generated wrappers in tensorflow/python/ops/gen_*.py and tensorflow/cc/ops/*.h. They are produced from the OpDefs.
• Compiler bridges (tensorflow/compiler/tf2xla/kernels/) re-implement many ops on top of XLA HLO. A TF op may have both a regular kernel and an XLA implementation.
• TFLite kernels are entirely separate — tensorflow/lite/kernels/ re-implements ~150 ops in a smaller library targeted at mobile.

Entry points for modification

• New op + kernel: register OpDef in tensorflow/core/ops/<family>_ops.cc, write kernel(s) in tensorflow/core/kernels/, add BUILD entries, write tests in *_test.cc.
• New dtype support: extend TF_CALL_* macro lists or add a new REGISTER_KERNEL_BUILDER line.
• Speeding up an existing op: drop in an MKL or oneDNN path under tensorflow/core/kernels/mkl/, or a fused GPU implementation in *_gpu.cu.cc.

Related

• core-runtime — the framework types kernels are written against.
• grappler — runs over graphs of these ops to fuse and rewrite.
• compilers/xla — the XLA path that lowers many of these ops to HLO.
• apps/tensorflow-lite — note that TFLite kernels are independent.