mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/gpu
History
George Necula 9f797990b5 Remove old backward compatibility mode for old PRGN custom call on GPU 
The backend support for the new custom call was added on June 28th, 2024 (#20997).

PiperOrigin-RevId: 723077990
2025-02-04 07:34:52 -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
Automated Code Change
2024-11-25 08:31:06 -08: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
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
hybrid_kernels.cc
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
hybrid_kernels.h
Add a GPU implementation of lax.linalg.eig.
2024-11-18 08:11:57 -08:00
hybrid.cc
Use thread-safe initialization of LAPACK kernels.
2025-01-13 02:51:38 -08: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
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
prng_kernels.cu.cc
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
prng_kernels.h
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
prng.cc
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
rnn_kernels.cc
[ROCm] Implement RNN support
2025-01-14 19:04:49 -06:00
rnn_kernels.h
Fix struct string encoding non-determinism in the RNN descriptor.
2025-01-09 12:57:09 +00: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 symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08:00
solver_interface.h
Port symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08:00
solver_kernels_ffi.cc
Port symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08:00
solver_kernels_ffi.h
Port symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08: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 symmetric tridiagonal reduction GPU kernel to FFI.
2024-12-09 12:41:23 -08: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
Integrate Triton up to [9732c047](9732c04701)
2024-12-03 10:55:56 -08: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
[ROCm] Implement RNN support
2025-01-14 19:04:49 -06:00