Eugene Zhulenev 1e03917c43 [xla:ffi] Use lazy decoding for Buffer<dtype,rank> ...

name                old cpu/op   new cpu/op   delta
BM_AnyBufferArgX1   11.0ns ± 3%  11.2ns ±10%   +1.76%  (p=0.000 n=67+69)
BM_AnyBufferArgX4   12.4ns ± 3%  12.4ns ± 4%   -0.31%  (p=0.006 n=69+69)
BM_BufferArgX1      12.5ns ± 1%  11.1ns ± 4%  -11.20%  (p=0.000 n=62+76)
BM_BufferArgX4      19.1ns ± 1%  14.4ns ± 4%  -24.84%  (p=0.000 n=64+73)
BM_BufferArgX8      36.0ns ± 5%  20.3ns ± 4%  -43.59%  (p=0.000 n=79+75)
BM_TupleOfI32Attrs  66.4ns ± 1%  66.4ns ± 2%   -0.03%  (p=0.000 n=66+72)

PiperOrigin-RevId: 650691450

2024-07-09 11:07:25 -07:00

..

blas_kernels.cc

Switch jaxlib to use nanobind instead of pybind11.

2023-08-24 16:07:56 -07:00

blas_kernels.h

Switch JAX to use the OpenXLA repository.

2023-03-13 18:38:26 +00:00

blas.cc

Move tsl/python to xla/tsl/python

2024-03-29 13:15:21 -07:00

BUILD

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07:00

cholesky_update_kernel.cc

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07:00

cholesky_update_kernel.cu.cc

add cholesky changes in bazel

2024-05-18 00:37:09 +00:00

cholesky_update_kernel.h

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07:00

gpu_kernel_helpers.cc

Fix cupti-related build failure under CUDA 11.

2023-09-27 14:33:30 -07:00

gpu_kernel_helpers.h

[CUDA] Add code to jax initialization that verifies that the CUDA libraries that are found are at least as new as the versions against which JAX was built.

2023-09-27 11:28:40 -07:00

gpu_kernels.cc

Fix error in custom call registration for some FFI functions

2024-06-29 12:18:34 -07:00

linalg.cc

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07:00

lu_pivot_kernels.cc

[xla:ffi] Use lazy decoding for Buffer<dtype,rank>

2024-07-09 11:07:25 -07:00

lu_pivot_kernels.cu.cc

Ported LuPivotsToPermutation to the typed XLA FFI

2024-05-02 08:12:05 -07:00

lu_pivot_kernels.h

Update lu_pivots_to_permutation to use FFI dimensions on GPU.

2024-06-27 09:27:15 -07:00

prng_kernels.cc

[xla:ffi] Use lazy decoding for Buffer<dtype,rank>

2024-07-09 11:07:25 -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

Ported threefry2x32 for GPU to the typed XLA FFI

2024-06-28 06:24:44 -07:00

prng.cc

Ported threefry2x32 for GPU to the typed XLA FFI

2024-06-28 06:24:44 -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_kernels.cc

[JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel.

2023-11-14 04:52:15 -08:00

solver_kernels.h

[shape_poly] Add support for dynamic shapes for eigh

2023-05-02 23:27:59 -07:00

solver.cc

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07: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

[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1

2024-01-19 05:55:41 -08: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]: Add missing hipStreamWaitEvent API call

2024-03-20 16:58:21 +00:00