企业官网建设流程全解析

CANN cann-samples 系统优化实战：MatMul 尾轮负载均衡与 Stream-K 多核切分优化

【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples

在昇腾 NPU 的 AI Core 多核矩阵乘法中，计算任务按基本块（Tile）分发到各核算数时，最后一轮往往会出现"剩余块少于核数"或"单核承载过大的 K 轴累加"两类负载不均问题。本文基于 CANN 开源样例仓库 cann-samples 的 系统优化目录，完整讲解其中两套典型的多核负载均衡策略——尾轮负载均衡（tail_rebalance）与 Stream-K（streamk）的设计原理、AscendC 源码实现要点与编译执行验证方法。读完后，你将能够理解多核 MatMul 任务分配中尾轮效率损失的成因，并掌握尾块拆分重分配与 K 轴切分两种方案的实现细节、适用场景与性能验证流程。

一、系统优化样例总览

系统优化目录 Samples/1_Features/system_optimization 收录了两类针对多核并行计算效率的优化样例，均以带优化特性的 MatMul 算子为载体，且都要求dav-3510架构（两个子目录的 CMakeLists 中均通过cann_sample_check_arch(dav-3510)声明）：

两者关注的是同一类矛盾的不同侧面：多核 MatMul 中"最后一轮"的计算效率。前者解决的是M/N 平面尾块分配不均，后者解决的是K 维串行累加过长与搬运模式冲突，可结合 matmul_story 等 MatMul 系列样例横向对比理解。

二、尾轮负载均衡（tail_rebalance）

2.1 背景与原理

在进行算子多核运算时，通常会遇到最后一轮迭代中剩余任务块（block）数量小于可用计算核数量的情况。此时仅部分核参与计算，其余核算数闲置，整体效率下降。尾轮负载均衡的做法是：将末轮未完全分配的基本块重切分，重新计算所需的总核数，并重新分配计算核心，使其均匀分布至各核，充分发挥计算核算力。重新分配后，端到端计算耗时减少，计算效率提升。

2.2 实现步骤一：计算尾块拆分维度

尾轮负载均衡首先要确定尾块拆分大小，即明确 M 维和 N 维上分别需要切分的块数。以下函数基于当前可用 AI Core 数计算最优尾块切分方案（样例源码见 CalcTailBasicBlock）：

// 计算尾块拆分维度 __aicore__ inline void CalcTailBasicBlock(uint64_t mTileNum, uint64_t nTileNum, uint64_t aicNum, uint64_t& tailMCnt, uint64_t& tailNCnt) { uint64_t mnCnt = mTileNum * nTileNum; uint64_t tailCnt = mnCnt - aicNum * (CeilDiv(mnCnt, aicNum) - 1); tailMCnt = 1UL; tailNCnt = 1UL; if (tailCnt != 0UL) { while ((tailMCnt + 1UL) * tailNCnt * tailCnt <= aicNum) { tailMCnt += 1UL; if (tailMCnt * (tailNCnt + 1UL) * tailCnt <= aicNum) { tailNCnt += 1UL; } } } }

从源码结构看，该函数采用贪心策略：在确保tailMCnt × tailNCnt × tailCnt ≤ aicNum（拆分后子块总数不超过可用核数）的前提下，尽可能扩大 M、N 两个维度的拆分数量，从而充分利用闲置核算数；若 M/N 总块数本身不足一核一轮（tailCnt == 0），则保持tailMCnt = tailNCnt = 1不做拆分。

2.3 实现步骤二：多核任务重分配

尾块拆分后需要重新计算总任务块数并调整多核分配策略，确保各核算数负载均衡（对应 任务分配循环）：

// 重新计算总块数：原始块数 + 尾块拆分后新增的块数 tileNum = tileNum + (tailCnt - 1) * perTailCnt; // Multi-core tile processing loop - distribute tiles across available cores for (uint64_t tileIdx = curBlockIdx; tileIdx < tileNum; tileIdx += blockNum) { // 多核任务分配循环 - 将任务块均匀分发至各核心 if (tileIdx / blockNum == (perCoreBlockNum - 1) && tailCnt > 1) { tileIdx = (perCoreBlockNum - 1) * blockNum + curBlockIdx / tailCnt; } // 单核计算逻辑 ... }

