CUDA 7.5 experimental __host__ __device__ lambdas

我玩了一下实验设备lambdas，它是在CUDA 7.5中引入的，Mark Harris在这篇博文中进行了推广。

对于下面的例子，我删除了很多不需要显示我的问题的东西(我的实际实现看起来更好一点…)。

我试着写一个foreach函数，根据模板参数对设备上的向量(每个元素1个线程)或主机(串行)进行操作。有了这个foreach函数，我可以很容易地实现BLAS函数。作为一个例子，我将一个标量赋值给向量的每个组件(我将完整的代码附加在最后):

template<bool onDevice> void assignScalar( size_t size, double* vector, double a )
{
    auto assign = [=] __host__ __device__ ( size_t index ) { vector[index] = a; };
    if( onDevice )
    {
        foreachDevice( size, assign );
    }
    else
    {
        foreachHost( size, assign );
    }
}

但是，由于__host__ __device__ lambda:

lambda的闭包类型("lambda ->void")不能在__global__函数模板实例化的模板实参类型中使用，除非lambda在__device__或__global__函数中定义

如果我从lambda表达式中删除__device__，我得到同样的错误，如果我删除__host__(仅__device__ lambda)，我没有得到编译错误，但在这种情况下，主机部分没有执行…

如果我将lambda分别定义为__host__或__device__，代码将按预期编译并工作。

template<bool onDevice> void assignScalar2( size_t size, double* vector, double a )
{
    if( onDevice )
    {
        auto assign = [=] __device__ ( size_t index ) { vector[index] = a; };
        foreachDevice( size, assign );
    }
    else
    {
        auto assign = [=] __host__ ( size_t index ) { vector[index] = a; };
        foreachHost( size, assign );
    }
}

然而，这引入了代码复制，实际上使使用lambdas的整个想法在本例中毫无用处。

是否有一种方法可以完成我想做的事情，或者这是实验功能中的错误?实际上，在编程指南的第一个示例中明确提到了定义__host__ __device__ lambda。即使对于那个更简单的例子(只是从lambda返回一个常量值)，我也找不到在主机和设备上同时使用lambda表达式的方法。

下面是完整的代码，使用选项-std=c++11 --expt-extended-lambda编译:

#include <iostream>
using namespace std;
template<typename Operation> void foreachHost( size_t size, Operation o )
{
    for( size_t i = 0; i < size; ++i )
    {
        o( i );
    }
}
template<typename Operation> __global__ void kernel_foreach( Operation o )
{
    size_t index = blockIdx.x * blockDim.x + threadIdx.x;
    o( index );
}
template<typename Operation> void foreachDevice( size_t size, Operation o )
{
    size_t blocksize = 32;
    size_t gridsize = size/32;
    kernel_foreach<<<gridsize,blocksize>>>( o );
}
__global__ void printFirstElementOnDevice( double* vector )
{
    printf( "dVector[0] = %fn", vector[0] );
}
void assignScalarHost( size_t size, double* vector, double a )
{
    auto assign = [=] ( size_t index ) { vector[index] = a; };
    foreachHost( size, assign );
}
void assignScalarDevice( size_t size, double* vector, double a )
{
    auto assign = [=] __device__ ( size_t index ) { vector[index] = a; };
    foreachDevice( size, assign );
}
// compile error:
template<bool onDevice> void assignScalar( size_t size, double* vector, double a )
{
    auto assign = [=]  __host__ __device__ ( size_t index ) { vector[index] = a; };
    if( onDevice )
    {
        foreachDevice( size, assign );
    }
    else
    {
        foreachHost( size, assign );
    }
}
// works:
template<bool onDevice> void assignScalar2( size_t size, double* vector, double a )
{
    if( onDevice )
    {
        auto assign = [=] __device__ ( size_t index ) { vector[index] = a; };
        foreachDevice( size, assign );
    }
    else
    {
        auto assign = [=] __host__ ( size_t index ) { vector[index] = a; };
        foreachHost( size, assign );
    }
}
int main()
{
    size_t SIZE = 32;
    double* hVector = new double[SIZE];
    double* dVector;
    cudaMalloc( &dVector, SIZE*sizeof(double) );
    // clear memory
    for( size_t i = 0; i < SIZE; ++i )
    {
        hVector[i] = 0;
    }
    cudaMemcpy( dVector, hVector, SIZE*sizeof(double), cudaMemcpyHostToDevice );
    assignScalarHost( SIZE, hVector, 1.0 );
    cout << "hVector[0] = " << hVector[0] << endl;
    assignScalarDevice( SIZE, dVector, 2.0 );
    printFirstElementOnDevice<<<1,1>>>( dVector );
    cudaDeviceSynchronize();
    assignScalar2<false>( SIZE, hVector, 3.0 );
    cout << "hVector[0] = " << hVector[0] << endl;
    assignScalar2<true>( SIZE, dVector, 4.0 );
    printFirstElementOnDevice<<<1,1>>>( dVector );
    cudaDeviceSynchronize();
//  assignScalar<false>( SIZE, hVector, 5.0 );
//  cout << "hVector[0] = " << hVector[0] << endl;
//
//  assignScalar<true>( SIZE, dVector, 6.0 );
//  printFirstElementOnDevice<<<1,1>>>( dVector );
//  cudaDeviceSynchronize();
    cudaError_t error = cudaGetLastError();
    if(error!=cudaSuccess)
    {
        cout << "ERROR: " << cudaGetErrorString(error);
    }
}

我使用CUDA 7.5的生产版本。

我尝试了第三个版本的assignScalar函数:

template<bool onDevice> void assignScalar3( size_t size, double* vector, double a )
{
#ifdef __CUDA_ARCH__
#define LAMBDA_HOST_DEVICE __device__
#else
#define LAMBDA_HOST_DEVICE __host__
#endif
    auto assign = [=] LAMBDA_HOST_DEVICE ( size_t index ) { vector[index] = a; };
    if( onDevice )
    {
        foreachDevice( size, assign );
    }
    else
    {
        foreachHost( size, assign );
    }
}

编译并运行没有错误，但是没有执行设备版本(assignScalar3<true>)。实际上，我认为__CUDA_ARCH__将始终是未定义的(因为函数不是__device__)，但我明确检查了有一个编译路径，它被定义。