CANN混合矩阵乘优化

hir.mix_matmul — 混合 CubeVector 矩阵乘【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills关键词mix_matmul, Mix Matrix Multiply, Post-Vector Function, Workspace, Communication概述hir.mix_matmul是 HIVM 方言中的混合矩阵乘操作在 Cube Core 执行矩阵乘法后支持在 Vector Core 上执行后处理函数post-vector function。这种 CubeVector 的混合模式允许在 tile 级别融合后处理操作如激活函数、类型转换等避免额外的 GM 读写开销。计算语义与 matmul 相同C descale * (A * B bias)但额外支持post_vector_func_insVector 后处理函数的输入参数workspace_ins工作空间缓冲区comm_params通信参数用于融合通信操作如 AllReducePython API 对应Triton 的tl.dot 后续 elementwise 操作在编译时可能被融合为 mix_matmul。IR 操作定义从 HIVMMacroOps.td 提取def MixMatmulOp : HIVM_GlobalMmadOpmix_matmul { let summary HIVM (Mix) Matrix Multiply Op with inputs from global memory; let arguments (ins AnyShaped:$a, AnyShaped:$b, VariadicAnyShaped:$postVecFuncIns, VariadicAnyShaped:$workspaceIns, OptionalAnyShaped:$tilingParams, OptionalAnyShaped:$commParams, OptionalAnyShaped:$bias, OptionalAnyShaped:$descale, OptionalAttrUnitAttr:$aTranspose, OptionalAttrUnitAttr:$bTranspose, OptionalAttrHIVM_DescaleModeAttr:$descaleMode, VariadicI64:$blockSizes, VariadicI64:$processSizes, OptionalI64:$swizzleOffset, OptionalI64:$swizzleDirection, OptionalI64:$epiloguePTiles, AnyShaped:$c); }参数说明输入操作数ins参数类型必选说明$aAnyShaped是矩阵 Am x k$bAnyShaped是矩阵 Bk x n$postVecFuncInsVariadicAnyShaped否Vector 后处理函数的输入参数$workspaceInsVariadicAnyShaped否工作空间缓冲区输入$tilingParamsAnyShaped否Tiling 参数$commParamsAnyShaped否通信相关参数拓扑、通信器、group 等$biasAnyShaped否Bias 向量$descaleAnyShaped否反量化缩放因子$cAnyShaped是矩阵 C输出输出操作数outs参数类型说明$resultVariadicAnyRankedTensor结果 Tensor属性属性类型必选说明$aTransposeUnitAttr否矩阵 A 转置加载$bTransposeUnitAttr否矩阵 B 转置加载$descaleModeHIVM_DescaleModeAttr否反量化模式post_vector_funcStrAttr否Vector 后处理函数名称通过属性指定I64 操作数参数类型必选说明$blockSizesVariadicI64否L1 层次 M/N/K 块大小$processSizesVariadicI64否L0 层次 M/N/K 块大小$swizzleOffsetI64否Swizzle 起始块编号$swizzleDirectionI64否Swizzle 方向$epiloguePTilesI64否Epilogue P tile 数量与 matmul 的差异特性matmulmix_matmul后处理融合不支持支持post_vector_func_ins工作空间不支持支持workspace_ins通信参数不支持支持comm_paramsOpNamematmulmix_matmulIR 示例基本用法hivm.hir.mix_matmul ins(%A_gm, %B_gm: memref16x16xf16, #hivm.address_spacegm, memref16x16xf16, #hivm.address_spacegm) outs(%res_gm : memref16x16xf16, #hivm.address_spacegm) tiling_params %tiling_params_gm : memref16xf16, #hivm.address_spacegm comm_params %comm_params_gm : memref16xi64, #hivm.address_spacegm带 Post-Vector Function 和 Workspacehivm.hir.set_ffts_base_addr %arg0 hivm.hir.mix_matmul {post_vector_func bishengir_gen_vector_epilogue_func} ins(%arg1, %arg2 : memref1024x1024xf16, #hivm.address_spacegm, memref1024x1024xf16, #hivm.address_spacegm) post_vector_func_ins(%arg3 : memref1024x1024xf16, #hivm.address_spacegm) workspace_ins(%arg4 : memref1024x1024xf16, #hivm.address_spacegm) outs(%arg5 : memref1024x1024xf16, #hivm.address_spacegm) block_sizes(%c128_i64, %c256_i64, %c256_i64 : i64, i64, i64) process_sizes(%c128_i64, %c256_i64, %c64_i64 : i64, i64, i64) swizzle_offset %c1_i64 : i64 swizzle_direction %c0_i64 : i64 epilogue_p_tiles %c4_i64 : i64IR 层约束与验证Core Type通过HIVMInferCoreTypeInterface推断通常为 Cube Core 执行矩阵乘法部分。Pipeline涉及 MTE2 和 MTE3 两个 Pipeline。Post-Vector Functionpost_vector_func属性指定 Vector 后处理函数名称postVecFuncIns提供其输入参数。WorkspaceworkspaceIns提供工作空间缓冲区用于中间计算结果存储。CommunicationcommParams用于融合通信操作如 AllReduce包含拓扑、通信器等信息。MIX Kernel使用 mix_matmul 的函数通常标记为hivm.func_core_type #hivm.func_core_typeAIV或MIX。常见问题Q: mix_matmul 的 post_vector_func 是什么A: 它是 Vector Core 上执行的后处理函数可以对矩阵乘结果进行激活函数、类型转换等操作。函数名通过post_vector_func属性指定输入通过postVecFuncIns传入。Q: workspace_ins 的用途A: 工作空间缓冲区用于存储中间计算结果例如在融合 AllReduce 时需要临时存储部分和。Q: 什么时候应该用 mix_matmul 而不是 matmulA: 当需要在矩阵乘后立即执行 Vector 后处理如激活、量化或需要融合通信操作时使用 mix_matmul 可以避免额外的 GM 读写。相关文档源码参考HIVMMacroOps.td测试用例ops.mlirmatmul 详解03-matmul.md【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考