Sharding into multiple IREE logical devices or queues with peered memory

This is a proposal on how to approach distribution of an MLIR program across multiple homogeneous physical devices that have peered memory. It is intended to improve memory usage under NUMA. Devices/Queues have direct peer-to-peer memory accesses during kernel execution. We want to explicitly shard tensor data and operations and assign them to devices/queues.

Tiling and distribution of work on multiple compute units treated as a single device and queue differs for this problem in that there each work item does not matter to what compute unit it is assigned. The compiler does not have to explicitly resolve the assignment. It is left up to the runtime. Also memory allocations does not have to distinguish where they are performed.

Here we use a program with Mesh dialect sharding annotations as an input. The original intent of the Mesh dialect is to utilize MPI-like collectives instead of placement operations like extracting and inserting tensor slices. We would like to expand its scope to support the latter case also. The algorithm that produces the sharding annotation is out-of-scope here and it is assumed that the program is fully annotated.

Below is an example of an annotated program that does matrix multiplication. It is sharded on a 1D process mesh of size 2.

```
2x3 3x2 2x2
+-----------+ +------------+ +-----------+
| process 0 | | fully | | process 0 |
+-----------+ X | replicated | = +-----------+
| process 1 | | | | process 1 |
+-----------+ +------------+ +-----------+
```

```mlir
mesh.mesh @mesh(shape = 2, peered_memory_axes = [0])

func.func @main(
%arg0 : tensor<2x3xf32>,
%arg1 : tensor<3x2xf32>,
%out_dps: tensor<2x2xf32>
) -> tensor<2x2xf32> {
%arg0_sharded = mesh.shard %arg0 to <@mesh, [[0]]> : tensor<2x3xf32>
%arg1_sharded = mesh.shard %arg1 to <@mesh, []> : tensor<3x2xf32>
%out_dps_sharded = mesh.shard %arg1 to <@mesh, [0]> : tensor<2x2xf32>

%res = linalg.matmul ins(%arg0_sharded, %arg1_sharded :
tensor<2x3xf32>, tensor<3x2xf32>)
outs(%out_dps_sharded : tensor<2x2xf32>) -> tensor<2x2xf32>

%res_sharded = mesh.shard %res to <@mesh, [[0]]> : tensor<2x2xf32>

return %res_sharded : tensor<2x2xf32>
}
```

Note that the mesh shape may be dynamic.

For example
```mlir
mesh.mesh @mesh(shape = ?x?x?, peered_memory_axes = [1, 2])
```
Here we have a 3D mesh with all dimensions being dynamic. It has peered memory across mesh axes 1 and 2. The mesh is partitioned into disjoint subsets where there is uniform memory only within a subset. Processes `(h, i, j)` and `(k, l, m)` have peered memory for `h = k` and distinct memory for `h != k`. In the immediate future we will be concerned with 1D meshes with peered memory.

-------------------------------------

As a first step we transform the program from annotations using `mesh.shard` operations to a distributed tensor form that utilizes the tensor encoding attribute.

```mlir
// Shard tensor axis 0 on mesh axis 0.
#shard_0_on_0 = #mesh.shard<@mesh, [[0]]>
// The tensor is replicated on all shards.
#fully_replicate = #mesh.shard<@mesh, []>

mesh.mesh @mesh(shape = 2, peered_memory_axes = [0])

func.func @main(
%arg0 : tensor<2x3xf32, #shard_0_on_0>,
%arg1 : tensor<3x2xf32, #fully_replicate>,
%out_dps: tensor<2x2xf32, #shard_0_on_0>
) -> tensor<2x2xf32, #shard_0_on_0> {
%res = linalg.matmul ins(%arg0, %arg1 :
tensor<2x3xf32, #shard_0_on_0>,
tensor<3x2xf32, #fully_replicate>)
outs(%out_dps : tensor<2x2xf32, #shard_0_on_0>)
-> tensor<2x2xf32, #shard_0_on_0>

return %res : tensor<2x2xf32, #shard_0_on_0>
}
```

The usage of the encoding attribute may clash with other uses like sparsity encoding. We access the mesh sharding attribute through a type interface that hides the encoding attribute detail to make future changes easier.

In this step we would insert resharding if required by the annotations. In general under peered memory when access uniformity is high we would like to avoid resharding and have the subsequent operation be able to consume the original sharding.

-------------------------------------

