PyTorch11/2 15:5484 分精选

PyTorch 将 isin 从 MPSGraph 迁移到 Metal(#202930)

viable/strict/1793620186: [MPS] Migrate isin to Metal (#202930)

阅读原文

推荐理由:我认为值得看:它给出算子级迁移的完整取舍与实测数据,11.14 倍几何平均加速可直接作为端侧推理优化参考。

PyTorch 合并 PR #202930,把 isin 算子从 MPSGraph 迁移到 Metal 实现。迁移原因之一是 MPSGraph 会生成 N×M 的 bool 张量,缓存时额外占用内存。新实现包含三个内核:isin_ 线性扫描,用于测试元素不超过 8 个的情况;isin_sorted_ 由主机排序测试元素、内核做二分查找;isin_mark_ 加 isin_gather_ 用于测试元素超过元素数 2 倍的情况,每个测试元素一个线程搜索并写 int32 标志,再由每个元素一个线程读取。作者称 MPS 未复用 CPU/CUDA 共享实现,是因为其 isin_sorting 调用 _unique 会与 CPU 同步,小输入性能降至 MPSGraph 的 0.14 倍。在 Apple M5 Pro、macOS 27.2 上,76 个用例的 Metal 实现相对 MPSGraph 几何平均加速 11.14 倍,最小 2.29 倍,最大 233.47 倍。

来源与依据

时间证据
11/2 15:54
收录于 11/2 21:30
交叉线索
1 个独立来源
内容可信度:高
证据链PyTorch