CUDA Thrust编译报错：__device__ lambda返回类型无法在主机代码查询

编译器错误信息 #

static_assert failed: 'Attempt to use an extended device lambda in a context that requires querying its return type in host code. Use a named function object, a host device lambda, or cuda::proclaim_return_type instead.'

出错代码 #

thrust::device_vector<float2> a;
thrust::device_vector<float> b;

float param1, param2;
float2 param3;

thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::make_tuple(a.begin(), b.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(a.end(), b.end())),
    [param1, param2, param3] __device__ (thrust::tuple<float2, float> const& tuple)
    {
        /* do something and return a float2 */
    },
    float2{},
    [] __device__ (float2 const& first, float2 const& second)
    {
        float2 result{};
        result.x = first.x + second.x;
        result.y = first.y + second.y;
        return result;
    });

用户疑问 #

作为CUDA/Thrust新手，不知道如何修改代码使其编译通过。另外，是否必须使用这个复杂的自定义归约函数？thrust::plus<float2>无法正常工作。

1. 修复编译错误 #

编译器报错的核心是：Thrust在主机端需要推导lambda的返回类型，但仅用__device__限定的lambda无法被主机端解析类型信息。以下三种方法可解决：

方法一：给lambda添加__host__ __device__限定符 #

让lambda同时支持主机和设备端执行，主机端就能正常推导返回类型：

thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::make_tuple(a.begin(), b.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(a.end(), b.end())),
    [param1, param2, param3] __host__ __device__ (thrust::tuple<float2, float> const& tuple)
    {
        /* do something and return a float2 */
    },
    float2{},
    [] __host__ __device__ (float2 const& first, float2 const& second)
    {
        float2 result{};
        result.x = first.x + second.x;
        result.y = first.y + second.y;
        return result;
    });

方法二：使用命名函数对象 #

如果不想让lambda在主机端执行，可定义带__device__限定的结构体函数对象：

struct TransformOp {
    float param1;
    float param2;
    float2 param3;

    __device__ float2 operator()(thrust::tuple<float2, float> const& tuple) const {
        /* do something and return a float2 */
    }
};

struct ReduceOp {
    __device__ float2 operator()(float2 const& first, float2 const& second) const {
        float2 result{};
        result.x = first.x + second.x;
        result.y = first.y + second.y;
        return result;
    }
};

// 调用transform_reduce
thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::make_tuple(a.begin(), b.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(a.end(), b.end())),
    TransformOp{param1, param2, param3},
    float2{},
    ReduceOp{});

方法三：用cuda::proclaim_return_type显式指定返回类型 #

通过模板强制告诉主机端lambda的返回类型，保留__device__限定：

thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::make_tuple(a.begin(), b.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(a.end(), b.end())),
    cuda::proclaim_return_type<float2>([param1, param2, param3] __device__ (thrust::tuple<float2, float> const& tuple) {
        /* do something and return a float2 */
    }),
    float2{},
    cuda::proclaim_return_type<float2>([] __device__ (float2 const& first, float2 const& second) {
        float2 result{};
        result.x = first.x + second.x;
        result.y = first.y + second.y;
        return result;
    }));

2. 解决thrust::plus<float2>失效问题 #

thrust::plus<float2>无法工作是因为CUDA的float2没有默认重载加法运算符。只需手动重载operator+，就能直接使用thrust::plus<float2>：

__host__ __device__ float2 operator+(float2 const& lhs, float2 const& rhs) {
    float2 res{};
    res.x = lhs.x + rhs.x;
    res.y = lhs.y + rhs.y;
    return res;
}

之后transform_reduce可简化为：

thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::make_tuple(a.begin(), b.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(a.end(), b.end())),
    [param1, param2, param3] __host__ __device__ (thrust::tuple<float2, float> const& tuple) {
        /* do something and return a float2 */
    },
    float2{},
    thrust::plus<float2>());

不用再写自定义归约函数。

内容的提问来源于stack exchange，提问作者0xbadf00d