如何用CuPy高效计算复数的绝对值平方？

当需要比较复数的模长（本质为$\sqrt{\text{real}^2 + \text{imag}^2}$）以找出最大绝对值时，只需比较绝对值的平方即可——省去耗时的sqrt()操作能有效提升计算速度。

我针对这个需求做了基准测试，对比不同实现方式的性能：仅计算绝对值的abs_only()版本耗时4381.648微秒，先计算绝对值存入临时变量再平方的abs_sq_temp()版本耗时4744.022微秒，GPU上的平方操作仅增加了362微秒。由此推测，如果能直接对复数的实部、虚部就地平方再相加，高效的绝对值平方计算总耗时应该能接近362微秒或约为其两倍，但我不清楚怎么用CuPy实现这种操作。

遗憾的是，NumPy与CuPy均未内置复数绝对值平方的计算函数。这类函数的实现逻辑大概率和np.absolute()、cp.absolute()一致，只是省去最后的平方根步骤，且速度会比.abs()更快。

基准测试代码 #

import cupy as cp
from cupyx.profiler import benchmark
import time
import numpy as np

# 生成大型复数数组
arr = cp.random.random(10000000) + 1j * cp.random.random(10000000)

def abs_only():
    return cp.absolute(arr)

def abs_sq():
    return cp.absolute(arr)**2

def abs_sq_temp():
    temp = cp.absolute(arr)
    return temp*temp

def conj():
    return arr*cp.conj(arr)

def real_imag():
    return cp.real(arr)**2 + cp.imag(arr)**2

# 执行基准测试
bench0 = benchmark(abs_only, n_repeat=20)
bench1 = benchmark(abs_sq, n_repeat=20)
bench2 = benchmark(abs_sq_temp, n_repeat=20)
bench3 = benchmark(conj, n_repeat=20)
bench4 = benchmark(real_imag, n_repeat=20)

print(bench0)
print(bench1)
print(bench2)
print(bench3)
print(bench4)

# CPU版本对比（NumPy）
arr2 = np.random.random(10000000) + 1j * np.random.random(10000000)
start_time = time.time()
plain_abs_numpy = np.abs(arr2)
time_plain_abs = time.time() - start_time
print(f"\nOutside the benchmark() function, CPU takes {time_plain_abs*1e6:.3f} us with np.abs()")

测试结果 #

abs_only            GPU-0:  4381.648 us   +/- 292.389 (min:  4261.536 / max:  5529.536) us
abs_sq              GPU-0: 21369.104 us   +/- 180.435 (min: 21300.129 / max: 22085.632) us
abs_sq_temp         GPU-0:  4744.022 us   +/- 38.082 (min:  4694.080 / max:  4829.056) us
conj                GPU-0:  7396.042 us   +/- 60.760 (min:  7289.728 / max:  7508.256) us
real_imag           GPU-0: 38408.486 us   +/- 211.628 (min: 38300.770 / max: 39266.209) us

Outside the benchmark() function, CPU takes 32956.123 us with np.abs()

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