PTO-ISA 指令详解:TCOLEXPANDADD 列广播加法(Column-wise Broadcast Add)的语义、汇编与内核实现
2026/9/20 1:29:43 网站建设 项目流程
  • 人工智能
  • 指令集
  • 算子库
  • CANN
  • Ascend

【免费下载链接】pto-isa

Parallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.

项目地址:https://gitcode.com/cann/pto-isa
点击查看免费下载

TCOLEXPANDADD 是 CANN PTO(Parallel Tile Operation)虚拟指令集 PTOISA.md 中 TCOLEXPAND 族的一员,用于将src1提供的"每列一个标量"广播到整列,并与src0对应元素执行加法,结果写入dst。本文从数学语义、两级汇编形式、C++ 内建接口、数据类型与布局约束出发,结合 A2/A3/A5 后端与 CPU 模拟器源码及 NPU 测试用例,完整还原该指令从声明、发射到向量单元执行的全链路,帮助算子开发者在自动模式与手动模式下正确使用并理解其底层行为。

指令示意图

指令语义与数学解释

TCOLEXPANDADD 属于按列广播的二元向量运算:它不是对整个 Tile 执行标量广播,而是把src1中的第j个标量值广播到第j列的所有行,再逐元素执行加法。

R = dst.GetValidRow()C = dst.GetValidCol()s_j为从src1中取出的第j列标量,则对于0 <= i < R0 <= j < C

$$ \mathrm{dst}{i,j} = \mathrm{src0}{i,j} + s_j $$

直观理解:

  • src0dst有效形状相同RC列),逐元素参与运算;
  • src1的有效形状只需覆盖C个值(通常为1 x CR_src1 x C,其中R_src1可为 1),其第j个元素沿列方向扩展;
  • 这与 TCOLEXPAND(仅做纯广播、无运算)形成对照:TCOLEXPANDADD 将"列扩展"与"二元加法"融合为单条指令,省去先广播再相加的两步操作。

从指令族的角度看,TCOLEXPAND 族在 include/pto/common/event.hpp 中按二元算子成族排列:TCOLEXPANDDIVTCOLEXPANDMULTCOLEXPANDADDTCOLEXPANDMAXTCOLEXPANDMINTCOLEXPANDSUBTCOLEXPANDEXPDIF。其中TCOLEXPANDADD在 include/pto/common/event.hpp 被映射到PIPE_V(向量计算单元),即该指令由 AIC(Vector 侧)执行。

汇编语法

同步形式

%dst = tcolexpandadd %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...>

AS Level 1(SSA)

%dst = pto.tcolexpandadd %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...>

AS Level 2(DPS)

pto.tcolexpandadd ins(%src0, %src1 : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>)

两级汇编的区别在于:

  • AS Level 1(SSA 形式):以!pto.tile<...>逻辑 Tile 为操作数,由编译器负责资源放置与调度,主要用于自动模式;
  • AS Level 2(DPS 形式):操作数变为显式绑定存储地址的!pto.tile_buf<...>ins(...)列出输入、outs(...)列出输出,是更贴近硬件资源分配的下沉形式。

C++ 内建接口

TCOLEXPANDADD 的 C++ 接口声明于 include/pto/common/pto_instr.hpp,公共包含头为<pto/pto-inst.hpp>

template <typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1, typename... WaitEvents> PTO_INST RecordEvent TCOLEXPANDADD(TileDataDst &dst, TileDataSrc0 &src0, TileDataSrc1 &src1, WaitEvents &... events);

接口特点:

  • 模板参数TileDataDst / TileDataSrc0 / TileDataSrc1决定了 Tile 的形状、数据类型与布局(RowMajor);
  • 变参WaitEvents &... events允许传入依赖事件:调用会先执行detail::PtoWaitEvents(events...)等待前置指令完成,再经MAP_INSTR_IMPL(TCOLEXPANDADD, dst, src0, src1)映射到后端实现;
  • 返回值RecordEvent记录了本次发射,可继续传给后续指令构成事件链,从而实现手动模式下跨管道(MTE2 → V → MTE3)的显式同步。

约束条件

数据类型

  • TileDataDst::DTypeTileDataSrc0::DTypeTileDataSrc1::DType必须一致,且属于以下集合:
    • halffloatint16int32:适用于 Atlas A2 训练系列产品 / Atlas A2 推理系列产品、Atlas A3 训练系列产品 / Atlas A3 推理系列产品,以及 Ascend 950PR / Ascend 950DT;
    • uint16uint32bfloat16_tint8uint8int64uint64:仅适用于 Ascend 950PR / Ascend 950DT(即 A5)。

该约束在源码层面对应双重检查:

  • 后端实现中的static_assert(如 include/pto/npu/a2a3/TColExpandBinOp.hpp)编译期校验 A2/A3 侧的数据类型集合;
  • CPU 模拟器在 include/pto/cpu/TColExpandOp.hpp 中通过IsColExpandAllowedTypeCheckColExtendTiles约束类型一致性(dstsrc0src1类型必须相同),并规定OP_ADD允许floathalfint64/uint64/int32/int16/uint32/uint16等整型(OP_EXPDIF除外)。

