尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

pyasc 纯 Cube 模式 Matmul 算子样例解析:C = A × B + Bias 的多核实现与 Tiling 原理

pyasc 纯 Cube 模式 Matmul 算子样例解析:C = A × B + Bias 的多核实现与 Tiling 原理 pyasc 纯 Cube 模式 Matmul 算子样例解析C A × B Bias 的多核实现与 Tiling 原理【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本技术指南以 examples/04_matmul_cube_only 样例为主体系统讲解在 CANN pyasc 中如何仅使用 AICCube 核实现矩阵乘算子涵盖 JIT 编译参数matmul_cube_onlyTrue的启用方式、Kernel 侧Matmul高阶 API 的调用流程、Host 侧MultiCoreMatmulTiling的多核切分原理以及 NPU 上板与仿真器两种运行验证方式。读完本文你将能独立编写并运行一个纯 Cube 模式的多核 Matmul 算子并理解其分块偏移与尾块处理的底层机制。一、概述什么是纯 Cube 模式在昇腾 AI 处理器的硬件架构中AI Core 通常包含负责矩阵乘法的 Cube 单元AIC和负责向量/标量运算的 AIV。常规的 Matmul 算子如 examples/03_matmul_mix会同时调度 AIC 与 AIVCube 计算 C 矩阵AIV 负责数据搬运、格式转换、Bias 叠加等辅助运算。而本样例展示的纯 Cube 模式只保留矩阵计算本身即C A × B BiasBias 可通过ENABLE_BIAS开关控制默认关闭此时仅 AIC 参与 Cube 矩阵计算AIV 空闲不参与任何计算。这一模式适用于只要矩阵乘结果、不需要任何后处理的场景可以最大限度降低调度开销也便于开发者聚焦 Matmul 高阶 API 本身的行为。启用方式非常简洁在 Kernel 核函数的 JIT 编译装饰器中设置matmul_cube_onlyTrueasc.jit(matmul_cube_onlyTrue, always_compileTrue) def matmul_kernel(a: asc.GlobalAddress, b: asc.GlobalAddress, c: asc.GlobalAddress, bias: asc.GlobalAddress, tiling: asc.adv.TCubeTiling, workspace: asc.GlobalAddress): ...对应源码见 matmul_cube_only.py。其中always_compileTrue表示每次都触发编译便于样例功能验证。二、运行环境要求类别要求AI 处理器Ascend 910B / 910CCANN 版本社区版 8.5.0.alpha001 及以上需要注意以下几点样例支持NPU 上板运行需要 NPU 硬件和仿真器模式不需要 NPU 硬件两种运行方式。仿真器模式的运行环境变量配置含LD_LIBRARY_PATH指向 simulator 库、以及接入 torch 时需LD_PRELOADlibruntime_camodel.so的说明请参考 docs/quick_start.md#envvar-config。PyTorch 和 torch_npu 的安装样例通过 torch 输入/输出 tensor 验证结果请参考 docs/quick_start.md#example-verification。纯 Cube 模式下USE_CORE_NUM需设置为 AI Core 数量仅 AIC实际启动核数为tiling.used_core_num。原因在于只有 AIC 参与计算AIV 不承担任何工作因此启动的核数与 AIC 核数一致而不是 AIC AIV 的总核数。三、样例规格本样例的输入输出规格如下参数名称输入/输出Shape数据类型格式是否转置a输入[128, 64]float16ND否b输入[64, 30720]float16ND否bias输入[1, 30720]float32ND—c输出[128, 30720]float32ND—这是一个典型的 瘦 M、长 N 矩阵乘M128、K64、N30720。N 方向 30720 的巨大宽度决定了多核切分主要发生在 N 方向见下文第五节。A、B 为 float16Cube 最常用的半精度输入C 与 Bias 为 float32累加精度。四、样例实现4.1 整体流程样例的数据流与计算流如下Global Memory (a_gm, b_gm, bias_gm) │ Matmul.set_tensor_a / set_tensor_b / set_bias ▼ Cube Unit (仅 AIC: C A × B Bias) │ Matmul.iterate_all ▼ Global Memory (c_gm) │ Matmul.end / pipe_barrier ▼ 完成4.2 Kernel 侧关键步骤完整源码见 matmul_cube_only.py。Kernel 核函数内的关键步骤依次为AIC 核判断—— 通过asc.ascend_is_aic()判断当前是否为 AIC 核仅 AIC 核执行矩阵乘计算。该接口在 python/asc/language/basic/common.py 中实现实际是生成一个asc.AscendIsAICOp的 IR 指令并返回 i1 类型的布尔值。创建 Matmul 对象—— 通过asc.adv.Matmul创建矩阵乘对象并通过asc.adv.MatmulType分别指定矩阵 A、B、C 和 Bias 的位置、格式和数据类型matmul asc.adv.Matmul( aasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype, IS_TRANS_A), basc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype, IS_TRANS_B), casc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, c_global.dtype), biasasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, bias_global.dtype), )其中MatmulType是包含position存储位置如 GM、format数据格式如 ND、dtype数据类型、is_trans是否转置和layout的不可变数据类定义见 python/asc/language/adv/matmul.py。本样例四个操作数全部位于 Global Memory 且为 ND 格式因此set_tensor_a/set_tensor_b后由 Cube 直接从 GM 取数计算不需要额外的 AIV 搬运。初始化 Matmul—— 调用asc.adv.register_matmul(pipe, workspace, matmul, tiling)完成初始化。从源码看该接口将 TPipe、workspace、Matmul 对象和 Tiling 结构体绑定生成asc.RegistMatmulObjOp指令见 python/asc/language/adv/matmul.py。设置矩阵数据和 Bias—— 通过matmul.set_tensor_a/set_tensor_b设置左右矩阵当tiling.is_bias为 True 时调用matmul.set_bias(bias_global)设置偏置通过matmul.set_tail(tail_m, tail_n, tiling.k_a)设置 M、N、K 三个方向的尾块大小if asc.get_block_idx() tiling.used_core_num: matmul.set_tensor_a(a_global, IS_TRANS_A) matmul.set_tensor_b(b_global, IS_TRANS_B) if tiling.is_bias: matmul.set_bias(bias_global) matmul.set_tail(tail_m, tail_n, tiling.k_a) matmul.iterate_all(c_global) matmul.end()注意set_bias的合法性校验当 A、B 均为 int8 时 Bias 仅支持 int32/int_其余情况支持 half/float16/float32见 python/asc/language/adv/matmul.py。执行与结束——matmul.iterate_all(c_global)一次性完成当前核负责的single_core_m × single_core_n大小 C 矩阵计算见 docs/python-api/language/generated/asc.language.adv.Matmul.iterate_all.mdmatmul.end()结束 Matmul 操作最后asc.pipe_barrier(asc.PipeID.PIPE_ALL)做全流水同步确保结果对后续消费者可见。4.3 核心接口速查接口用途asc.adv.Matmul矩阵乘高阶 API封装 Cube 计算单元asc.adv.MatmulType定义矩阵乘操作数的存储位置、数据格式和数据类型asc.adv.register_matmul初始化 Matmul 对象绑定 TPipe、workspace 和 Tiling 参数matmul.set_tensor_a/set_tensor_b设置矩阵乘的左/右操作数matmul.set_bias设置矩阵乘的 Bias 偏置操作数matmul.set_tail设置 M、N、K 方向的尾块大小处理非对齐场景matmul.iterate_all单次调用完成所有分块的矩阵乘计算结果写入目标 Tensormatmul.end结束矩阵乘操作asc.ascend_is_aic判断当前核是否为 AICCube Coreasc.pipe_barrier阻塞相同流水具有数据依赖的相同流水之间需要插入此同步asc.lib.host.MultiCoreMatmulTilingHost 侧 Tiling API获取 Matmul 分块参数关于set_tail需要补充说明它在不改变 Tiling 的情况下重新设置本次计算的 singleCoreM/singleCoreN/singleCoreK以元素为单位对应 Ascend C 原型void SetTail(int tailM -1, int tailN -1, int tailK -1)默认值 -1 表示不调整。其典型用途正如本样例注释所言如果是尾核需要调整 single_core_m/single_core_n/single_core_k详见 docs/python-api/language/generated/asc.language.adv.Matmul.set_tail.md。五、分块、多核与 Tiling 逻辑这是本样例最核心的工程细节分三层展开5.1 多核切分纯 Cube 模式下只有矩阵计算USE_CORE_NUM用于设置启动多少个 CubeAIC实例执行。例如 Ascend 910B1 平台有 24 个 Cube 核建议设置为 24样例中USE_CORE_NUM 24见 matmul_cube_only.py。每个核根据block_idx计算自己在 M、N 方向的索引和偏移量。具体切分逻辑在calc_offsets函数中block_idx asc.get_block_idx() m_single_blocks tiling.m.ceildiv(tiling.single_core_m) n_index block_idx // m_single_blocks m_index block_idx % m_single_blocks即M 方向被切分成ceil(M / single_core_m)个块核按M 方向块优先的次序编号——m_index block_idx % m_single_blocks、n_index block_idx // m_single_blocks。对 M128、N30720 的规格M 方向块数很少绝大多数核都分布在 N 方向的不同子块上。5.2 各操作数的偏移计算偏移计算同样在calc_offsets中完成见 matmul_cube_only.pyoffset_a m_index * tiling.k_a * tiling.single_core_m # A 按行方向取 m_index 个子块 offset_b n_index * tiling.single_core_n # B 按列方向取 n_index 个子块 offset_c m_index * tiling.n * tiling.single_core_m n_index * tiling.single_core_n offset_bias n_index * tiling.single_core_nA 矩阵是 [M, K] 的 ND 布局第m_index个子块的行首偏移为m_index * K * single_core_mK 即tiling.k_a。B 矩阵是 [K, N] 的 ND 布局第n_index个子块的列首偏移为n_index * single_core_n。C 矩阵是 [M, N] 的 ND 布局偏移由行方向m_index * N * single_core_m与列方向n_index * single_core_n共同决定。Bias 是 [1, N] 的行向量只需按 N 方向取第n_index个子块的偏移。代码同时处理了is_trans_a / is_trans_b两种转置情形转置时偏移公式不同。5.3 尾块Tail处理由于 M、N 不一定能被single_core_m × single_core_n整除最后一个子块的核需要计算实际剩余的尾块大小tail_m tiling.m - m_index * tiling.single_core_m if tail_m 0 or tail_m tiling.single_core_m: tail_m tiling.single_core_m tail_n tiling.n - n_index * tiling.single_core_n if tail_n 0 or tail_n tiling.single_core_n: tail_n tiling.single_core_n若剩余量不大于 0该核实际不承担有效数据或不小于完整块非尾核则回退为完整的single_core_m/single_core_n。随后通过matmul.set_tail(tail_m, tail_n, tiling.k_a)通知 Matmul 对象本次实际计算范围。5.4 Host 侧 Tiling 生成Tiling 在 Host 侧通过MultiCoreMatmulTiling生成见generate_tiling函数matmul_cube_only.pymatmul_tiling host.MultiCoreMatmulTiling(host.get_ascendc_platform()) matmul_tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16, IS_TRANS_A) matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16, IS_TRANS_B) matmul_tiling.set_c_type(host.TPosition.GM, 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(USE_CORE_NUM) matmul_tiling.set_org_shape(m, n, k) matmul_tiling.set_shape(m, n, k) matmul_tiling.enable_bias(ENABLE_BIAS) matmul_tiling.set_buffer_space(-1, -1, -1) tiling asc.adv.TCubeTiling() matmul_tiling.get_tiling(tiling)关键点解读set_dim(USE_CORE_NUM)设置参与多核 Matmul 计算的核数对应 Ascend C 原型int32_t SetDim(int32_t dim)返回 -1 表示设置失败、0 表示成功见 docs/python-api/lib/generated/asc.lib.host.MultiCoreMatmulTiling.set_dim.md。set_org_shape/set_shape分别设置原始形状与实际计算形状本样例两者一致均为 m/n/k。set_buffer_space(-1, -1, -1)表示 L1 / L0C / UB 缓冲空间交由 Tiling 自动决策。最终get_tiling(tiling)将计算结果写入asc.adv.TCubeTiling结构体。若返回 -1 表示 Tiling 计算失败详见 docs/python-api/lib/generated/asc.lib.host.MatmulApiTiling.get_tiling.md失败原因可在日志级别 WARNING 下搜索关键字 MatmulApi Tiling 定位。5.5 分块计算规则M 方向按single_core_m切分、N 方向按single_core_n切分每个核负责single_core_m × single_core_n大小的结果子矩阵。切分参数由 Host Tiling 依据硬件 Cube 规格L0A/L0B/L0C 容量自动计算并写入TCubeTilingKernel 侧只需读取tiling.single_core_m、tiling.single_core_n、tiling.used_core_num、tiling.k_a等字段即可完成自身工作量的确定。六、编译执行与功能验证环境配置含下载源码、安装 pyasc、配置仿真器环境变量等请参考 docs/quick_start.md#envready。完成环境配置后在样例目录执行如下命令进行功能验证cd pyasc/examples/04_matmul_cube_only python3 matmul_cube_only.py -r [RUN_MODE] -v [SOC_VERSION]运行前请根据实际平台的 Cube 核数修改 matmul_cube_only.py 中的USE_CORE_NUM参数。脚本参数说明RUN_MODE编译执行方式可选ModelNPU 仿真或NPUNPU 上板。SOC_VERSION昇腾 AI 处理器型号。如果无法确定具体型号可在安装昇腾 AI 处理器的服务器上执行npu-smi info查询在查询到的 Name 前增加Ascend信息例如 Name 对应取值为xxxyy则实际配置的[SOC_VERSION]为Ascendxxxyy。运行示例Ascend910B1请替换为实际的 AI 处理器型号# 仿真器模式 python3 matmul_cube_only.py -r Model -v Ascend910B1 # NPU 上板模式 python3 matmul_cube_only.py -r NPU -v Ascend910B1参数解析与校验逻辑见 matmul_cube_only.py-r缺省为Model-v缺省为 None仿真器模式下默认Ascend910B1环境NPU 上板模式下自动检测Backend仅支持Model/NPU两个枚举值定义见 python/asc/runtime/config.pyPlatform枚举则定义了受支持的昇腾芯片型号。6.1 结果正确性验证样例通过 PyTorch 的 float32 精度矩阵乘结果做对照验证matmul_cube_only_custom见 matmul_cube_only.pyif ENABLE_BIAS: matmul (torch.matmul(a.to(torch.float32), b.to(torch.float32)).to(torch.float32) bias).to(torch.float32) else: matmul torch.matmul(a.to(torch.float32), b.to(torch.float32)).to(torch.float32) ... assert torch.allclose(c, matmul, rtol1e-3, atol1e-3)A、B 的 float16 输入先被提升为 float32 计算参考结果避免累积精度差异Kernel 输出与参考结果通过torch.allclose以rtol1e-3, atol1e-3容差比对。执行成功后输出[INFO] start process sample matmul_cube_only. [INFO] Sample matmul_cube_only run success.七、从样例到源码纯 Cube 模式的实现原理除了样例本身pyasc 源码为该模式提供了清晰的实现证据ascend_is_aic的 IR 语义在 python/asc/language/basic/common.py 中该函数生成AscendIsAICOp指令编译期将展开为对当前核类型的运行时判断。这是仅 AIC 参与计算的编程入口非 AIC 核直接跳过整个 Matmul 代码路径。Matmul 高阶 API 的指令映射python/asc/language/adv/matmul.py 是asc.adv.Matmul的全部实现。可以看到set_tensor_a/set_tensor_b/set_bias/set_tail/iterate_all/end分别映射为MatmulSetTensorAOp、MatmulSetBiasOp、MatmulSetTailOp、MatmulIterateAllOp、MatmulEndOp等 ASC IR 指令iterate_all支持en_atomic0~3、sync、en_sequential_write、wait_iterate_all等扩展参数见 python/asc/language/adv/matmul.py输出到 GM 时默认en_sequential_writeFalse。这一层与 Ascend C 的 Matmul 高阶 API 一一对应可推断纯 Cube 模式最终生成的底层执行序列即为GM 取数 → Cube 计算 → GM 写回的最小化流水。与混合模式的差异对比 examples/03_matmul_mix 的混合模式样例可以发现纯 Cube 模式省去了 AIV 侧的 Bias 叠加、格式转换等后处理逻辑Kernel 内只有if asc.ascend_is_aic()保护下的 Cube 计算路径这正是其名字cube_only的来源。八、总结本样例是理解 pyasc Matmul 高阶 API 与多核 Tiling 的最小可运行范本。核心要点可归纳为四条模式开关asc.jit(matmul_cube_onlyTrue)即可让 Kernel 进入纯 Cube 模式仅 AIC 计算核数设定纯 Cube 模式下USE_CORE_NUM应等于 AIC 核数如 910B1 为 24实际启动核数取tiling.used_core_num分块与偏移M/N 方向分别按single_core_m/single_core_n切分核通过block_idx计算行列索引与各操作数偏移尾核用set_tail修正计算范围验证闭环Host 侧MultiCoreMatmulTiling生成 TilingKernel 侧执行计算最终与 torch float32 参考结果做allclose容差对比。在此基础上读者可以自行修改ENABLE_BIAS、USE_CORE_NUM、矩阵 Shape 与数据类型进一步观察纯 Cube 模式在不同规格下的分块行为或参考 examples 目录下的其他样例如 03 混合模式、05 Matmul LeakyReLU对比 AIC/AIV 协同与纯 AIC 两种调度方式的差异。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表