面向华为 CANN 的临时内存管理：TBuf 在算子开发中的实践指南

在基于华为昇腾处理器的算子开发体系中，Ascend C 提供了高度抽象化的算子编写框架，包括矢量计算接口、内存管理接口、队列机制（Queue）等模块。其中，TBuf（Temporary Buffer）作为临时内存管理的关键组件，被广泛用于核函数内部的中间结果存储。

对于许多算子而言，单次计算往往并非“一步到位”，会经历多轮数据类型转换、中间变量生成、精度调整等操作。如果将所有中间结果都以 LocalTensor 的形式临时申请，会带来显著的编程复杂度和内存开销。TBuf 的设计正是为了解决这一痛点：以统一、可控的方式管理核函数执行期间的临时内存。

本文将从 TBuf 的设计理念讲起，进而结合一个典型示例——在不支持 bfloat16 加法的硬件环境上实现 Add 算子，展示 TBuf 在实际算子开发中的关键作用。

---

一、为什么需要 TBuf？算子中的临时变量问题

大多数 Ascend C 的算子实现都遵循以下计算流程：

1. CopyIn：从全局内存 GM 读取数据到 L1/UB LocalTensor
2. Compute：执行矢量计算流程
3. CopyOut：将结果写回全局内存 GM

其中 Compute 阶段最复杂。

例如在一个需要进行数据类型转换的场景中：

• 输入为 bfloat16，不支持直接进行 Add
• 中间需要转换为 float
• 计算完毕后再转回 bfloat16

此时就涉及多个中间结果：

• xCast → float
• yCast → float
• resultFloat → float
• output → bfloat16

这些临时变量的生命周期都只在 Compute 内部，适合通过 TBuf 分配管理。

相比直接创建 LocalTensor：

• TBuf 可以预先一次分配一块连续空间
• 反复 Get() 时无额外分配成本
• 自动适配不同类型的 LocalTensor 视图
• 由 TPipe 统一管理其生命周期、更适合流水线执行

因此，TBuf 成为 Ascend C 编程范式中最常用的临时空间管理方式。

---

二、示例背景：bfloat16 Add 算子为何需要临时内存？

在某些硬件型号（如 Atlas A2/Atlas 800I A2）上，Add 矢量接口 不支持 bfloat16 源操作数的直接求和。因此，Add 算子的执行流程需要进一步扩展：

1. 输入：bfloat16 → float（Cast）
2. 执行 Add：float + float → float
3. 输出：float → bfloat16（Cast）

整个过程涉及三处临时空间：

• tmp0：x cast to float
• tmp1：y cast to float
• tmp2：add result (float)

如果这些中间空间由开发者直接以 LocalTensor 手工分配，会使代码显著冗长。而借助 TBuf，我们只需在初始化阶段申请内存，Compute 中即可快速获取。

---

三、算子整体规格：Add（bfloat16 输入）

为了便于理解，这里总结一下算子的输入输出规格。

名称	Shape	数据类型	Format
x（输入）	(1, 2048)	bfloat16	ND
y（输入）	(1, 2048)	bfloat16	ND
z（输出）	(1, 2048)	bfloat16	ND

该算子的执行核函数被定义为：

add_custom

所依赖的主要接口包括：

• DataCopy：内存搬移
• Cast：类型转换
• Add：基础矢量算术
• EnQue / DeQue：队列调度
• TBuf：临时内存管理

---

四、TBuf 在算子类中的声明

下面是 Add 算子核心类的结构（节选）：

class KernelAdd {
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN,  BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;

    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuf0, tmpBuf1, tmpBuf2;

    AscendC::GlobalTensor<bfloat16_t> xGm; 
    AscendC::GlobalTensor<bfloat16_t> yGm;
    AscendC::GlobalTensor<bfloat16_t> zGm;
};

我们重点关注三条：

TBuf<VECCALC> tmpBuf0, tmpBuf1, tmpBuf2;

每一个 TBuf 对象本质上代表一块可重复使用的临时内存区域。在本算子中：

• tmpBuf0：存放 x 转 float
• tmpBuf1：存放 y 转 float
• tmpBuf2：存放 float add 结果

位置选择为：

TPosition::VECCALC

表示内存位于矢量计算空间内（通常是内存层级较高的 UB/L1）。

---

