3、CuPy索引与切片：整数索引、布尔索引、花式索引、切片性能对比

索引和切片，是任何数组操作的基础。在 CPU 上你早就玩熟了，对吧？但到了 GPU 上，事情会有点不一样。我刚开始用 CuPy 的时候，直接把 NumPy 的索引习惯搬过来，结果性能惨不忍睹。后来才明白——GPU 的内存访问模式对索引方式极其敏感。选对了，速度飞起；选错了，还不如 CPU。

这一节，我把 CuPy 里四种主流索引手法掰开揉碎，一个个讲清楚。每种都附上我踩过的坑和实际项目的调优经验。

整数索引：最直接的定位

整数索引就是指定行和列的具体下标。比如 arr[2, 3] 取第 3 行第 4 列的元素。CuPy 支持跟 NumPy 一模一样的写法：

import cupy as cp

arr = cp.arange(24).reshape(4, 6)
print(arr[2, 3])   # 输出 15

这里有个小陷阱。单元素索引在 GPU 上会触发一次 kernel launch，开销其实不小。如果你需要频繁访问单个元素，我建议一次性读成 host 数组再做操作。我在做流式数据处理时就吃过这个亏——每次迭代取一个标量，慢到怀疑人生。

个人习惯：能用向量化操作就别用标量索引。实在要取单个值，用 .get() 转成 NumPy 标量再处理，避免反复触发 GPU 通信。

布尔索引：条件筛选的利器

布尔索引用一个 True/False 数组做掩码，选出符合条件的元素。CuPy 在这块的执行效率非常高，因为它能直接用 GPU 的并行比较指令：

arr = cp.random.randn(1000)
mask = arr > 0.5
positive_large = arr[mask]   # 布尔索引
print(positive_large.shape)

说实话，我项目里有一半的数据过滤场景都用布尔索引。它比写循环快上两个数量级。不过注意一点：布尔索引返回的是副本，不是视图。你改返回的数组不会影响原数组。

我曾经在一个在线学习系统里用布尔索引做样本筛选，每次只取异常值做更新。arr[mask] 会分配新显存，如果 mask 很稀疏（比如只有 1% 的 True），那浪费就大了。这时候我换成了 cp.where 配合 cp.take，显存占用更可控。

花式索引：任意顺序选元素

花式索引允许你用一个整数数组指定要取的下标。CuPy 跟 NumPy 一样支持：

arr = cp.arange(12).reshape(3, 4)
indices = cp.array([0, 2])
rows = arr[indices]   # 取第0行和第2行
print(rows)

这里有个知识点：花式索引在 CuPy 中会触发 gather 操作，本质上是一个数据重排。如果下标数组无序，GPU 的显存访问可能不连续，导致带宽利用率下降。我测试过，按有序下标（如 [0,1,2,3]）和乱序下标（如 [3,1,0,2]）的性能差能到 30%。原因是 GPU 的 warp 在访问连续地址时能合并内存事务。

避坑指南：花式索引返回的永远是副本，无法通过改变返回数组来影响原数组。如果你需要原地修改元素，用 cp.scatter_add 或手动写 kernel，别用花式索引赋值。

切片：连续块的极致优化

切片是 CuPy 里最友好的操作。因为它直接对应 GPU 的连续内存区域，不需要额外复制，返回的是视图：

arr = cp.ones((1024, 1024), dtype=cp.float32)
block = arr[100:200, 50:150]   # 视图，零拷贝
block[:, :] = 0                # 直接修改原数组

切片的高效来自于 GPU 的内存合并访问——所有线程按顺序读连续地址，带宽几乎打满。我实测过，对一个大数组做切片赋值，比用布尔索引快 5~8 倍。所以只要是规则连续的区域，我首选切片。

但要注意，切片如果带步长（如 arr[::2, ::3]），性能会下降，因为地址不再连续。我在生成图像特征图时就碰到过——步长切片导致内存访问不规整，显存带宽上不去。后来改成了 cp.reshape 和 cp.take 的组合，速度反而更快。

性能对比：选对工具，事半功倍

下面这张表是我在相同数据量（1000 万元素，float32）下测试的对比结果。测试环境是 RTX 3090，CuPy 12.0。

索引方式返回类型内存布局相对耗时适用场景
整数索引（单元素）标量随机访问极慢仅调试用
布尔索引（稀疏 mask）副本条件筛选中等过滤异常、筛选子集
花式索引（有序下标）副本连续较快重排、打乱数据
花式索引（乱序下标）副本随机慢尽量避免
切片（连续无步长）视图连续极快截取子块、分块处理
切片（带大步长）视图跨步较慢降采样、隔点取样

你看，切片的无步长版本几乎是零开销，这是 GPU 最喜欢的方式。布尔索引和有序花式索引也还不错。最该避开的，就是单元素索引和乱序花式索引——它们会让 GPU 内存控制器忙不过来。

说到这里，我把四种索引的访问关系画成一张图，方便你对比理解。

CuPy索引与切片分类图 CuPy 索引与切片分类 GPU 数组 整数索引 单元素/标量 布尔索引 条件掩码筛选 花式索引 任意下标选择 切片（连续块） 性能总结 切片(无步长)>花式索引(有序)>布尔索引>花式索引(乱序)>整数索引(单元素)

实战建议与避坑

光知道语法还不够，我把自己踩过的坑总结一下，你照着做基本不会翻车：

  • 批量操作优先：别在 for 循环里写 arr[i]，用切片或布尔索引一次搞定。GPU 喜欢并行，讨厌串行。
  • 连续优先：能用切片切连续块，就别用花式索引打散。比如 arr[100:200] 比 arr[[100,101,...,199]] 快 10 倍以上。
  • 注意分配：布尔索引和花式索引会分配新显存。如果你的显存很紧，考虑用 cp.asnumpy 临时搬到 host 或者改用视图切片。
  • 调试辅助：如果索引结果不对，先用 .get() 转到 NumPy 上验证。GPU 上调试不方便，养成先检查逻辑的习惯。
  • 避免乱序花式索引：除非你真的需要随机重排，否则尽量让下标有序。比如排序之后再用花式索引，性能会好很多。

我的经验：在实现 K-Means 的样本分配步骤时，我用布尔索引取代了 for 循环判断类别，速度从 200ms 降到 12ms。本质就是把串行比较变成了并行 mask。索引方式选对，威力巨大。

好了，CuPy 索引与切片的精髓就在这里。你只需要记住：连续优先、批量优先、视图优先。下一次写 GPU 代码前，先想想这三点，性能基本不会太差。

← 上一章 📖 返回目录 下一章 →