PTO TCOLEXPANDSUB 指令详解：列广播减法（Column-wise Broadcast Subtract）的语义、三级汇编语法与双平台源码实现-平芜编程栈

人工智能
指令集
算子库
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

点击查看免费下载

TCOLEXPANDSUB 是 CANN PTO（Parallel Tile Operation）虚拟指令集中用于列广播减法的核心向量指令：它将每行中的每个元素减去一个"每列一个标量"的向量，常与 TCOLEXPAND 系列指令配合用于逐列统计后的去均值（center）等 tile 级运算。本文以 docs/isa/TCOLEXPANDSUB.md 为骨架，结合include/pto/下的 A2A3/A5 后端实现与tests/中的 ST 测试用例，完整讲解其数学语义、汇编语法（PTO Assembly / AS Level 1 SSA / AS Level 2 DPS）、C++ 内建接口、类型与布局约束、64 位元素模拟机制以及底层向量指令调用链，读者可在掌握语法后直接编写可运行的 PTO kernel。

指令概述

TCOLEXPANDSUB 的语义是列广播减法（column-wise broadcast subtract）：src0是一个R × C的 Tile，src1是一个长度为C（每列一个标量）的向量，指令把src1中第j列的值s_j广播到每一行，然后计算dst[i][j] = src0[i][j] - s_j。

其指令示意图如下（来源：docs/figures/isa/TCOLEXPANDSUB.svg）：

它与同族指令的关系：PTO 提供了一整组TCOLEXPAND系列二元运算指令（ADD/DIV/MAX/MIN/MUL/SUB/EXPDIF 等），底层共享同一套"列广播二元运算"模板（见下文实现原理小节），其中减法即 TCOLEXPANDSUB。TCOLEXPAND（单操作数）负责把源列首元素广播到整个目标列，而TCOLEXPANDSUB等二元指令则在广播的同时完成逐元素运算，常被组合用于归一化类算子。

数学语义

设R = dst.GetValidRow()（有效行数）、C = dst.GetValidCol()（有效列数），s_j为从src1中取出的第j列的标量（每列一个值）。指令的计算公式为：

$$ \mathrm{dst}{i,j} = \mathrm{src0}{i,j} - s_j \qquad (0 \le i < R,\ 0 \le j < C) $$

要点：

广播发生在列维度上：s_j不随行号i变化，同一列的所有行共享同一个减数；
src1只提供C个有效值（一行、C列），实现层面会为每一行重复读取这一行向量；
结果dst的形状与src0一致，为R × C。

汇编语法

TCOLEXPANDSUB 的汇编书写分为三个层级：PTO 汇编（同步形式）、AS Level 1（SSA 形式）与 AS Level 2（DPS 形式）。

PTO 汇编（同步形式）

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

AS Level 1（SSA 形式）

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

AS Level 2（DPS 形式）

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

三个层级的差异体现在操作数描述粒度上：PTO 汇编与 AS Level 1 使用 SSA 值（%dst/%src0/%src1）与!pto.tile<...>类型描述；AS Level 2 则显式区分输入（ins）与输出（outs），并落到!pto.tile_buf<...>的 buffer 类型，对应更接近硬件资源视图的表示。同族指令（TCOLEXPANDADD 等）的语法结构完全一致，仅指令名不同。

C++ 内建接口

指令通过 C++ 内建函数暴露，声明位于 include/pto/common/pto_instr.hpp。公共包含头为<pto/pto-inst.hpp>，内部声明位于pto/common/pto_instr.hpp。

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

接口要点（可从源码 pto_instr.hpp 中 TCOLEXPANDSUB 的实现印证）：

返回RecordEvent，可用于指令间事件同步；
变参模板WaitEvents...支持传入事件对象，实现"先等待、再执行"的依赖语义；
内部实现先调用detail::PtoWaitEvents(events...)等待传入事件，再通过MAP_INSTR_IMPL(TCOLEXPANDSUB, dst, src0, src1)宏分派到平台相关实现；
三个 Tile 操作数模板参数分别对应dst、src0、src1，与汇编形式中的操作数顺序一致。

