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_opsexposes 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.