Adam Paszke 2db03ba54b [Pallas:MGPU] Add support for grid dims in GPUMesh ...

Of course no communication can happen across grid dimensions (unlike over the WG dim),
but we need to be able to launch multiple blocks somehow.

PiperOrigin-RevId: 688488660

2024-10-22 04:10:46 -07:00

..

mosaic

[Pallas TPU] Use new MLIR op names

2024-10-18 16:14:27 -07:00

mosaic_gpu

[Pallas:MGPU] Add support for grid dims in GPUMesh

2024-10-22 04:10:46 -07:00

triton

Reverts ebb75db8a523150c48376d15391f84380a2bb110

2024-10-22 03:29:32 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

core.py

[Pallas:MGPU] Allow allocating transformed refs in run_scoped

2024-10-22 01:38:46 -07:00

pallas_call.py

[Pallas:MGPU] Allow allocating transformed refs in run_scoped

2024-10-22 01:38:46 -07:00

primitives.py

[Pallas:MGPU] Allow allocating transformed refs in run_scoped

2024-10-22 01:38:46 -07:00

utils.py

[Pallas TPU] Add lowering for 64 bit

2024-09-19 16:42:45 +01:00