jax
b0a7e8af - [Mosaic GPU] Implement canonicalization for `TiledLayout`s.

Commit
254 days ago
[Mosaic GPU] Implement canonicalization for `TiledLayout`s. Layout canonicalization gets rid of unnecessary tiles and dimensions in `Tiling`s---such that (1) any tiling operation partitions at least one dimension into more than `1` tile, and (2) the leading dimensions of each tile are not `1` (if canonicalizing a tile in this way leads to an empty tile, then the tile is given shape `(1,)`---which is still a meaningful (final) tile). Canonicalization allows simplifying the definition of a few of our predefined layouts---`WGMMA`, `WGMMA_ROW`, `TCGEN05`, `TCGEN05_ROW`, and `TCGEN05_COL`. PiperOrigin-RevId: 772933245
Author
Committer
Parents
Loading