cublasXt 矩阵乘法在 C++ 中成功，在 Python 中失败

Question

我正在尝试在 Ubuntu Linux 16.04 上使用 Python 2.7.14 中的 ctypess 将 CUDA 9.0 中的 cublasXt*gemm 函数包装起来。这些函数接受主机内存中的数组作为它们的一些参数。我已经能够在 C++ 中成功使用它们，如下所示：

#include <iostream>
#include <cstdlib>
#include "cublasXt.h"
#include "cuda_runtime_api.h"

void rand_mat(float* &x, int m, int n) {
    x = new float[m*n];
    for (int i=0; i<m; ++i) {
        for (int j=0; j<n; ++j) {
            x[i*n+j] = ((float)rand())/RAND_MAX;
        }
    }
}

int main(void) {
    cublasXtHandle_t handle;
    cublasXtCreate(&handle);

    int devices[1] = {0};
    if (cublasXtDeviceSelect(handle, 1, devices) !=
        CUBLAS_STATUS_SUCCESS) {
        std::cout << "initialization failed" << std::endl; 
        return 1;
    }

    float *a, *b, *c;
    int m = 4, n = 4, k = 4;

    rand_mat(a, m, k);
    rand_mat(b, k, n);
    rand_mat(c, m, n);

    float alpha = 1.0;
    float beta = 0.0;

    if (cublasXtSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                      m, n, k, &alpha, a, m, b, k, &beta, c, m) != 
           CUBLAS_STATUS_SUCCESS) {
        std::cout << "matrix multiply failed" << std::endl; 
        return 1;
    }
    delete a; delete b; delete c;
    cublasXtDestroy(handle);
}

但是，当我尝试按如下方式将它们包装在 Python 中时，我在 cublasXt*gemm 调用中遇到了段错误：

import ctypes
import numpy as np

_libcublas = ctypes.cdll.LoadLibrary('libcublas.so')
_libcublas.cublasXtCreate.restype = int
_libcublas.cublasXtCreate.argtypes = [ctypes.c_void_p]
_libcublas.cublasXtDestroy.restype = int
_libcublas.cublasXtDestroy.argtypes = [ctypes.c_void_p]
_libcublas.cublasXtDeviceSelect.restype = int
_libcublas.cublasXtDeviceSelect.argtypes = [ctypes.c_void_p,
                                            ctypes.c_int,
                                            ctypes.c_void_p]
_libcublas.cublasXtSgemm.restype = int
_libcublas.cublasXtSgemm.argtypes = [ctypes.c_void_p,
                                     ctypes.c_int,
                                     ctypes.c_int,
                                     ctypes.c_int,
                                     ctypes.c_int,
                                     ctypes.c_int,
                                     ctypes.c_void_p,
                                     ctypes.c_void_p,
                                     ctypes.c_int,
                                     ctypes.c_void_p,
                                     ctypes.c_int,
                                     ctypes.c_void_p,
                                     ctypes.c_void_p,
                                     ctypes.c_int]

handle = ctypes.c_void_p()
_libcublas.cublasXtCreate(ctypes.byref(handle))
deviceId = np.array([0], np.int32)
status = _libcublas.cublasXtDeviceSelect(handle, 1,
                                         deviceId.ctypes.data)
if status:
    raise RuntimeError

a = np.random.rand(4, 4).astype(np.float32)
b = np.random.rand(4, 4).astype(np.float32)
c = np.zeros((4, 4), np.float32)

status = _libcublas.cublasXtSgemm(handle, 0, 0, 4, 4, 4,
                                  ctypes.byref(ctypes.c_float(1.0)),
                                  a.ctypes.data, 4, b.ctypes.data, 4, 
                                  ctypes.byref(ctypes.c_float(0.0)),
                                  c.ctypes.data, 4)
if status:
    raise RuntimeError
print 'success? ', np.allclose(np.dot(a.T, b.T).T, c_gpu.get())
_libcublas.cublasXtDestroy(handle)

奇怪的是，如果我稍微修改上面的 Python 包装器以接受我已转移到 GPU 的pycuda.gpuarray.GPUArray 矩阵，它们就会起作用。关于为什么在将主机内存传递给函数时我只在 Python 中遇到段错误的任何想法？

Answer 1

这些Xt<t>gemm 函数的CUBLAS 文档中似乎存在错误。至少从 CUDA 8 开始，参数m,n,k,lda,ldb,ldc 都是size_t 类型。这可以通过查看头文件cublasXt.h 来发现。

您的包装器的以下修改对我来说似乎可以正常工作：

$ cat t1340.py
import ctypes
import numpy as np

_libcublas = ctypes.cdll.LoadLibrary('libcublas.so')
_libcublas.cublasXtCreate.restype = int
_libcublas.cublasXtCreate.argtypes = [ctypes.c_void_p]
_libcublas.cublasXtDestroy.restype = int
_libcublas.cublasXtDestroy.argtypes = [ctypes.c_void_p]
_libcublas.cublasXtDeviceSelect.restype = int
_libcublas.cublasXtDeviceSelect.argtypes = [ctypes.c_void_p,
                                            ctypes.c_int,
                                            ctypes.c_void_p]
_libcublas.cublasXtSgemm.restype = int
_libcublas.cublasXtSgemm.argtypes = [ctypes.c_void_p,
                                     ctypes.c_int,
                                     ctypes.c_int,
                                     ctypes.c_size_t,
                                     ctypes.c_size_t,
                                     ctypes.c_size_t,
                                     ctypes.c_void_p,
                                     ctypes.c_void_p,
                                     ctypes.c_size_t,
                                     ctypes.c_void_p,
                                     ctypes.c_size_t,
                                     ctypes.c_void_p,
                                     ctypes.c_void_p,
                                     ctypes.c_size_t]

handle = ctypes.c_void_p()
_libcublas.cublasXtCreate(ctypes.byref(handle))
deviceId = np.array([0], np.int32)
status = _libcublas.cublasXtDeviceSelect(handle, 1,
                                         deviceId.ctypes.data)
if status:
    raise RuntimeError

a = np.random.rand(4, 4).astype(np.float32)
b = np.random.rand(4, 4).astype(np.float32)
c = np.zeros((4, 4), np.float32)
alpha = ctypes.c_float(1.0)
beta = ctypes.c_float(0.0)

status = _libcublas.cublasXtSgemm(handle, 0, 0, 4, 4, 4,
                                 ctypes.byref(alpha),
                                 a.ctypes.data, 4, b.ctypes.data, 4,
                                 ctypes.byref(beta),
                                 c.ctypes.data, 4)
if status:
    raise RuntimeError
print 'success? ', np.allclose(np.dot(a.T, b.T).T, c)
_libcublas.cublasXtDestroy(handle)
$ python t1340.py
success?  True
$

列举我所做的更改：

1. 将m、n、k、lda、ldb、ldc 的cublasXtSgemm 参数从cublasXtSgemm 更改为ldc 参数从@9876543439@ 到@9876
2. 为您的 alpha 和 beta 参数提供了显式变量；这可能无关紧要
3. 在您的 np.allclose 函数中，将 c_gpu.get 更改为 c

以上内容在 CUDA 8 和 CUDA 9 上进行了测试。我已向 NVIDIA 提交了一个内部错误以更新文档（即使当前的 CUDA 9 文档也不能反映头文件的当前状态。）