CANN pyasc MatmulLeakyRelu 融合算子实战Matmul 高阶 API 与 LeakyReLU 的单 Kernel 融合实现【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本篇文章围绕 pyasc 仓库中examples/05_matmul_leakyrelu样例完整讲解如何在昇腾 AI 处理器上使用 Matmul 高阶 API 与 LeakyReLU 基础 API在单次 Kernel 调用中完成C LeakyReLU(A × B Bias)的矩阵乘与激活融合计算。读完本文你将掌握 Matmul 结果输出到 VECCALC 的融合写法、with matmul.iterate()分块迭代范式、带DataCopyParams参数的数据搬出方式以及 MultiCoreMatmulTiling 的 Tiling 配置与多核划分原理并可直接运行该样例进行功能验证。概述为什么要做 Matmul 与 LeakyReLU 融合本样例实现的是 MatmulLeakyRelu 融合算子计算公式为C A * B Bias C C 0 ? C : C * alpha即先在 Cube 计算单元AIC完成矩阵乘A × B并叠加 Bias 偏置再在矢量计算单元AIV对结果逐元素执行 LeakyReLU 激活非负元素保持不变负元素乘以斜率alpha。如果不做融合矩阵乘结果需要从 Local Memory 写回 Global Memory再由下一个独立的激活算子读回中间结果的往返会引入额外的内存带宽开销。而通过 pyasc 提供的 Matmul 高阶 API 与 LeakyReLU 基础 API 组合矩阵乘结果可以直接保留在片上由矢量计算单元就地消费避免中间结果的 Global Memory 回传从而提升性能。这也是examples/05_matmul_leakyrelu样例相对于独立的 Matmul 样例如 examples/03_matmul_mix的核心差异点。运行环境要求类别要求AI 处理器Ascend 910B / 910CCANN 版本社区版 8.5.0.alpha001 及以上注意事项样例支持 NPU 上板运行需要 NPU 硬件和仿真器模式不需要 NPU 硬件两种运行方式。仿真器模式运行方式请参考 运行环境变量配置 完成配置。PyTorch 和 torch_npu 的安装请参考 样例运行验证。样例规格参数名称输入/输出Shape数据类型格式是否转置a输入[1024, 256]float16ND否b输入[256, 640]float16ND否bias输入[1, 640]float32ND—alpha输入标量—float32——c输出[1024, 640]float32ND—从规格可以看出融合场景的一个典型特征矩阵乘的输入 A、B 为 float16而输出 C 为 float32与 Bias 一致。这对应源码中set_c_type(host.TPosition.VECCALC, host.CubeFormat.ND, host.DataType.DT_FLOAT)的配置矩阵乘输出的高精度结果直接在片上交给激活函数处理避免精度在往返 Global Memory 的过程中受损。样例实现整体流程Global Memory (a_gm, b_gm, bias_gm) │ Matmul.set_tensor_a / set_tensor_b / set_bias ▼ Cube Unit (矩阵乘: C_temp A × B Bias) │ Matmul.get_tensor_c (输出到 VECCALC) ▼ Vector Compute (LeakyReLU) │ TQue.enque → TQue.deque ▼ Global Memory (c_gm) │ data_copy (带 DataCopyParams) ▼ 结果校验该流程在示例代码 matmul_leakyrelu.py 中完整落地calc_offsets负责按核索引计算各张量在 Global Memory 中的偏移matmul_leakyrelu_kernel为被asc.jit(always_compileTrue)装饰的 Kernel 主体generate_tiling负责在 Host 侧生成 Tiling 信息。关键步骤创建 Matmul 对象— 矩阵 C 的输出位置设置为TPosition.VECCALC使得矩阵乘结果保留在 Local Memory 中供后续激活函数直接消费。迭代计算与融合— 使用with matmul.iterate() as count逐块迭代。matmul.get_tensor_c(relu_out_local, en_sequential_writeTrue)获取当前块的矩阵乘结果asc.leaky_relu就地计算激活结果通过 TQue 传递给 copy_out 阶段。带参数的数据搬出— 矩阵 C 以base_m × base_n的块为单位输出使用asc.DataCopyParams指定 block_count、block_len、src_stride、dst_stride配合asc.data_copy(repeat_paramsparams)将分块结果写回 Global Memory。结束计算—matmul.end()结束矩阵乘操作asc.pipe_barrier同步。Kernel 内的融合实现逐行解读以下为 Kernel 主体中与融合直接相关的核心片段见 matmul_leakyrelu.pypipe asc.TPipe() relu_out_queue asc.TQue(asc.TPosition.VECOUT, 1) pipe.init_buffer(querelu_out_queue, num1, lensize) matmul asc.adv.Matmul( aasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype), basc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype), casc.adv.MatmulType(asc.TPosition.VECCALC, asc.CubeFormat.ND, c_global.dtype), biasasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, bias_global.dtype), ) asc.adv.register_matmul(pipe, workspace, matmul, tiling) matmul.set_tensor_a(a_global) matmul.set_tensor_b(b_global) matmul.set_bias(bias_global) with matmul.iterate() as count: relu_out_local relu_out_queue.alloc_tensor(c.dtype) matmul.get_tensor_c(relu_out_local, en_sequential_writeTrue) asc.leaky_relu(relu_out_local, relu_out_local, alpha, counttiling.base_m * tiling.base_n) relu_out_queue.enque(relu_out_local) relu_out_local relu_out_queue.deque(c.dtype) ... asc.data_copy(c_global[start_offset:], relu_out_local, repeat_paramsparams) relu_out_queue.free_tensor(relu_out_local) matmul.end() asc.pipe_barrier(asc.PipeID.PIPE_ALL)各环节的底层实现与设计意图如下融合的载体是TPosition.VECCALC。从源码看MatmulType 是一个包含position、format、dtype、is_trans、layout五个字段的冻结数据类其中c...TPosition.VECCALC就是告诉矩阵乘引擎C 矩阵的中间结果不要落回 Global Memory而是写到矢量计算单元可访问的 VECCALC 空间。这是融合能成立的结构前提。register_matmul完成对象绑定。register_matmul 接收 pipe、workspace、matmul 对象与可选的 tiling底层创建create_asc_RegistMatmulObjOp把 TPipe、workspace 空间和 Tiling 参数绑定到 Matmul 对象上之后才能调用各set_*与迭代接口。with matmul.iterate() as count是分块迭代入口。iterate 返回一个MatmulIterator每调用一次迭代即进入一次 with 循环体就计算出一块baseM × baseN的 C 矩阵count给出当前迭代序号用于后续推导数据搬出的目标偏移。这种迭代式写法天然支持与矢量计算的流水线融合。asc.leaky_relu是就地激活的关键。leaky_relu定义在 vec_binary_scalar.py对应底层算子LeakyRelu按元素执行 Leaky ReLU 操作支持dst与src为同一 Tensor 的就地in-place写法因此这里asc.leaky_relu(relu_out_local, relu_out_local, alpha, ...)直接在矩阵乘结果上完成激活无需额外分配临时空间。counttiling.base_m * tiling.base_n指定本次激活处理的元素个数。TQue 完成跨阶段数据流转。激活结果通过relu_out_queue.enque投递、deque取出队列位置为TPosition.VECOUT深度为 1。enque/deque 之间构成数据依赖关系配合matmul.end()后的asc.pipe_barrier(asc.PipeID.PIPE_ALL)实现 Cube 与 Vector 两条流水之间的同步。核心接口接口用途asc.adv.Matmul矩阵乘高阶 API支持结果输出到 VECCALC 供融合算子消费asc.adv.MatmulType定义矩阵乘操作数的存储位置、数据格式和数据类型asc.adv.register_matmul初始化 Matmul 对象绑定 TPipe、workspace 和 Tiling 参数matmul.set_tensor_a/set_tensor_b设置矩阵乘的左/右操作数matmul.set_bias设置矩阵乘的 Bias 偏置matmul.iterate每调用一次计算出一块 baseM × baseN 的 C 矩阵支持与矢量计算融合matmul.get_tensor_c获取当前迭代块的矩阵乘结果 Tensorasc.leaky_reluLeakyReLU 激活函数dst src 0 ? src : src * alphaasc.data_copy(repeat_params)带 DataCopyParams 的数据搬运按 block 步长将 VECCALC 结果写回 GMasc.DataCopyParams数据搬运参数block_count、block_len、src_stride、dst_stridematmul.end结束矩阵乘操作asc.pipe_barrier阻塞相同流水具有数据依赖的相同流水之间需要插入此同步补充说明两个接口的底层细节便于读者放心使用set_bias有明确的数据类型约束见 matmul.py当 A、B 为 int8 时Bias 必须为 int32否则 Bias 支持 float16 / float32。本样例 A、B 为 float16Bias 取 float32符合约束。DataCopyParams 的四个参数均有默认值block_count1, block_len0, src_stride0, dst_stride0在样例会显式构造。其字段含义为block_count为搬运的 block 数量对应base_mblock_len为每个 block 的有效数据长度按 C016 元素对齐即base_n * sizeof(dtype) // DEFAULT_C0_SIZEsrc_stride/dst_stride为源/目的侧相邻 block 间的步长按 16 字节为单位。分块、多核、Tiling 逻辑本样例的 Tiling 生成与多核划分在 Host 侧函数generate_tiling中完成见 matmul_leakyrelu.pymatmul_tiling host.MultiCoreMatmulTiling(host.get_ascendc_platform()) matmul_tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16, False) matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16, False) matmul_tiling.set_c_type(host.TPosition.VECCALC, host.CubeFormat.ND, host.DataType.DT_FLOAT) matmul_tiling.set_bias_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT) matmul_tiling.set_dim(2) matmul_tiling.set_org_shape(m, n, k) matmul_tiling.set_shape(m, n, k) matmul_tiling.enable_bias(True) matmul_tiling.set_traverse(host.MatrixTraverse.FIRSTM) matmul_tiling.set_fix_split(256, 128, -1) matmul_tiling.set_buffer_space(-1, -1, -1) tiling asc.adv.TCubeTiling() matmul_tiling.get_tiling(tiling)结合 wrappers.py 中的声明各配置项的作用如下平台配置—MultiCoreMatmulTiling构造时传入host.get_ascendc_platform()获取当前平台信息set_dim(2)指定 2 个核参与计算set_fix_split(256, 128, -1)设置固定的 M、N 方向分块大小M 方向 256、N 方向 128K 方向 -1 表示由框架自动决定。Tiling 生成—set_a_type/set_b_type/set_c_type/set_bias_type分别设置 A/B/C/Bias 的类型信息存储位置、格式、数据类型其中 C 的位置必须与 Kernel 内 Matmul 对象一致均为VECCALCset_org_shape(m, n, k)与set_shape(m, n, k)设置原始 shape 与当前 shape本样例 shape 为[1024, 640, 256]注意传入顺序为 m、n、kenable_bias(True)开启 Biasset_buffer_space(-1, -1, -1)让各缓存空间大小由框架自动推导。遍历方式—set_traverse(host.MatrixTraverse.FIRSTM)设置遍历方式为先 M 后 N。结合 Kernel 中round_m tiling.single_core_m // tiling.base_m与start_offset count % round_m * tiling.base_m * tiling.n count // round_m * tiling.base_n可以验证迭代序号count对round_m取余得到 M 方向的轮次、整除得到 N 方向的轮次与 FIRSTM先遍历 M 方向的语义一致。多核切分—base_m × base_n为每次迭代的计算粒度single_core_m // base_m个 M 方向轮次和 N 方向轮次共同决定每个核的总迭代次数。每个核负责的子块位置由 calc_offsets 依据get_block_idx()计算m_index block_idx % m_single_blocks、n_index block_idx // m_single_blocks进而得到 A/B/C/Bias 各自的偏移量。启动核数— 启动配置为matmul_leakyrelu_kernel[tiling.used_core_num // 2, rt.current_stream()]见 matmul_leakyrelu.py。由于本样例同时使用 CubeAIC与 VectorAIV两条流水一个 AI Core 上需要同时容纳两种任务的调度因此实际启动核数取used_core_num // 2。这一点与仓库中 MIX 模式样例的说明一致见 examples/03_matmul_mix/README.md 中MIX 模式下 USE_CORE_NUM 需设置为 AI Core 数量的 2 倍启动核数为 USE_CORE_NUM // 2。编译执行环境配置请参考 quick_start.md。完成环境配置后执行如下命令可进行功能验证cd pyasc/examples/05_matmul_leakyrelu python3 matmul_leakyrelu.py -r [RUN_MODE] -v [SOC_VERSION]其中脚本参数说明如下RUN_MODE编译执行方式可选择 NPU 仿真、NPU 上板对应参数分别为Model/NPU。SOC_VERSION昇腾 AI 处理器型号。如果无法确定具体的 SOC_VERSION则在安装昇腾 AI 处理器的服务器执行npu-smi info命令进行查询在查询到的 Name 前增加 Ascend 信息。例如 Name 对应取值为 xxxyy实际配置的 SOC_VERSION 值为 Ascendxxxyy。示例如下Ascend910B1 请替换为实际的 AI 处理器型号# 仿真器模式 python3 matmul_leakyrelu.py -r Model -v Ascend910B1 # NPU 上板模式 python3 matmul_leakyrelu.py -r NPU -v Ascend910B1脚本入口的校验逻辑见 matmul_leakyrelu.py会检查-r取值必须属于config.Backend的成员Model/NPU-v取值必须属于config.Platform非法取值会抛出带支持列表的ValueError。Kernel 运行完成后脚本使用 PyTorch 参考实现进行精度对比matmul (torch.matmul(a.to(torch.float32), b.to(torch.float32)).to(torch.float32) bias).to(torch.float32) assert torch.allclose(c, torch.where(matmul 0, matmul, matmul * alpha), rtol1e-3, atol1e-3)即先将 A、B 升为 float32 做矩阵乘再加 Bias再用torch.where按 LeakyReLU 公式生成参考结果与算子输出做allclose对比容差为 rtol1e-3、atol1e-3。执行成功后输出[INFO] start process sample matmul_leakyrelu. [INFO] Sample matmul_leakyrelu run success.小结examples/05_matmul_leakyrelu是一个典型的矩阵乘 激活融合算子样例它通过把 Matmul 的 C 矩阵输出位置配置为TPosition.VECCALC让矩阵乘结果停留在片上 Local Memory再以with matmul.iterate()的逐块迭代方式调用asc.leaky_relu就地完成激活最后用带DataCopyParams的data_copy按块写回 Global Memory从而在单次 Kernel 调用内完成整条计算链避免了中间结果的内存往返。如果你希望进一步了解不带矢量融合的纯 Cube 矩阵乘参考 examples/04_matmul_cube_only或更复杂的算子融合如 examples/07_swiglu可以在本样例基础上对照阅读体会不同计算单元之间的数据流转方式。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考