什么是CUDA数学库等效的arg()?

Question

c++标准库有一个叫做std::arg()的函数。它返回其复数参数的相位角(或角分量)，以弧度表示。

我需要在CUDA内核中调用std::arg()之类的东西。在CUDA数学库或CUDA复杂数学库中是否有等效的函数?

Answer 1

对注释进行总结和稍微扩展:

• CUDA数学库没有直接支持arg()函数的复数值。
• cuComplex.h是非常初级的。正如@njuffa所说:它只有CUBLAS和CUFFT所要求的。
• 您可以，并且应该使用atan2(y, x)(或std::atan2(y, x))，其中y是虚分量，x是实分量。

那么，推力呢?毕竟，thrust::complex<T>有arg()方法吗?

嗯，它是这样实现的:

template <typename T>
__host__ __device__
T arg(const complex<T>& z)
{
// Find `atan2` by ADL.
using std::atan2;
return atan2(z.imag(), z.real());
}

(参见GitHub)…所以，atan2()也在那里。