Clean up some unused GPU linear algebra kernels.

This change removes the legacy `csrlsvqr` and `sytrd` custom calls from jaxlib. These were never covered by the export compatibility policy, and their FFI counterparts have been targeted by JAX for several releases.

PiperOrigin-RevId: 766298494

Author: dfm

jaxlib/cuda/BUILD

@@ -155,24 +155,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "cusolver_kernels",
-    srcs = ["//jaxlib/gpu:solver_kernels.cc"],
-    hdrs = ["//jaxlib/gpu:solver_kernels.h"],
-    deps = [
-        ":cuda_gpu_kernel_helpers",
-        ":cuda_solver_handle_pool",
-        ":cuda_vendor",
-        "//jaxlib:kernel_helpers",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@local_config_cuda//cuda:cuda_headers",
-        "@xla//xla/service:custom_call_status",
-        "@xla//xla/tsl/cuda:cudart",
-        "@xla//xla/tsl/cuda:cusolver",
-    ],
-)
-
 cc_library(
     name = "cusolver_interface",
     srcs = ["//jaxlib/gpu:solver_interface.cc"],
@@ -223,21 +205,14 @@ nanobind_extension(
     features = ["-use_header_modules"],
     module_name = "_solver",
     deps = [
-        ":cuda_gpu_kernel_helpers",
-        ":cuda_solver_handle_pool",
         ":cuda_vendor",
-        ":cusolver_kernels",
         ":cusolver_kernels_ffi",
         "//jaxlib:kernel_nanobind_helpers",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
         "@local_config_cuda//cuda:cuda_headers",
         "@nanobind",
         "@xla//xla/tsl/cuda:cublas",
         "@xla//xla/tsl/cuda:cudart",
         "@xla//xla/tsl/cuda:cusolver",
-        "@xla//xla/tsl/python/lib/core:numpy",
     ],
 )
 
@@ -472,7 +447,6 @@ cc_library(
         ":cuda_prng_kernels",
         ":cuda_vendor",
         ":cudnn_rnn_kernels",
-        ":cusolver_kernels",
         ":cusolver_kernels_ffi",
         ":cusparse_kernels",
         ":triton_kernels",

jaxlib/gpu/BUILD

@@ -59,8 +59,6 @@ exports_files(srcs = [
     "solver_handle_pool.h",
     "solver_interface.cc",
     "solver_interface.h",
-    "solver_kernels.cc",
-    "solver_kernels.h",
     "solver_kernels_ffi.cc",
     "solver_kernels_ffi.h",
     "sparse.cc",

jaxlib/gpu/gpu_kernels.cc

@@ -19,7 +19,6 @@ limitations under the License.
 #include "jaxlib/gpu/linalg_kernels.h"
 #include "jaxlib/gpu/prng_kernels.h"
 #include "jaxlib/gpu/rnn_kernels.h"
-#include "jaxlib/gpu/solver_kernels.h"
 #include "jaxlib/gpu/solver_kernels_ffi.h"
 #include "jaxlib/gpu/sparse_kernels.h"
 #include "jaxlib/gpu/triton_kernels.h"
@@ -40,14 +39,12 @@ XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_syrk_ffi", "CUDA",
                          SyrkFfi);
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_geqrf_ffi", "CUDA",
                          GeqrfFfi);
-XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cusolver_csrlsvqr", Csrlsvqr, "CUDA");
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_csrlsvqr_ffi", "CUDA",
                          CsrlsvqrFfi);
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_orgqr_ffi", "CUDA",
                          OrgqrFfi);
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_syevd_ffi", "CUDA",
                          SyevdFfi);
-XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cusolver_sytrd", Sytrd, "CUDA");
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_sytrd_ffi", "CUDA",
                          SytrdFfi);
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_gesvd_ffi", "CUDA",

jaxlib/gpu/solver.cc

@@ -13,101 +13,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <stdexcept>
-#include <utility>
-
-#include "absl/container/flat_hash_map.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_format.h"
 #include "nanobind/nanobind.h"
-#include "nanobind/stl/pair.h"  // IWYU pragma: keep
-#include "jaxlib/gpu/gpu_kernel_helpers.h"
-#include "jaxlib/gpu/solver_handle_pool.h"
-#include "jaxlib/gpu/solver_kernels.h"
 #include "jaxlib/gpu/solver_kernels_ffi.h"
 #include "jaxlib/gpu/vendor.h"
 #include "jaxlib/kernel_nanobind_helpers.h"
