Phase 3, Section 4: Collective Coercions

Phase 3: Sharding (the distribution axis). Section 4 of 4: Collective Coercions.

1. Type-Changing Coercions

When the GSPMD propagation pass encounters a mismatch between the derived sharding of a tensor and the required sharding for a downstream operation, the type-checker automatically inserts a type-changing coercion.

At the type level, these are identity functions that change the refinement type. At the operational level, they lower to collective communication primitives across the device mesh.

2. Collective Mappings

2.1 All-Gather

Transition: @sharded(axis=d) -> @replicated Description: Every device broadcasts its slice of the tensor to all other devices, concatenating them along axis=d. The result is a full copy of the tensor on every device.

2.2 Reduce-Scatter

Transition: @partial -> @sharded(axis=d) Description: Devices hold partial sums of a tensor. They reduce (e.g., sum) the partials across the mesh and scatter the resulting complete chunks such that each device only holds a slice along axis=d.

2.3 All-Reduce

Transition: @partial -> @replicated Description: Devices hold partial sums. They reduce the partials and broadcast the complete result to all devices. (Logically equivalent to a Reduce-Scatter followed by an All-Gather).

2.4 All-To-All (Resharding)

Transition: @sharded(axis=a) -> @sharded(axis=b) Description: The tensor is sharded along dimension a, but the downstream operation requires it to be sharded along dimension b. Devices exchange chunks to transpose the sharding axis.

3. Lowering to MLIR

The backend emitter (emitter.ml) is responsible for translating these typed coercions into concrete MLIR.

Since standard MLIR dialects do not have built-in NCCL primitives by default, Kina will emit custom calls or stub functions for the collectives:

// Example lowering of All-Gather
%replicated = func.call @kina_collective_all_gather(%sharded, %axis) : (tensor<...>, i32) -> tensor<...>

In the Simulator Runtime, these kina_collective_* calls will be intercepted and executed via shared-memory thread barriers to mock network communication.