CANN/ops-tensor StreamK矩阵乘后处理块

Block Epilogue StreamK【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明StreamK 矩阵乘后处理 Block运行在 AIV 核。从 workspace 读取 AIC 计算的中间结果执行 K 轴切分累加Add、类型转换Cast、ReLU 激活并输出到 GM。继承自Block Epilogue 基础框架特殊约束AIV 核专用仅运行在 AIVVector核与 AICCube核协同工作AIC 核执行矩阵乘计算结果输出到 workspaceAIV 核从 workspace 读取中间结果执行后处理Workspace 必需必须提供 workspace 用于存储 AIC 计算的中间结果workspace 数据类型WorkspaceType通常为 float大小tailMNTileNum × skKTileNum × BLOCK_BASE_M × BLOCK_BASE_N × sizeof(WorkspaceType)输出类型支持两种输出类型float → float无需 Cast直接输出float → half/bfloat16执行 Cast 转换后输出ReLU 支持float 输出支持 ReLU非 bfloat16_thalf 输出支持 ReLUbfloat16_t 输出不支持 ReLU硬件限制L0C2Out 模式支持两种 Fixpipe 输出模式ON_THE_FLYCompact 模式紧凑布局ND_FIXPIPE_1_2Normal 模式stride 对齐到 32BAIV 并行度SK 模式aivMte2Num_ GetTaskRation()多个 AIV 并行处理一个 AICDP 模式aivMte2Num_ BLOCK_CUBE单个 AIV 处理数据对齐burstLen对齐到 32BBLOCK_SIZEsrcGapworkspace 中 K 轴切分之间的间隔dstGapGM 输出中 N 轴的间隔n_ - curNL1InAiv特殊静态常量常量说明BLOCK_BASE_MBlock 基础 M 维度256BLOCK_BASE_NBlock 基础 N 维度256NUM_TWO数值常量2MAIN_WINDOW主窗口大小4UB2GM_SRCGAP_UNITUB2GM srcGap 对齐单位32特殊类型别名类型说明WorkspaceTypeWorkspace 数据类型模板参数OutType输出数据类型模板参数DispatchPolicy调度策略模板参数特殊数据结构Argumentsstruct Arguments { GM_ADDR cGmAddr{nullptr}; // C 矩阵 GM 地址 GM_ADDR workspaceGmAddr{nullptr}; // Workspace GM 地址 };Paramsusing Params Arguments;AivParamsstruct AivParams { uint64_t indexParams 0; // AIV 处理的 tile 索引 uint64_t mCntIndex 0; // M 轴 tile 索引 uint64_t nCntIndex 0; // N 轴 tile 索引 uint64_t kCntIndex 0; // K 轴切分索引 uint64_t curML1InAiv 0; // 当前 tile M 维度 uint64_t curNL1InAiv 0; // 当前 tile N 维度 uint64_t curAlignedNInAiv 0; // 对齐后的 N 维度ND_FIXPIPE_1_2 模式 };CopyGm2UbParamsstruct CopyGm2UbParams { uint64_t offsetWorkspaceGM 0; // workspace 起始偏移 uint64_t kCnt 0; // K 轴切分数量 uint64_t mBurstOri 0; // 原始 M burst 数量 uint64_t mBurst 0; // 当前 AIV 处理的 M burst 数量 uint64_t burstLen 0; // 每次 burst 的数据长度 uint64_t srcGap 0; // workspace 中 K 切分之间的间隔 };CopyUb2GmParamsstruct CopyUb2GmParams { uint64_t offsetCGm 0; // GM 输出偏移 uint64_t mLength 0; // M 维度长度 uint64_t burstLen 0; // burst 长度N 维度 uint64_t dstGap 0; // GM 输出 N 轴间隔 uint64_t srcGap 0; // UB 中 N 轴间隔 };特殊成员方法构造函数__aicore__ inline BlockEpilogueStreamK()功能构造 BlockEpilogueStreamK 对象。析构函数__aicore__ inline ~BlockEpilogueStreamK()功能析构 BlockEpilogueStreamK 对象。Init函数__aicore__ inline void Init( Params const params, // 参数cGmAddr, workspaceGmAddr BlockShape blockShapeInAiv, // 问题规模 (m, n, k) BlockShape tileL1ShapeInAiv, // L1 tile 形状 (mL1, nL1, kL1) BlockCoord coordInAiv, // tile 坐标 (mCnt, nCnt, kCnt) uint64_t usedCoreNum, // 使用的核数 bool checkIsSkScene) // 是否为 SK 模式功能初始化 BlockEpilogueStreamK 组件。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | params | Params | 包含 cGmAddr 和 workspaceGmAddr | | blockShapeInAiv | BlockShape | 问题规模(m, n, k)| | tileL1ShapeInAiv | BlockShape | L1 tile 形状(mL1, nL1, kL1)| | coordInAiv | BlockCoord | tile 坐标(mCnt, nCnt, kCnt)| | usedCoreNum | uint64_t | 使用的核数 | | checkIsSkScene | bool | true SK 模式false DP 模式 |执行流程设置问题规模m_,n_,mL1_,nL1_设置 tile 坐标mCnt_,nCnt_,kCnt_设置 AIV 并行度aivMte2Num_设置 GM buffercGlobal_,workspaceGlobal_ICache 预加载2 个 cache lineDP 模式同步设置并等待MTE3_MTE2标志Run函数__aicore__ inline void Run()功能执行 StreamK 后处理。 执行流程更新索引UpdateAivBasicIndex()更新 AIV tile 索引更新块参数UpdateAivBasicBlock()更新当前 tile 形状AIV 并行循环按aivMte2Num_循环更新参数UpdateAivParams(index)搬运 workspace → UBDataCopyPadfloat等待 MTE2 完成WaitFlagMTE2_VK 轴切分累加循环执行AddReLU 激活可选Relu类型转换可选Castfloat → half/bfloat16等待 V 完成WaitFlagV_MTE3搬运 UB → GMDataCopyPadOutTypeoperator函数__aicore__ inline void operator()()功能执行后处理调用Run()。UpdateAivBasicIndex函数__aicore__ inline void UpdateAivBasicIndex()功能更新 AIV tile 索引mCntIndex, nCntIndex, kCntIndex。 算法说明计算newBlockIdxGetBlockIdx() / (GetTaskRation() * kCnt_)计算kCntIndexGetBlockIdx() % (GetTaskRation() * kCnt_)计算cGmIndex基于 DPSK 混合策略的索引扫描窗口计算mainWindow min(4, mCnt_)Z 型扫描偶数行正向奇数行反向UpdateAivBasicBlock函数__aicore__ inline void UpdateAivBasicBlock()功能更新当前 tile 形状curML1InAiv, curNL1InAiv, curAlignedNInAiv。 算法说明curML1InAiv尾块使用剩余尺寸否则使用mL1_curNL1InAiv尾块使用剩余尺寸否则使用nL1_curAlignedNInAivND_FIXPIPE_1_2 模式对齐到 32BUpdateAivParams函数__aicore__ inline void UpdateAivParams(uint64_t index)功能更新 AIV 数据搬运参数copyGm2UbParams, copyUb2GmParams。 算法说明mBurstBaseCeilAlign(curML1InAiv / (kCnt_ * GetTaskRation()), 32 / curAlignedNInAiv)mBurstCntCeilDiv(curML1InAiv, mBurstBase)mBurstTailcurML1InAiv - (mBurstCnt - 1) * mBurstBaseoffsetWorkspaceGM(indexParams * kCnt_ * BLOCK_BASE_M * BLOCK_BASE_N) (kCntIndex * mBurstBase mBurst * index) * curAlignedNInAivoffsetCGmnCntIndex * nL1_ mCntIndex * mL1_ * n_ (kCntIndex * mBurstBase mBurst * index) * n_burstLenCeilAlign(mBurst * curAlignedNInAiv, 32)srcGapBLOCK_BASE_M * BLOCK_BASE_N - burstLendstGapn_ - curNL1InAiv事件同步事件用途MTE3_MTE2DP 模式下 AIC 与 AIV 同步MTE2_Vworkspace → UB 完成同步V_MTE3UB → GM 完成同步说明DP 模式下需要等待 AIC 完成写入 GMSK 模式下直接从 workspace 读取。调用示例组件组装using WorkspaceType float; using OutType half; using DispatchPolicy Blaze::Gemm::MatmulMultiBlockWithStreamKBlaze::Gemm::MatMulL0C2Out::ON_THE_FLY; using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueStreamK WorkspaceType, OutType, DispatchPolicy;参数准备using Params typename BlockEpilogue::Params; Params params { cGmAddr, // C 矩阵 GM 地址 workspaceGmAddr // Workspace GM 地址 };组件初始化BlockEpilogue blockEpilogue; BlockShape blockShape{m, n, k}; BlockShape tileL1Shape{mL1, nL1, kL1}; BlockCoord coord{mCnt, nCnt, kCnt}; uint64_t usedCoreNum 8; bool checkIsSkScene true; // SK 模式 blockEpilogue.Init(params, blockShape, tileL1Shape, coord, usedCoreNum, checkIsSkScene);组件执行blockEpilogue(); // 或 blockEpilogue.Run();数据流存储层次Workspace (float) → UB → Add 汇聚 → Cast → ReLU → GM (OutType)执行流程UpdateAivBasicIndex (更新 tile 索引) ↓ UpdateAivBasicBlock (更新 tile 形状) ↓ AIV 并行循环aivMte2Num_ 次 ↓ UpdateAivParams (更新搬运参数) ↓ DataCopyPad: Workspace → UB (float) ↓ WaitFlagMTE2_V ↓ Add 循环K 轴切分累加 ↓ ReLU可选 ↓ Cast可选float → half/bfloat16 ↓ WaitFlagV_MTE3 ↓ DataCopyPad: UB → GM (OutType)Workspace 布局offsetWorkspaceGM indexParams * kCnt_ * BLOCK_BASE_M * BLOCK_BASE_N (kCntIndex * mBurstBase mBurst * index) * curAlignedNInAiv说明每个 tile 的 K 轴切分结果按顺序存储在 workspace 中间隔为BLOCK_BASE_M * BLOCK_BASE_N。GM 输出布局offsetCGm nCntIndex * nL1_ mCntIndex * mL1_ * n_ (kCntIndex * mBurstBase mBurst * index) * n_说明按 Z 型扫描顺序输出到 GM。性能优化建议AIV 并行度配置SK 模式aivMte2Num_ GetTaskRation()多 AIV 并行处理一个 AICDP 模式aivMte2Num_ BLOCK_CUBE单 AIV 处理建议SK 模式下增加 AIV 数量提升并行度K 轴切分累加累加循环for (i 1; i kCnt; i) { Add(ub, ub, ub[i * burstLen], burstLen); }首次 K 切分数据在ub[0]后续切分在ub[burstLen]开始类型转换优化float → half使用Cast(ubDst, ubSrc, RoundMode::CAST_RINT, burstLen)ReLU 位置float 输出在 Cast 前half 输出在 Cast 后L0C2Out 模式选择ON_THE_FLYCompact 模式srcGap计算alignedN - actualNND_FIXPIPE_1_2Normal 模式curAlignedNInAiv对齐到 32B数据对齐burstLen对齐到 32BsrcGap确保 K 轴切分数据正确间隔dstGap确保 GM 输出 N 轴正确间隔适用场景StreamK KernelAIC AIV 双核协同K 轴切分汇聚多个 K 切分结果累加类型转换float → half/bfloat16ReLU 激活可选 ReLU 后处理【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考