-#include "xla/tsl/python/lib/core/numpy.h"
 
 namespace jax {
 namespace JAX_GPU_NAMESPACE {
 namespace {
 
 namespace nb = nanobind;
 
-// Converts a NumPy dtype to a Type.
-SolverType DtypeToSolverType(const dtype& np_type) {
-  static auto* types =
-      new absl::flat_hash_map<std::pair<char, int>, SolverType>({
-          {{'f', 4}, SolverType::F32},
-          {{'f', 8}, SolverType::F64},
-          {{'c', 8}, SolverType::C64},
-          {{'c', 16}, SolverType::C128},
-      });
-  auto it = types->find({np_type.kind(), np_type.itemsize()});
-  if (it == types->end()) {
-    nb::str repr = nb::repr(np_type);
-    throw std::invalid_argument(
-        absl::StrFormat("Unsupported dtype %s", repr.c_str()));
-  }
-  return it->second;
-}
-
-#ifdef JAX_GPU_CUDA
-
-// csrlsvqr: Linear system solve via Sparse QR
-
-// Returns a descriptor for a csrlsvqr operation.
-nb::bytes BuildCsrlsvqrDescriptor(const dtype& dtype, int n, int nnzA,
-                                  int reorder, double tol) {
-  SolverType type = DtypeToSolverType(dtype);
-  return PackDescriptor(CsrlsvqrDescriptor{type, n, nnzA, reorder, tol});
-}
-
-#endif  // JAX_GPU_CUDA
-
-// Returns the workspace size and a descriptor for a geqrf operation.
-std::pair<int, nb::bytes> BuildSytrdDescriptor(const dtype& dtype, bool lower,
-                                               int b, int n) {
-  SolverType type = DtypeToSolverType(dtype);
-  auto h = SolverHandlePool::Borrow(/*stream=*/nullptr);
-  JAX_THROW_IF_ERROR(h.status());
-  auto& handle = *h;
-  int lwork;
-  gpusolverFillMode_t uplo =
-      lower ? GPUSOLVER_FILL_MODE_LOWER : GPUSOLVER_FILL_MODE_UPPER;
-  switch (type) {
-    case SolverType::F32:
-      JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpusolverDnSsytrd_bufferSize(
-          handle.get(), uplo, n, /*A=*/nullptr, /*lda=*/n, /*D=*/nullptr,
-          /*E=*/nullptr, /*tau=*/nullptr, &lwork)));
-      break;
-    case SolverType::F64:
-      JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpusolverDnDsytrd_bufferSize(
-          handle.get(), uplo, n, /*A=*/nullptr, /*lda=*/n, /*D=*/nullptr,
-          /*E=*/nullptr, /*tau=*/nullptr, &lwork)));
-      break;
-    case SolverType::C64:
-      JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpusolverDnChetrd_bufferSize(
-          handle.get(), uplo, n, /*A=*/nullptr, /*lda=*/n, /*D=*/nullptr,
-          /*E=*/nullptr, /*tau=*/nullptr, &lwork)));
-      break;
-    case SolverType::C128:
-      JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpusolverDnZhetrd_bufferSize(
-          handle.get(), uplo, n, /*A=*/nullptr, /*lda=*/n, /*D=*/nullptr,
-          /*E=*/nullptr, /*tau=*/nullptr, &lwork)));
-      break;
-  }
-  return {lwork, PackDescriptor(SytrdDescriptor{type, uplo, b, n, n, lwork})};
-}
-
 nb::dict Registrations() {
   nb::dict dict;
-  dict[JAX_GPU_PREFIX "solver_sytrd"] = EncapsulateFunction(Sytrd);
-
-#ifdef JAX_GPU_CUDA
-  dict["cusolver_csrlsvqr"] = EncapsulateFunction(Csrlsvqr);
-#endif  // JAX_GPU_CUDA
 
   dict[JAX_GPU_PREFIX "solver_getrf_ffi"] = EncapsulateFfiHandler(GetrfFfi);
   dict[JAX_GPU_PREFIX "solver_geqrf_ffi"] = EncapsulateFfiHandler(GeqrfFfi);
@@ -127,12 +45,7 @@ nb::dict Registrations() {
 }
 
 NB_MODULE(_solver, m) {
-  tsl::ImportNumpy();
   m.def("registrations", &Registrations);
-  m.def("build_sytrd_descriptor", &BuildSytrdDescriptor);
-#ifdef JAX_GPU_CUDA
-  m.def("build_csrlsvqr_descriptor", &BuildCsrlsvqrDescriptor);
-#endif  // JAX_GPU_CUDA
 }
 
 }  // namespace