Kernel Matmul Basic【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明基础矩阵乘 Kernel仅支持 AIC 计算无 AIV 参与不支持 workspace。适用于小矩阵、简单计算场景集成 BlockScheduler 调度、BlockMmad 计算和 BlockEpilogueEmpty 后处理组件。继承自Kernel Matmul 基础框架特殊约束BlockEpilogue 限制仅支持Block::BlockEpilogueEmpty不支持任何后处理操作。using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueEmpty;计算模式仅在 AIC 核函数中执行不支持 AIV 计算AIV 核直接返回。if ASCEND_IS_AIV { return; // AIV 核直接返回不执行任何计算 }Workspace 不支持不支持 workspace无法存储中间结果适用于完整 tile 计算场景。非连续输入不支持带 stride 的非连续输入仅支持连续 ND/NZ layout。FP32 大 K不支持 FP32 大 K 场景K 轴切分受硬件限制K 值过大时需使用 StreamK Kernel。特殊成员方法构造函数__aicore__ inline KernelMatmulBasic()功能构造 KernelMatmulBasic 对象。析构函数__aicore__ inline ~KernelMatmulBasic()功能析构 KernelMatmulBasic 对象。UnsetHf32函数__aicore__ inline void UnsetHf32(bool isHf32)功能关闭 HF32 模式。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | isHf32 | bool | 是否启用 HF32 模式 |说明当isHf32 true时调用AscendC::SetHF32Mode(0)关闭 HF32 模式Basic Kernel 在计算完成后自动调用此函数清理 HF32 状态HF32 模式设置流程// Kernel 开始时 if (isHf32) { AscendC::SetHF32Mode(1); AscendC::SetHF32TransMode(1); } // Kernel 结束时 UnsetHf32(isHf32); // 关闭 HF32 模式MM Layout TransformSetMMLayoutTransform(true); // 调用 BlockMmad 前设置为列主序适配 Fixpipe // ... BlockMmad 计算 ... SetMMLayoutTransform(false); // 计算后关闭说明MM Layout Transform 用于适配 Fixpipe 输出格式确保 L0C 数据正确搬出到 GM。L2 Cache 配置// 根据 scheduler 参数禁用 A/B 的 L2 Cache if (bs.GetBL2CacheDisable()) { gmB.SetL2CacheHint(AscendC::Te::CacheMode::CACHE_MODE_DISABLE); } if (bs.GetAL2CacheDisable()) { gmA.SetL2CacheHint(AscendC::Te::CacheMode::CACHE_MODE_DISABLE); }说明可选禁用 A/B 矩阵的 L2 Cache避免大矩阵场景下的缓存污染。Tile 循环策略Basic 特有的循环策略for (int64_t tileIdx curBlockIdx; tileIdx tileNum; tileIdx blockNum) { // 每个 block 处理 tileIdx curBlockIdx, curBlockIdx blockNum, curBlockIdx 2*blockNum, ... // 多 block 并行处理不同 tile // 每个 tile 独立调用 BlockMmad 执行矩阵乘 }说明简单 stride 策略block 按固定 strideblockNum遍历 tile无 workspace每个 tile 完整计算不依赖中间结果无 AIC-AIV 同步单核计算无需跨核同步BlockScheduler 功能Basic 特有功能说明GetTileNum获取总 tile 数量GetTileL1Shape获取 L1 tile 形状GetTileL0Shape获取 L0 tile 形状GetBlockNum计算实际需要的 block 数量GetBlockCoord获取 tile 的坐标GetBlockShape获取 tile 的形状Gethf32Flag获取 HF32 模式标志GetL1BuferNum获取 L1 缓冲数量GetL0cDB获取 L0C 双缓冲标志GetBL2CacheDisable获取 B 矩阵 L2 Cache 禁用标志GetAL2CacheDisable获取 A 矩阵 L2 Cache 禁用标志调用示例Kernel 组装与调用// 1. 类型定义 using AType half; // A 矩阵数据类型 using BType half; // B 矩阵数据类型 using CType half; // C 矩阵输出类型 using BiasType half; // Bias 数据类型可选 using LayoutA AscendC::Te::NZLayoutPtn; // A 矩阵布局NZ/ND using LayoutB AscendC::Te::NZLayoutPtn; // B 矩阵布局NZ/ND using LayoutC AscendC::Te::NDLayoutPtn; // C 矩阵布局ND using LayoutBias LayoutC; // Bias 布局 // 2. ProblemShape 定义 // 形状(m, n, k, batch)batch0 或 1 表示单 batch using ProblemShape AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t; // 3. BlockScheduler 组装 // FullLoadMode: 0非全载默认, 1A全载, 2B全载 constexpr int64_t FULL_LOAD_MODE 0; using BlockScheduler Blaze::Gemm::Block::BlockSchedulerMatmulBasic ProblemShape, FULL_LOAD_MODE; // 4. BlockMmad 组装 // DispatchPolicy: 调度策略FusedOpType: 融合操作类型 constexpr uint64_t FUSED_OP_TYPE 0; using DispatchPolicy Blaze::Gemm::MatmulMultiBlockBasicFULL_LOAD_MODE, FUSED_OP_TYPE; using BlockMmad Blaze::Gemm::Block::BlockMmad DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias; // 5. BlockEpilogue 组装 // Basic Kernel 仅支持 Empty不支持后处理 using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueEmpty; // 6. Kernel 组装 using MatmulKernel Blaze::Gemm::Kernel::KernelMatmulBasic ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler; // 7. Params 构造 using Params typename MatmulKernel::Params; Params params; // --- ProblemShape 参数 --- params.problemShape {m, n, k, batch}; // (M, N, K, Batch) // --- BlockMmad 参数 --- params.mmadParams.aGmAddr aGM; // A 矩阵 GM 地址 params.mmadParams.bGmAddr bGM; // B 矩阵 GM 地址 params.mmadParams.cGmAddr cGM; // C 矩阵 GM 地址 params.mmadParams.biasGmAddr biasGM; // Bias GM 地址可选nullptr 表示无 bias // --- BlockScheduler 参数 --- // L1 tile 形状决定每个 tile 的 M/N/K 轴尺寸 params.schedulerParams.mL1 256; // M 轴 L1 tile 尺寸 params.schedulerParams.nL1 256; // N 轴 L1 tile 尺寸 params.schedulerParams.kL1 128; // K 轴 L1 tile 尺寸 // L0 base 形状决定每次 Mmad 计算的 M/N/K 轴尺寸 params.schedulerParams.baseM 128; // M 轴 L0 base 尺寸 params.schedulerParams.baseN 128; // N 轴 L0 base 尺寸 params.schedulerParams.baseK 64; // K 轴 L0 base 尺寸 // 尾块切分Batch1 场景提升尾块并行度 params.schedulerParams.mTailCnt 2; // M 轴尾块切分数量建议 1~4 params.schedulerParams.nTailCnt 2; // N 轴尾块切分数量建议 1~4 // L1 尾块切分矩阵不能被 tile 整除时使用 params.schedulerParams.mBaseTailSplitCnt 1; // M 轴 L1 尾块切分数量建议 1 params.schedulerParams.nBaseTailSplitCnt 1; // N 轴 L1 尾块切分数量建议 1 params.schedulerParams.mTailMain 1; // M 轴 L1 尾块主尺寸 params.schedulerParams.nTailMain 1; // N 轴 L1 尾块主尺寸 // 双缓冲配置提升数据搬运与计算并行度 params.schedulerParams.l1BufferNum 2; // L1 缓冲数量1单缓冲, 2双缓冲 params.schedulerParams.l0cDB 2; // L0C 双缓冲1单缓冲, 2双缓冲 params.schedulerParams.ubDB 2; // UB 双缓冲1单缓冲, 2双缓冲 // HF32 模式可选用于特定精度场景 params.schedulerParams.isHf32 0; // HF32 模式标志0关闭 // L2Cache 配置可选控制 A/B 矩阵 L2Cache 行为 params.schedulerParams.l2CacheDisable Blaze::Gemm::L2CacheMode::L2_CACHE_DEFAULT; // L2Cache 使能默认 // 其他选项 // A_L2_CACHE_DISABLE 禁用 A 矩阵 L2Cache // B_L2_CACHE_DISABLE 禁用 B 矩阵 L2Cache // ALL_L2_CACHE_DISABLE 禁用所有 L2Cache // 非连续场景参数连续 ND 格式不需要设置 params.schedulerParams.sliceM 0; // M 轴 slice 尺寸非连续场景 params.schedulerParams.srcNdStride 0; // stride非连续场景 params.schedulerParams.innerBatch 1; // transpose 内轴 batch // --- BlockEpilogue 参数 --- // Empty Epilogue 无需设置参数 params.epilogueParams {}; // 8. Kernel 调用 MatmulKernel mm; mm(params); // 执行矩阵乘计算参数详解ProblemShape 参数参数类型说明示例mint64_tM 轴尺寸1024nint64_tN 轴尺寸1024kint64_tK 轴尺寸512batchint64_tBatch 数量0 或 1 为单 batch1BlockMmad 参数参数类型说明示例aGmAddrGM_ADDRA 矩阵 GM 地址aGMbGmAddrGM_ADDRB 矩阵 GM 地址bGMcGmAddrGM_ADDRC 矩阵 GM 地址cGMbiasGmAddrGM_ADDRBias GM 地址nullptr 表示无 biasbiasGM 或 nullptrBlockScheduler 参数详见 BlockSchedulerMatmulBasic 参数详解常用配置示例小矩阵场景params.schedulerParams.mL1 128; params.schedulerParams.nL1 128; params.schedulerParams.kL1 64; params.schedulerParams.baseM 64; params.schedulerParams.baseN 64; params.schedulerParams.baseK 32; params.schedulerParams.l1BufferNum 1; // 单缓冲 params.schedulerParams.l0cDB 1; params.schedulerParams.ubDB 1;大矩阵场景params.schedulerParams.mL1 256; params.schedulerParams.nL1 256; params.schedulerParams.kL1 128; params.schedulerParams.baseM 128; params.schedulerParams.baseN 128; params.schedulerParams.baseK 64; params.schedulerParams.l1BufferNum 2; // 双缓冲 params.schedulerParams.l0cDB 2; params.schedulerParams.ubDB 2; params.schedulerParams.l2CacheDisable Blaze::Gemm::L2CacheMode::ALL_L2_CACHE_DISABLE;尾块优化场景Batch1params.schedulerParams.mTailCnt 4; // 尾块切为 4×4 16 份 params.schedulerParams.nTailCnt 4; // 16 个 Block 并行处理尾块数据流存储层次GM (A/B/Bias) → BlockScheduler (tile 切分) → L1 (双缓冲) → L0A/L0B (双缓冲) → L0C (双缓冲) → GM (C)Kernel 执行流程BlockScheduler 初始化 ↓ 获取 tile 数量和形状 (L1/L0 tile) ↓ HF32 模式设置可选 ↓ BlockMmad 初始化 (设置缓冲策略) ↓ 创建 GM Tensor (ND/NZ layout) ↓ 配置 L2 Cache (可选禁用) ↓ 遍历 tile → BlockMmad 执行 (每个 tile 独立计算) ↓ 清理 (关闭 HF32/MM Layout Transform)性能优化建议Basic 特有L2 Cache 配置大矩阵场景建议禁用 L2 Cache 避免缓存污染小矩阵场景可保留 L2 Cache 提升数据复用Bias 预处理常量 bias可预先处理减少运行时开销动态 bias保持实时加载NZ 格式优化权重矩阵B优先使用 NZ 格式提升 L1/L0 搬运效率激活矩阵A使用 ND 格式即可K 轴切分kL1和baseK应根据数据局部性和复用率优化避免 K 轴切分过于细碎导致搬运开销增加适用场景小矩阵(m × n × k) 较小时Basic Kernel 更高效简单计算无复杂后处理需求时Basic Kernel 足够单核场景仅需 AIC 计算时Basic Kernel 更简洁【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考