jax
53b46ec4 - [Mosaic GPU] Add layout inference rule for `MemRefReshapeOp`.

Commit
56 days ago
[Mosaic GPU] Add layout inference rule for `MemRefReshapeOp`. We currently only support inferring tiled layouts for reshape that can be expressed as `dest_shape = CollapseShape(src_shape, ...)` or `dest_shape = ExpandShape(src_shape, ...)` but not a combination of both. We defaults to untiled layout for such cases. This is aligned with LANE semantics where we rejects unfolding dimensions when commuting an `UntilingTransform` with a `ReshapeTransform`. PiperOrigin-RevId: 953269111
Author
Parents
Loading