CANN/catlass列广播乘法API

TileOneBlkColumnBroadcastMul

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass

代码位置

[TOC]

功能说明

TileOneBlkColumnBroadcastMul实现 epilogue 阶段的列广播乘法操作。将形状 (m, 1) 的列向量在 block 内广播到 (m, n) 后与输入相乘。broadcast 粒度为一个 block（BYTE_PER_BLK字节），即 column 上的每 1 个元素广播到 1 个完整的 block。

• 适用范围：所有架构（无架构特化）
• 风格：非 TLA

模板原型

template < class ArchTag_, // 架构标签 class ComputeType_, // 计算数据类型 class TileShape_ // Tile 形状 > struct TileOneBlkColumnBroadcastMul;

调用接口

void operator()( AscendC::LocalTensor<ElementCompute> const &ubOut, AscendC::LocalTensor<ElementCompute> const &ubIn0, AscendC::LocalTensor<ElementCompute> const &ubIn1 // (m, eleNumPerBlk) 形状 )

通过AscendC::Mul+BinaryRepeatParams（src1RepStride = 0,src1BlkStride = 1）实现列广播。

调用示例

#include "catlass/epilogue/tile/tile_broadcast_mul.hpp" using namespace Catlass::Epilogue::Tile; using ComputeType = Gemm::GemmType<half, layout::RowMajor>; using TileShape = Shape<128, 256>; using ColumnBroadcastMul = TileOneBlkColumnBroadcastMul<Arch::AtlasA2, ComputeType, TileShape>; AscendC::LocalTensor<half> ubOut, ubIn0, ubIn1; ColumnBroadcastMul op; op(ubOut, ubIn0, ubIn1);

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass