PyTorch中方括号索引与index_select的区别及CUDA非法内存报错问题

方括号索引与index_select的实现差异 #

• 功能适配逻辑：a[b]属于PyTorch高级索引实现，支持多维度混合索引（切片、布尔张量、整数张量的组合使用），内部会先完成多索引广播、维度对齐的通用处理，再生成输出张量；index_select是单维度张量索引的专用轻量实现，仅支持指定单个维度、传入1D整数张量完成索引，无额外的通用索引兼容处理逻辑。
• 反向传播实现：高级索引的反向传播依赖IndexBackward算子，为了处理重复索引的梯度累加需求，内部会调用CUDA侧的merge_sort核函数对索引做分组排序；index_select的反向传播逻辑更轻量化，直接通过IndexSelectBackward算子实现，无需对索引做排序处理，仅需按索引位置回填梯度即可。

报错底层原因 #

该问题是PyTorch 1.9.0与PyTorch3D 0.6.0搭配使用时的已知兼容性问题：

• Pulsar渲染器的底层CUDA内核调用会修改当前CUDA上下文的状态，触发PyTorch 1.9.0版本IndexBackward中merge_sort核函数的边界漏洞：该版本的merge_sort核函数在该异常CUDA上下文状态下会错误计算显存地址偏移，导致非法显存访问，触发你遇到的cudaErrorIllegalAddress报错。
• 由于index_select的反向实现完全不依赖merge_sort核函数，因此替换索引方式后问题会完全消失。该bug和业务逻辑无关，只要调用过Pulsar渲染器修改了CUDA上下文，后续所有高级索引的反向传播步骤都会触发报错，和你描述的复现特征完全匹配。

临时解决方案 #

如果不想修改原有索引写法，可将PyTorch版本升级至1.10.0及以上，搭配对应兼容版本的PyTorch3D即可修复该内核漏洞。

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