CANN/catlass矩阵乘API

CATLASS Gemm API【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlassCATLASS针对NPU上不同层级上执行的矩阵乘累加MMAD操作提供了一个统一的编程模型。CATLASS的Gemm API对应于以下分层由高到低分别是CATLASS Gemm模型CATLASS基于上述的分层结构实现了经典“三层嵌套循环”的矩阵乘算法。以下伪代码描述了针对像mmad这样的单核内同步矩阵乘法指令的Matmul 内核的模型。整个算子被称为“Matmul”这是伪代码仅用于说明哪些层次的部分对应于矩阵乘的内部或外部循环。// Catlass::Gemm::Kernel::BasicMatmul: BlockTileM 和 BlockTileN 循环 // 在AICores上并行 for (int block_m 0; block_m MatmulM; block_m BlockTileM) { for (int block_n 0; block_n MatmulN; block_n BlockTileN) { // Catlass::Gemm::Block::BlockMmad: 在k-tile上迭代的主循环 // 在这个阶段没有循环展开 for (int k_tile 0; k_tile MatmulK; k_tile) { // Tile级内循环: (m,k) x (k,n) (m,n) // TileMmad使用硬件指令 AscendC::Mmad for (int tile_mma_m 0; tile_mma_m m; tile_mma_m) { for (int tile_mma_n 0; tile_mma_n n; tile_mma_n) { for (int tile_mma_k 0; tile_mma_k k; tile_mma_k) { mmad.call(c, a, b); } // tile_mma_k } // tile_mma_n } // tile_mma_m } // k_tile mainloop } // block_n } // block_m前两重嵌套的for循环对应于多个AICore上的并行性。代码上并不将它们显式地表达为两重for循环而是通过BlockIdx来区分不同核心处理的数据块。在两重嵌套的for循环内部将全局内存分片然后将分片搬运到更“局部”的内存如L1 Buffer或L0 Buffer并执行MMAD计算。这些分片拷贝和分片MMAD计算的迭代通常是完全静态的并且完全展开。CATLASS Gemm组件CATLASS使用以下组件表达上述循环嵌套这些组件针对数据类型、数据排布和数学指令进行特化。API 层级API 类 和/或 函数 名称DeviceCatlass::Gemm::Device::DeviceGemmKernelCatlass::Gemm::Kernel::BasicMatmulBlockCatlass::Gemm::Block::BlockMmadCatlass::Epilogue::Block::BlockEpilogueTile (MMAD and Copy)TileMmadandTileCopyBasicAscendC::MmadandAscendC::DataCopy在CATLASS中我们通过首先在Kernel层组合Block主循环和Block后处理然后用主机侧适配器包装它们来组装内核。用户使用这些组件组装内核时需要通过以下顺序实例化。组装所需的Block主循环和Block后处理。将Blocks组合在一起构建成Kernel。用Device层适配器包装Kernel。这个顺序也反映在CATLASS的示例中examples/00_basic_matmul如下文摘录所示。// 第一步: 创建所需的特化Block层mmad // 参数 using DispatchPolicy Gemm::MmadAtlasA2Pingpongtrue; using L1TileShape GemmShape128, 256, 256; using L0TileShape GemmShape128, 256, 64; using AType Gemm::GemmTypeElementA, LayoutA; using BType Gemm::GemmTypeElementB, LayoutB; using CType Gemm::GemmTypeElementC, LayoutC; using BlockMmad Gemm::Block::BlockMmadDispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType; // 第二步指定Block层的后处理类型可选 using BlockEpilogue void; // 第三步指定计算时的数据走位方式 using BlockScheduler typename Gemm::Block::GemmIdentityBlockSwizzle; // 第四步在kernel层将mmad和后处理组合到一起 using MatmulKernel Gemm::Kernel::BasicMatmulBlockMmad, BlockEpilogue, BlockScheduler; // 第五步将kernel放入device适配器中host侧处理kernel使用 using Matmul Catlass::Gemm::Device::DeviceGemmMatmulKernel;Block APIBlock API包括“矩阵乘累加”和后处理。它负责实现上述三层嵌套循环伪代码中的k_tile循环。Block在昇腾NPU的SPMD编程模型中指一个Process是逻辑核的概念。它提供的硬件特性包括异步内存拷贝例如从全局内存到L1 Buffer适用于L0 Buffer中的分片粒度数据的MMAD指令多核间、核内不同流水线间的同步操作用以确保异步操作之间的数据依赖关系得到满足。Block使用TileMma和TileCopyAPI见下文来执行分片粒度的数据拷贝和MMAD运算。Block中的不同硬件流水线例如MTE1、MTE2或FixPipe提供不同能力不同的硬件流水线间需要共享数据并协调对共享数据的访问。例如MTE2将数据从全局内存拷贝到L1 Buffer后需要让MTE1知道输入已准备好。这里Kernel::层负责对Block::层接口进行调用Block::层接口负责独立的C矩阵分块计算。Block MmadCatlass::Gemm::Block::BlockMmad是矩阵乘Block层级累加MMAD主循环的主要接口。BlockMmad类定义在头文件中 include/catlass/gemm/block/block_mmad.hpp.namespace Catlass::Gemm::Block { //////////////////////////////////////////////////////////////////// template class DispatchPolicy, class L1TileShape, class L0TileShape, class AType, class BType, class CType, class BiasType void, class TileCopy Gemm::Tile::TileCopytypename DispatchPolicy::ArchTag, AType, BType, CType, BiasType, class TileMmad Gemm::Tile::TileMmadtypename DispatchPolicy::ArchTag, AType, BType, BiasType struct BlockMmad {}; //////////////////////////////////////////////////////////////////// } // namespace Catlass::Gemm::BlockDispatchPolicy是Block层重要的参数之一下一节会详细介绍。L1TileShape和L0TileShape对应L1 Buffer和L0 Buffer上使用的基本块大小后续详细介绍。AType、BType、CType、BiasType是GemmType的实例其中包含了全局内存上A、B、C矩阵和Bias向量的数据类型和数据排布。TileCopy是Tile::TileCopy的实例包含了不同访存层级间的块粒度数据拷贝如全局内存到L1 BufferL1 Buffer到L0 Buffer等。TileMmad是Tile::TileMmad的实例完成L0上基本块粒度的矩阵乘累加运算。Block Dispatch PoliciesBlockMmad的实现不是通用的。相反它们必须针对每个算法和NPU架构特例化。用户可以通过选择与该特例化匹配的模板参数来调度到BlockMmad的特例化。 CATLASS采用基于标签的调度策略类型来特例化Block层Mmad实现并为其提供调优能力。以下给出了一个Dispatch Policy的样例对应AtlasA2的架构下采用L1 Buffer上pingpong Buffer启用unitflag优化// 2-Buffer in L1 Buffer , // unitflag enable struct MmadAtlasA2Pingpong { using ArchTag Arch::AtlasA2; static constexpr uint32_t STAGES 2; static constexpr bool ENABLE_UNIT_FLAG true; };STAGES参数使用户可以方便地调整多buffer场景的buffer片数ENABLE_UNIT_FLAG参数用于表示是否启用Mmad运算与L0C结果拷贝到全局内存的细粒度并行。采用Dispatch Policy的设计还有如下优点它避免了代码重复主循环可以被多个不同的内核使用。它使编写通用代码更容易因为主要类型名称BlockMmad在任何实现中都不会改变。它提供了一个清晰、单一的扩展点供用户插入针对他们自己调度策略特化实现的新的、定制的主循环。TileShapeL1TileShape和L0TileShape对应L1 Buffer和L0 Buffer上使用的基本块大小由m,n,k表示Epilogue尾处理实现了涉及输出矩阵的逐元素操作。用户可以提供自定义的尾处理或者使用标准尾处理之一。这些尾处理位于目录include/catlass/epilogue/block/中包括像Catlass::Epilogue::Block::BlockEpilogue这样的类。CATLASS提供的尾处理不在include/catlass/gemm目录下也不在Catlass::Gemm命名空间中因为它们可以用于除Gemm之外的其他计算。Kernel APIKernel对应了所有Block在NPU上执行逻辑的集合。Kernel层BasicMatmul承担以下功能对包含的不同Block的逻辑进行组合加入必要的同步逻辑。不同Block和处理全局内存上数据的对应关系Swizzling。将输入数据在Block粒度分片。Kernel层API是设备侧调用的入口也是融合连续矩阵乘、尾处理或其他操作的组合点。Kernel API 入口在Catlass::Gemm::Kernel::BasicMatmul, 位于头文件 include/catlass/gemm/kernel/basic_matmul.hpp.BasicMatmul是一个无状态的设备侧内核实现的矩阵乘运算由两部分组成Block MmadBlock Epiloguenamespace Catlass::Gemm::Kernel { template class BlockMmad_, class BlockEpilogue_, class BlockScheduler_ class BasicMatmul; } // namespace Catlass::Gemm::Kernel注无状态指调用者管理着内核的状态例如上述描述的设备API。内核仅接收输入和输出参数 (Params).其中Block Mmad代表局部数据块上的矩阵乘加运算Block Epilogue代表mmad之后的运算例如C : beta * C alpha * A * B中的beta * C。Device APIDevice层是Host侧调用的入口在这一层屏蔽调用Device侧函数的差异。用户定义完Kernel结构之后放入Device层模板便可以执行算子。using BlockMmad Gemm::Block::BlockMmadDispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType; using BlockEpilogue void; using BlockScheduler typename Gemm::Block::GemmIdentityBlockSwizzle; // kernel using MatmulKernel Gemm::Kernel::BasicMatmulBlockMmad, BlockEpilogue, BlockScheduler; // device using Matmul Gemm::Device::DeviceGemmMatmulKernel; Matmul matmulOp; // args为维度信息等参数的封装结构体 matmulOp(args, workspace, stream);Tile MMAD and CopyTile粒度的MMAD和Copy是对基础API的MMAD和数据拷贝接口的组合这一层的目的是构建可组合的NPU微内核这些微内核由硬件加速的数学运算和数据拷贝操作组成每个操作都有其数据类型和排布。Tile粒度的MMAD和Copy提供了不同硬件上完成相同计算或数据拷贝语义的统一API。用户可以在本文档顶部的三重嵌套循环伪代码的“内层”循环中使用Tile::TileMmad()或Tile::CopyGmToL1()、Tile::CopyL0CToGm()等来调用这些操作。我们将这个API层级称为“Tile”因为它使用基础API提供的原子能力去构建更大粒度的操作作为一个可重用组件它就像将单独的瓷砖拼接在一起构建成马赛克的图案。Basic APIBasic层级API封装了实际的硬件指令调用这些指令加速了MMAD或数据拷贝操作对应CATLASS的基础API实现了对硬件能力的抽象开放了芯片能力保证了完备性和兼容性其中ISASI类API不保证跨硬件版本兼容。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考