其中perCoreBlockNum是拆分前每核平均分得的轮次数，perTailCnt是最后一轮核上剩余的尾块数。当循环走到"拆分前的最后一轮"（tileIdx / blockNum == perCoreBlockNum - 1）且确实存在尾块拆分（tailCnt > 1）时，把所有核的tileIdx重映射进尾块区，并用curBlockIdx / tailCnt决定本核领取哪个子块，实现末轮任务的均匀分发。

2.4 实现步骤三：尾块坐标重映射

处理到末轮拆分的尾块时，需要重新计算当前核算数负责的子块在 M 维和 N 维上的起始坐标与尺寸，确保数据切分正确（对应 尾块拆分处理逻辑）：

// 判断是否为尾块拆分场景：最后一个核心且存在尾块拆分 if (tileIdx / blockNum == (perCoreBlockNum - 1) && tailCnt > 1) { // 计算拆分后每个子块在M维度和N维度上的尺寸 int64_t splitBlkM = tool::CeilDiv(curM, tailMCnt); int64_t splitBlkN = tool::CeilDiv(curN, tailNCnt); // 计算当前核心在尾块拆分中的索引位置 int64_t mSplitIdx = (curBlockIdx % tailCnt) % tailMCnt; int64_t nSplitIdx = (curBlockIdx % tailCnt) / tailMCnt; // 计算子块在M维度和N维度上的起始偏移 int64_t mSplitOffset = mSplitIdx * splitBlkM; int64_t nSplitOffset = nSplitIdx * splitBlkN; //跳过那些超出原始维度范围的无效子块 if (mSplitOffset >= curM || nSplitOffset >= curN) { continue; } // 更新当前子块的实际尺寸（边界处理） curM = (curM - mSplitOffset) < splitBlkM ? (curM - mSplitOffset) : splitBlkM; curN = (curN - nSplitOffset) < splitBlkN ? (curN - nSplitOffset) : splitBlkN; // 重新设置张量GM地址，指向当前子块对应的数据区域 tensorAGmBlock = tensorAgm(AscendC::Te::MakeCoord(mTileIdx * baseM + mSplitOffset, 0L), AscendC::Te::MakeShape(curM, k)); tensorBGmBlock = tensorBgm(AscendC::Te::MakeCoord(0L, nTileIdx * baseN + nSplitOffset), AscendC::Te::MakeShape(k, curN)); tensorCGmBlock = tensorCgm(AscendC::Te::MakeCoord(mTileIdx * baseM + mSplitOffset, nTileIdx * baseN + nSplitOffset), AscendC::Te::MakeShape(curM, curN)); // 根据实际输出尺寸重新设置L0C布局 layoutL0C = AscendC::Te::MakeL0CLayout(curM, curN); // L0C layout for output tensorL0C = AscendC::Te::MakeTensor(AscendC::Te::MakeL0CmemPtr<float>(l0cOffset), layoutL0C); }

关键点在于：向上取整（CeilDiv）切分子块会产生越界子块，需通过mSplitOffset >= curM || nSplitOffset >= curN直接continue跳过无效子块；对有效子块再按"剩余不足一块时取剩余"的规则收缩curM/curN，最后用收缩后的尺寸重建 A/B/C 三个 GM 张量视图和 L0C 输出布局，保证张量地址映射正确。

该样例的完整内核（tail_rebalance/main.asc）在尾轮负载均衡之外还叠加了 SWAT 扫描（Swath）策略——通过WINDOW_LEN滑动窗口把线性 tile 索引映射成二维网格并做蛇形翻转（见 SWAT 映射代码），以改善 L1 复用；计算主体采用Mutex锁保护的 L1/L0 双缓冲 ping-pong 流水线与Mmad累加，数据类型为 bfloat16。

2.5 性能结果对比

