commit d4b6d74202d7ed6014bcba6d4fe97e7d15ae519e
author Quinn Dawkins <quinn.dawkins@gmail.com> Thu Feb 22 18:02:35 2024 -0500
committer GitHub <noreply@github.com> Thu Feb 22 23:02:35 2024 +0000
tree 769798a661212ff51cd0841269c9bc38c48d6d49
parent 0353d1234f17b4aacaf2324a59240e068c7189d5

[VectorExt] Add support for projecting nested layouts (#16528)

This adds two new fields to the layout to track which ids in the basis
are correspond to the subgroups_per_workgroup and threads_per_outer. The
reason these masks are required is because when performing a
rank-reducing projection, we lose necessary degrees of freedom
(dimensions) that can indicate how the data in the underlying vector is
replicated across threads. The simplest example is reducing a 2x2 vector
with a grid of 2x2 threads. In this case there are two ways the reduced
data can be replicated across threads.

```
vector<2> = s0, s1, s0, s1
vector<2> = s0, s0, s1, s1
```

however we only have one valid layout based on the constraints on the
thread basis/thread counts. Assuming all other dimensions are 1, it is
required that thread_count == vector size == 2. Similarly, the total
number of threads in the basis must be equal to the flat total number of
threads (in this case 4). Since there is only one valid layout, there is
no way to differentiate between these two distributed cases.

The active_id masks fix this by further decoupling the thread basis
(i.e. how to delinearize the flat thread id to a set of ids used by the
layout) from the actual vector shape. Handling projections then is
simply a matter of masking off the ids of the basis according to the
projected dims.