mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 00:36:04 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/pallas
History
Benjamin Chetioui c7199fe8a5 [Pallas/Mosaic GPU] Enable progressive lowering for integer addition. 
The helpers `_fragmented_array_to_ir` and `_fragmented_array_from_ir` in
`dialect_lowering.py` have been modified, such that a fragmented array's
signedness no longer appears in its IR representation.

This is because signedness is a reflection of how we make use of the value,
and not an inherent property of it. The appropriate signedness value to use
to reload a fragmented array from IR must be provided by the caller.

PiperOrigin-RevId: 726030853
2025-02-12 06:29:25 -08:00
..
mosaic
[JAX/Pallas] Add has_side_effect parameter to CompilerParams to stop CSE of operations.
2025-02-11 13:33:01 -08:00
mosaic_gpu
[Pallas/Mosaic GPU] Enable progressive lowering for integer addition.
2025-02-12 06:29:25 -08:00
triton
[pallas:triton] Lift dot_general restriction on minimal tile size for a.
2025-02-11 06:27:16 -08:00
__init__.py
[Pallas] Upstream pallas to JAX
2023-08-01 16:43:13 -07:00
BUILD
[Pallas] Refactor Pallas HLO interpret mode to a standalone file.
2025-01-27 17:52:27 -08:00
core.py
Merge pull request #26423 from gnecula:debug_info_jaxpr_7
2025-02-10 12:58:26 -08:00
cost_estimate.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
helpers.py
[Pallas TPU] Add helpers for writing collectives
2025-02-04 15:39:10 -08:00
hlo_interpreter.py
[better_errors] Continue adding debug info to Jaxprs (step 3)
2025-02-06 16:26:49 +02:00
pallas_call.py
[Pallas/Mosaic GPU] Add initial support for warpgroup semantics in lowering.
2025-02-12 01:47:49 -08:00
primitives.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
utils.py
[Pallas TPU] Add helpers for writing collectives
2025-02-04 15:39:10 -08:00