以基础 MatMul 算子为例，在相同输入规模（M=2560, K=1024, N=2560）下通过 Profiling 工具采集硬件流水线执行状态。使用尾轮负载均衡策略优化后，尾轮计算时间显著缩短。样例文档给出的 msprof 分解数据为：

[Profile Breakdowm] +-------------------+------------+---------+------------+----------+----------+-------------+----------------+ || candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | +===================+============+=========+============+==========+==========+=============+================+ || tail_rebalance | 82.135 | 41.781 | 1.863 | 10.539 | 33.148 | 2.132 | 2.500 | +-------------------+------------+---------+------------+----------+----------+-------------+----------------+

与相同输入规模下的基础 matmul 算子相比：

[Profile Breakdowm] +-----------+------------+---------+------------+----------+----------+-------------+----------------+ || candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | +===========+============+=========+============+==========+==========+=============+================+ || matmul | 86.870 | 43.804 | 1.850 | 12.997 | 51.857 | 2.970 | 2.200 | +-----------+------------+---------+------------+----------+----------+-------------+----------------+

可以看到，由于尾轮计算效率提升，整体计算时间从 86.87us 缩短到 82.14us，其中 MTE2（GM→L1 搬运）耗时从 51.857us 降到 33.148us，收益最为明显。

适用场景：矩阵维度不是基本块大小的整数倍，导致末轮存在不完整分配、尾块占比较高且核算数较多的场景。

三、Stream-K 多核切分（streamk）

3.1 背景：为什么切 M/N 会破坏搬运模式

当矩阵规模可被单核承载时，采用负载均衡策略将任务分散到多核通常能提升效率；但若沿 M 方向和 N 方向切分基本块，会破坏原有最优的数据搬运模式。具体而言，MTE2 总的数据搬运量可表示为 $M \cdot K \cdot (N / \text{baseN}) + N \cdot K \cdot (M / \text{baseM})$，M 和 N 方向切分基本块会使数据重复搬运开销增加。Stream-K 策略的出发点正是：不改变原有 M/N 切分策略，而是把任务沿 K 方向切分为多份，在不同核算数上并行计算，从而避免破坏数据搬运策略。

3.2 原理：AIC 计算与 AIV 累加的数据并行

各核算数对自己那份 K 分片独立做矩阵乘后，将中间部分和搬出到外部内存（GM workspace），再由统一缓冲区（UB）逐次累加得到最终结果。更进一步，DPSK（Data-Parallel Streamk）策略将最后一轮的 AIC 计算提前执行：最后一轮原本 AIC 完全空闲、需要等待 AIV 累加结束；提前后，AIV 的累加计算与最后一轮 AIC 的计算互不阻塞，形成 AIC/AIV 数据并行，整体流水时序提前。

3.3 实现步骤一：AIC 与 AIV 逻辑分离

样例内核通过__global__ __aicore__ __mix__(1, 2)声明同一核函数同时包含 AIC（Cube 侧）与 AIV（Vector 侧）两类计算单元的逻辑，比例为 1 个 AIC 对应 2 个 AIV（见 内核入口声明）。AIC 侧负责设置核间同步标志CrossCoreSetFlag，AIV 侧等待CrossCoreWaitFlag就绪后执行向量累加（见 AIC 侧同步逻辑）：

