[gpu] Use syevBatched for eigh

Use syevBatched for eigh

Author: jaro-sevcik

jaxlib/gpu/solver_kernels_ffi.cc

@@ -423,16 +423,21 @@ ffi::Error Syevd64Impl(int64_t batch, int64_t n, gpuStream_t stream,
       params_cleanup(
           params, [](gpusolverDnParams_t p) { gpusolverDnDestroyParams(p); });
 
+  int64_t matrix_size = n * n * ffi::ByteWidth(dataType);
+  int64_t batch_step = std::numeric_limits<int>::max() / matrix_size;
+  if (batch_step >= 32 * 1024) {
+    batch_step = 32 * 1024;
+  }
   size_t workspaceInBytesOnDevice, workspaceInBytesOnHost;
-  JAX_FFI_RETURN_IF_GPU_ERROR(gpusolverDnXsyevd_bufferSize(
+  JAX_FFI_RETURN_IF_GPU_ERROR(gpusolverDnXsyevBatched_bufferSize(
       handle.get(), params, jobz, uplo, n, aType, /*a=*/nullptr, n, wType,
       /*w=*/nullptr, aType, &workspaceInBytesOnDevice,
-      &workspaceInBytesOnHost));
+      &workspaceInBytesOnHost, std::min(batch, batch_step)));
 
   auto maybe_workspace = scratch.Allocate(workspaceInBytesOnDevice);
   if (!maybe_workspace.has_value()) {
     return ffi::Error(ffi::ErrorCode::kResourceExhausted,
-                      "Unable to allocate device workspace for syevd");
+                      "Unable to allocate device workspace for syevBatched");
   }
   auto workspaceOnDevice = maybe_workspace.value();
   auto workspaceOnHost =
@@ -447,17 +452,18 @@ ffi::Error Syevd64Impl(int64_t batch, int64_t n, gpuStream_t stream,
         out_data, a_data, a.size_bytes(), gpuMemcpyDeviceToDevice, stream));
   }
 
-  size_t out_step = n * n * ffi::ByteWidth(dataType);
-  size_t w_step = n * ffi::ByteWidth(ffi::ToReal(dataType));
+  size_t out_step = n * n * ffi::ByteWidth(dataType) * batch_step;
+  size_t w_step = n * ffi::ByteWidth(ffi::ToReal(dataType)) * batch_step;
 
-  for (auto i = 0; i < batch; ++i) {
-    JAX_FFI_RETURN_IF_GPU_ERROR(gpusolverDnXsyevd(
+  for (int64_t i = 0; i < batch; i += batch_step) {
+    size_t batch_size = static_cast<size_t>(std::min(batch_step, batch - i));
+    JAX_FFI_RETURN_IF_GPU_ERROR(gpusolverDnXsyevBatched(
         handle.get(), params, jobz, uplo, n, aType, out_data, n, wType, w_data,
         aType, workspaceOnDevice, workspaceInBytesOnDevice,
-        workspaceOnHost.get(), workspaceInBytesOnHost, info_data));
+        workspaceOnHost.get(), workspaceInBytesOnHost, info_data, batch_size));
     out_data += out_step;
     w_data += w_step;
-    ++info_data;
+    info_data += batch_step;
   }
 
   return ffi::Error::Success();
@@ -576,18 +582,23 @@ ffi::Error SyevdDispatch(gpuStream_t stream, ffi::ScratchAllocator scratch,
       CheckShape(out->dimensions(), {batch, rows, cols}, "out", "syevd"));
   FFI_RETURN_IF_ERROR(CheckShape(w->dimensions(), {batch, cols}, "w", "syevd"));
   FFI_RETURN_IF_ERROR(CheckShape(info->dimensions(), batch, "info", "syevd"));
-  if (algorithm == SyevdAlgorithm::kJacobi ||
-      (algorithm == SyevdAlgorithm::kDefault && cols <= 32)) {
+#if JAX_GPU_HAVE_64_BIT
+  if (algorithm == SyevdAlgorithm::kJacobi) {
     SOLVER_DISPATCH_IMPL(SyevdjImpl, batch, cols, stream, scratch, lower, a,
-                         out, w, info);
+                          out, w, info);
   } else {
-#if JAX_GPU_HAVE_64_BIT
     return Syevd64Impl(batch, cols, stream, scratch, lower, a, out, w, info);
+  }
 #else
+  if (algorithm == SyevdAlgorithm::kJacobi ||
+     (algorithm == SyevdAlgorithm::kDefault && cols <= 32)) {
+    SOLVER_DISPATCH_IMPL(SyevdjImpl, batch, cols, stream, scratch, lower, a,
+                         out, w, info);
+  } else {
     SOLVER_DISPATCH_IMPL(SyevdImpl, batch, cols, stream, scratch, lower, a, out,
                          w, info);
-#endif
   }
+#endif
   return ffi::Error::InvalidArgument(absl::StrFormat(
       "Unsupported dtype %s in syevd", absl::FormatStreamed(dataType)));
 }

jaxlib/gpu/vendor.h

@@ -412,6 +412,8 @@ typedef cusolverDnParams_t gpusolverDnParams_t;
 
 #define gpusolverDnXsyevd_bufferSize cusolverDnXsyevd_bufferSize
 #define gpusolverDnXsyevd cusolverDnXsyevd
+#define gpusolverDnXsyevBatched_bufferSize cusolverDnXsyevBatched_bufferSize
+#define gpusolverDnXsyevBatched cusolverDnXsyevBatched
 #define gpusolverDnXgesvd_bufferSize cusolverDnXgesvd_bufferSize
 #define gpusolverDnXgesvd cusolverDnXgesvd