[Python] cusolverDnDgesvdj for the computation of singular values only under Python

I have set up the code below for the computation of the SVD of a matrix. The code uses cuSOLVER's cusolverDnDgesvdj.

import ctypes
import numpy as np
import pycuda.autoinit # implicit context creation
import pycuda.driver as cuda

# -----------------------------------------------------------------------------
# 0) C types
# -----------------------------------------------------------------------------
c_int = ctypes.c_int
c_double = ctypes.c_double
c_voidp = ctypes.c_void_p

# -----------------------------------------------------------------------------
# 1) Load cuSOLVER and CUDA driver
# -----------------------------------------------------------------------------
cusolver = ctypes.cdll.LoadLibrary("libcusolver.so")
cuda_drv = ctypes.cdll.LoadLibrary("libcudart.so")

# -----------------------------------------------------------------------------
# 2) Declare cuSOLVER / CUDA-driver function prototypes
# -----------------------------------------------------------------------------
# cusolverDnCreate / Destroy
cusolver.cusolverDnCreate.argtypes = [ctypes.POINTER(c_voidp)]
cusolver.cusolverDnCreate.restype = c_int
cusolver.cusolverDnDestroy.argtypes = [c_voidp]
cusolver.cusolverDnDestroy.restype = c_int

# Set stream on handle
cusolver.cusolverDnSetStream.argtypes = [c_voidp, ctypes.c_void_p] # Use c_void_p for stream handle
cusolver.cusolverDnSetStream.restype = c_int

# gesvdjInfo create/destroy
cusolver.cusolverDnCreateGesvdjInfo.argtypes = [ctypes.POINTER(c_voidp)]
cusolver.cusolverDnCreateGesvdjInfo.restype = c_int
cusolver.cusolverDnDestroyGesvdjInfo.argtypes = [c_voidp]
cusolver.cusolverDnDestroyGesvdjInfo.restype = c_int

# set tolerance / max sweeps
cusolver.cusolverDnXgesvdjSetTolerance.argtypes = [c_voidp, c_double]
cusolver.cusolverDnXgesvdjSetTolerance.restype = c_int
cusolver.cusolverDnXgesvdjSetMaxSweeps.argtypes = [c_voidp, c_int]
cusolver.cusolverDnXgesvdjSetMaxSweeps.restype = c_int

# query bufferSize (13 args!)
cusolver.cusolverDnDgesvdj_bufferSize.argtypes = [
c_voidp, # handle
c_int, # jobz
c_int, # econ
c_int, # m
c_int, # n
c_voidp, # A
c_int, # lda
c_voidp, # S
c_voidp, # U
c_int, # ldu
c_voidp, # V
c_int, # ldv
ctypes.POINTER(c_int), # &lwork
c_voidp # gesvdjInfo
]
cusolver.cusolverDnDgesvdj_bufferSize.restype = c_int

# run SVD
cusolver.cusolverDnDgesvdj.argtypes = [
c_voidp, c_int, c_int, c_int, c_int,
c_voidp, c_int, c_voidp, c_voidp, c_int,
c_voidp, c_int, c_voidp, c_int, c_voidp, c_voidp
]
cusolver.cusolverDnDgesvdj.restype = c_int

# get sweeps & residual
cusolver.cusolverDnXgesvdjGetSweeps.argtypes = [c_voidp, c_voidp, ctypes.POINTER(c_int)]
cusolver.cusolverDnXgesvdjGetSweeps.restype = c_int
cusolver.cusolverDnXgesvdjGetResidual.argtypes = [c_voidp, c_voidp, ctypes.POINTER(c_double)]
cusolver.cusolverDnXgesvdjGetResidual.restype = c_int

# -----------------------------------------------------------------------------
# 3) Problem parameters & host data
# -----------------------------------------------------------------------------
m, n = 3, 2
lda, ldu, ldv = m, m, n
minmn = min(m, n)

# A in column-major “matlab” layout:
h_A = np.array([1.0, 4.0, 2.0, # col 0
2.0, 5.0, 1.0], # col 1
dtype=np.float64)

# containers for results
h_U = np.zeros((ldu*m,), dtype=np.float64)
h_V = np.zeros((ldv*n,), dtype=np.float64)
h_S = np.zeros((minmn,), dtype=np.float64)
h_info = np.zeros((1,), dtype=np.int32)

# exact singulars for later check
h_S_exact = np.array([7.065283497082729, 1.040081297712078], dtype=np.float64)

