mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,443 Commits 504 Branches 263 Tags
Commit Graph

29 Commits

Author SHA1 Message Date
Tzu-Wei Sung
5179642eb5 [Mosaic] Rename dep name. 
PiperOrigin-RevId: 732985217
 2025-03-03 11:01:25 -08:00
William S. Moses
8262987a1c Fix build dependencies 
PiperOrigin-RevId: 731330542
 2025-02-26 08:38:31 -08:00
jax authors
b510127a13 Internal compatibility change 
PiperOrigin-RevId: 729428478
 2025-02-21 01:21:56 -08:00
jax authors
eaceac3bf9 [Pallas] Reductions with replicated axes. 
PiperOrigin-RevId: 727292293
 2025-02-15 07:41:16 -08:00
Tzu-Wei Sung
d4758b6d5e [Mosaic][NFC] Factor out xla-array related utils in a separate file. 
Also added tests.

PiperOrigin-RevId: 721424194
 2025-01-30 09:49:41 -08:00
Sergei Lebedev
4221f109d1 [mosaic] Extracted serialization pass traversal logic into a reusable function 
I will use it to implement Mosaic GPU serialization pass in a follow up.

PiperOrigin-RevId: 716156650
 2025-01-16 02:58:06 -08:00
Tzu-Wei Sung
bf94389b08 [Mosaic] Use tpu::CreateMask for getX32VmaskByPaddingEnd. 
It was cmp + iota before.

PiperOrigin-RevId: 713240888
 2025-01-08 04:18:53 -08:00
Tzu-Wei Sung
57b21541a2 [Mosaic] NFC: Pull out vreg related functions to util. 
These functions are related to vreg manipulation and are used in different rules.

PiperOrigin-RevId: 711484002
 2025-01-02 11:50:19 -08:00
Sergei Lebedev
8987867faa [mosaic_gpu] Include Mosaic GPU dialect fiels into jaxlib 2024-12-23 13:46:25 +00:00
Sergei Lebedev
a14e6968bf [mosaic] Migrated the serialization pass from codegen to pass_boilerplate.h 
This prepares teh generalization of the serialization pass to handle both
Mosaic TPU and GPU.

PiperOrigin-RevId: 705628923
 2024-12-12 14:19:36 -08:00
Naums Mogers
0d36b0b433 [Mosaic] Add target core type parameter to tpu.sem_signal 
Adds the optional core type parameter to `tpu.sem_signal` for cross-core signalling.
If the target core type is not provided, the target core type is assumed to be that of the core issuing the signal.
The issuing core type is determined based on the core type annotation of the parent function; if the annotation is not provided, the issuing core type is assumed to be TensorCore.

PiperOrigin-RevId: 698129842
 2024-11-19 13:40:13 -08:00
Praveen Batra
8296f6e0ba [Mosaic] Add extension files for infer/apply vector layout. 
PiperOrigin-RevId: 691868278
 2024-10-31 11:08:37 -07:00
Peter Hawkins
6d1f51e63d Clean up BUILD files. 
PiperOrigin-RevId: 667604964
 2024-08-26 09:11:17 -07:00
Jevin Jiang
aa16485457 [XLA:Mosaic] Support memref shapecast. 
This cl supports memref shapecast:
1. if tile is (1, 128), we support shapecast on any dim.
2. if shapecast on sublane dim, we only support tile aligned shape.
3. if shapecast on non-tiling dim, we support any shapecast.
4. all other cases would be considered as invalid memref shapecast.

PiperOrigin-RevId: 651924552
 2024-07-12 17:05:03 -07:00
Kyle Lucke
84d748f43c Stop using xla/statusor.h now that it just contains an alias for absl::Status. 
In some situations, this meant also changing unrelated files to directly include tsl/platform/statusor.h to get the definitions for TF_ASSIGN_OR_RETURN, etc., where they were getting transitively included for free.

PiperOrigin-RevId: 645169743
 2024-06-20 15:09:40 -07:00
jax authors
92d892b425 Adds rewrite patterns for arith and math operations with bf16 operands/results that are not supported by the underlying hardware. 
PiperOrigin-RevId: 635865752
 2024-05-21 11:08:19 -07:00
jax authors
8c2425e571 Adds rewrite patterns to LinalgVectorizationPass to eliminate transfer_read and transfer_write ops. 
PiperOrigin-RevId: 628500668
 2024-04-26 13:51:04 -07:00
Tomás Longeri
027c24e602 [Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout. 
PiperOrigin-RevId: 597332393
 2024-01-10 13:00:21 -08:00
Blake Hechtman
1cba270f14 [XLA:MOSAIC] implement downcast from s32 to s8 correctly 
PiperOrigin-RevId: 589830744
 2023-12-11 08:20:53 -08:00
Tomás Longeri
250486f13e [Mosaic] Expose C API for apply-vector-layout's assemble, disassemble, relayout and applyLayoutOp 
PiperOrigin-RevId: 580405900
 2023-11-07 22:32:02 -08:00
Tomás Longeri
1c1dd7c8c7 [Mosaic] Expose C API for VectorLayout, VRegDataBounds 
This is in preparation for Python bindings

PiperOrigin-RevId: 579355000
 2023-11-03 18:24:16 -07:00
Tomás Longeri
a3a5926225 [Mosaic] apply_vector_layout rewrite (6): arith.constant 
PiperOrigin-RevId: 567407089
 2023-09-21 13:54:16 -07:00
Tomás Longeri
838f59e576 [Mosaic] apply_vector_layout C++ rewrite (4) Elementwise ops 
PiperOrigin-RevId: 565255860
 2023-09-13 22:19:10 -07:00
Tomás Longeri
a6eed40f24 [MOSAIC] apply_vector_layout C++ rewrite (3) applyLayoutOp and relayout 
PiperOrigin-RevId: 563556815
 2023-09-07 15:08:30 -07:00
Tomás Longeri
d02b59e410 [MOSAIC] apply_vector_layout C++ rewrite (1) VectorLayout functions 
PiperOrigin-RevId: 561237760
 2023-08-29 22:57:13 -07:00
Jevin Jiang
046bcc0ad9 [Mosaic] Add missing headers in linalg vectorization. 
PiperOrigin-RevId: 560865251
 2023-08-28 17:39:47 -07:00
Richard Levasseur
f891cbf64b Load Python rules from rules_python 
PiperOrigin-RevId: 559789250
 2023-08-24 10:22:57 -07:00
Sharad Vikram
3baa6e7a89 Enable building jaxlib w/ Mosaic 
PiperOrigin-RevId: 551159246
 2023-07-26 03:59:30 -07:00
Sharad Vikram
3d556b7a19 Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic 
PiperOrigin-RevId: 549801858
 2023-07-20 18:28:51 -07:00