Provide access to local tiles in distributed matrix?

[Benjamin Brock]

Today we discussed adding execution policies to the C++ API, as well as other facilities that might be necessary for a distributed API.

One of the issues that came up is whether---and how---to provide access to a process' local tiles in a distributed matrix.

One option would be to explicitly expose a tile grid, while another would be to expose a range of submatrices, along with the particular coordinates of the submatrix in the global matrix. The second option would support any distributed matrix distributions as long as they had rectangular, non-overlapping tiles. (The tile shapes could be non-uniform.) This could even potentially support non-rectangular tiles if we find a way to represent the coordinates of non-rectangular submatrices.

Whether to expose this, and what exactly we'd expect users to do with this, is still up for discussion.

distributed_matrix<float> m = ...;


auto local_tile = ...;

// Local multiply
grb::multiply(local_tile, ...);

// Distributed multiply
grb::multiply(m, ...);

Tile grid

0 1 0 1
2 3 2 3
0 1 0 1
2 3 2 3

auto g = m.grid_shape();

// (4, 4)
fmt::print("{}\n", g);

for (int i = 0; i < m.grid_shape()[0]; i++) {
  for (int j = 0; j < m.grid_shape()[1]; j++) {
    if (m.tile({i, j}).rank() == my_rank()) {
      algorithm(m.tile({i, j}), ...);
    }
  }
}

if (m.tile({0, 2}).rank() == my_rank()) {
  ...;
}

for (auto&& [submatrix_coordinates, tile] : m.my_local_tiles()) {
  // I know where my submatrix lies based on `submatrix_coordinates`
  algorithm(tile, ...);
}