After converting to distributed tensors we shard the operations, where we describe what each process is doing. One approach is to utilize a `mesh.for_all` operation. All processes share the same code. It is similar to the `flow.dispatch.workgroups`.

```mlir
#shard_0_on_0 = #mesh.shard<@mesh, [[0]]>
#fully_replicate = #mesh.shard<@mesh, [[]]>

mesh.mesh @mesh(shape = 2, peered_memory_axes = [0])

func.func @main(
%arg0 : tensor<2x3xf32, #shard_0_on_0>,
%arg1 : tensor<3x2xf32, #fully_replicate>,
%out_dps: tensor<2x2xf32, #shard_0_on_0>
) -> tensor<2x2xf32, #shard_0_on_0> {

%res = mesh.for_all on @mesh ->
tensor<2x3xf32, #shard_0_on_0> {
%c2 = arith.constant 2 : index
%proc_idx = mesh.process_multi_index on @mesh : index
%slice_offset = arith.muli %proc_idx, %c2 : index
%arg0_slice = tensor.extract_slice %arg0[%slice_offset, 0][1, 3][1, 1] :
tensor<2x3xf32, #shard_0_on_0> to tensor<1x3xf32>
%out_dps_slice = tensor.extract_slice %out_dps[%slice_offset, 0][1, 2][1, 1] :
tensor<2x2xf32, #shard_0_on_0> to tensor<1x2xf32>

%res_slice = linalg.matmul ins(%arg0_slice, %arg1 :
tensor<1x3xf32>,
tensor<3x2xf32, #fully_replicate>)
outs(%out_dps_slice : tensor<2x2xf32, #shard_0_on_0>)
-> tensor<1x2xf32>
// We return the shard for this process.
return %res_slice : tensor<1x2xf32>
}

return %res : tensor<2x2xf32, #shard_0_on_0>
}
```

-------------------------------------

Instead of a new `mesh.for_all` operation we can introduce `mesh.for_one`. It describes what each process is doing separately. With this approach we combine it with `scf.forall` to describe the operations for the entire mesh.

```mlir
#shard_0_on_0 = #mesh.shard<@mesh, [[0]]>
#fully_replicate = #mesh.shard<@mesh, [[]]>

mesh.mesh @mesh(shape = 2, peered_memory_axes = [0])

func.func @main(
%arg0 : tensor<2x3xf32, #shard_0_on_0>,
%arg1 : tensor<3x2xf32, #fully_replicate>,
%out_dps: tensor<2x2xf32, #shard_0_on_0>
) -> tensor<2x2xf32, #shard_0_on_0> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%mesh_size = mesh.mesh_shape @mesh : index
%res = scf.forall (%proc_idx) in (%mesh_size)
shared_outs(%out = %out_dps) -> (tensor<2x2xf32, #shard_0_on_0>) {
scf.forall.in_parallel {
%res = mesh.for_one %proc_idx on @mesh {
%c2 = arith.constant 2 : index
%slice_offset = arith.muli %proc_idx, %c2 : index
%arg0_slice = tensor.extract_slice %arg0[%slice_offset, 0][1, 3][1, 1] :
tensor<2x3xf32, #shard_0_on_0> to tensor<1x3xf32>
%out_slice = tensor.extract_slice %out[%slice_offset, 0][1, 2][1, 1] :
tensor<2x2xf32, #shard_0_on_0> to tensor<1x2xf32>

%res_slice = linalg.matmul ins(%arg0_slice, %arg1 :
tensor<1x3xf32>,
tensor<3x2xf32, #fully_replicate>)
outs(%out_slice : tensor<2x2xf32, #shard_0_on_0>)
-> tensor<1x2xf32>

tensor.parallel_insert_slice %res_slice into
%out[%slice_offset, 0][1, 2][1, 1] :
tensor<1x2xf32> into tensor<2x2xf32, #shard_0_on_0>
}
}
}

return %res : tensor<2x2xf32, #shard_0_on_0>
}
```

-------------------------------------

After all operations have been sharded we need convert each distributed tensor into a list of tensors.

In the case of a dynamic mesh we need to use a dynamic list. We can utilize `tensor`. This dynamic nature will require propagation all the way to the runtime. The VM would need to handle dynamic lists of buffers.

In the case of a static mesh we can use the values that correspond to shards directly. In this case we have to also unroll all loops that run operations on the mesh. Bellow is an example of doing that.

