CANN/cannbot-skills Triton条件优化-北京尧图网络

Where 条件优化单位置 False 替换为 get_element insert_slice【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills触发条件当 Agent 发现 Triton kernel 中存在tl.where语句且其 condition 满足以下所有条件时应考虑应用此优化condition 形式为X ! y其中X是关于地址偏移的向量如base tl.arange(0, BLOCK_SIZE)y是标量或常量仅有一个位置为 FalseX y仅在一个索引处成立即y在X的取值范围内仅出现一次A、B 和 X 的 shape 相同向量维度一致替换逻辑清晰原tl.where(condition, A, B)表示当X ! y时选 A否则选 B直观理解当y在[base, base BLOCK_SIZE)范围内时X y仅在偏移量_offs y - base处成立因此where只在一个位置选择 B其他位置都选 A。核心知识为什么需要优化tl.where在 NPU 上会导致离散访存退化为 scalar 操作。当 condition 是逐元素比较时Vector 单元需要对每个元素分别判断条件并选择结果无法充分利用向量化内存访问严重影响性能。当 condition 仅在一个位置为 False 时where的语义等价于在 A 的某个位置替换为 B 的对应值。这个操作可以用更高效的get_elementinsert_slice组合来实现get_element从 B 中提取单个元素标量操作但只执行一次insert_slice将单个元素写入 A 的指定位置向量化操作只修改一个位置关键前提前提说明condition 仅一个位置为 False这是此优化的核心前提多个位置为 False 时不适用y 可能在范围外当y不在[base, base BLOCK_SIZE)时condition 恒为 True不需要任何操作dtype 一致tl.full创建的 tensor 必须与 A 的数据类型一致extension 回退机制get_element和insert_slice来自triton.language.extra.cann.extension。如果该模块不可用需要回退到tlTriton 核心库也支持这两个操作try: import triton.language.extra.cann.extension as extension except Exception: extension tl代码模式优化前tl.where 退化为 scalar 计算X base tl.arange(0, BLOCK_SIZE) A tl.load(A_ptr X) B tl.load(B_ptr X) condition (X ! y) A tl.where(condition, A, B)问题tl.where(condition, A, B)中 condition 是逐元素比较导致离散访存退化为 scalar 操作。优化后get_element insert_slicetry: import triton.language.extra.cann.extension as extension except Exception: extension tl X base tl.arange(0, BLOCK_SIZE) A tl.load(A_ptr X) B tl.load(B_ptr X) if base y base BLOCK_SIZE: _offs y - base _val extension.get_element(B, (_offs,)) _tensor tl.full((1,), _val, dtypeA.dtype) A extension.insert_slice( A, _tensor, offsets(_offs,), sizes(1,), strides(1,) )改进消除了tl.where的离散访存改为单次get_element提取单次insert_slice写入保持向量化访存。关键代码解析步骤代码作用1if base y base BLOCK_SIZE:判断 y 是否在当前 block 范围内2_offs y - base计算 y 在 block 内的偏移3extension.get_element(B, (_offs,))从 B 中提取_offs位置的值4tl.full((1,), _val, dtypeA.dtype)创建形状为 (1,) 的 tensor用于 insert_slice5extension.insert_slice(A, _tensor, ...)将 _tensor 插入 A 的_offs位置if 条件不满足时当y不在[base, base BLOCK_SIZE)范围内时X ! y恒为 Truetl.where的结果就是 A 本身不需要任何操作也不用保留原有的tl.where逻辑。910_95 特别注意Reg-based 架构下的 where 行为910_95 属于 Reg-based 架构支持 SIMT VF 模式。在 SIMT 模式下tl.where的退化行为与 910BMem-based 架构不同架构where 退化行为910B (Mem-based)where 退化为标量循环性能损失 10x-100x910_95 (Reg-based, SIMD VF)where 同样退化为标量操作性能损失严重910_95 (Reg-based, SIMT VF)通用算术操作在 SIMT VF 模式下会降级where 场景更复杂结论无论哪种架构单位置 False 的 where 都应替换为get_elementinsert_slice。get_element / insert_slice 的数据类型支持910_95 上get_element和insert_slice支持的数据类型int8int16int32uint8uint16uint32uint64int64fp16fp32bf16boolYYYYYYYYYYYN注意bool类型不支持这两个操作。如果 A/B 是 bool 类型不能使用此优化。insert_slice 参数约束insert_slice的参数必须满足以下约束参数约束offsets长度必须与张量维度数相同sizes长度必须与张量维度数相同且sub的 shape 必须与sizes匹配strides长度必须与张量维度数相同插入区域不能超出ful张量的边界对于 1D 场景本优化最常见的场景参数为offsets(_offs,)- 1D 偏移sizes(1,)- 插入 1 个元素strides(1,)- 步长为 1多维场景扩展当 A/B 是多维 tensor 时get_element和insert_slice的索引需要对应每个维度if base_m y_m base_m BLOCK_M and base_n y_n base_n BLOCK_N: _offs_m y_m - base_m _offs_n y_n - base_n _val extension.get_element(B, (_offs_m, _offs_n)) _tensor tl.full((1, 1), _val, dtypeA.dtype) A extension.insert_slice( A, _tensor, offsets(_offs_m, _offs_n), sizes(1, 1), strides(1, 1) )相关文档标量降级规避指南 - where 退化与标量降级的关系内存访问模式优化指南 - 离散访存 vs 连续访存硬件速查手册 - 910_95 Reg-based 架构差异get_element API 文档 - get_element 参数规格与数据类型支持insert_slice API 文档 - insert_slice 参数规格与约束vec_ops.py - get_element / insert_slice 实现源码【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

CANN/cannbot-skills Triton条件优化

相关推荐

Cursor Rules 配置必须今天完成！VS Code即将在v1.90移除legacy rule loader——仅剩72小时迁移窗口期

WarcraftHelper：魔兽争霸3终极优化工具，让你重拾经典游戏完美体验

SUID样式系统揭秘：从styled到sx属性的完整样式解决方案 [特殊字符]

更多资讯

Agent技能库设计实战：打造可复用、可观测的智能体能力体系

脑机接口技术：医疗与消费应用现状与挑战

2026汕头电气检测机构排名 TOP5 CMA 资质机构提供防爆设备检测+防爆安全检测联系方式推荐

后端开发避坑指南：奸人世家高频面试题与实战拆解

Mac本地部署Qwen Coder：从零开始玩转开源AI编程模型

PS通道抠图全解：原理、实操与复杂图像无痕处理

今日更新

2k显示屏性能优化踩坑：版本升级后API全变了，这份源码解析救了我

3步搞定美眉图实战项目，告别官方文档抓不住重点

3步搞定黄金大劫案项目搭建从入门到精通

本周精选

深入解析Transformer多头注意力机制与工程优化

OpenClaw 的 Skills 跑学习任务，模型通道改到 TaoToken 通道行不行？

ChatGPT报错Oops, an error occurred! 全链路排查指南

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？