Tile 形状与布局(编译期约束)

  • TileDataDst::isRowMajor必须为true,即只支持RowMajor(行主序)布局;CPU 侧同样要求TileDst/TileSrc0/TileSrc1均为 RowMajor(见 include/pto/cpu/TColExpandOp.hpp 的static_assert)。

src1 的形状约束

  • src1预期提供每列一个标量,即其有效形状必须覆盖C个值(通常1 x C),而不是与src0相同的R x C
  • 从 include/pto/npu/a2a3/TColExpandBinOp.hpp 的实现看,当src1有效形状与dst相同时(src1eqdst),其行步长退化为与dst一致,仍只使用每列首个标量参与广播;而TCOLEXPANDOP_IMPL还会自动比较src0dst的有效形状(src0eqdst),决定以src0还是src1作为主操作数进行发射,兼顾了调用方传入顺序的灵活性。

布局与分形约束

  • 确切的布局/分形约束是目标特定的,需参见include/pto/npu/*/TColExpand*.hpp下的后端头文件(如 include/pto/npu/a2a3/TColExpandBinOp.hpp、include/pto/npu/a5/TColExpandBinOp.hpp)。

64 位元素类型(A5 特有)

int64/uint64仅在 Ascend 950PR / Ascend 950DT 上受支持,且存在两条关键规则:

  1. 寄存器对模拟:A5 没有原生的 64 位向量 ALU,指令通过一对 32 位寄存器(分别保存每个元素的低 32 位与高 32 位)模拟实现。对应实现见 include/pto/npu/a5/TColExpandAdd.hpp:在PTO_NPU_ARCH_A5 / PTO_NPU_ARCH_A6下调用Int64BinaryCalcRegs<Int64Op::Add, T>(dstLow, dstHigh, src0Low, src0High, src1Low, src1High, preg)完成高低字分别参与运算的 64 位加法;
  2. 解交织布局:每列标量操作数src1与全尺寸操作数使用相同的解交织(de-interleaved)布局读取,保证高低字配对正确。

计算结果为精确的 64 位二进制补码值。Tile 对齐遵循 64 位元素的通用规则:RowMajor 的 Tile 要求Cols % 4 == 0

底层实现解析

A2 / A3 后端:基于重复步长(repeat stride)的 vadd

在 include/pto/npu/a2a3/TColExpandAdd.hpp 中,ColExpandAddOp将运算落到硬件vadd指令上:

vadd(dst, src0, src1, repeats, 1, 1, 1, 8, 8, 8); // 广播模式:src1 的 repeat 步长为 0(不前进) vadd(dst, src0, src1, repeats, 1, 1, 1, dstRepeatStride, src0RepeatStride, 0); // 通用步长版本

其关键思想是:src1的 repeat 步长传0,使向量硬件在连续 repeat 中反复读取同一份标量数据,从而实现"列广播"而不需要显式拷贝展开。

调度逻辑位于 include/pto/npu/a2a3/TColExpandBinOp.hpp:

  • NormMode(常规模式):当 Tile 的Cols == ValidCol(列无 padding)或Rows == 1时,以ElementsPerRepeat(每个 repeat 可承载的元素数)为粒度拆分循环,分整段与余数两阶段设置连续掩码SetContMaskByDType<T>后发射;
  • CountMode(计数模式):其余情况(存在列 padding)退化为逐行循环,每行通过SetVectorCount(validCol)限定有效列数后单独发射,避免 padding 区被误写。

发射前还会依据blockSizeElem = BLOCK_BYTE_SIZE / sizeof(DType)elementsPerRepeat = REPEAT_BYTE / sizeof(DType)(见 include/pto/npu/a2a3/TColExpandBinOp.hpp)换算出行步长与 repeat 步长,保证不同数据类型下都能对齐硬件最小粒度。

A5 后端:掩码寄存器与 64 位模拟

A5 后端(include/pto/npu/a5/TColExpandAdd.hpp)使用寄存器张量(RegTensor)与掩码寄存器:

vadd(reg_dst, reg_src0, reg_src1, preg, MODE_ZEROING);

preg掩码配合MODE_ZEROING处理列尾部无效元素;64 位类型则走上述高低字寄存器对模拟路径。

CPU 模拟器:并行按列广播

CPU 侧实现位于 include/pto/cpu/TColExpandOp.hpp:TColExpand_Op先做类型/布局编译期检查,随后通过cpu::parallel_for_1d按列并行,每列取出src1的标量值src1Val,再对整列逐行执行ElementOpCal<T, OP_ADD>::apply完成广播加法。该实现同时是单元测试与功能仿真的参考实现,保证指令语义在 CPU 上可验证、可调试。

使用示例:自动模式与手动模式

自动模式(Auto Mode)

自动模式下由编译器/运行时负责 Tile 的资源放置与调度,用户只表达数据流:

# Auto mode: compiler/runtime-managed placement and scheduling. %dst = pto.tcolexpandadd %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...>

手动模式(Manual Mode)

手动模式下资源必须显式绑定后再发射指令,可通过pto.tassign将虚拟操作数绑定到具体存储地址:

# Manual mode: resources must be bound explicitly before issuing the instruction. # Optional for tile operands: # pto.tassign %arg0, @tile(0x1000) # pto.tassign %arg1, @tile(0x2000) %dst = pto.tcolexpandadd %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...>

PTO 汇编形式

%dst = tcolexpandadd %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...> # AS Level 2 (DPS) pto.tcolexpandadd ins(%src0, %src1 : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>)

内核实测示例(NPU 测试用例)

仓库在 A2/A3、A5、kirin9030、kirinDev0000、kirinX90 等多个平台均提供了 TCOLEXPANDADD 的系统测试用例,例如 tests/npu/a5/src/st/testcase/tcolexpandadd/tcolexpandadd_kernel.cpp:

template <typename T, uint32_t dstRow, uint32_t dstCol, uint32_t src1Row, uint32_t src1Col> __global__ AICORE void runCOLEXPANDADD(__gm__ T __out__* out, __gm__ T __in__* src0, __gm__ T __in__* src1) { using DynShapeDim5 = Shape<1, 1, 1, src1Row, src1Col>; using DynStridDim5 = pto::Stride<1, 1, 1, src1Col, 1>; using GlobalData = GlobalTensor<T, DynShapeDim5, DynStridDim5>; using TileData = Tile<TileType::Vec, T, src1Row, src1Col, BLayout::RowMajor, -1, -1>; using DstDynShapeDim5 = Shape<1, 1, 1, dstRow, dstCol>; using DstDynStridDim5 = pto::Stride<1, 1, 1, dstCol, 1>; using DstGlobalData = GlobalTensor<T, DstDynShapeDim5, DstDynStridDim5>; using DstTileData = Tile<TileType::Vec, T, dstRow, dstCol, BLayout::RowMajor, -1, -1>; DstTileData src0Tile(dstRow, dstCol); TileData src1Tile(src1Row, src1Col); DstTileData dstTile(dstRow, dstCol); TASSIGN(src0Tile, 0x0); TASSIGN(src1Tile, 0x10000); TASSIGN(dstTile, 0x20000); int offset = 0; DstGlobalData src0Global(src0 + offset); GlobalData src1Global(src1 + offset); DstGlobalData dstGlobal(out + offset); TLOAD(dstTile, dstGlobal); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); #ifndef __PTO_AUTO__ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); #endif TCOLEXPANDADD(dstTile, src0Tile, src1Tile); #ifndef __PTO_AUTO__ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); #endif TSTORE(dstGlobal, dstTile); #ifndef __PTO_AUTO__ set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); pipe_barrier(PIPE_ALL); #endif out = dstGlobal.data(); }

该用例展示了完整的"装载(TLOAD)→ 运算(TCOLEXPANDADD)→ 存储(TSTORE)"流水:TASSIGN显式绑定三个 Tile 的存储基址(0x0/0x10000/0x20000),src1Tile的形状为src1Row x src1Col(测试中取1 x C,即每列一个标量),dstTilesrc0Tile形状为dstRow x dstCol__PTO_AUTO__宏区分自动模式与手动模式:非自动模式下需用set_flag / wait_flag显式同步 MTE2 → V → MTE3 → S 管道事件。

测试用例覆盖的形状与数据类型组合包括:

  • float16x12832x321x128
  • aclFloat16half):4x25610x64(后者验证非对齐列数下的掩码/计数路径);
  • int32_t16x32int16_t16x64
  • int64_tuint64_t16x32(A5 的 64 位寄存器对模拟路径)。

数据生成脚本见 tests/npu/a5/src/st/testcase/tcolexpandadd/gen_data.py,A2/A3 对应用例见 tests/npu/a2a3/src/st/testcase/tcolexpandadd/,CPU 仿真测试见 tests/cpu/st/testcase/tcolexpandop/。

进一步阅读

  • 同族指令参考:TCOLEXPAND、TCOLEXPANDDIV、TCOLEXPANDMUL、TCOLEXPANDSUB、TCOLEXPANDMAX、TCOLEXPANDMIN、TCOLEXPANDEXPDIF,以及与行广播对应的 TROWEXPANDADD;
  • 虚拟指令集整体说明:PTO-Virtual-ISA-Manual.md 与 PTOISA.md;
  • 后端实现:include/pto/npu/a2a3/TColExpandBinOp.hpp、include/pto/npu/a2a3/TColExpandAdd.hpp、include/pto/npu/a5/TColExpandAdd.hpp;
  • CPU 参考实现:include/pto/cpu/TColExpandOp.hpp;
  • 编程入门教程:docs/coding/tutorials/README.md 与 docs/coding/ProgrammingModel.md。
  • 人工智能
  • 指令集
  • 算子库
  • CANN
  • Ascend

【免费下载链接】pto-isa

Parallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.

项目地址:https://gitcode.com/cann/pto-isa
点击查看免费下载

相关推荐

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询