CUDA 中 SFINAE 的坑

发表于： 更新于：
字数： 538 阅读：≈ 2分钟

TL;DR

如果想要写一些 traits，而且涉及到了函数和变量（而不是类型），则需要同时加上 __device__ 和 __host__ 以确保它在 device 侧和 host 侧都可见。否则会因为 SFINAE 不报错（实际上可能是在 device 侧找不到函数或变量），而且不知道为什么结果不对。

描述

项目用的语言标准是 C++17，atomicAdd 是 CUDA 中的一族函数，能对支持类型进行原子操作，其参数为一个地址和一个值，我想要判断 atomicAdd 是否有某个基本数据类型对应的重载版本。

template <typename T, typename U = std::decay_t<T>>
__device__ __host__ auto atomicAddTest(T&&)
    -> std::void_t<decltype(atomicAdd((U*)(nullptr), std::declval<U>()))>;

__device__ __host__ auto atomicAddTest(...) -> int;

template <typename T>
__device__ __host__ static inline constexpr auto atomicAddable =
    std::is_same_v<decltype(atomicAddTest(std::declval<T>())), void>;

static_assert(atomicAddable<float>);
static_assert(atomicAddable<unsigned>);

看起来很正常对不对？但是如果想要同时在 host 侧和 device 侧来使用 atomicAddable 变量，但是又忘记给这些函数加 __device__ __host__，就会有 calling device function from host 的错误，然后因为 SFINAE 不报错，下面的程序实际使用时总是走了不一样的分支，半天不知道原因在哪里！后来加了 static_assert，同时暂时性地将 SFINAE 保护的部分放到函数体里面去，才通过 LSP 了解到具体的出错原因。

接下来编译又出现一个问题：

a __device__ variable "atomicAddable [with T=DType]" cannot be directly read in a host function

所以还是把变量改成 Trait 类算了：

template <typename T, typename U = std::decay_t<T>>
__device__ __host__ auto atomicAddTest(T&&)
    -> std::void_t<decltype(atomicAdd((U*)(nullptr), std::declval<U>()))>;

__device__ __host__ auto atomicAddTest(...) -> int;

template <typename T, typename U = void>
struct AtomicAddable : std::false_type {};

template <typename T>
struct AtomicAddable<
    T,
    std::enable_if_t<std::is_same_v<decltype(atomicAddTest(std::declval<T>())), void>>>
    : std::true_type {};

static_assert(AtomicAddable<float>::value);
static_assert(AtomicAddable<int>::value);
static_assert(AtomicAddable<unsigned>::value);

在 C++20 有了 concept 之后，这种情况可能会好很多，因为 concept 的结果是右值，而且不用写好几个函数迂回地描述一件事情，也就不会存在这种到底是 host 侧还是 device 侧的问题。