# -----------------------------------------------------------------------------
# 4) Create handle, stream, gesvdjInfo
# -----------------------------------------------------------------------------
handle = c_voidp()
assert cusolver.cusolverDnCreate(ctypes.byref(handle)) == 0

# Create a PyCUDA stream
stream = cuda.Stream()
# Pass the underlying CUDA stream handle to cusolver
cusolver.cusolverDnSetStream(handle, ctypes.c_void_p(stream.handle))



params = c_voidp()
assert cusolver.cusolverDnCreateGesvdjInfo(ctypes.byref(params)) == 0

# tune
tol = c_double(1e-7)
sweeps = c_int(15)
cusolver.cusolverDnXgesvdjSetTolerance(params, tol)
cusolver.cusolverDnXgesvdjSetMaxSweeps(params, sweeps)

# -----------------------------------------------------------------------------
# 5) Allocate + copy device memory
# -----------------------------------------------------------------------------
d_A = cuda.mem_alloc(h_A.nbytes)
d_S = cuda.mem_alloc(h_S.nbytes)
d_U = cuda.mem_alloc(h_U.nbytes)
d_V = cuda.mem_alloc(h_V.nbytes)
d_info = cuda.mem_alloc(h_info.nbytes)

# Use the PyCUDA stream object in PyCUDA functions
cuda.memcpy_htod_async(d_A, h_A, stream)
# (we’ll copy results back later)

# -----------------------------------------------------------------------------
# 6) Query workspace size
# -----------------------------------------------------------------------------
lwork = c_int()
status = cusolver.cusolverDnDgesvdj_bufferSize(
handle,
1, # jobz = CUSOLVER_EIG_MODE_VECTOR
0, # econ = full
m, n,
int(d_A), lda, # Cast d_A to int
int(d_S), # Cast d_S to int
int(d_U), ldu, # Cast d_U to int
int(d_V), ldv, # Cast d_V to int
ctypes.byref(lwork),
params
)
assert status == 0

d_work = cuda.mem_alloc(lwork.value * ctypes.sizeof(c_double))

# -----------------------------------------------------------------------------
# 7) Run batched SVD (batch = 1 here)
# -----------------------------------------------------------------------------
status = cusolver.cusolverDnDgesvdj(
handle,
1, # jobz
0, # econ
m, n,
int(d_A), lda, # Cast d_A to int
int(d_S), # Cast d_S to int
int(d_U), ldu, # Cast d_U to int
int(d_V), ldv, # Cast d_V to int
int(d_work), lwork, # Cast d_work to int
int(d_info), # Cast d_info to int
params
)
assert status == 0

# -----------------------------------------------------------------------------
# 8) Copy results back to host & sync
# -----------------------------------------------------------------------------
# Use the PyCUDA stream object in PyCUDA functions
cuda.memcpy_dtoh_async(h_U, d_U, stream)
cuda.memcpy_dtoh_async(h_V, d_V, stream)
cuda.memcpy_dtoh_async(h_S, d_S, stream)
cuda.memcpy_dtoh_async(h_info, d_info, stream)
# Synchronize the PyCUDA stream
stream.synchronize()

print("info =", int(h_info[0]))
print("S =", h_S)
print("S_exact=", h_S_exact)

# get internal stats
sweeps_out = c_int()
resid_out = c_double()
cusolver.cusolverDnXgesvdjGetSweeps(handle, params, ctypes.byref(sweeps_out))
cusolver.cusolverDnXgesvdjGetResidual(handle, params, ctypes.byref(resid_out))
print("executed sweeps =", sweeps_out.value)
print("residual =", resid_out.value)

# -----------------------------------------------------------------------------
# 9) Cleanup
# -----------------------------------------------------------------------------
# cuda_drv.cudaStreamDestroy(stream) # No need to destroy PyCUDA stream via ctypes
cusolver.cusolverDnDestroyGesvdjInfo(params)
cusolver.cusolverDnDestroy(handle)


The code works.

Now, if I change the second parameter of cusolver.cusolverDnDgesvdj_bufferSize and cusolver.cusolverDnDgesvdj to 0 to compute the singular values only, the code still works.

However, if I try to remove the pointers to d_U and d_V, the code does not work anymore. I have tried to replace the pointers with 0 and with None, obtaining every time a different status error.

How should I consistently modify the calls to cusolver.cusolverDnDgesvdj_bufferSize and cusolver.cusolverDnDgesvdj to avoid the need of defining d_U and d_V?

I'm running the code under Google Colab.

Continue reading...