自动模式与手动模式

自动模式（Auto Mode）：资源放置与调度由编译器/运行时托管，直接书写 SSA 指令即可：

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

手动模式（Manual Mode）：发射指令前需先显式绑定资源，可通过pto.tassign将参数绑定到指定 tile 地址（可选，当指令包含 tile 操作数时）：

# 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.tcolexpandsub %src0, %src1 : !pto.tile<...>, !pto.tile<...> -> !pto.tile<...>

两种模式仅影响资源绑定与调度方式，指令语义不变。

约束（Constraints）

使用 TCOLEXPANDSUB 必须满足以下约束：

数据类型：TileDataDst::DType与TileDataSrc1::DType必须是以下类型之一：
- 通用类型（适用于 A2、A3 与 A5 平台）：half、float、int16、int32；
- A5 专属扩展类型：uint16、uint32、bfloat16_t、int8、uint8、int64、uint64。
布局约束（编译期）：TileDataDst::isRowMajor必须为真，即目标 Tile 必须是行主序布局。
src1形状：src1预期提供每列一个标量，即其有效形状必须覆盖C个值（通常为1 × C的 Vec Tile）。
平台特定约束：确切的布局/分形（fractal）约束是目标平台相关的，参见 include/pto/npu/a2a3/TColExpand*.hpp 与 include/pto/npu/a5/TColExpand*.hpp 下的后端头文件。

这些约束在源码中有对应的编译期检查。例如 a2a3/TColExpandBinOp.hpp 中的TCOLEXPANDOP_IMPL通过static_assert限定int32_t、int、int16_t、half、float16_t、float、float32_t等类型，并断言TileData::isRowMajor，否则直接编译失败（错误提示"Fix: TCOLEXPANDOP Invalid data type."与"Fix: TCOLEXPANDOP not supported Layout type"）；a5/TColExpandBinOp.hpp 则额外允许int64_t、uint64_t、uint32_t、uint16_t、bfloat16_t、int8_t、uint8_t，与文档中的 A5 扩展类型一致。

64 位元素类型（A5 专属）

int64/uint64仅在 A5（Ascend 950PR/Ascend 950DT）上受支持，其实现有两个关键特性：

模拟执行：A5 没有原生的 64 位向量运算单元，指令通过一对 32 位寄存器（分别保存每个元素的低 32 位与高 32 位）模拟实现；每列标量操作数与全尺寸操作数使用相同的解交织（de-interleaved）布局读取。
精度与对齐：计算结果为精确的 64 位补码值；Tile 对齐遵循 64 位元素的通用规则——RowMajor 的 Tile 要求Cols % 4 == 0（物理列数需为 4 的倍数，有效列数不要求对齐）。

源码级实现原理

TCOLEXPANDSUB 在不同平台上有不同的底层向量指令路径，均通过模板统一分派，理解实现有助于把握性能特征与边界条件。

A2/A3 平台：基于 vsub 的重复计算

A2/A3 后端实现在 include/pto/npu/a2a3/TColExpandSub.hpp：

template <typename T> struct ColExpandSubOp { PTO_INTERNAL static void ColExpandBinInstr(__ubuf__ T* dst, __ubuf__ T* src0, __ubuf__ T* src1, uint8_t repeats) { vsub(dst, src0, src1, repeats, 1, 1, 1, 8, 8, 8); } ... };

其核心是向量减指令vsub，并定义了两组算子（ColExpandSubOp与ColExpandSubOp2，后者交换src0/src1顺序，用于处理操作数形状互换的情形）。真正的广播循环在 a2a3/TColExpandBinOp.hpp 的TCOLEXPANDOP_IMPL中完成：

