mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib
History
Sergei Lebedev 7929cd8410 [pallas:triton] The lowering now uses PTX instead of Triton IR 
This change improves the stability and backward compatibility of Pallas Triton
calls, because unlike PTX, the Triton dialect has no stability guarantees
and does change in practice.

See #25196.

A few notes

* Pallas Triton no longer delegates compilation to PTX to XLA:GPU. Instead,
  compilation is done via a new PjRt extension, which uses its own compilation
  pipeline mirrored after the one in the Triton Python bindings.
* The implementation of the old custom call used by Pallas Triton is
  deprecated and will be removed after 6 months as per
  [compatibility guarantees] [*]

[*]: https://jax.readthedocs.io/en/latest/export/export.html#compatibility-guarantees

PiperOrigin-RevId: 722773884
2025-02-03 13:21:40 -08:00
..
cpu
Use thread-safe initialization of LAPACK kernels.
2025-01-13 02:51:38 -08:00
cuda
Remove CUDA rpaths from jaxlib build.
2025-01-17 17:09:30 -08:00
gpu
Merge pull request #25755 from ROCm:ci_rnn_final-upstream
2025-01-15 10:40:54 -08:00
mlir
[mosaic_gpu] Fixed mosaic_gpu-serde pass registration
2025-01-24 06:35:54 -08:00
mosaic
[Mosaic TPU] Support dynamic gather along axis 0 or 1 for 32-bit vreg-sized vector.
2025-01-31 18:47:25 -08:00
rocm
Merge pull request #25755 from ROCm:ci_rnn_final-upstream
2025-01-15 10:40:54 -08:00
tools
[ROCm] Implement RNN support
2025-01-14 19:04:49 -06:00
triton
Fix triton capi_objects target to depend on MLIR CAPIIRObjects bazel
2024-09-06 01:06:27 +00:00
absl_status_casters.h
Migrate JAX MLIR Python dialect extensions to nanobind.
2024-12-13 07:08:28 -08:00
BUILD
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08:00
cpu_feature_guard.c
Added PyUnstable_Module_SetGIL to PyInit_cpu_feature_guard
2024-08-20 02:32:37 +02:00
cuda_plugin_extension.cc
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
ffi_helpers.h
Simplify logic in jaxlib FFI_ASSIGN_OR_RETURN macro, and fix gcc build.
2024-09-17 11:23:25 -07:00
gpu_common_utils.py
Improve the gpu lowering error message if users forget link the gpu library.
2023-09-11 16:14:18 -07:00
gpu_linalg.py
Remove forward compatibility checks from cholesky_update lowering.
2024-09-20 08:32:25 -07:00
gpu_plugin_extension.cc
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08:00
gpu_plugin_extension.h
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
gpu_prng.py
Remove some dead code from gpu_prng.py
2024-10-29 09:29:56 -04:00
gpu_rnn.py
[ROCm] Implement RNN support
2025-01-14 19:04:49 -06:00
gpu_solver.py
Activate FFI implementation of tridiagonal reduction on GPU.
2025-01-10 09:28:15 -08:00
gpu_sparse.py
[ROCm] Bring-up pjrt support
2024-06-17 16:49:22 +00:00
gpu_triton.py
[ROCm]: Add get_arch_details for triton kernel call
2024-08-12 06:16:27 +00:00
handle_pool.h
[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.
2023-05-08 15:37:04 -07:00
hlo_helpers.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
init.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
jax.bzl
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08:00
kernel_helpers.h
Migrate JAX MLIR Python dialect extensions to nanobind.
2024-12-13 07:08:28 -08:00
kernel_nanobind_helpers.h
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
lapack.py
Activate Triangular Solve to XLA's FFI
2024-12-11 02:22:37 -08:00
pass_boilerplate.h
[mosaic] Migrated the serialization pass from codegen to pass_boilerplate.h
2024-12-12 14:19:36 -08:00
README.md
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
rocm_plugin_extension.cc
Deduplicate some GPU plugin definition code.
2024-12-10 11:32:06 -08:00
setup.py
Increase the minimum SciPy version to 1.11.1.
2025-01-07 16:10:45 -08:00
symlink_files.bzl
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
utils.cc
Add a unittest test extension that runs test cases in parallel using threads.
2025-01-08 09:11:47 -08:00

README.md

jaxlib: support library for JAX

jaxlib is the support library for JAX. While JAX itself is a pure Python package, jaxlib contains the binary (C/C++) parts of the library, including Python bindings, the XLA compiler, the PJRT runtime, and a handful of handwritten kernels. For more information, including installation and build instructions, refer to main JAX README: https://github.com/jax-ml/jax/.