4. CuPy 数学运算：算术运算、广播机制、通用函数（ufunc）、聚合运算

聊到 GPU 加速，核心就是算得快。怎么个快法？

说白了，就是把循环扔给成千上万个核心并行跑。CuPy 的数学运算，就是帮你干这事的主力。

我个人习惯，上手新库先看它的运算体系是否顺手。CuPy 这块做得挺自然，跟 NumPy 几乎镜像，但底层跑在 CUDA 上。

4.1 算术运算：逐元素操作

最基本的加减乘除，在 CuPy 里就是逐元素操作。你想想看，两个同样形状的数组，对应位置算一下。

我在项目中遇到过几次，有人把 CPU 上的 for 循环直接搬到 GPU 上写，那效率……嗯，惨不忍睹。用 CuPy 的算术运算，一行顶你十行 for。

import cupy as cp

a = cp.array([1, 2, 3, 4])
b = cp.array([10, 20, 30, 40])

c_add = a + b
c_sub = a - b
c_mul = a * b
c_div = b / a

print(c_add)  # [11 22 33 44]
print(c_mul)  # [10 40 90 160]

这些运算，CuPy 内部会生成 CUDA kernel，每个元素一个线程。数据量越大，优势越明显。

4.2 广播机制：不同形状也能算

运算时形状不匹配怎么办？CuPy 沿用了 NumPy 的广播机制。规则很简单：从尾部维度开始对齐，要么相等，要么其中一个为 1。

举个例子，一个 3x4 的矩阵加上一个长度为 4 的向量。向量会在行方向自动复制 3 份。

matrix = cp.ones((3, 4))
vector = cp.array([1, 2, 3, 4])

result = matrix + vector
print(result)
# [[2. 3. 4. 5.]
#  [2. 3. 4. 5.]
#  [2. 3. 4. 5.]]

广播规则：

• 从后往前比较各维度的尺寸。
• 如果两个维度相等，继续比较前一个维度。
• 如果其中一个维度为 1，它会扩展到另一个维度的大小。
• 如果两个维度不等且都不为 1，直接报错。

4.3 通用函数（ufunc）：向量化的核心

通用函数，简称 ufunc。它是对数组每个元素执行同一操作的函数。CuPy 的 ufunc 直接编译为 CUDA 核函数。

我习惯把 ufunc 分为几类：数学运算、三角函数、指数对数、类型转换。写模型训练时用得最多的是 cp.exp、cp.log 和 cp.sqrt。

x = cp.array([0.0, 1.0, 2.0, 3.0])
y = cp.exp(x)       # 指数
z = cp.log(x + 1)   # 对数
w = cp.sqrt(x)      # 平方根

print(y)
print(z)
print(w)

CuPy 还支持创建自定义 ufunc，用 cp.ElementwiseKernel 或 cp.RawKernel。这块稍微进阶，但当你需要特殊的逐元素操作时，非常有用。

4.4 聚合运算：从数组到标量

聚合运算，就是沿着某个轴把数据“压”成一个更小的结果。求和、均值、最大值、最小值，是最常用的。

CuPy 的聚合运算会利用 GPU 的 reduction 操作，并行归约。我刚开始用的时候，一度怀疑结果算错了，因为太快了。

data = cp.random.randn(1000000).astype(cp.float32)

total = cp.sum(data)        # 所有元素求和
mean_val = cp.mean(data)    # 均值
max_val = cp.max(data)      # 最大值
min_val = cp.min(data)      # 最小值
std_val = cp.std(data)      # 标准差

print(f"Sum: {total:.4f}")
print(f"Mean: {mean_val:.4f}")
print(f"Max: {max_val:.4f}")
print(f"Std: {std_val:.4f}")

沿着轴聚合：

matrix = cp.random.randn(4, 5)

row_sum = cp.sum(matrix, axis=1)   # 每一行求和，结果形状 (4,)
col_sum = cp.sum(matrix, axis=0)   # 每一列求和，结果形状 (5,)

print("行求和:", row_sum)
print("列求和:", col_sum)

CuPy 支持几乎所有的 NumPy 聚合函数：cp.prod、cp.cumsum、cp.argmax、cp.argmin。用法完全一致。

最后提一句：CuPy 的数学运算，大部分是“非就地”的，即原数组不变，返回新数组。如果你要节省显存，可以用 out 参数或原地操作符（如 +=）。

a = cp.array([1, 2, 3])
b = cp.array([4, 5, 6])

# 输出到已有数组，避免重新分配内存
result = cp.empty_like(a)
cp.add(a, b, out=result)

# 原地操作
a += b   # a 被修改
print(a) # [5 7 9]