ttnn.broadcast

ttnn.broadcast(input_tensor: ttnn.Tensor, sender_coord: ttnn.MeshCoordinate, *, cluster_axis: int = None, num_links: int = None, memory_config: ttnn.MemoryConfig = the input tensor memory config, topology: ttnn.Topology = ttnn.Topology.Linear, subdevice_id: ttnn.SubDeviceId | None) → ttnn.Tensor

Performs a broadcast operation from a sender device to all other mesh devices across a cluster axis.

Parameters:

  • input_tensor (ttnn.Tensor)Input tensor. The data residing on sender_coord is the data that gets broadcast.

  • sender_coord (ttnn.MeshCoordinate)Coordinate of the sender device in the mesh.

Mesh Tensor Programming Guide : https://github.com/tenstorrent/tt-metal/blob/main/tech_reports/Programming_Mesh_of_Devices/Programming_Mesh_of_Devices_with_TT-NN.md

Keyword Arguments:

  • cluster_axis (int, optional)The mesh axis to broadcast across. Defaults to None, in which case the tensor must have a line (1D) topology.

  • num_links (int, optional)Number of links to use for the broadcast operation. Defaults to None, for which the number of links is determined automatically.

  • memory_config (ttnn.MemoryConfig, optional)Output memory configuration. Defaults to the input tensor memory config.

  • topology (ttnn.Topology, optional)Fabric topology. Valid options are Ring and Linear. Defaults to ttnn.Topology.Linear.

  • subdevice_id (ttnn.SubDeviceId, optional)Subdevice id for worker cores.

Returns:

ttnn.TensorThe output tensor, same shape as the input, holding the sender device’s data on every device along the cluster axis.

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B, FLOAT32

TILE, ROW_MAJOR

Broadcast is a data-movement collective and does not restrict the input dtype; the output preserves the input dtype.

Memory Support:

  • Interleaved: DRAM and L1

  • Sharded: WIDTH_SHARDED, HEIGHT_SHARDED, BLOCK_SHARDED (each may reside in DRAM or L1; the op places no buffer-type restriction)

Example

torch_input = torch.randn([1, 1, 32, 256], dtype=torch.bfloat16)
tt_input = ttnn.from_torch(
    torch_input,
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=mesh_device,
    mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
)

# Copy the data on the sender coordinate to every device along the cluster axis.
output = ttnn.broadcast(tt_input, ttnn.MeshCoordinate(0, 0), cluster_axis=1)
logger.info(output.shape)  # [1, 1, 32, 256]