Phase 3, Section 2: Mesh and Sharding Types

Phase 3: Sharding (the distribution axis). Section 2 of 4: Mesh and Sharding Types.

1. The Device Mesh Kind

To support sharding at the type level, the compiler must understand the physical or logical topology of the available hardware. This is modeled as a Mesh.

1.1 Mesh Declaration

In Phase 3, we restrict meshes to a 1D logical array of devices for simplicity. A mesh is declared at the program level and assigned an identifier.

let mesh = @mesh([4]) in

This binds mesh to a 1D array of 4 devices.

2. Sharding Annotations

Tensors are given a sharding refinement, which dictates how their data is distributed across a specific mesh.

2.1 Replicated

A tensor is @replicated if every device in the mesh holds an identical, full copy of the tensor's data.

Syntax:

let w = tensor<F32, [1024, 1024], @replicated(mesh)> { ... }

2.2 Sharded

A tensor is @sharded if it is partitioned along one of its dimensions across the devices in the mesh.

Syntax:

let x = tensor<F32, [1024, 1024], @sharded(mesh, axis=0)> { ... }

Here, axis=0 indicates that the rows of x are evenly divided among the 4 devices in mesh (each device gets a [256, 1024] slice).

3. Type Calculus Rules

The typed IR is extended such that a tensor_type is parameterized by its sharding refinement: Tensor[dtype, shape, placement, sharding]

3.1 Well-Formedness

A sharding refinement @sharded(mesh, axis=d) is well-formed only if:

1. d is a valid dimension index (0 <= d < rank).
2. The size of dimension d is divisible by the size of the mesh (for Phase 3, we assume perfectly divisible static shapes).

3.2 Default Placements

If a tensor originates on the @host, it has no sharding refinement. When a tensor is moved to a @device or mesh, it must receive a sharding refinement, which can be explicitly annotated or inferred by the GSPMD propagation pass.