PyPTO 运行时动态配置接口 set_operation_optionscombine_axis 尾轴 broadcast inline 优化实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoPyPTOParallel Tensor/Tile Operation 编程范式的pypto.experimental.set_operation_options是编译框架提供的运行时动态配置管理功能核心接口它将原本静态写在tile_fwk_config.json中的operation.combine_axis参数转变为动态、可编程的指令。读完本文你将掌握该接口的函数原型、参数语义、作用域规则与底层实现链路从 Python API 到 C Pass 再到代码生成并能在 jit/loop 内核中以逐函数粒度控制尾轴 broadcast inline 优化。产品支持情况set_operation_options接口当前支持以下硬件平台以仓库文档标注为准Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持功能说明从静态配置到动态指令在 PyPTO 编译框架中大量编译与代码生成参数原本通过配置文件 tile_fwk_config.json 统一管理。其中operation段默认开启了尾轴合并优化operation: { combine_axis : true },配置文件方案是全局静态的一旦下发会影响当前编译范围内所有函数的代码生成行为粒度粗、无法按函数差异化控制。set_operation_options的出现正是为了解决这一问题——它把配置文件中的参数动态化、可编程化允许开发者在 jit/loop 作用域内按需开启或关闭特定优化实现配置与代码同源、作用范围局部可控的精细化管理。从源码结构看PyPTO 的选项体系在 config.py 中统一组织为codegen、host、pass、runtime、verify、debug、operation等命名空间见PREFIX_MAPoperation.前缀对应操作级优化选项set_operation_options正是operation命名空间的写入入口。函数原型set_operation_options(*, combine_axis: Optional[bool] None)接口位于pypto.experimental命名空间通过 python/pypto/experimental/init.py 从 operation.py 导出。注意参数前有*即必须以关键字参数形式传参不支持位置传参。参数说明参数名输入/输出说明combine_axis输入含义在代码生成阶段实现尾轴 broadcast inline。说明双目运算(32,1) (32,128)时不需要将(32,1)先 broadcast 到(32,128)而是将(32,1)通过 brcb 指令扩展到(32,8)再进行(32,8) (32,128)的计算。前提是(32,1)必须是连续的。类型bool取值范围{True, False}combine_axis 参数详解combine_axis是一个布尔开关取True时强制启用代码生成阶段的尾轴 axis 合并axis combine / broadcast inline优化。其核心收益在于减少不必要的显式 broadcast 计算默认路径关闭(32,1)输入需先被显式 broadcast 扩展为(32,128)再参与(32,128) (32,128)的完整计算中间多出一份临时内存与搬运开销。开启后将(32,1)通过 brcbbroadcast copy指令仅扩展为(32,8)然后执行(32,8) (32,128)。由于向量单元计算时尾轴是逐行展开的(32,8)的窄行数据与(32,128)的宽行数据可以在指令级完成 inline 广播消除了显式 broadcast 带来的冗余访存与计算量。该优化在代码生成端的具体实现可参见 codegen_op.cpp 中的NeedCombineAxis、CombineAxisShape、CombineAxisOffset三个方法当操作数带有inputCombineAxis/outputCombineAxis属性且判定需要合并时代码生成器会同时改写操作数的 shape 与 offset含动态 offset使 brcb 扩展后的窄行能与宽行对齐计算向量二元运算代码生成 codegen_vector_binary.cpp 会将isInputForceCombineAxis作为参数传递给底层TAdd等指令宏见 aicore_runtime.h 中的RT_OPERATION_OP_ADD宏定义。返回值说明voidSet 方法无返回值。设置操作成功即生效无需接收返回值。约束说明使用set_operation_options(combine_axis...)时必须满足以下约束使用场景约束——尾轴连续尾轴 broadcast 输入尾轴必须是连续的否则功能失效。具体而言如果前序节点是尾轴 reducereduce 接口能够保证连续性如果前序节点是COPY_IN从 GM 拷入需要在前端保证数据在 GM 中连续。从源码看这一约束源于ForceCombineAxisForAxisCombine的判定逻辑codegen_preproc.cpp只有当操作数内存类型不是MEM_DEVICE_DDR、且rawshape.back() 1时才会被标记为可合并输入输出侧还需经过ReduceNeedCombineAxis对 reduce 类算子的特判。非连续内存布局下强行合并会产生错误的地址计算。类型安全约束必须确保传入的combine_axis的类型与参数定义的类型完全一致即bool否则可能导致未定义行为或运行时错误。作用范围约束——局部生效参数设置是局部的只会影响当前 jit/loop 内的编译过程若在某个作用域内未设置则继承上层 jit/loop 作用域中的设置。这一作用域语义与 PyPTO 整体的 Options 作用域机制一致——set_options会结束旧作用域并开启新作用域见 config.py 中set_options的 docstring 及_Options上下文管理器实现。调用示例最直接的用法是在 jit 内核函数内部开启尾轴合并import pypto pypto.frontend.jit def kernel(x: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32), y: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32), out: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32)): pypto.experimental.set_operation_options(combine_axisTrue) out[...] x y真实场景Sinkhorn 迭代内核仓库测试 test_mhc_sinkhorn.py 展示了 combine_axis 在规约-广播交替密集场景下的典型用法。该内核计算 Sinkhorn 归一化涉及大量amax/sum(..., keepdimTrue)后紧跟广播减法/除法的模式pypto.frontend.jit( runtime_options{run_mode: pypto.RunMode.SIM, device_sched_mode: 1}, pass_options{vec_nbuffer_setting: {-2: 1, -1: 2}}, ) def mhc_sinkhorn_kernel_sim(...): pypto.experimental.set_operation_options(combine_axisTrue) ... row_max pypto.amax(comb_flag, 1, True) # 尾轴为 1 的 reduce 输出 comb_flag pypto.exp(comb_flag - row_max) # 广播减法 row_sum pypto.sum(comb_flag, 1, True) comb_flag comb_flag / (row_sum eps) # 广播除法 ...这里amax/sum的keepdimTrue输出形状尾轴为 1且 reduce 接口保证连续随后紧跟广播二元运算——正是combine_axis优化的典型收益场景reduce 结果不再显式广播到全宽而是以窄行 brcb 后 inline 参与计算。类似的用法还出现在 flash attention、MHA、interleave rope、gated delta rule 等内核实现中例如 flash_attention_mha_impl.py、chunked_gated_delta_rule_impl.py。底层实现链路解析了解实现链路有助于判断何时该用、何时不该用该接口。从源码可以梳理出combine_axis从 Python 层到指令层的完整路径1. Python 层选项写入operation.py 中的实现非常精简def set_operation_options(*, combine_axis: Optional[bool] None): options_dict {k: v for k, v in locals().items() if v is not None} set_options(operation_optionsoptions_dict)它只收集非 None的关键字参数随后调用 config.py 的set_options(operation_options...)。_Options.prepare_options()会根据PREFIX_MAP将operation_options展开为operation.combine_axis形式最终通过pypto_impl.SetScope(...)写入 C 侧当前作用域。2. C 侧配置读取与下发C 侧通过 config_manager_ng.h 中的KEY_COMBINE_AXIS combine_axis标识该键并在 program.cpp 中将其读入函数参数配置paramConfigs.combineAxis configScope-GetOperationConfigbool(KEY_COMBINE_AXIS);paramConfigs.combineAxis定义于 function.h随后作为该函数的编译参数贯穿整个 Pass 流水线。3. Pass 阶段AxisCombine 与强制标记AxisCombine passaxis_combine.cpp在RunOnFunction开头检查function.paramConfigs_.combineAxis为false时直接跳过整个 axis 合并流程日志输出AxisCombine is skipped.。该 pass 负责为 broadcast 类操作数插入对齐逻辑例如通过OP_BRCB创建对齐张量。CodegenPreproc passcodegen_preproc.cpp中的ForceCombineAxisForAxisCombine是核心标记逻辑遍历 AIV 核vector算子与 UB copy 算子对满足尾轴为 1、内存非 GM条件的输入输出分别设置inputCombineAxis/outputCombineAxis属性对OP_EXPAND等算子还会调用FixExpandDimForAxisCombine修正 expand 维度对 reduce 输出使用ReduceNeedCombineAxis做特判OP_ROWSUMLINE等场景不可合并。PadLocalBufferpad_local_buffer.cpp也会读取combineAxis_在COMBINE_AXIS模式下调整本地 buffer 的 padding 策略保证窄行数据布局可被向量单元直接消费。4. 代码生成与指令发射最终在 codegen_op.cpp 中NeedCombineAxis依据操作属性判断是否需要合并CombineAxisShape/CombineAxisOffset重写参与计算的 shape 与偏移binary_tiled.h 中isCombineAxisEnabled function.paramConfigs_.combineAxis isInWhiteList明确标注尾轴last axis在 combine_axis 开启时走 brcinline 路径——这与文档中(32,1) 通过 brcb 指令扩展到 (32,8)的描述完全对应。验证与调试仓库中针对该接口提供了自动化测试可用于验证配置读写链路的正确性test_config_options.py 中的test_operation_option是直接的功能验证def test_operation_option(): set_operation_options(combine_axisTrue) option get_operation_options() assert option[combine_axis]该测试同时印证了配套的读取接口pypto.experimental.get_operation_options()实现在 operation.py 中通过get_current_scope().get_operation_options()返回当前作用域内operation.前缀的全部选项。如需恢复默认配置可调用pypto.reset_options()见 config.py 的test_reset_option用法。除set_operation_options外PyPTO 还提供成体系的动态配置接口均位于pypto.experimentalset_runtime_options/get_runtime_options运行期选项、set_pass_options/get_pass_optionsPass 级选项如vec_nbuffer_setting、cube_l1_reuse_setting、sg_set_scope等它们共用同一套作用域与 Options 管理机制可组合使用实现内核级精细调优。总结set_operation_options(combine_axisTrue)是 PyPTO 中开启尾轴 broadcast inline 优化的推荐入口它把全局静态的tile_fwk_config.json配置转化为 jit/loop 作用域内可编程、可继承的动态指令配合 reduce 等保证尾轴连续的前序算子可在不改变语义的前提下减少显式 broadcast 的访存与计算开销。使用时要牢记三条约束尾轴输入必须连续、参数类型必须为 bool、作用范围仅限当前作用域并向上继承。底层由 AxisCombine / CodegenPreproc 等 Pass 与 codegen 阶段的 brcb inline 发射共同完成仓库中的单元测试与 Sinkhorn、FlashAttention 等内核实现可作为进一步参考的实证样例。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyPTO 运行时动态配置接口 set_operation_optionscombine_axis 尾轴 broadcast inline 优化实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoPyPTOParallel Tensor/Tile Operation 编程范式的pypto.experimental.set_operation_options是编译框架提供的运行时动态配置管理功能核心接口它将原本静态写在tile_fwk_config.json中的operation.combine_axis参数转变为动态、可编程的指令。读完本文你将掌握该接口的函数原型、参数语义、作用域规则与底层实现链路从 Python API 到 C Pass 再到代码生成并能在 jit/loop 内核中以逐函数粒度控制尾轴 broadcast inline 优化。产品支持情况set_operation_options接口当前支持以下硬件平台以仓库文档标注为准Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持功能说明从静态配置到动态指令在 PyPTO 编译框架中大量编译与代码生成参数原本通过配置文件 tile_fwk_config.json 统一管理。其中operation段默认开启了尾轴合并优化operation: { combine_axis : true },配置文件方案是全局静态的一旦下发会影响当前编译范围内所有函数的代码生成行为粒度粗、无法按函数差异化控制。set_operation_options的出现正是为了解决这一问题——它把配置文件中的参数动态化、可编程化允许开发者在 jit/loop 作用域内按需开启或关闭特定优化实现配置与代码同源、作用范围局部可控的精细化管理。从源码结构看PyPTO 的选项体系在 config.py 中统一组织为codegen、host、pass、runtime、verify、debug、operation等命名空间见PREFIX_MAPoperation.前缀对应操作级优化选项set_operation_options正是operation命名空间的写入入口。函数原型set_operation_options(*, combine_axis: Optional[bool] None)接口位于pypto.experimental命名空间通过 python/pypto/experimental/init.py 从 operation.py 导出。注意参数前有*即必须以关键字参数形式传参不支持位置传参。参数说明参数名输入/输出说明combine_axis输入含义在代码生成阶段实现尾轴 broadcast inline。说明双目运算(32,1) (32,128)时不需要将(32,1)先 broadcast 到(32,128)而是将(32,1)通过 brcb 指令扩展到(32,8)再进行(32,8) (32,128)的计算。前提是(32,1)必须是连续的。类型bool取值范围{True, False}combine_axis 参数详解combine_axis是一个布尔开关取True时强制启用代码生成阶段的尾轴 axis 合并axis combine / broadcast inline优化。其核心收益在于减少不必要的显式 broadcast 计算默认路径关闭(32,1)输入需先被显式 broadcast 扩展为(32,128)再参与(32,128) (32,128)的完整计算中间多出一份临时内存与搬运开销。开启后将(32,1)通过 brcbbroadcast copy指令仅扩展为(32,8)然后执行(32,8) (32,128)。由于向量单元计算时尾轴是逐行展开的(32,8)的窄行数据与(32,128)的宽行数据可以在指令级完成 inline 广播消除了显式 broadcast 带来的冗余访存与计算量。该优化在代码生成端的具体实现可参见 codegen_op.cpp 中的NeedCombineAxis、CombineAxisShape、CombineAxisOffset三个方法当操作数带有inputCombineAxis/outputCombineAxis属性且判定需要合并时代码生成器会同时改写操作数的 shape 与 offset含动态 offset使 brcb 扩展后的窄行能与宽行对齐计算向量二元运算代码生成 codegen_vector_binary.cpp 会将isInputForceCombineAxis作为参数传递给底层TAdd等指令宏见 aicore_runtime.h 中的RT_OPERATION_OP_ADD宏定义。返回值说明voidSet 方法无返回值。设置操作成功即生效无需接收返回值。约束说明使用set_operation_options(combine_axis...)时必须满足以下约束使用场景约束——尾轴连续尾轴 broadcast 输入尾轴必须是连续的否则功能失效。具体而言如果前序节点是尾轴 reducereduce 接口能够保证连续性如果前序节点是COPY_IN从 GM 拷入需要在前端保证数据在 GM 中连续。从源码看这一约束源于ForceCombineAxisForAxisCombine的判定逻辑codegen_preproc.cpp只有当操作数内存类型不是MEM_DEVICE_DDR、且rawshape.back() 1时才会被标记为可合并输入输出侧还需经过ReduceNeedCombineAxis对 reduce 类算子的特判。非连续内存布局下强行合并会产生错误的地址计算。类型安全约束必须确保传入的combine_axis的类型与参数定义的类型完全一致即bool否则可能导致未定义行为或运行时错误。作用范围约束——局部生效参数设置是局部的只会影响当前 jit/loop 内的编译过程若在某个作用域内未设置则继承上层 jit/loop 作用域中的设置。这一作用域语义与 PyPTO 整体的 Options 作用域机制一致——set_options会结束旧作用域并开启新作用域见 config.py 中set_options的 docstring 及_Options上下文管理器实现。调用示例最直接的用法是在 jit 内核函数内部开启尾轴合并import pypto pypto.frontend.jit def kernel(x: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32), y: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32), out: pypto.Tensor([pypto.DYNAMIC, 128], pypto.DT_FP32)): pypto.experimental.set_operation_options(combine_axisTrue) out[...] x y真实场景Sinkhorn 迭代内核仓库测试 test_mhc_sinkhorn.py 展示了 combine_axis 在规约-广播交替密集场景下的典型用法。该内核计算 Sinkhorn 归一化涉及大量amax/sum(..., keepdimTrue)后紧跟广播减法/除法的模式pypto.frontend.jit( runtime_options{run_mode: pypto.RunMode.SIM, device_sched_mode: 1}, pass_options{vec_nbuffer_setting: {-2: 1, -1: 2}}, ) def mhc_sinkhorn_kernel_sim(...): pypto.experimental.set_operation_options(combine_axisTrue) ... row_max pypto.amax(comb_flag, 1, True) # 尾轴为 1 的 reduce 输出 comb_flag pypto.exp(comb_flag - row_max) # 广播减法 row_sum pypto.sum(comb_flag, 1, True) comb_flag comb_flag / (row_sum eps) # 广播除法 ...这里amax/sum的keepdimTrue输出形状尾轴为 1且 reduce 接口保证连续随后紧跟广播二元运算——正是combine_axis优化的典型收益场景reduce 结果不再显式广播到全宽而是以窄行 brcb 后 inline 参与计算。类似的用法还出现在 flash attention、MHA、interleave rope、gated delta rule 等内核实现中例如 flash_attention_mha_impl.py、chunked_gated_delta_rule_impl.py。底层实现链路解析了解实现链路有助于判断何时该用、何时不该用该接口。从源码可以梳理出combine_axis从 Python 层到指令层的完整路径1. Python 层选项写入operation.py 中的实现非常精简def set_operation_options(*, combine_axis: Optional[bool] None): options_dict {k: v for k, v in locals().items() if v is not None} set_options(operation_optionsoptions_dict)它只收集非 None的关键字参数随后调用 config.py 的set_options(operation_options...)。_Options.prepare_options()会根据PREFIX_MAP将operation_options展开为operation.combine_axis形式最终通过pypto_impl.SetScope(...)写入 C 侧当前作用域。2. C 侧配置读取与下发C 侧通过 config_manager_ng.h 中的KEY_COMBINE_AXIS combine_axis标识该键并在 program.cpp 中将其读入函数参数配置paramConfigs.combineAxis configScope-GetOperationConfigbool(KEY_COMBINE_AXIS);paramConfigs.combineAxis定义于 function.h随后作为该函数的编译参数贯穿整个 Pass 流水线。3. Pass 阶段AxisCombine 与强制标记AxisCombine passaxis_combine.cpp在RunOnFunction开头检查function.paramConfigs_.combineAxis为false时直接跳过整个 axis 合并流程日志输出AxisCombine is skipped.。该 pass 负责为 broadcast 类操作数插入对齐逻辑例如通过OP_BRCB创建对齐张量。CodegenPreproc passcodegen_preproc.cpp中的ForceCombineAxisForAxisCombine是核心标记逻辑遍历 AIV 核vector算子与 UB copy 算子对满足尾轴为 1、内存非 GM条件的输入输出分别设置inputCombineAxis/outputCombineAxis属性对OP_EXPAND等算子还会调用FixExpandDimForAxisCombine修正 expand 维度对 reduce 输出使用ReduceNeedCombineAxis做特判OP_ROWSUMLINE等场景不可合并。PadLocalBufferpad_local_buffer.cpp也会读取combineAxis_在COMBINE_AXIS模式下调整本地 buffer 的 padding 策略保证窄行数据布局可被向量单元直接消费。4. 代码生成与指令发射最终在 codegen_op.cpp 中NeedCombineAxis依据操作属性判断是否需要合并CombineAxisShape/CombineAxisOffset重写参与计算的 shape 与偏移binary_tiled.h 中isCombineAxisEnabled function.paramConfigs_.combineAxis isInWhiteList明确标注尾轴last axis在 combine_axis 开启时走 brcinline 路径——这与文档中(32,1) 通过 brcb 指令扩展到 (32,8)的描述完全对应。验证与调试仓库中针对该接口提供了自动化测试可用于验证配置读写链路的正确性test_config_options.py 中的test_operation_option是直接的功能验证def test_operation_option(): set_operation_options(combine_axisTrue) option get_operation_options() assert option[combine_axis]该测试同时印证了配套的读取接口pypto.experimental.get_operation_options()实现在 operation.py 中通过get_current_scope().get_operation_options()返回当前作用域内operation.前缀的全部选项。如需恢复默认配置可调用pypto.reset_options()见 config.py 的test_reset_option用法。除set_operation_options外PyPTO 还提供成体系的动态配置接口均位于pypto.experimentalset_runtime_options/get_runtime_options运行期选项、set_pass_options/get_pass_optionsPass 级选项如vec_nbuffer_setting、cube_l1_reuse_setting、sg_set_scope等它们共用同一套作用域与 Options 管理机制可组合使用实现内核级精细调优。总结set_operation_options(combine_axisTrue)是 PyPTO 中开启尾轴 broadcast inline 优化的推荐入口它把全局静态的tile_fwk_config.json配置转化为 jit/loop 作用域内可编程、可继承的动态指令配合 reduce 等保证尾轴连续的前序算子可在不改变语义的前提下减少显式 broadcast 的访存与计算开销。使用时要牢记三条约束尾轴输入必须连续、参数类型必须为 bool、作用范围仅限当前作用域并向上继承。底层由 AxisCombine / CodegenPreproc 等 Pass 与 codegen 阶段的 brcb inline 发射共同完成仓库中的单元测试与 Sinkhorn、FlashAttention 等内核实现可作为进一步参考的实证样例。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
