c++ - 如何在模板函数中使用CUBLAS库？

CUBLAS对于每种数据类型都有一个separate function，但我想从模板中调用CUBLAS，例如:

template <typename T> foo(...) {
    ...
    cublas<S/D/C/Z>geam(..., const T* A, ...);
    ...
}

如何触发正确的函数调用？

最佳答案

我为具有相同函数名称的不同类型编写了cublas包装函数。

inline cublasStatus_t cublasGgeam(cublasHandle_t handle,
        cublasOperation_t transa, cublasOperation_t transb,
        int m, int n,
        const float *alpha,
        const float *A, int lda,
        const float *beta,
        const float *B, int ldb,
        float *C, int ldc)
{
    return cublasSgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc);
}

inline cublasStatus_t cublasGgeam(cublasHandle_t handle,
        cublasOperation_t transa, cublasOperation_t transb,
        int m, int n,
        const double *alpha,
        const double *A, int lda,
        const double *beta,
        const double *B, int ldb,
        double *C, int ldc)
{
    return cublasDgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc);
}

之后，您可以为具有相同函数名称的任何类型调用geam()。 C++编译器将根据参数类型选择合适的函数。在你的情况下，它应该像

template <typename T> foo(...) {
    ...
    cublasGgeam(..., A, ...);
    ...
}

尽管您必须为包装函数编写一长串列表，但这是编译时重载，完全没有运行时成本。

关于c++ - 如何在模板函数中使用CUBLAS库？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/16402087/