2024-10-08 11:13:04 +01:00
|
|
|
``jax.experimental.pallas.mosaic_gpu`` module
|
|
|
|
|
=============================================
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.experimental.pallas.mosaic_gpu
|
|
|
|
|
|
|
|
|
|
Classes
|
|
|
|
|
-------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
Barrier
|
|
|
|
|
GPUBlockSpec
|
|
|
|
|
GPUCompilerParams
|
|
|
|
|
GPUMemorySpace
|
2024-10-17 08:07:31 -07:00
|
|
|
Layout
|
2024-10-11 05:10:46 -07:00
|
|
|
SwizzleTransform
|
2024-10-08 11:13:04 +01:00
|
|
|
TilingTransform
|
|
|
|
|
TransposeTransform
|
|
|
|
|
WGMMAAccumulatorRef
|
|
|
|
|
|
|
|
|
|
Functions
|
|
|
|
|
---------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
2024-10-16 04:04:07 -07:00
|
|
|
barrier_arrive
|
|
|
|
|
barrier_wait
|
2024-10-31 02:20:33 -07:00
|
|
|
commit_smem
|
2024-10-08 11:13:04 +01:00
|
|
|
copy_gmem_to_smem
|
|
|
|
|
copy_smem_to_gmem
|
2024-10-28 08:25:08 -07:00
|
|
|
emit_pipeline
|
2024-10-17 08:07:31 -07:00
|
|
|
layout_cast
|
2024-10-16 04:04:07 -07:00
|
|
|
set_max_registers
|
2024-10-08 11:13:04 +01:00
|
|
|
wait_smem_to_gmem
|
|
|
|
|
wgmma
|
|
|
|
|
wgmma_wait
|
|
|
|
|
|
|
|
|
|
Aliases
|
|
|
|
|
-------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
ACC
|
|
|
|
|
GMEM
|
|
|
|
|
SMEM
|
