跨 g++-nvcc 边界(包括内核)使用函数模板代码

Using function-templated code across the g++-nvcc boundary (including kernels)

假设我使用 NVIDIA CUDA 的 nvcc 编译器编译以下内容:

template<typename T, typename Operator>
__global__ void fooKernel(T t1, T t2)  {
    Operator op;
    doSomethingWith(t1, t2);
}

template<typename T>
__device__ __host__ void T bar(T t1, T t2)  {
    return t1 + t2;
}

template<typename T, typename Operator>
void foo(T t1, T t2)  {
    fooKernel<<<2, 2>>>(t1, t2);
}

// explicit instantiation
template decltype(foo<int, bar<int>>) foo<int, bar<int>);

现在,我希望我的 gcc、非 nvcc 代码调用 foo():

...

template<typename T, typename Operator> void foo(T t1, T t2);


foo<int, bar<int>> (123, 456);
...

我在用 CUDA 编译的 .o/.a/.so 文件中有适当的 (?) 实例化。

我能做到吗?

这里的问题是模板代码通常在使用的地方实例化,这不起作用,因为 foo 包含一个内核调用,g++ 无法解析。您显式实例化模板并为主机编译器转发声明它的方法是正确的。下面是如何做到这一点。我稍微修正了你的代码并将其分成 3 个文件:

  1. gpu.cu
  2. gpu.cuh
  3. cpu.cpp

gpu.cuh

此文件包含供 gpu.cu 使用的模板代码。我在您的 foo() 函数中添加了一些用途以确保其正常工作。

#pragma once
#include <cuda_runtime.h>

template <typename T>
struct bar {
    __device__ __host__ T operator()(T t1, T t2)
    {
        return t1 + t2;
    }
};

template <template <typename> class Operator, typename T>
__global__ void fooKernel(T t1, T t2, T* t3)
{
    Operator<T> op;
    *t3 = op(t1, t2);
}

template <template <typename> class Operator, typename T>
T foo(T t1, T t2)
{
    T* t3_d;
    T t3_h;
    cudaMalloc(&t3_d, sizeof(*t3_d));
    fooKernel<Operator><<<1, 1>>>(t1, t2, t3_d);
    cudaMemcpy(&t3_h, t3_d, sizeof(*t3_d), cudaMemcpyDeviceToHost);
    cudaFree(t3_d);
    return t3_h;
}

gpu.cu

此文件仅实例化 foo() 函数以确保它可用于 linking:

#include "gpu.cuh"

template int foo<bar>(int, int);

cpu.cpp

在这个纯 C++ 源文件中,我们需要确保我们没有获得模板实例化,因为那样会产生编译错误。相反,我们只转发声明结构 bar 和函数 foo。代码如下所示:

#include <cstdio>

template <template <typename> class Operator, typename T>
T foo(T t1, T t2);

template <typename T>
struct bar;

int main()
{
    printf("%d \n", foo<bar>(3, 4));
}

生成文件

这会将所有代码放入可执行文件中:

.PHONY: clean all
all: main

clean:
        rm -f *.o main

main: gpu.o cpu.o
        g++ -L/usr/local/cuda/lib64 $^ -lcudart -o $@

gpu.o: gpu.cu
        nvcc -c -arch=sm_20 $< -o $@

cpu.o: cpu.cpp
        g++ -c $< -o $@

设备代码由 nvcc 编译,主机代码由 g++ 编译,并且全部由 g++ 编辑 link。在 运行 后,您会看到漂亮的结果:

7

这里要记住的关键是内核启动和内核定义必须在 nvcc 编译的 .cu 文件中。为了将来参考,我还将在 separation of linking and compilation with CUDA.

上将此 link 留在这里