CANN pyascMatmulApiTiling.set_traverse接口详解固定 M 轴/N 轴优先的 Matmul 计算方向【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyascasc.lib.host.MatmulApiTiling.set_traverse是 CANN pyasc 在 Python 侧提供的 Host 端 Matmul Tiling 设置接口用于固定 Matmul 分块计算时大矩阵的输出遍历方向M 轴优先还是 N 轴优先。本文以 关联 API 文档 为主线结合 pybind11 绑定源码 与仓库内 Matmul 示例完整讲解接口签名、参数语义、返回值、底层实现与真实算子中的调用方式帮助读者在编写 Matmul 类算子的 Host 侧 Tiling 时正确控制计算方向。接口背景为什么需要设置 Matmul 计算方向在昇腾 AI 处理器上编写 Matmul 类算子时通常需要两份代码Kernel 侧的矩阵乘计算逻辑以及 Host 侧的 Tiling 逻辑。Tiling 负责根据输入矩阵的形状、数据类型、片上 Buffer 容量等信息计算出 Kernel 执行所需的切分参数如 baseM/baseN/baseK、单核计算范围等。CANN pyasc 在asc.lib.host模块中提供了一组 Matmul Tiling API其定位与 模块总览 中的描述一致用户只需要传入 A/B/C 矩阵的 Position 位置、Format 格式和 DType 数据类型等信息调用 API 接口即可获取 Kernel 侧Init中TCubeTiling结构体中的相关参数。set_traverse正是这组 API 中用于控制矩阵遍历顺序的接口一次迭代计算出[baseM, baseN]大小的 C 矩阵分片后下一次迭代的输出分片是优先沿 M 方向移动还是优先沿 N 方向移动。这一顺序直接影响数据搬运的局部性与多核切分的结果是影响 Matmul 性能的细节参数之一。接口签名与 Ascend C 函数原型对应set_traverse是MatmulApiTiling/MultiCoreMatmulTiling/BatchMatmulTiling的共有接口Python 侧签名如下set_traverse(self: libhost.MatmulApiTilingBase, traverse: libhost.MatrixTraverse) - int它对应的 Ascend C 函数原型为int32_t SetTraverse(MatrixTraverse traverse)从绑定实现看python/asc/lib/host/bindings/MatmulApiTiling.cpp#L1147-L1148 通过 pybind11 将 C 的SetTraverse(traverse)直接暴露为 Python 方法.def( set_traverse, [](MatmulApiTilingBase self, MatrixTraverse traverse) { return self.SetTraverse(traverse); }, traverse_a, ...)这意味着 Python 侧的方法名、参数顺序与 C 侧完全一致返回值也原样透传。参数说明MatrixTraverse 枚举的三种取值参数traverse用于设置固定的 Matmul 计算方向类型为MatrixTraverse。该枚举在 python/asc/lib/host/bindings/Enums.cpp#L89-L92 中注册为 Python 枚举共三个取值枚举值说明MatrixTraverse.NOSET不设置固定遍历方向由 Tiling 函数内部根据形状与 Buffer 情况自行决策是多数接口如MatmulConfigParams构造函数、set_matmul_config_params的默认值MatrixTraverse.FIRSTMM 轴优先C 矩阵分片沿 M 方向优先迭代MatrixTraverse.FIRSTNN 轴优先C 矩阵分片沿 N 方向优先迭代在 Python 代码中通过host.MatrixTraverse.FIRSTM/host.MatrixTraverse.FIRSTN引用。从源码结构看NOSET主要用于表示未显式设置的默认状态而set_traverse文档明确限定可选值为FIRSTM/FIRSTN两者。traverse 的深层语义参考同模块set_matmul_config_params接口对traverse参数的解释见 绑定源码 docstring可以更准确理解该参数的含义traverseMatmul 做矩阵运算的循环迭代顺序即一次迭代计算出[baseM, baseN]大小的 C 矩阵分片后自动偏移到下一次迭代输出的 C 矩阵位置的偏移顺序。也就是说Matmul 计算被拆分为若干个[baseM, baseN]的输出分片set_traverse决定这些分片在 C 矩阵上的游走路径FIRSTM表示沿 M 方向先把当前 N 列走完再切到下一组 NFIRSTN表示沿 N 方向优先。选择与数据排布如 NZ 格式、A/B 矩阵的搬运方式匹配的方向可以减少缓存失效、提升搬运效率。返回值说明与错误处理接口返回int-1设置失败0设置成功。与模块内其他set_*系列接口set_a_type、set_shape、set_buffer_space等的返回值约定完全一致。设置失败时建议检查传入的traverse取值是否合法必须是MatrixTraverse枚举成员以及调用时机是否正确。调用示例完整可运行以下示例完整继承自 API 文档展示了在配置完 A/B/C/Bias 矩阵类型、形状与 Buffer 空间之后如何设置 M 轴优先的遍历方向并获取 Tiling 结果import asc.lib.host as host ascendc_platform host.get_ascendc_platform() tiling host.MatmulApiTiling(ascendc_platform) tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16) tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16) tiling.set_c_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT) tiling.set_bias_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT) tiling.set_shape(1024, 1024, 1024) tiling.set_org_shape(1024, 1024, 1024) tiling.set_bias(True) tiling.set_traverse(host.MatrixTraverse.FIRSTM) # 固定 M 轴优先 tiling.set_buffer_space(-1, -1, -1) tiling_data host.TCubeTiling() ret tiling.get_tiling(tiling_data) # ret 不为 -1 表示 Tiling 计算成功代码要点说明host.get_ascendc_platform()根据当前运行时 SOC 版本获取平台对象作为MatmulApiTiling构造参数实现见 host/init.pyset_shape/set_org_shape分别设置计算形状与原始完整形状单位为元素个数set_buffer_space(-1, -1, -1)使用默认的片上 Buffer 空间L1/L0C/UB 均取 AI 处理器默认大小get_tiling将计算结果写入TCubeTiling结构体供 Kernel 侧使用若需改为 N 轴优先仅需将set_traverse参数替换为host.MatrixTraverse.FIRSTN其余配置保持不变。仓库内的真实调用案例set_traverse在仓库的 Matmul 系列示例中均有实际使用可作为参照examples/05_matmul_leakyrelu/matmul_leakyrelu.py#L107在生成 Tiling 时调用matmul_tiling.set_traverse(host.MatrixTraverse.FIRSTM)其 README 明确说明该行用于设置遍历方式为先 M 后 Nexamples/09_linear/linear.py#L326线性层GEMM 类算子Tiling 生成中同样使用mm.set_traverse(host.MatrixTraverse.FIRSTM)examples/09_linear/ascendc/linear.asc#L138C 侧对应的 Host Tiling 实现调用mmTiling.SetTraverse(matmul_tiling::MatrixTraverse::FIRSTM);与 Python 侧形成一一对应examples/10_fused_infer_attention/ascendc/fused_infer_attention.asc#L1066 与 L1092融合注意力算子的 QK 与 PV 两次矩阵乘均固定为FIRSTM遍历方向。这些案例表明无论是 Python 侧通过asc.lib.host计算 Tiling还是 C 侧直接调用SetTraverse仓库内的 Matmul 类算子普遍选择FIRSTMM 轴优先作为固定遍历方向。使用约束与建议调用时机set_traverse属于配置类接口与set_shape、set_buffer_space等一样需要在get_tiling之前调用才会生效在get_tiling之后再调用不会影响已生成的 Tiling 结果。与set_matmul_config_params的关系set_matmul_config_params的traverse参数默认值为MatrixTraverse::NOSET见 绑定源码即不强制遍历方向而set_traverse用于显式固定方向。若两种方式同时使用需要保证语义一致避免 Tiling 与 Kernel 侧预期不符。与 Kernel 侧保持一致遍历方向会体现在 Tiling 参数如 C 分片的偏移计算中Kernel 侧按 Tiling 结果执行因此无需额外在 Kernel 中重复设置但若自行编写 Kernel 侧的循环顺序需要与 Tiling 侧的方向匹配。性能调优参考当输出 C 矩阵分片的遍历路径与数据在 L2/UB 中的排布吻合时数据复用率更高。实际算子开发中可通过对比FIRSTM与FIRSTN两档配置下的执行耗时来选择仓库示例的默认选择可作为起点。总结asc.lib.host.MatmulApiTiling.set_traverse是 Host 侧 Matmul Tiling 配置链路上控制计算方向的关键一环它通过MatrixTraverse枚举FIRSTM/FIRSTN固定 M 轴或 N 轴优先的输出分片遍历顺序并以-1/0的返回值告知设置是否成功。该接口由 CSetTraverse经 pybind11 直接绑定而来MatmulApiTiling.cpp#L1146-L1184与set_matmul_config_params中的traverse参数、C 侧SetTraverse调用见 linear.asc#L138共同构成完整的遍历方向控制手段。在编写或调优 Matmul 类算子时将set_traverse与 set_shape、set_org_shape、set_buffer_space 等接口配合使用即可获得完整的、可供 Kernel 直接消费的TCubeTiling参数。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考