template <typename T> __global__ __aicore__ __mix__(1, 2) void MatmulKernel( GM_ADDR aGm, GM_ADDR bGm, GM_ADDR cGm, GM_ADDR workspaceGm, uint32_t m, uint32_t k, uint32_t n) { // -------------------- AIC 逻辑（AI Core） -------------------- if ASCEND_IS_AIC { // 计算实际参与运算的核心数：取任务总块数与可用核心数的最小值 uint64_t usedCoreNum = tileNum < blockNum ? tileNum : blockNum; // 若当前核心索引超出有效核心范围，则该核心不参与实际计算 if (curBlockIdx >= usedCoreNum) { AscendC::CrossCoreSetFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_FIX>(tool::AIC_SYNC_AIV_FLAG); AscendC::CrossCoreSetFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_FIX>(tool::AIC_SYNC_AIV_FLAG + tool::FLAG_ID_MAX); return; } // 后续AIC计算逻辑代码 // 若当前分片索引超出总分片数，同样设置同步标志后提前返回 if (tileIdx + blockNum >= tileNum) { AscendC::CrossCoreSetFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_FIX>(tool::AIC_SYNC_AIV_FLAG); AscendC::CrossCoreSetFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_FIX>(tool::AIC_SYNC_AIV_FLAG + tool::FLAG_ID_MAX); } } // -------------------- AIV 逻辑（AI Vector Core） -------------------- if ASCEND_IS_AIV { // 判断当前AIV核心是否已完成所有预期的循环轮次（lastLoopTotalCnt * 任务分配比） if (curBlockIdx >= lastLoopTotalCnt * AscendC::GetTaskRatio()) { AscendC::CrossCoreWaitFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_MTE3>(tool::AIC_SYNC_AIV_FLAG); AscendC::SyncAll(); return; } // 常规AIV计算路径：等待AIC标志就绪，再进行向量计算 AscendC::CrossCoreWaitFlag<tool::AIC_SYNC_AIV_MODE_4, PIPE_MTE3>(tool::AIC_SYNC_AIV_FLAG); AscendC::SyncAll(); // 后续AIV计算逻辑代码 } }

从源码结构看，AIC 侧对"超额核算数"和"最后一轮"都设置了同步标志后再提前返回，AIV 侧先CrossCoreWaitFlag再SyncAll，避免在 workspace 未写完时提前累加或空等核算数造成死锁——这就是文档强调的"越界保护防止死锁"。

3.4 实现步骤二：Workspace 中转机制

Workspace 是位于全局内存（GM）中的中转缓冲区，用于 AIC 与 AIV 之间的数据交互。每个 AIC 核算数完成自身 K 轴累加后，通过CopyL0C2GM把 L0C 中的部分累加结果写入各自独立开辟的 workspace 空间，按(mTileNum, nTileNum, skKTileNum)三维分块组织、每块大小为BLOCK_BASE_M × BLOCK_BASE_N（样例中均为 256），确保互不覆盖（见 AIC 侧 workspace 写入 与 L0C2GM 拷贝）：

// -------------------- AIC 逻辑（AI Core） -------------------- if ASCEND_IS_AIC { // 将 workspace 指针转换为全局内存浮点指针，便于后续地址计算 __gm__ float* workspaceGmAddr = reinterpret_cast<__gm__ float*>(workspaceGm); // 计算 K 维度上的分块数（每个 block 在 K 轴上被切分的份数） uint64_t skKTileNum = blockNum / (mTileNum * nTileNum); // 循环遍历当前核心负责的所有 tile（步长为 blockNum，实现负载均衡） for (uint64_t tileIdx = curBlockIdx; tileIdx < tileNum; tileIdx += blockNum) { // 计算当前 tile 在 K 维度上的分块索引 uint64_t kTileIdx = (tileIdx % blockNum) % skKTileNum; // 计算当前 tile 在 workspace 中的偏移量 // 布局逻辑：(mTileIndex, nTileIndex, kTileIndex) 三维映射到线性地址 int64_t offsetWorkspace = (((tileIdx % blockNum) / skKTileNum) * skKTileNum + kTileIdx) * tool::BLOCK_BASE_M * tool::BLOCK_BASE_N; // 构建 workspace 张量对象（GM 内存视图） auto gmWorkSpace = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr(workspaceGmAddr + offsetWorkspace), layoutWorkspace); for (uint64_t iter0 = 0; iter0 < kL1TileNum; ++iter0) { // L0 层内部迭代（实际的计算/搬运操作） for (uint16_t iter1 = 0; iter1 < kL0IterNum; ++iter1) { // 矩阵乘计算核心逻辑 } // 当 K 维度的最后一轮迭代完成时，将 L0C 中的累加结果搬运到 workspace if (iter0 + 1 == kL1TileNum) { // 创建 L0C 到 GM 的拷贝操作 auto CopyL0C2GM = AscendC::Te::MakeCopy(AscendC::Te::CopyL0C2GM{}); // 执行拷贝：将 L0C 中的数据写入 workspace 指定偏移位置 // FINAL_ACCUMULATION 表示这是最终累加结果，需要写回 AscendC::Te::Copy( CopyL0C2GM, gmWorkSpace, tensorL0C, AscendC::Te::FixpipeParams{tool::FINAL_ACCUMULATION}); } } } } // -------------------- AIV 逻辑（AI Vector Core） -------------------- if ASCEND_IS_AIV { // 计算当前 AIV 核心需要读取的 workspace 起始地址偏移量 // 公式含义： // - newBlockIdx * skKTileNum * BLOCK_BASE_M * BLOCK_BASE_N: 按 block 索引到对应分区 // - kTileIdx * mBurstBase * curN: 按 K 分块和 burst 维度进一步定位 // - copyGm2UbParams_.mBurst * index: 按 burst 索引计算具体数据块偏移 copyGm2UbParams_.offsetWorkspaceGM = newBlockIdx * skKTileNum * tool::BLOCK_BASE_M * tool::BLOCK_BASE_N + (kTileIdx * mBurstBase + copyGm2UbParams_.mBurst * index) * curN; // 计算其他搬运参数（如搬运长度、burst 配置等，原代码省略） // ... 参数计算逻辑 ... // 执行 GM 到 UB（统一缓冲区）的数据搬运 // 将 workspace 中指定偏移的数据搬运到 ubAddTensor（UB 上的张量） DataCopyPad<float>( ubAddTensor, // 目标：UB 上的张量 workspaceGlobal_[copyGm2UbParams_.offsetWorkspaceGM], // 源：workspace 中的指定位置 dataCopyExtParams, // 搬运扩展参数（长度、步长等） {false, 0, 0, 0}); // 对齐/填充参数 // 后续 AIV 向量计算逻辑 // ... } }

AIV 侧再从 workspace 读取中间结果搬运至本地 UB 逐次累加（源码中通过DataCopyPad搬入、Add逐份相加、Cast转回 bfloat16 后经 MTE3 写回 C，见 AIV 累加主循环），保证累加结果的确定性。按 AIC:AIV 配比划分累加份额：例如配比为 2 时，可将每两份 K 切分计算得到的 AIC 结果进一步划分，划分数量等于"切 K 数量 × AIC 与 AIV 配比"。

主机侧 workspace 的分配逻辑同样直观（main 函数内存分配）：

uint8_t *dWorkSpace = nullptr; size_t sizeWorkSpace = numBlocks * tool::BASIC_BLOCK_SIZE_256 * tool::BASIC_BLOCK_SIZE_256 * tool::DATA_SIZE_FP32 + tool::RPC_WORKSIZE * tool::MB_SIZE;

即"核算数 × 256 × 256 个 fp32 元素"的中转区（每核算数一块 256×256 的部分和区），再加 20MB 的 RPC 预留空间，用ACL_MEM_MALLOC_HUGE_ONLY独占大页内存分配。

3.5 实现步骤三：分核与坐标重设

切 K 后需将线性 tile 索引重新映射为(mTileIdx, nTileIdx, kTileIdx)三维坐标并处理尾块边界（见 SK 场景判定与 skKTileNum 计算）：

// (M,N)块数较少时，增大K轴切分，提高并行度 if(tileNum <= blockNum / 2) { skKTileNum = blockNum / tileNum; // K轴切分数 = 总核数 / 块数 skKSingleCore = CeilDiv(k, skKTileNum); // 每核处理K长度 } // 块数能整除核算数：无需切K，走纯数据并行 else if ((tileNum % blockNum) == 0) { skKTileNum = 1; skKSingleCore = k; } // (M,N)块数较多时，基于尾块计算切分，并向上取整保证整除 else { skKTileNum = blockNum / (tileNum % blockNum); skKSingleCore = CeilDiv(k, skKTileNum); skKTileNum = CeilDiv(k, skKSingleCore); // 反推实际切分数 } // 计算尾块（不足一个完整 block）的 (M, N) 分块数量 int64_t tailMNTileNum = tileNum < blockNum ? tileNum : tileNum % blockNum; uint64_t totalMNTileNumInDP = tileNum - tailMNTileNum; tileNum = totalMNTileNumInDP + tailMNTileNum * skKTileNum; int64_t tailSKTotalTileNum = tailMNTileNum * skKTileNum;

从源码结构看，切分策略分三档：M/N 块数不足核数一半时全力切 K；能整除核算数时退化为纯数据并行（DP）；其余情况只对尾块切 K。随后线性索引按如下规则解算三维坐标（见 坐标重映射）：

• SK 场景（尾块区）：kTileIdx = (tmpTileIdx % usedCoreNum) % curKTileNum，(M,N)块索引取商并加上totalMNTileNumInDP偏移，落在尾部块区间；
• DP 场景（完整块）：kTileIdx = 0，(M,N)索引直接为tmpTileIdx / curKTileNum；
• 尾块边界尺寸：curM/curN取最后一块的实际剩余量，K 分片长度curSK在最后一个分片处收缩为k - (curKTileNum - 1) * skKSingleCore。

是否处于 SK 场景由辅助函数CheckIsSkScene判定（判定函数），即CeilDiv(tileIdx + 1, blockNum) == CeilDiv(tileNum, blockNum)——当前 tile 落在最后一个分配轮次即为 SK 尾块区。此外，源码在 DP+SK 混合调度下还做了一次"倒数第二轮后推一个批次、最后一轮前推一个批次"的索引扰动（tmpTileIdx ± usedCoreNum，见 SK Preload 调度），把尾轮的 AIC 计算提前，配合 3.2 节的 AIC/AIV 流水重叠。AIV 侧用GetBlockIdx() / (GetTaskRatio() * skKTileNum)与取模反解出一致的(block, kTile)坐标（AIV 坐标设置），确保 AIC/AIV 两侧对 workspace 分区的寻址规则完全一致。

3.6 性能结果对比

样例文档给出的 msprof 分解数据为（相同输入规模）：

[Profile Breakdowm] +-----------+------------+---------+------------+----------+----------+-------------+----------------+ || candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | +===========+============+=========+============+==========+==========+=============+================+ || streamk | 23.423 | 12.776 | 1.954 | 4.009 | 11.701 | 9.391 | 4.000 | +-----------+------------+---------+------------+----------+----------+-------------+----------------+

与相同输入规模下的基础开 db（双缓冲）的 matmul 算子相比：

[Profile Breakdowm] +-----------+------------+---------+------------+----------+----------+-------------+----------------+ || candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | +===========+============+=========+============+==========+==========+=============+================+ || n_buffer | 28.455 | 18.479 | 2.148 | 5.984 | 16.761 | 0.950 | 2.900 | +-----------+------------+---------+------------+----------+----------+-------------+----------------+

由样例的仿真流水图可以看出，通过切分 K 维度并分配至多核计算，整体计算时间从 28.46us 缩短到 23.42us，提前了流水时序。

适用场景：

• 多核负载不均：各核算数因任务分配不均导致部分核算数空闲、整体利用率偏低时，Stream-K 通过在 K 维度上细粒度切分并将子任务均匀分配到各核算数，提升多核利用率和计算吞吐量；
• 大 K 场景：当矩阵 K 维度较大（如 K ≥ 8192）时，单核独立承载完整计算效率低，Stream-K 可将负载切分到多核算数并行处理，充分利用多核资源实现加速。

四、两种策略的选型对比

两者可以看作对"最后一轮"问题的两级处理：尾轮负载均衡先解决 M/N 尾块核算数闲置，Stream-K 再对剩余难以均衡的尾块沿 K 维切分。若 M/N 块数能整除核算数，Stream-K 源码会自动退化为纯数据并行路径，不会引入额外切分开销。

五、编译、执行与性能测试

5.1 编译样例

从项目根目录启动构建，参考 README.md。以 tail_rebalance 为例，在仓库根目录下完成编译和安装后，进入当前样例目录：

cmake -S . -B build -DNPU_ARCH=dav-3510 cmake --build build --parallel cmake --install build --prefix ./build_out cd ./build_out/1_Features/system_optimization/tail_rebalance/

如需单独编译当前样例：

cmake --build build --target tail_rebalance cp ./Samples/1_Features/system_optimization/tail_rebalance/scripts/* ./build/Samples/1_Features/system_optimization/tail_rebalance/ cd ./build/Samples/1_Features/system_optimization/tail_rebalance/

streamk 的 target 名为streamk，路径为1_Features/system_optimization/streamk，命令形式相同。两个样例的可执行文件由 各自的 CMakeLists.txt 定义，编译时以--npu-arch=${NPU_ARCH} -O3构建.asc源文件，链接platform、tiling_api与cann_samples::tensor_api库，并把scripts/下的辅助脚本一并安装到输出目录。

5.2 运行样例

使用可执行文件直接执行算子用例，需要指定矩阵乘维度，并随机生成输入数据：

./tail_rebalance 1024 2048 1024

streamk 样例的调用方式一致：

./streamk 1024 2048 1024

运行成功后，终端将打印如下类似信息（host 端会先调gen_data.py生成 bf16 输入，再调verify_result.py与 CPU 参考结果比对，流程见 host 主函数）：

Data generated successfully! [verify] shape(1024, 1024), elements=1048576 - summary (large matrix, full tensors omitted) abs_err: max=2.560000e+02, mean=6.103516-03, rmse=1.250000e+00 rel_err: max=6.410256e-03 count(|abs_err| > 0.001): 108 / 1048576 cpu golden (top-left 4x4): tensor([[40448., 41728., 41472., 41984.], [39680., 40704., 40448., 40960.], [40192., 41472., 41472., 41984.], [40960., 41984., 41728., 42240.]], dtype=torch.bfloat16) npu out (top-left 4x4): tensor([[40448., 41728., 41472., 41984.], [39680., 40704., 40448., 40960.], [40192., 41472., 41472., 41984.], [40960., 41984., 41728., 42240.]], dtype=torch.bfloat16) max abs diff: 256.0 point error count(>0.1): 0/1048576 ratio error count(>0.001): 25/1048576, error ratio: 0.0000024 [PASS] NPU results are consistent with CPU.

如果存在精度问题，则会打印错误数据，并显示如下结果：

[ERROR] NPU results differ from CPU.

5.3 测试性能

运行性能测试脚本，指定矩阵乘法的维度后执行：

python3 profile_matmul.py 1024 2048 1024

profile_matmul.py 的工作方式是从脚本目录自动发现所有可执行文件作为候选算法，逐个用msprof包裹运行，从op_summary_*.csv中解析 Task Duration、aic_mac_time、aic_scalar_time、aic_mte1/mte2_time、aic_fixpipe_time与aic_icache_miss_rate等指标，最后按 kernel 耗时排序打印[Recommended Algorithm Ranking]与[Profile Breakdown]表格。打印如下执行结果即证明性能测试成功：

[Profile Breakdowm] +-------------------+------------+---------+------------+----------+----------+-------------+----------------+ || candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | +===================+============+=========+============+==========+==========+=============+================+ || tail_rebalance | 82.135 | 41.781 | 1.863 | 10.539 | 33.148 | 2.132 | 2.500 | +-------------------+------------+---------+------------+----------+----------+-------------+----------------+

streamk 样例输出形式相同（candidate 为streamk，kernel 耗时 23.423us，见 3.6 节对比数据）。

5.4 支持架构

两个样例均支持 NPU ARCH 3510（dav-3510），构建时通过-DNPU_ARCH=dav-3510指定；在其它架构上编译会因cann_sample_check_arch检查被跳过，实际运行也需具备对应的昇腾环境与 msprof profiling 工具。

【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples