jax
4864c71c - [Mosaic GPU][NFC] Add `SMEMTiling` layout.

Commit
213 days ago
[Mosaic GPU][NFC] Add `SMEMTiling` layout. This is sufficient to express all SMEM transforms that we infer, because: - We support at most one tiling transformation. - If we do have a tiling transformation, we will automatically add a swizzle with the max possible swizzle, so there's no need to have swizzle transforms as part of the equation system. - We don't support any other transforms (e.g. transpose) at this level. PiperOrigin-RevId: 810884001
Parents
Loading