mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/rocm
History
Peter Hawkins 93b839ace4 Use input-output aliasing for jaxlib GPU custom calls. 
Previously we had no way to tell XLA that inputs and outputs of GPU custom calls must alias. This now works in XLA:GPU so we can just ask XLA to enforce the aliasing we need.

PiperOrigin-RevId: 479642543
2022-10-07 12:22:04 -07:00
..
BUILD.bazel
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_gpu_kernel_helpers.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_gpu_kernel_helpers.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_linalg.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_lu_pivot_kernels.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_lu_pivot_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_lu_pivot_kernels.hip.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_prng_kernels.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_prng_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_prng_kernels.hip.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hip_prng.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hipblas_kernels.cc
Use input-output aliasing for jaxlib GPU custom calls.
2022-10-07 12:22:04 -07:00
hipblas_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hipblas.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hipsolver_kernels.cc
Use input-output aliasing for jaxlib GPU custom calls.
2022-10-07 12:22:04 -07:00
hipsolver_kernels.h
[ROCm] Upgrade to ROCm 5.3 and associated enhancements
2022-10-01 04:45:26 -07:00
hipsolver.cc
[ROCm] Upgrade to ROCm 5.3 and associated enhancements
2022-10-01 04:45:26 -07:00
hipsparse_kernels.cc
Use input-output aliasing for jaxlib GPU custom calls.
2022-10-07 12:22:04 -07:00
hipsparse_kernels.h
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
hipsparse.cc
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00