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/gpu
History
Dan Foreman-Mackey 30d54ec6ff Refactor FFI shape inference functions to include dimension check. 
Previously we always had two steps when extracting the batch size: (1) check the buffer has enough dimensions, (2) get the shape. And, in a few cases, this first check was missing. Now these steps are combined into one function that returns a StatusOr.

As part of this, I needed to fix our implementation of the `ASSIGN_OR_RETURN` macro to properly handle parentheses.

PiperOrigin-RevId: 664803225
2024-08-19 07:41:28 -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
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.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
Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.
2024-08-14 09:20:40 -07:00
linalg_kernels.cc
Update GPU implementation of lu_pivots_to_permutation to infer the permutation size directly from the input dimensions, instead of using an input parameter.
2024-08-08 07:35:47 -07:00
linalg_kernels.cu.cc
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
linalg_kernels.h
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
linalg.cc
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
2024-07-10 12:09:12 -07:00
make_batch_pointers.cu.cc
Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.
2024-08-16 04:37:09 -07:00
make_batch_pointers.h
Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.
2024-08-16 04:37:09 -07: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_kernels_ffi.cc
Refactor FFI shape inference functions to include dimension check.
2024-08-19 07:41:28 -07:00
solver_kernels_ffi.h
Port Getrf GPU kernel from custom call to FFI.
2024-08-01 15:02:25 -07:00
solver_kernels.cc
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
solver_kernels.h
Move jaxlib GPU handlers to separate build target.
2024-08-01 12:30:04 -07:00
solver.cc
[ROCM] Fix BUILD.bazel library source paths
2024-08-07 09:18:20 -05:00
sparse_kernels.cc
Add batching rules to jax.lax.linalg.tridiagonal_solve.
2023-08-10 16:25:59 -07:00
sparse_kernels.h
Add batching rules to jax.lax.linalg.tridiagonal_solve.
2023-08-10 16:25:59 -07:00
sparse.cc
Move tsl/python to xla/tsl/python
2024-03-29 13:15:21 -07:00
triton_kernels.cc
[Jax/Triton] Skip benchmarking while autotuning for configs that cannot be launched.
2024-04-18 14:35:51 -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] Fix BUILD.bazel library source paths
2024-08-07 09:18:20 -05:00