使用模板实现 CUDA 设备阵列

Question

我正在尝试实现 Thrust 设备向量的固定大小版本。我编写了一些初始版本，但遇到了一个奇怪的模板错误。

代码如下：

#include <iostream>
#include <array>

enum class memcpy_t {
    host_to_host,
    host_to_device,
    device_to_host,
    device_to_device
};

template <typename T, std::size_t N>
struct cuda_allocator {
    using pointer = T*;

    static void allocate(T *dev_mem) {
        cudaMalloc(&dev_mem, N * sizeof(T)); 
    }

    static void deallocate(T *dev_mem) {
        cudaFree(dev_mem); 
    }

    template <memcpy_t ct>
    static void copy (T *dst, T *src) {
        switch(ct) {
        case memcpy_t::host_to_host:
            cudaMemcpy(dst, src, N * sizeof(T), cudaMemcpyHostToHost);
            break;
        case memcpy_t::host_to_device:
            cudaMemcpy(dst, src, N * sizeof(T), cudaMemcpyHostToDevice);
            break;
        case memcpy_t::device_to_host:
            cudaMemcpy(dst, src, N * sizeof(T), cudaMemcpyDeviceToHost);
            break;
        case memcpy_t::device_to_device:
            cudaMemcpy(dst, src, N * sizeof(T), cudaMemcpyDeviceToDevice);
            break;
        default:
            break;
        }
    }
};

template <typename T, std::size_t N>
struct gpu_array {
    using allocator = cuda_allocator<T, N>;
    using pointer = typename allocator::pointer;
    using value_type = T;
    using iterator = T*;
    using const_iterator = T const*;

    gpu_array() {
        allocator::allocate(data);
    }

    gpu_array(std::array<T, N> host_arr) {
        allocator::allocate(data);
        allocator::copy<memcpy_t::host_to_device>(data, host_arr.begin());
    }

    gpu_array& operator=(gpu_array const& o) {
        allocator::allocate(data);
        allocator::copy<memcpy_t::device_to_device>(data, o.begin());
    }

    operator std::array<T, N>() {
        std::array<T, N> res;
        allocator::copy<memcpy_t::device_to_host>(res.begin(), data);
        return res;
    }

    ~gpu_array() {
        allocator::deallocate(data);
    }

    __device__ iterator begin() { return data; }
    __device__ iterator end() { return data + N; }
    __device__ const_iterator begin() const { return data; }
    __device__ const_iterator end() const { return data + N; }

private:
    T* data;
};

template <typename T, std::size_t N>
__global__ void add_kernel(gpu_array<T,N> &r,
                           gpu_array<T,N> const&a1,
                           gpu_array<T,N> const&a2) {
    int i = blockIdx.x*blockDim.x + threadIdx.x;
    r.begin()[i] = a1.begin()[i] + a2.begin()[i];
}

template <typename T, std::size_t N>
gpu_array<T, N> operator+(gpu_array<T,N> const&a1,
                          gpu_array<T,N> const&a2)
{
    gpu_array<T, N> res;
    add_kernel<<<(N+255)/256, 256>>>(res, a1, a2);
    return res;
}

const int N = 1<<20;

int main() {
    std::array<float, N> x,y;

    for (int i = 0; i < N; i++) {
        x[i] = 1.0f;
        y[i] = 2.0f;
    } 

    gpu_array<float, N> dx{x};
    gpu_array<float, N> dy{y};

    std::array<float, N> res = dx + dy;

    for(const auto& elem : res) {
        std::cout << elem << ", ";
    }
}

可能还有很多其他错误，但我遇到了一个奇怪的错误。 nvcc 给我以下错误：

error: no match for 'operator<' (operand types are '<unresolved overloaded function    type>' and 'memcpy_t')
allocator::copy<memcpy_t::host_to_device>(data, host_arr.begin());

出于某种原因，它是否将我的枚举类模板参数视为operator<？顺便说一句，这是使用选项-arch=sm_70 -std=c++14 编译的。我没有受过关于 C++ 和 CUDA 如何交互的良好教育，所以我无法解决这个问题。

Answer 1

这有点让人摸不着头脑，但这里的根本问题是根据 C++ 标准的语法有缺陷。它是生成错误的主机编译器，据我所知，它这样做是完全正确的。有关所有血腥细节，请参阅here。

使用copy 特化的代码应如下所示：

gpu_array(std::array<T, N> host_arr) {
    allocator::allocate(data);
    allocator::template copy<memcpy_t::host_to_device>(data, host_arr.begin());
}

gpu_array& operator=(gpu_array const& o) {
    allocator::allocate(data);
    allocator::template copy<memcpy_t::device_to_device>(data, o.begin());
}

operator std::array<T, N>() {
    std::array<T, N> res;
    allocator::template copy<memcpy_t::device_to_host>(res.begin(), data);
    return res;
}

这可能是有史以来看起来最奇怪的语法，但它是使编译器将< 视为模板标记而不是运算符所必需的。修复代码中的任何地方，这个特定的编译器错误应该消失。