Runtime, kernels and threading

TensorFlow · Expert · 14 min

Tuning CPU inference, debugging placement and writing custom kernels all require knowing how the runtime executes ops.

Every TensorFlow operation is an op registered in a global registry with a name, typed inputs and attributes (AddV2, MatMul, Conv2D). Each op has one or more kernels: implementations for a device type and dtype. When an op runs, the placer chooses a device that has a kernel (preferring GPUs), and the executor dispatches to that kernel. The full low-level op set is exposed as tf.raw_ops, which is what high-level Python functions ultimately call.

In eager mode, each Python call executes one op immediately. In graph mode (tf.function), the graph is optimised by Grappler, partitioned across devices with send/receive ops inserted between them, and run by an executor that schedules ready ops in parallel.

CPU parallelism uses two thread pools:

  • intra-op: threads used inside a single op (a big matmul);
  • inter-op: how many independent ops run concurrently.

Set them with tf.config.threading.set_intra_op_parallelism_threads(n) and set_inter_op_parallelism_threads(n) before the runtime initialises; tuning them matters for CPU serving with many concurrent requests.

When no op does what you need, you can write a custom op in C++ (REGISTER_OP plus kernels), build it into a shared library and load it with tf.load_op_library. Prefer composing existing ops first: custom ops complicate builds, XLA support and portability.

Example

import tensorflow as tf

# Thread pools must be configured before the runtime starts executing ops.
tf.config.threading.set_intra_op_parallelism_threads(4)   # threads inside one op
tf.config.threading.set_inter_op_parallelism_threads(2)   # ops run concurrently

print(tf.raw_ops.AddV2(x=tf.constant(1.0), y=tf.constant(2.0)).numpy())   # 3.0

tf.debugging.set_log_device_placement(True)      # log each op's chosen device
a = tf.constant([[1.0, 2.0]])
b = tf.matmul(a, a, transpose_b=True)            # logs: Executing op MatMul in device ...
print(b.numpy())                                 # [[5.]]
tf.debugging.set_log_device_placement(False)

print([d.device_type for d in tf.config.list_logical_devices()])   # e.g. ['CPU', 'GPU']

tf.config.run_functions_eagerly(True)            # debugging: tf.functions run op by op
tf.config.run_functions_eagerly(False)

Remember

  • Ops are registered interfaces; kernels implement them per device and dtype
  • Intra-op vs inter-op thread pools control CPU parallelism
  • tf.raw_ops exposes every op; custom C++ ops are a last resort

Watch out for

  • Setting thread-pool sizes after the first op has run, which raises a RuntimeError.

Check yourself

  1. An op exists but has no GPU kernel for the requested dtype. What does TensorFlow typically do with soft placement on?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: www.tensorflow.org/guide/intro_to_graphs

Builds on

Unlocks

Nothing builds on this yet.