Benjamin Chetioui 1e36cbe597 [Mosaic GPU] Raise a NotImplementedError if swizzle=16. ...

Unswizzled MMAs don't lower correctly, and are not currently intended to be
supported.

PiperOrigin-RevId: 737981373

2025-03-18 06:29:13 -07:00

..

examples

[Mosaic GPU] Add support for .cta_group::2 MMA with n=512 on Blackwell

2025-03-11 09:53:20 -07:00

__init__.py

[Mosaic GPU] Introduce an initial transform inference pass.

2025-03-13 23:22:59 -07:00

core.py

[Mosaic GPU] Add support for .cta_group::2 MMA with n=512 on Blackwell

2025-03-11 09:53:20 -07:00

dialect_lowering.py

[pallas:mosaic_gpu] jnp.reduce_sum now works for >1D arrays

2025-03-17 05:32:07 -07:00

fragmented_array.py

[Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts

2025-03-18 05:38:49 -07:00

inference_utils.py

[Mosaic GPU][NFC] Simplify implementation for in_{layout,transforms}_for_operand utils.

2025-03-14 03:52:10 -07:00

launch_context.py

[pallas:mosaic_gpu] copy_smem_to_gmem now allows skipping cp.async.commit_group

2025-03-07 07:43:54 -08:00

layout_inference.py

[pallas:mosaic_gpu] jnp.reduce_sum now works for >1D arrays

2025-03-17 05:32:07 -07:00

layouts.py

[Mosaic GPU] Replace WGMMAFragLayout with TiledLayout in the mlir dialect and use it in layout inference.

2025-03-11 13:50:42 -07:00

mma_utils.py

[Mosaic GPU] Remove support for large tiles on Blackwell

2025-03-05 10:34:53 -08:00

profiler.py

[mosaic_gpu] Warmup the kernel when doing CUPTI based profiling

2025-02-23 10:20:37 -08:00

tcgen05.py

[Mosaic GPU] Raise a NotImplementedError if swizzle=16.

2025-03-18 06:29:13 -07:00

transform_inference.py

[Mosaic GPU] Add transform inference rule for mgpu.slice_smem.

2025-03-18 04:53:54 -07:00

utils.py

[Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts

2025-03-18 05:38:49 -07:00

wgmma.py

[Mosaic GPU] Raise a NotImplementedError if swizzle=16.

2025-03-18 06:29:13 -07:00