TPU support

TPU device, ops, topology, and the runtime glue that makes tf.distribute.TPUStrategy work. Lives in tensorflow/core/tpu/ (C++) and tensorflow/python/tpu/ (Python).

Purpose

• Expose Google's Tensor Processing Units as a Device to the TF runtime.
• Provide TPU-specific ops (TPUReplicateMetadata, TPUReplicatedInput, TPUExecute, TPUOrdinalSelector, etc.).
• Drive TPU compilation through XLA (most TPU code lowers to HLO and runs through XLA).
• Implement the topology (mesh) and pod-level coordination.

Directory layout

tensorflow/core/tpu/
├── kernels/                  # TPU-specific op kernels
├── graph_rewrite/            # Graph rewrites that prepare TF graphs for TPU execution
├── ops/                      # TPU op registrations
├── tpu_compilation_cache_*.{h,cc}
├── tpu_executor_*.{h,cc}     # TPU stream executor
├── tpu_node_context.{h,cc}
├── tpu_on_demand_compiler.{h,cc}
├── tpu_topology.{h,cc}
└── ...

tensorflow/python/tpu/
├── tpu.py                    # Public Python entry: tf.tpu.*
├── tpu_strategy_util.py
├── tpu_embedding_v3*         # TPU embedding API
├── topology.py
├── device_assignment.py
├── ops/                      # TPU op wrappers
└── ...

Key abstractions

Type / function	File	Description
tf.distribute.TPUStrategy	tensorflow/python/distribute/tpu_strategy.py	Strategy for running on TPU pods.
tf.tpu.experimental.initialize_tpu_system	tensorflow/python/tpu/tpu.py	Resets and initializes TPU runtime; required at start.
TPUClusterResolver	tensorflow/python/distribute/cluster_resolver/tpu/	Discovers TPU cluster topology (Cloud TPU, GKE).
TPUExecute op	tensorflow/core/tpu/kernels/tpu_execute_op.cc	Runs an XLA-compiled TPU program.
Topology / DeviceAssignment	tensorflow/python/tpu/topology.py, device_assignment.py	Mesh shape and replica→core mapping.
TPU embedding ops	tensorflow/core/tpu/kernels/tpu_embedding_*	High-throughput sparse embedding lookups on TPU.

How a TPU model runs

sequenceDiagram
    participant User
    participant Strategy as TPUStrategy
    participant Resolver as TPUClusterResolver
    participant Compiler as XLA TPU compiler
    participant TPU as TPU device

    User->>Resolver: TPUClusterResolver(cloud_tpu)
    User->>Strategy: tf.distribute.TPUStrategy(resolver)
    User->>User: with strategy.scope(): build model
    User->>Strategy: strategy.run(train_step, ...)
    Strategy->>Compiler: lower train_step subgraph to HLO (via tf2xla)
    Compiler-->>Strategy: TPU executable
    Strategy->>TPU: TPUExecute(executable, inputs)
    TPU-->>Strategy: outputs

The defining feature: almost everything that runs on a TPU is XLA-compiled. The TF→HLO bridge (tensorflow/compiler/tf2xla/) lowers the train_step graph to HLO, and the XLA TPU compiler turns that into a TPU executable. There is no fall-through to interpreted op execution on a TPU.

TPU embedding

TPU embedding is its own subsystem because embedding lookups on a TPU need to fan out across the mesh and use specialised hardware. The tpu_embedding_v3* Python API and the C++ kernels under tensorflow/core/tpu/kernels/ implement large-scale sparse-feature embedding tables that can exceed a single chip's HBM.

SPMD vs replicated

TPU programs run in two main modes:

• Replicated (via TPUStrategy.run): the same XLA program runs on every replica with replica-specific inputs. Cross-replica communication via CrossReplicaSum/CollectivePermute.
• SPMD (via DTensor or XLA SPMD partitioner): a single XLA program is partitioned across devices automatically; the partitioner inserts the right collectives.

Both lower through tensorflow/compiler/tf2xla/ to HLO, then through XLA's TPU backend.

Pod-level orchestration

For TPU pods (multi-host TPU clusters), TF uses the gRPC-based TPUClusterResolver plus the coordination service to bring up workers, agree on topology, and synchronise checkpointing.

Integration points

• XLA bridge — tensorflow/compiler/tf2xla/ produces the HLO for TPU.
• tf.distribute — TPUStrategy uses TPU collectives and topology.
• tf.function — wrapping a step in @tf.function(jit_compile=True) is the explicit way to force TPU compilation; TPUStrategy does this implicitly.
• DTensor — DTensor's TPU mesh support uses the same TPU topology code.

Entry points for modification

• New TPU kernel — tensorflow/core/tpu/kernels/. Most TPU "kernels" are actually graph rewrites or HLO emitters; pure runtime kernels are rare on TPU because XLA owns execution.
• Graph rewrites that prepare a TF graph for TPU compilation — tensorflow/core/tpu/graph_rewrite/.
• Python TPU API — tensorflow/python/tpu/tpu.py.
• TPU embedding — tensorflow/python/tpu/tpu_embedding_v3* and the matching C++ kernels.

Related

• compilers/xla — TPU runs everything through XLA.
• distribution-strategy — TPUStrategy is one of the main strategies.
• distributed-runtime — pod orchestration.