首先比较src0/src1与dst的有效形状是否一致（src0eqdst/src1eqdst），从而决定使用哪个算子、哪个操作数作为"每列标量"（TColExpandBinOp.hpp#L107-L120）；
若目标 tile 连续（Cols == ValidCol或Rows == 1），走TColExpandBinaryNormMode，以一次向量指令配合 repeat 步长（repeat stride）完成多行广播；否则走TColExpandBinaryCountMode，逐行循环并配合SetVectorCount设置向量计数（TColExpandBinOp.hpp#L76-L82）。

A5 平台：寄存器张量 + 谓词掩码，含 64 位模拟

A5 后端实现在 include/pto/npu/a5/TColExpandSub.hpp，其标量路径为：

template <typename T> struct ColExpandSubOp { PTO_INTERNAL static void ColExpandBinaryInstr( RegTensor<T>& reg_dst, RegTensor<T>& reg_src0, RegTensor<T>& reg_src1, MaskReg& preg) { vsub(reg_dst, reg_src0, reg_src1, preg, MODE_ZEROING); } ... };

差异点在于 A5 使用RegTensor（寄存器张量）+MaskReg（谓词掩码）+MODE_ZEROING的编程模型，并在 a5/TColExpandBinOp.hpp 中提供1D/2D与PostUpdate/NoPostUpdate多种实现变体（VFImplKind可选，默认走VFIMPL_DEFAULT，连续 tile 使用 PostUpdate 版本以复用地址自增特性）。

64 位路径则由Int64ColExpandBinary承担（a5/TColExpandBinOp.hpp#L155-L219）：每个 64 位元素拆为(low, high)两个 32 位寄存器，通过Int64LoadBounded按列偏移读取，Int64BinaryCalcRegs<Int64Op::Sub, T>完成低/高字对的减法运算，最后用vintlv交织回写，并配合pintlv_b32生成的低/高掩码分别存储。对于非 A5/A6 目标，该函数仅有声明式 stub，编译期不可用。

使用示例

C++ kernel 示例（A5 平台测试用例）

以下示例取自 tests/npu/a5/src/st/testcase/tcolexpandsub/tcolexpandsub_kernel.cpp，展示了完整的数据搬运—计算—存回流程：

#include <pto/pto-inst.hpp> using namespace pto; template <typename T, uint32_t dstRow, uint32_t dstCol, uint32_t src1Row, uint32_t src1Col> __global__ AICORE void runCOLEXPANDSUB(__gm__ T __out__* out, __gm__ T __in__* src0, __gm__ T __in__* src1) { using TileData = Tile<TileType::Vec, T, src1Row, src1Col, BLayout::RowMajor, -1, -1>; 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); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); // 手动模式下的流水线同步：等待 MTE2 搬运完成后向量单元再计算 set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); TCOLEXPANDSUB(dstTile, src0Tile, src1Tile); // 计算完成后等待，再触发 MTE3 存回 set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); TSTORE(dstGlobal, dstTile); }

该用例验证了src1形状为1 × C（src1Row == 1）的典型用法，并覆盖多种数据类型与形状组合，例如（tcolexpandsub_kernel.cpp#L74-L83）：

float：6×128、18×32；
half（aclFloat16）：10×256、12×64；
int32_t：16×32；int16_t：16×64；
64 位类型int64_t、uint64_t：16×32（对应 A5 的 64 位模拟路径）。

CPU 参考实现

CPU 侧的参考测试位于 tests/cpu/st/testcase/tcolexpandop/tcolexpandop_kernel.cpp，其中LaunchTCOLEXPANDSUB以 lambda 方式调用TCOLEXPANDSUB(dst, src0, src1)，src1声明为Tile<TileType::Vec, T, 1, iCol, BLayout::RowMajor, -1, -1>（即单行向量），可用于在没有 NPU 的环境下做语义正确性对照。

PTO 汇编完整形式

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