```mlir
mesh.mesh @mesh(shape = 2, peered_memory_axes = [0])

func.func @main(
%arg0_shard0 : tensor<1x3xf32>,
%arg0_shard1 : tensor<1x3xf32>,
%arg1_shard0 : tensor<3x2xf32>,
%arg1_shard1 : tensor<3x2xf32>,
%out_dps_shard0: tensor<1x2xf32>,
%out_dps_shard1: tensor<1x2xf32>,
) -> (tensor<1x2xf32>, tensor<1x2xf32>) {
// Shard 0.
%res_shard0 = mesh.for_one 0 on @mesh {
%res = linalg.matmul ins(%arg0_shard0, %arg1_shard0 :
tensor<1x3xf32>, tensor<3x2xf32>)
outs(%out_dps_shard0 : tensor<1x2xf32>)
-> tensor<1x2xf32>
return %res : tensor<1x2xf32>
}

// Shard 1.
%res_shard1 = mesh.for_one 1 on @mesh {
%res = linalg.matmul ins(%arg0_shard1, %arg1_shard1 :
tensor<1x3xf32>, tensor<3x2xf32>)
outs(%out_dps_shard2 : tensor<1x2xf32>)
-> tensor<1x2xf32>
return %res : tensor<1x2xf32>
}

return %res_shard0, %res_shard0 : tensor<1x2xf32>, tensor<1x2xf32>
}
```

This example is manicured for illustration. In real-world programs a mesh process may have to ingest operands that require extracting slices that fall on the boundary of tensor shards. If we can reliably fuse `mesh.for_one` operations to avoid copying, operand resharding can be handled as a separate concern.

At this point we are ready to convert the remaining `mesh.for_one` ops into IREE specific descriptions of device/queue affinity.

-------------------------------------

We could shard after dispatch region formation, but before transforming into workgroups. It is expected that the user may have provided some sharding annotations throughout the program. Then we form the dispatch regions, complete the sharding annotations and lastly partition into Mesh processes. The user-provided sharding annotations may hinder fusion of the dispatch regions. Can we place sharding annotations inside dispatch regions and what does that even mean? This approach also requires that sharding propagation handles dispatch regions. This drives the requirement for solving the more general problem of propagating through a composite operation.

The other option is to shard before forming the dispatch regions. In this scenario during dispatch region formation we will have to handle a lot of tensor placement (insert/extract) operations due to resharding. It would also have to handle device/queue placement.

We could expand the scope of `flow.dispatch.workgroups` to handle additional dimensions that are mapped to devices/queues.

-------------------------------------

We need to provide the user functionality to shard/unshard function arguments/results. This should be coupled with function variants that reserve the original signature and handle all sharding inside.

```mlir
func.func @main_shard_arg0(%arg0 : tensor<2x3xf32>)
-> tensor<2x3xf32, #shard_0_on_0> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%mesh_size = mesh.mesh_shape @mesh : index
%out_dps = tensor.empty : tensor<2x3xf32, #shard_0_on_0>
%res = scf.forall (%proc_idx) in (%mesh_size)
shared_outs(%out = %out_dps) -> (tensor<2x3xf32, #shard_0_on_0>) {
scf.forall.in_parallel {
%res = mesh.for_one %proc_idx on @mesh {
%c2 = arith.constant 2 : index
%slice_offset = arith.muli %proc_idx, %c2 : index
%slice = tensor.extract_slice %arg0[%slice_offset, 0][1, 3][1, 1] :
tensor<2x3xf32> to tensor<1x3xf32>

tensor.parallel_insert_slice %slice into
%out[%slice_offset, 0][1, 3][1, 1] :
tensor<1x3xf32> into tensor<2x3xf32, #shard_0_on_0>
}
}
}

return %res : tensor<2x3xf32, #shard_0_on_0>
}

func.func @main(%arg0 : tensor<2x3xf32>, ...) ... {
%arg0_sharded = func.call @main_shard_arg0(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32, #shard_0_on_0>
...
%res = func.call @main_shrded(%arg0_sharded, ...) ...
// unshard %res and return
}
```

The resulting IR can be handled as the rest of the program to convert each distributed tensor into a list of tensors.

How to stream the tensor shards from the host directly into the destination slices without having to allocate the original unsharded tensor?

----------------------

There is a potential optimization for some operations that have high computational complexity like matrix multiplication. When tiled, in order to avoid read multiple times the same tile from a slow memory region in the kernel, each such tile can be copied closer to the compute unit. We don't want to do that before starting computation, but to interleave it into the kernel.