五、初始化阶段：为 TBuf 分配内存

在 Init() 中，需要为每个 TBuf 分配空间：

pipe.InitBuffer(tmpBuf0, TOTAL_LENGTH * sizeof(float));
pipe.InitBuffer(tmpBuf1, TOTAL_LENGTH * sizeof(float));
pipe.InitBuffer(tmpBuf2, TOTAL_LENGTH * sizeof(float));

特点：

1. 只分配一次，生命周期贯穿整个核函数执行
2. 计算过程中反复使用，避免大量临时分配
3. 支持不同类型 LocalTensor 的动态视图

内存是按 float 大小申请的，因为临时结果均为 float 类型。

---

六、Compute：TBuf 临时内存使用的完整流程解析

Compute 是 TBuf 发挥作用的核心部分。流程如下：

1. 从 Queue 中取出输入 Tensor

auto xLocal = inQueueX.DeQue<bfloat16_t>();
auto yLocal = inQueueY.DeQue<bfloat16_t>();
auto zLocal = outQueueZ.AllocTensor<bfloat16_t>();

xLocal 和 yLocal 为 bfloat16 类型的局部 Tensor。

---

2. 从 TBuf 获取临时 float Tensor

auto tmpTensor0 = tmpBuf0.Get<float>();
auto tmpTensor1 = tmpBuf1.Get<float>();
auto tmpTensor2 = tmpBuf2.Get<float>();

每次 Get() 并不会重新分配空间，而是：

• 在 TBuf 的固定内存上创建一个基于 float 的 LocalTensor 视图
• 类型只是视图层面的 reinterpret cast，不影响底层内存结构

这是 TBuf 的最大优势之一。

---

3. 输入 Cast：bfloat16 → float

Cast(tmpTensor0, xLocal, RoundMode::CAST_NONE, TOTAL_LENGTH);
Cast(tmpTensor1, yLocal, RoundMode::CAST_NONE, TOTAL_LENGTH);

临时结果放入 tmpTensor0 / tmpTensor1。

---

4. Add 矢量计算（float 加法）

Add(tmpTensor2, tmpTensor0, tmpTensor1, TOTAL_LENGTH);

执行完后，tmpTensor2 中已包含 float 结果。

---

5. 输出 Cast：float → bfloat16

Cast(zLocal, tmpTensor2, RoundMode::CAST_RINT, TOTAL_LENGTH);

将结果转换回输出所需的 bfloat16。

---

6. 入队 & 释放资源

outQueueZ.EnQue<bfloat16_t>(zLocal);
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);

这一步将结果放入输出队列，并释放输入局部变量，让下一轮计算可以重复利用 LocalTensor。

---

七、TBuf 使用的几个关键特性总结

1. 临时空间复用

TBuf 通常在 Init 中一次性分配空间，其后在 compute 中反复使用，从而大幅减少动态分配带来的指令开销。

2. 按类型构建视图

同一块 Buffer 可以通过 Get() 转成不同类型的 LocalTensor，只要底层大小能容纳即可。

3. 生命周期由算子类控制

避免多线程/多 pipeline 情况下的内存混乱。

4. 完全适配 Ascend C 流水线模型

TBuf 与 Queue、LocalTensor、TPipe 等协同构成完整的资源调度体系。

---

八、示例意义：TBuf 在算子设计中的典型作用

虽然本文示例聚焦在 Add 算子的 bfloat16 类型转换问题，但这一模式具有普适性：

• 所有需要“中间计算结果”的算子都可以使用 TBuf
• 如卷积前的形状展开、LayerNorm 的方差/均值缓存、Softmax 的 max/sum 缓存等

TBuf 提供了一套高性能的临时内存解决方案，是 Ascend C 算子开发不可或缺的基础组件。

---

九、结语

在华为 CANN 的算子开发体系中，TBuf 不仅仅是一个“临时内存容器”，它基于 Ascend C 的计算模型，为开发者提供了：

• 高效的内存管理
• 清晰的生命周期控制
• 与矢量流水线完美契合的资源调度能力

无论你在开发基础矢量算子，还是构建更复杂的 AI 算法算子，TBuf 都会成为不可或缺的重要工具。理解并熟练使用 TBuf，将显著提升算子实现的性能、代码结构与开发效率。

训练营简介

2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro