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
Peter Hawkins a852710a09 Merge CUDA and ROCM kernel code in jaxlib. 
The code for both CUDA and ROCM is almost identical, so with a small shim library to handle the differences we can share almost everything.

PiperOrigin-RevId: 483666051
2022-10-25 07:23:34 -07:00
..
blas_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
blas_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
blas.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
BUILD
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
gpu_kernel_helpers.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
gpu_kernel_helpers.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
gpu_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
linalg.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
lu_pivot_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
lu_pivot_kernels.cu.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
lu_pivot_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.cu.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
prng.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
solver_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
solver_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
solver.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
sparse_kernels.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
sparse_kernels.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
sparse.cc
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00
vendor.h
Merge CUDA and ROCM kernel code in jaxlib.
2022-10-25 07:23:34 -07:00