CANN/cannbot-skills UB缓冲区管理指南

UB 缓冲区管理指南

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

TBuf/TQue 选择、Double Buffer 流水线并行、批量搬运模式。

目录

1. TBuf vs TQue 选择
2. TQue 详解
3. TBuf 详解
4. Double Buffer 流水线并行
5. 批量搬运 + 逐行计算模式

TBuf vs TQue 选择

TQue 详解

模板参数

template <TPosition pos, int32_t depth, auto mask = 0> class TQue;

depth 参数关键说明

注意：depth与 Double Buffer 无关。Double Buffer 由InitBuffer的num参数控制。

// ✅ 非连续入队（普通场景）：depth=1 即可 AscendC::TQue<AscendC::TPosition::VECIN, 1> que; pipe->InitBuffer(que, 1, size); auto tensor = que.AllocTensor<T>(); que.EnQue(tensor); tensor = que.DeQue<T>(); que.FreeTensor(tensor);

Double Buffer 配置

Double Buffer 是在InitBuffer的num参数中设置，与模板参数depth无关。

// ✅ 开启 Double Buffer：在 InitBuffer 中设置 num=2 AscendC::TQue<AscendC::TPosition::VECIN, 1> que; // 模板 depth=1 即可 pipe->InitBuffer(que, 2, size); // num=2 开启 Double Buffer // ✅ 关闭 Double Buffer AscendC::TQue<AscendC::TPosition::VECIN, 1> que; pipe->InitBuffer(que, 1, size); // num=1 单 Buffer

TQue Buffer 数量限制

注意：

• 不开启 Double Buffer（num=1）：最多可申请 8 个 TQue
• 开启 Double Buffer（num=2）：每个 TQue 占用 2 个 buffer，最多只能申请 4 个 TQue

// 开启 Double Buffer 时，最多只能申请 4 个 TQue pipe->InitBuffer(que0, 2, size); // ✅ pipe->InitBuffer(que1, 2, size); // ✅ pipe->InitBuffer(que2, 2, size); // ✅ pipe->InitBuffer(que3, 2, size); // ✅ pipe->InitBuffer(que4, 2, size); // ❌ 超过限制

TQue 正确用法

// TQue：需要队列管理（MTE 搬运相关） // 模板 depth=1 即可，Double Buffer 在 InitBuffer 的 num 参数中设置 AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueX; pipe->InitBuffer(inQueueX, 2, bufferSize); // num=2 开启 Double Buffer AscendC::LocalTensor<half> x = inQueueX.AllocTensor<half>(); AscendC::DataCopyPad(x, xGm, {1, size * sizeof(half), 0, 0}, {false, 0, 0, 0}); inQueueX.EnQue(x); // ... AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>(); inQueueX.FreeTensor(xLocal);

TBuf 详解

特性

// TBuf：纯计算缓冲区 AscendC::TBuf<AscendC::TPosition::VECCALC> workBuf; pipe->InitBuffer(workBuf, bufferSize); // ✅ 使用 Get<T>() 获取 Tensor，无需释放 AscendC::LocalTensor<float> work = workBuf.Get<float>(); // ... 计算逻辑 ... // 无需 FreeTensor

Double Buffer 流水线并行

核心认知

Double Buffer 不是"用2块内存计算"，而是"用2块内存做搬入/搬出，使 MTE2/MTE3 与 Vector 计算并行"。

本质：内存搬运与计算并行，掩盖搬运延迟。

硬件原理

• MTE2：搬运工，GM → UB
• Vector：加工员，计算
• MTE3：搬运工，UB → GM

时间线对比

无 Double Buffer（串行）：

Row 0: [MTE2][Vector][MTE3] Row 1: [MTE2][Vector][MTE3]

有 Double Buffer（并行）：

Row 0: [MTE2-B0][Vector-B0][MTE3-B0] Row 1: [MTE2-B1][Vector-B1][MTE3-B1] ↑ MTE2与Vector并行！

实现原则

正确用法

// 1. Init: num=2 开启 Double Buffer pipe->InitBuffer(inQueueX, 2, tileSize * sizeof(T)); pipe->InitBuffer(outQueueY, 2, tileSize * sizeof(T)); pipe->InitBuffer(workBuf, workSize * sizeof(T)); // 2. Process: 单循环结构，TQue 自动轮转 for (int i = 0; i < totalTiles; i++) { CopyIn(i); // MTE2 异步搬运 Compute(i); // Vector 计算 CopyOut(i); // MTE3 异步搬出 } // 3. CopyIn void CopyIn(int i) { LocalTensor<T> x = inQueueX.AllocTensor<T>(); DataCopyPad(x, xGm[i * tileSize], {1, (uint32_t)(tileSize * sizeof(T)), 0, 0}, {false, 0, 0, 0}); inQueueX.EnQue(x); } // 4. Compute void Compute(int i) { LocalTensor<T> x = inQueueX.DeQue<T>(); LocalTensor<T> y = outQueueY.AllocTensor<T>(); Add(y, x, constTensor, tileSize); outQueueY.EnQue(y); inQueueX.FreeTensor(x); } // 5. CopyOut void CopyOut(int i) { LocalTensor<T> y = outQueueY.DeQue<T>(); DataCopyPad(yGm[i * tileSize], y, {1, (uint32_t)(tileSize * sizeof(T)), 0, 0}); outQueueY.FreeTensor(y); }

为什么能并行？

常见误区

批量搬运 + 逐行计算模式

适用场景

处理多行数据时，批量搬运减少 MTE2/MTE3 调用次数，充分利用带宽。

模式结构

CopyInBatch(N行) → 逐行计算(N次) → CopyOutBatch(N行)

代码模板

__aicore__ inline void ProcessBatch() { uint32_t totalRowsToProcess = endRow - startRow; if (totalRowsToProcess == 0) return; for (uint32_t tile = 0; tile < tilesPerCore; tile++) { uint32_t startLocalRow = tile * tileRows; // 边界检查：防止 uint32_t 下溢 if (startLocalRow >= totalRowsToProcess) break; uint32_t remaining = totalRowsToProcess - startLocalRow; uint32_t rowsThisTile = (remaining < tileRows) ? remaining : tileRows; CopyInBatch(startLocalRow, rowsThisTile); ComputeBatch(rowsThisTile); CopyOutBatch(startLocalRow, rowsThisTile); } }

Host 侧 Tiling 计算

// A2/A3 UB = 192KB constexpr uint64_t UB_SIZE = 192 * 1024; constexpr uint32_t MAX_BLOCK_COUNT = 4095; // DataCopyPad blockCount 限制 // bytesPerTileRow: double buffer (in*2 + out*2) uint32_t bytesPerTileRow = paddedColsT * typeSizeBytes * 4; // tileRows uint32_t tileRows = (UB_SIZE - overheadBytes) / bytesPerTileRow; tileRows = std::max(1u, std::min(tileRows, MAX_BLOCK_COUNT));

注意事项

1. tileRows 限制：DataCopyPad 的blockCount最大 4095
2. 尾核处理：startLocalRow >= totalRowsToProcess时提前退出
3. stride 计算：UB 侧 stride 单位是 32 字节块，GM 侧是字节

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills