Dimitar (Mitko) Asenov d2bf034c47 [Mosaic GPU] Test the wgmma_op lowering when a is in registers. ...

I had to add support for wgmma layout in vector_load. Not sure if this is useful outside the test.

PiperOrigin-RevId: 735384104

2025-03-10 08:25:43 -07:00

..

examples

[Mosaic GPU] Remove support for large tiles on Blackwell

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

__init__.py

[Mosaic GPU] Remove old jaxlib version guards.

2025-02-12 08:49:40 -08:00

core.py

[Mosaic GPU] Define TMEMLayout without referring to the PTX guide

2025-02-14 02:06:17 -08:00

dialect_lowering.py

[Mosaic GPU] Test the wgmma_op lowering when a is in registers.

2025-03-10 08:25:43 -07:00

fragmented_array.py

[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics.

2025-03-10 02:14:39 -07:00

inference_utils.py

[Mosaic GPU][NFC] Move some functions to a new file called inference_utils.py.

2025-02-28 03:02:59 -08: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] Added support for math functions to the WG lowering

2025-03-10 05:08:19 -07:00

layouts.py

[Mosaic GPU][NFC] Move some functions to a new file called inference_utils.py.

2025-02-28 03:02:59 -08: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] Remove support for large tiles on Blackwell

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

utils.py

[Mosaic GPU] Add a new layout to help with transposing WGMMA results

2025-03-07 07:42:01 -08:00

wgmma.py

[Mosaic GPU] Add support for small tiles for (WG)MMA LHS

2025-03-05 09:01:20 -08:00