如何在模板函数中使用CUBLAS库？ [英] How to use CUBLAS library within a template function?

问题描述

CUBLAS对每种类型的数据都有单独的功能，但我想从一个模板内调用CUBLAS，例如：

CUBLAS has a separate function for each type of data, but I want to call CUBLAS from within a template, e.g.:

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

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

How do I trigger the correct function call?

推荐答案

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

I wrote cublas wrapper functions for different types with same function name.

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 ++编译器将根据参数的类型选择正确的函数。在你的情况下，它应该像

After that, you can call geam() for any type with the same function name. C++ compiler will choose the right function by the type of the parameters. In you case it should be like

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

这是一个完成时间重载运行时成本，虽然你必须写一个长列表的包装函数。

This is a comple-time overload and no runtime cost at all, although you have to write a long list for wrapper functions.

这篇关于如何在模板函数中使用CUBLAS库？的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！