[mosaic] Indexing into 3rd minor-most+ dimensionwith sub 32-bit dtypes

[oliverdutton-iso]

Could indexing into 3rd minor-most+ dimensions with sub 32-bit dtypes be supported?

[I know indexing on refs are supported, but I was hoping to use jax.vjp on something directly inside a mosaic kernel but this blocks me]

MosaicError: INTERNAL: Mosaic failed to compile TPU kernel: Only 32-bit types supported

The MLIR operation involved:
  %69 = "vector.extract"(%65) <{static_position = array<i64: 0>}> : (vector<7x128x128xbf16>) -> vector<128x128xbf16>

[justinjfu]

I've forwarded this request to the internal Mosaic team.

It seems like this operation should be easy to support since it just involves selecting the correct tiles out of the source array.

I'm not sure if you have control over whatever is generating this op, but as a workaround it's also possible to store to a scratch ref in VMEM and immediately load out the slice you want. The compiler should elide the memory transfer.

[oliverdutton-iso]

Ooo, I didn't know a load/store from scratch would get elided. I thought any write to Refs in Mosaic TPU was treated as if it was a barrier before any ops defined later (even if they don't depend on the ref written to)
e.g. for

def kernel(a_ref, b_ref, c_ref):
   a_ref[...] *= 5
   b_ref[...]  = c_ref[...]

can't be rearranged of the c->b before the a*=5
and assumed same rules for scratch refs?

[apaszke]

This should be fixed now!