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
rocm_jax/jaxlib/gpu
History
Dan Foreman-Mackey 21e98b5ce4 Fix overflow error in GPU batched linear algebra kernels. 
As reported in https://github.com/jax-ml/jax/issues/24843, our LU decomposition on GPU hits overflow errors when the batch size approaches int32 max. This was caused by an issue in how we were constructing the batched pointers used by cuBLAS.

PiperOrigin-RevId: 695694648
2024-11-12 05:33:49 -08:00
..
blas_handle_pool.cc
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
blas_handle_pool.h
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
blas_kernels.cc
Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.
2024-08-16 04:37:09 -07:00
blas_kernels.h
Switch JAX to use the OpenXLA repository.
2023-03-13 18:38:26 +00:00
blas.cc
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -07:00
BUILD
cleanup: remove api_version from BUILD files
2024-10-08 09:44:15 -07:00
gpu_kernel_helpers.cc
Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.
2024-08-16 04:37:09 -07:00
gpu_kernel_helpers.h
Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.
2024-08-16 04:37:09 -07:00
gpu_kernels.cc
Update FFI target name for syrk operation to be consistent with other kernels.
2024-09-06 13:21:38 -07:00
linalg_kernels.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
linalg_kernels.cu.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
linalg_kernels.h
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
linalg.cc
Port the GPU Cholesky update custom call to the FFI.
2024-08-20 05:46:03 -07:00
make_batch_pointers.cu.cc
Fix overflow error in GPU batched linear algebra kernels.
2024-11-12 05:33:49 -08:00
make_batch_pointers.h
Fix overflow error in GPU batched linear algebra kernels.
2024-11-12 05:33:49 -08:00
prng_kernels.cc
Move FFI helper macros from jaxlib/cpu/lapack_kernels.cc to a jaxlib/ffi_helpers.h.
2024-07-10 15:09:45 -07:00
prng_kernels.cu.cc
Ported threefry2x32 for GPU to the typed XLA FFI
2024-06-28 06:24:44 -07:00
prng_kernels.h
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
prng.cc
Remove forward compatibility mode for old PRGN custom call on GPU
2024-07-31 08:10:17 -07:00
rnn_kernels.cc
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
rnn_kernels.h
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
rnn.cc
Run LSTM test using FP32 math (as opposed to TF32)
2023-09-19 14:45:14 -04:00
solver_handle_pool.cc
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
solver_handle_pool.h
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
solver_interface.cc
Port GPU kernels for SVD to the FFI.
2024-09-20 07:34:50 -07:00
solver_interface.h
Port GPU kernels for SVD to the FFI.
2024-09-20 07:34:50 -07:00
solver_kernels_ffi.cc
(follow-up of PR #23852) add missing typename keyword to work with gcc
2024-11-07 23:55:38 -06:00
solver_kernels_ffi.h
Port GPU kernels for SVD to the FFI.
2024-09-20 07:34:50 -07:00
solver_kernels.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
solver_kernels.h
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
solver.cc
Port GPU kernels for SVD to the FFI.
2024-09-20 07:34:50 -07:00
sparse_kernels.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
sparse_kernels.h
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
sparse.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
triton_kernels.cc
Fix a number of minor problems in the ROCM build.
2024-08-26 17:04:01 -07:00
triton_kernels.h
[jax_triton] Only use side stream to do autotuning when doing graph capture
2024-02-02 10:48:26 -08:00
triton_utils.cc
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
triton_utils.h
[jax_triton] Add user-specified name field to serialized format.
2023-08-16 02:53:51 -07:00
triton.cc
[ROCm]: Add get_arch_details for triton kernel call
2024-08-12 06:16:27 +00:00
triton.proto
[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1
2024-01-19 05:55:41 -08:00
vendor.h
Update Eigh kernel on GPU to use 64-bit interface when it is available.
2024-10-10 12:59:37 -07:00