mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 19:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/gpu
History
Peter Hawkins 3f91b4b43a Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/ 
Move the common jaxlib/gpu_plugin_extension into jaxlib/gpu/

Cleanup only, no functional changes intended.

PiperOrigin-RevId: 738183402
2025-03-18 16:29:37 -07: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
Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/
2025-03-18 16:29:37 -07:00
ffi_wrapper.h
Update the sparse GPU kernels in jaxlib to use the FFI.
2025-02-06 11:45:57 -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 the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
gpu_plugin_extension.cc
Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/
2025-03-18 16:29:37 -07:00
gpu_plugin_extension.h
Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/
2025-03-18 16:29:37 -07:00
hybrid_kernels.cc
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
hybrid_kernels.h
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
hybrid.cc
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
linalg_kernels.cc
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
linalg_kernels.cu.cc
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
linalg_kernels.h
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08:00
linalg.cc
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
2025-02-07 08:48:15 -08: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
Fix integer overflow in workspace size computations for experimental.rnn.*.
2025-03-12 08:22:04 -07:00
rnn_kernels.h
Fix integer overflow in workspace size computations for experimental.rnn.*.
2025-03-12 08:22:04 -07:00
rnn.cc
Fix integer overflow in workspace size computations for experimental.rnn.*.
2025-03-12 08:22:04 -07: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
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
solver_interface.h
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
solver_kernels_ffi.cc
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
solver_kernels_ffi.h
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00: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
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
sparse_kernels.cc
Update the sparse GPU kernels in jaxlib to use the FFI.
2025-02-06 11:45:57 -08:00
sparse_kernels.h
Update the sparse GPU kernels in jaxlib to use the FFI.
2025-02-06 11:45:57 -08:00
sparse.cc
Update the sparse GPU kernels in jaxlib to use the FFI.
2025-02-06 11:45:57 -08:00
triton_kernels.cc
Reverts 74b4d868e3751c1b4efa315ff8cf771faeb0b663
2025-03-13 14:59:09 -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
[ROCm] Implement RNN support
2025-01-14 19:04:49 -06:00