ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PyPTO 的 SIMT 原子操作指南:atomic_max 原子最大值归约实现与源码解析

PyPTO 的 SIMT 原子操作指南:atomic_max 原子最大值归约实现与源码解析 人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载导读pypto_pro.language.simt.atomic_max是 CANN PyPTO 的 SIMT单指令多线程编程范式中用于无锁并发聚合的核心原语它以原子方式比较目的操作数target的旧值与源操作数value并把较大者写回目的操作数。本指南围绕该接口的官方 API 文档展开覆盖产品支持情况、函数原型、参数与数据类型约束、返回值语义及完整调用示例并深入 IR 构建与解析实现 与 端到端测试帮助读者在 256 线程并行求取最大分数的典型场景中正确、安全地使用原子最大值操作。产品支持情况atomic_max属于 SIMT 原子操作家族其硬件支持范围与当前 AI 加速器的 SIMT 执行能力直接相关官方文档明确如下产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持从仓库测试代码看atomic_max的端到端用例通过pytest.mark.soc(950)标记并断言设备名包含Ascend950与文档中仅 Ascend 950 系列支持的结论一致参见 test_atomic_max.py 中的_require_a5()辅助函数。功能说明atomic_max以原子方式执行如下读-改-写序列读取目的操作数target的当前值将其与源操作数value比较取较大者将较大值写回target。整个过程对并发线程不可分割多个线程同时调用时不会产生数据竞争。值得注意的是对于DT_FP16与DT_BF16两种半精度数据类型该操作只更新目的操作数不返回旧值这一语义差异会直接影响返回值的使用方式详见返回值说明一节。函数原型pypto_pro.language.simt.atomic_max( target: Scalar, value: Scalar, ) - Optional[Scalar]该接口位于 SIMT 命名空间pl.simt.*下。在仓库中_simt_api.py提供了 Python 侧的类型声明与 docstring用于支持 IDE 跳转到定义、导入期拼写检查与静态类型校验这些声明函数并不会在运行时被真正调用——在 PyPTO kernel 内部AST 解析器会在 Python 真正执行前拦截每一个pl.simt.xxx调用并翻译为 IRkernel 外调用会抛出RuntimeError。参数说明参数输入/输出说明target输入目的操作数Scalar类型。必须直接传入 Tile 或 Tensor 的单元素下标访问表达式例如ub_tile[0, 0]或gm_tensor[0, 0]。-UB Tile必须位于 UBUnified Buffer使用 NDN 维描述格式支持DT_INT32、DT_UINT32、DT_FP16、DT_BF16、DT_FP32。-GM Tensor必须为 ND 格式支持DT_INT32、DT_UINT32、DT_FP16、DT_BF16、DT_FP32、DT_INT64、DT_UINT64。value输入源操作数Scalar类型表示用于比较的值。数据类型必须与target一致数值字面量按target的数据类型处理整数目的操作数不接受浮点字面量。从源码看 target 的校验逻辑必须直接传入下标表达式不是文档的泛泛之谈而是解析器的硬性约束。在 simt_ops.py 的_parse_atomic_call中target参数必须是ast.Subscript节点否则抛出InvalidVal并提示 target must be a direct Tile or Tensor subscript——这意味着不能先把元素读取成标量再传入例如current dst[0, 0]; pl.simt.atomic_max(current, value)会被拒绝若下标中包含切片dst[0:1, 0:1]会抛出NotSupported即原子操作只允许精确的整数下标选中单个元素容器类型必须是TileType或TensorType随后从容器类型取出dtype用于后续操作数校验。对应解析器单测见 test_simt_atomic.pyscalar_alias用例验证间接标量目标被拒绝slice_target用例验证切片目标被拒绝。value 的数据类型与字面量规则_parse_atomic_operandsimt_ops.py实现了按 target 数据类型上下文类型化字面量的规则若value是数值字面量且target为整数类型而字面量是浮点数则抛出InvalidType提示 requires an integer value合法字面量会被转换成与target相同的标量常量非字面量表达式如另一个标量则正常解析其数据类型必须在后端与target一致否则触发CommonExternal报错 operand 0 dtype must match target dtype。约束说明atomic_max只能在由pypto_pro.language.vector_function(modesimt)定义的SIMT 入口函数或辅助函数中调用。从解析器实现看_parse_atomic_call会检查当前函数类型是否为SimtVF或SimtCallee否则抛出InvalidOperation并提示应将逻辑移入pl.vector_function(modesimt)定义的函数体内同时要求恰好传入target与value两个位置参数不支持关键字参数。此外SIMT 入口函数必须通过simt_functhreads语法在with pl.section_vector():代码块内启动对应 IRsimt.launch从源码推断atomic_max作为 SIMT 上下文相关的操作同样遵循这一调度约束——在 调用示例 中可以看到这一标准结构。返回值说明atomic_max的返回值语义与target的数据类型强相关target数据类型为DT_FP16或DT_BF16时无返回值只能作为独立语句调用不能把结果赋给变量或作为表达式的一部分target数据类型为其他受支持类型如DT_INT32、DT_UINT32、DT_FP32、DT_INT64、DT_UINT64时返回更新前的target旧值返回值类型与target一致。这一行为在解析器单测中被显式验证test_half_precision_atomic_add_max_min_are_void_on_ub_tile与对应的 GM Tensor 用例test_simt_atomic.py断言半精度原子调用的 IR 类型为NoneTypetest_half_precision_atomic_result_cannot_be_returned则验证将半精度原子结果直接return会触发InvalidValmust return None or one scalar value。端到端测试test_atomic_max_returns_old_value_and_preserves_other_elementstest_atomic_max.py则验证了 INT32 场景下返回值确实是更新前的旧值且不改变target之外的其他元素。调用示例官方示例多线程求最大值归约以下示例来自官方 API 文档256 个线程并发对max_score执行原子最大值归约最终max_score[0, 0]收敛为scores中所有元素的最大值import pypto_pro.language as pl pl.vector_function(modesimt, max_threads256) def reduce_max( scores: pl.Tensor[[1, 256], pl.DT_INT32], max_score: pl.Tensor[[1, 1], pl.DT_INT32], ) - None: tid pl.simt.linear_thread_idx() pl.simt.atomic_max(max_score[0, 0], scores[0, tid]) pl.jit() def atomic_max_kernel( scores: pl.Tensor[[1, 256], pl.DT_INT32], max_score: pl.Tensor[[1, 1], pl.DT_INT32], ): with pl.section_vector(): reduce_max256示例要点拆解pl.vector_function(modesimt, max_threads256)将reduce_max声明为 SIMT 入口函数max_threads256声明线程规模上限pl.simt.linear_thread_idx()返回当前 block 内按 x 优先展开的扁平线程号对应 IRsimt.linear_thread_idx每个线程负责一个元素scores[0, tid]所有线程把各自的分数原子性地与max_score[0, 0]取最大值256 次原子更新后该元素即为全局最大值外层pl.jit()kernel 必须通过with pl.section_vector():代码块承载 SIMT 启动reduce_max256即simt.launch的源码形式。完整可运行变体全数据类型覆盖与旧值返回仓库端到端测试提供了更贴近工程实践的完整代码可在 Ascend 950 上直接运行。其 SIMT 函数同时覆盖 UB Tile 与 GM Tensor 的全部受支持数据类型test_atomic_max.pypl.vector_function(modesimt, max_threads64) def atomic_max_ub_all_dtypes( int32_tile, uint32_tile, fp16_tile, bf16_tile, fp32_tile, ): pl.simt.atomic_max(int32_tile[0, 0], 7) pl.simt.atomic_max(uint32_tile[0, 0], 7) pl.simt.atomic_max(fp16_tile[0, 0], 7.0) pl.simt.atomic_max(bf16_tile[0, 0], 7.0) pl.simt.atomic_max(fp32_tile[0, 0], 7.0)注意这里整型 Tile 使用整数字面量7浮点 Tile 使用浮点字面量7.0正对应参数说明中整数目的操作数不接受浮点字面量的规则。GM Tensor 场景下还额外支持int64与uint64pl.vector_function(modesimt, max_threads64) def atomic_max_gm_all_dtypes( int32_state, uint32_state, fp16_state, bf16_state, fp32_state, int64_state, uint64_state, ): pl.simt.atomic_max(int32_state[0, 0], 7) pl.simt.atomic_max(uint32_state[0, 0], 7) pl.simt.atomic_max(fp16_state[0, 0], 7.0) pl.simt.atomic_max(bf16_state[0, 0], 7.0) pl.simt.atomic_max(fp32_state[0, 0], 7.0) pl.simt.atomic_max(int64_state[0, 0], 7) pl.simt.atomic_max(uint64_state[0, 0], 7)如果希望读取更新前的旧值例如用于记录最大值的历史状态仅对非半精度类型可用且必须将返回值显式赋给标量目标pl.vector_function(modesimt, max_threads1) def atomic_max_return_value_gm( state: pl.Tensor[[1, 64], pl.DT_INT32], old_values: pl.Tensor[[1, 64], pl.DT_INT32], ): old_values[0, 0] pl.simt.atomic_max(state[0, 0], 12)对应的 jit kernel 通过with pl.section_vector():启动并在运行后调用torch.npu.synchronize()等待 NPU 同步再断言state[0, 0] 12、old_values[0, 0] 10同时验证target之外的其他元素保持初始值123不变test_atomic_max.py。UB Tile 的使用前提使用 UB Tile 作为target时需在 kernel 内先以pl.make_tile创建指定target_memorypl.MemorySpace.Vec的 Tile 并明确地址参见 test_atomic_max.py 中的addr0x0000等参数随后通过pl.load从 GM 加载、用pl.system.sync_src/pl.system.sync_dst完成 MTE2↔Vector 流水线同步最后以pl.store写回才能保证原子操作看到正确的数据。实现原理从 Python 调用到 SIMT 原子 IR理解atomic_max的底层路径有助于排查使用问题整条链路在仓库中清晰可查API 声明层pl.simt.atomic_max的 Python 签名与 docstring 定义在 _simt_api.pyAST 解析层kernel 编译时simt_ops.py 中注册的_parse_atomic_max解析器会调用通用流程_parse_atomic_call完成上文所述的下标形式、切片、容器类型与字面量类型校验IR 构建层校验通过后atomic_maxbuildersimt_ops.py将调用规约为simt.atomic_max操作操作数为[container, offset, value]其中container是 Tile/Tensor 容器、offset是解析后的单元素下标——这也是解析器单测中断言block.getval not in function_ir的原因原子操作直接作用于容器元素地址不会先执行一次普通的元素读取test_simt_atomic.pyIR 类型层半精度目标时调用结果类型为NoneType其余类型为与 target 一致的标量类型这一信息被后续代码生成阶段用于决定是否生成返回旧值的指令序列。与其他 SIMT 原子操作的关系atomic_max并非孤立接口而是 PyPTO SIMT 原子操作家族的一员。官方文档在 atomic/index.md 中列出了全部 11 个成员atomic_add、atomic_sub、atomic_exch、atomic_max、atomic_min、atomic_inc、atomic_dec、atomic_cas、atomic_and、atomic_or、atomic_xor覆盖数值累加、极值归约、计数器、比较交换与按位运算等场景。与最常用的 atomic_add 文档 对比可以发现atomic_max与其在产品支持、参数约束、返回值语义半精度无返回值、其余类型返回旧值上完全一致区别仅在于更新规则取大值 vs 累加。这 11 个操作在 simt_ops.py 中共享同一个_parse_atomic_call解析框架仅在操作名与操作数个数上有所不同如atomic_cas额外携带compare操作数这也是它们行为规则高度统一的原因。若你的场景需要求最小值可平移本文示例直接改用pl.simt.atomic_min其约束与用法完全对称。赞分享人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐CANN PyPTO SIMT 原子操作 atomic_exch原子交换的接口语义、数据类型约束与源码实现CANN PyPTO SIMT 原子操作 atomic_exch原子交换的接口语义、数据类型约束与源码实现 导读 pypto_pro.language.sim人工智能编译器模型编译高性能计算深度学习CANNCANN PyPTO 张量归约算子详解pypto.Tensor.amax 指定维度求最大值的原理、约束与实战CANN PyPTO 张量归约算子详解pypto.Tensor.amax 指定维度求最大值的原理、约束与实战 导读 pypto.Tensor.amax 是 C人工智能编译器模型编译高性能计算深度学习CANNPyPTO Tensor.softmax 算子详解维度归一化的 API 用法、数值原理与源码实现PyPTO Tensor.softmax 算子详解维度归一化的 API 用法、数值原理与源码实现 本文围绕 CANN PyPTO 张量库中的 Tensor.s人工智能编译器模型编译高性能计算深度学习CANN上一篇5个Dify工作流实战技巧如何用开源项目加速AI应用开发下一篇React Native View Shot 项目常见问题解决方案创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表