企业官网建设流程全解析

• 人工智能
• 指令集
• 算子库
• CANN
• Ascend

TPARTMAX 是 CANN pto-isa 虚拟指令集（Parallel Tile Operation，面向 Ascend 平台的 tile 级操作指令集）中用于在目标有效区域内执行逐元素最大值选择的核心指令：当两个源 Tile 在某个位置都有效时取二者最大值，仅一个有效时直接透传该输入，从而支持 shape 不完全一致的 Tile 之间进行安全的归并计算。本文以 TPARTMAX_zh.md 为骨架，结合仓库内 CPU 模拟、Atlas A2/A3 与 Ascend 950 系列（A5）的真实实现和 ST 测试用例，完整讲解该指令的数学语义、三级汇编语法、C++ 内建接口、平台约束、底层实现原理与可复制的编程示例，帮助你正确地在自定义算子中完成带边界掩码的最大值融合。

指令概述

TPARTMAX 在目标有效区域（destination valid region）内执行逐元素最大值选择。其核心设计目标是处理部分有效（partial validity）的场景：在深度学习算子中，两个参与融合的输入往往 shape 并不完全一致（例如 softmax 后的 mask、padding 后的张量、不同 tile 切分策略产生的碎片）。TPARTMAX 允许：

• 两个输入在该位置均有效：结果为max(src0, src1)；
• 仅 src0 有效：结果直接复制src0的值；
• 仅 src1 有效：结果直接复制src1的值；
• 其余有效区域不匹配的情况由具体实现定义（在 CPU 参考实现中表现为写入 0，见下文源码分析）。

它属于 TPART 系列指令（同系列还有 TPARTADD、TPARTMIN、TPARTMUL、TPARTARGMAX/ARGMIN 等，均定义于 docs/isa/），用于在两个输入有效区域不完全对齐时仍能完成二元/归约运算，避免开发者手工编写逐元素掩码逻辑。

数学语义

对目标有效区域内的每个元素(i, j)，指令结果可形式化描述为：

$$ \mathrm{dst}{i,j} = \begin{cases} \max(\mathrm{src0}{i,j}, \mathrm{src1}{i,j}) & \text{若两个输入在 } (i,j) \text{ 处均有定义} \ \mathrm{src0}{i,j} & \text{若仅src0在 } (i,j) \text{ 处有定义} \ \mathrm{src1}_{i,j} & \text{若仅src1在 } (i,j) \text{ 处有定义} \end{cases} $$

从 CPU 参考实现可以最直观地印证这三分支语义。在 include/pto/cpu/TPartOp.hpp 的TPartInstr中，逐元素遍历时用两个布尔量判断输入有效范围：

bool InSrc0 = i < Src0ValidRow && j < Src0ValidCol; bool InSrc1 = i < Src1ValidRow && j < Src1ValidCol; if (InSrc0 && InSrc1) { InstrOp::PartInstr(dst, src0, src1, DstOffset, Src0Offset, Src1Offset); } else if (InSrc0 && !InSrc1) { dst[DstOffset] = src0[Src0Offset]; } else if (!InSrc0 && InSrc1) { dst[DstOffset] = src1[Src1Offset]; } else { dst[DstOffset] = 0; }

其中PartInstr在 include/pto/cpu/TPartMax.hpp 中被具体化为dst[DstOffset] = std::max(src0[Src0Offset], src1[Src1Offset])。注意最后一种"双无效"分支写入 0——这正是文档中"其余有效区域不匹配的情况由具体实现定义"在 CPU 参考实现中的落地行为。

汇编语法

同步形式（PTO 汇编）

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

AS Level 1（SSA 形式）

SSA 形式携带完整类型签名，两个输入与一个输出均为 tile 类型：

%dst = pto.tpartmax %src0, %src1 : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

AS Level 2（DPS 形式）

DPS（Data Parallel Semantics）形式显式区分ins（输入）与outs（输出），操作数类型为!pto.tile_buf<...>：

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

三级形式对应 pto-isa 的分层汇编设计：PTO 汇编面向手写内核的简洁表达，Level 1 SSA 供编译器中间表示使用，Level 2 DPS 则明确了数据依赖关系，便于调度器做流水线优化。

C++ 内建接口

指令声明于 include/pto/common/pto_instr.hpp，公共包含头为 include/pto/pto-inst.hpp：

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

关键点：

• 模板参数：TileDataDst、TileDataSrc0、TileDataSrc1为 Tile 数据类型（携带元素类型、行列、布局等编译期信息）；可变参数WaitEvents支持传入事件用于同步等待。
• 返回RecordEvent：指令发射后返回事件记录，可配合事件机制实现异步流水与同步，这也是 pto-isa 中绝大多数指令的统一内建形态（可参考 docs/coding/Event.md 了解事件机制）。
• 内部实现入口：TPARTMAX最终调用各平台对应的TPARTMAX_IMPL（CPU、A2/A3、A5 各有实现，见下文）。

约束与平台差异

通用约束

• dst、src0和src1的元素类型必须一致；
• 目标有效区域定义结果的计算范围；
• 若dst的有效区域为零（行列任一为 0），指令直接返回；
• 支持的部分有效区域模式要求至少一个源 Tile 的有效区域与dst完全一致，另一个源 Tile 的有效区域在两个维度上都不能超过dst；
• 上述范围之外的有效区域组合，行为均由具体实现定义。

这部分约束与 CPU 实现 include/pto/cpu/TPartOp.hpp 的TPartCheck完全对应：它在编译期通过static_assert校验元素类型合法性（int32_t/int16_t/half/bfloat16_t/float/int8_t/uint8_t/uint16_t/int64_t/uint64_t/uint32_t），并在运行期对DstValidRow == 0 || DstValidCol == 0直接提前返回。

Atlas A2/A3 训练/推理系列产品实现检查

在 include/pto/npu/a2a3/TPartMax.hpp 中通过static_assert强制：

• 支持元素类型：int32_t、int、int16_t、half、float16_t、float、float32_t；
• dst、src0、src1必须全部为行主序（isRowMajor），否则编译期报错"TPARTMAX not supported BLayout type."。

Ascend 950PR / Ascend 950DT 实现检查

在 include/pto/npu/a5/TPartMax.hpp 中支持的位宽更广：

• 支持元素类型：int8_t、uint8_t、int16_t、uint16_t、int32_t、uint32_t、int64_t、uint64_t、half、bfloat16_t、float。

A5 实现对int64_t/uint64_t走独立的Int64Part<Int64Op::Max, ...>路径（见 include/pto/npu/a5/TPartMax.hpp），其余类型走统一的带谓词寄存器的vmax路径。CPU 参考实现的类型集合与 A5 一致，方便先用 CPU 模拟验证再上昇腾硬件。

底层实现原理：有效区域如何被翻译为硬件指令

CPU 参考实现

CPU 实现（include/pto/cpu/TPartMax.hpp + include/pto/cpu/TPartOp.hpp）是逐元素标量模拟，逻辑清晰、适合阅读与调试，也是 tests/cpu/st/testcase/tpartmax/ 测试的验证基准。

Atlas A2/A3 实现

A2/A3 的实现把"部分有效"翻译为硬件向量指令的组合，核心在 include/pto/npu/a2a3/TPartOp.hpp 的TPartInstr，按src1与dst的有效区域关系分三种情况处理：

• src1ValidRow < dstValidRow：先在src1有效的行区间内用TPartOps做向量化vmax，剩余行直接用TPartCopyInstr将src0对应行复制到dst（此时相当于"仅 src0 有效"，直接透传）；
• src1ValidCol < dstValidCol：先将整块src0复制到dst，再在src1有效的列区间内用vmax(dst, src1, dst)原地取最大值，中间通过pipe_barrier(PIPE_V)保证先拷贝后计算的顺序；
• src1ValidRow == dstValidRow && src1ValidCol == dstValidCol（完全相等）：直接对整个有效区域做vmax。

指令本体 include/pto/npu/a2a3/TPartMax.hpp 的PartMaxOp最终调用昇腾向量指令vmax(dst, src0, src1, repeats, 1, 1, 1, ...)，其中elementsPerRepeat与blockSizeElem由REPEAT_BYTE / sizeof(T)和BLOCK_BYTE_SIZE / sizeof(T)推导（见 include/pto/npu/a2a3/TPartMax.hpp）。TPartOps还会根据行 stride 是否超过REPEAT_STRIDE_MAX以及行列规模自动在计数模式（PartCountMode）与常规模式（PartNormMode）之间选择，以兼顾非连续布局与访存效率。

Ascend 950 系列（A5）实现

A5 实现（include/pto/npu/a5/TPartMax.hpp）通过TPartBinOps的TPARTOP_IMPL调度，核心运算为带谓词寄存器的向量指令：

vmax(dst, src0, src1, preg, MODE_ZEROING);

MODE_ZEROING表明掩码未覆盖的位置写 0，与文档中"具体实现定义"的行为保持一致；int64/uint64位宽则拆分为Int64Part专门处理。从源码结构看，A5 的 TPart 系列（TPARTADD/TPARTMIN/TPARTMAX 等）共享同一套TPartBinOps基础设施，指令差异仅体现在BinInstr内调用的底层向量指令不同。

编程示例

自动模式（Auto）

自动模式下，Tile 的放置与调度由编译器/运行时负责，用户只需声明 Tile 并调用指令：

#include <pto/pto-inst.hpp> using namespace pto; void example_auto() { using TileT = Tile<TileType::Vec, float, 16, 16>; TileT src0, src1, dst; TPARTMAX(dst, src0, src1); }

Tile<TileType::Vec, float, 16, 16>表示一个 16×16 的 float 向量 Tile，TPARTMAX按模板实参推导出三者的类型。

手动模式（Manual）

手动模式下需先用TASSIGN显式绑定 Tile 在缓冲区中的地址（地址以字节为单位）：

#include <pto/pto-inst.hpp> using namespace pto; void example_manual() { using TileT = Tile<TileType::Vec, float, 16, 16>; TileT src0, src1, dst; TASSIGN(src0, 0x1000); TASSIGN(src1, 0x2000); TASSIGN(dst, 0x3000); TPARTMAX(dst, src0, src1); }

带部分有效区域的完整内核示例

仓库 ST 测试给出了更贴近实战的完整用法。以 tests/cpu/st/testcase/tpartmax/tpartmax_kernel.cpp 为例：src0为完整的 64×64 Tile，src1只有 32×32 的有效区域（kValidRows1 = kValidCols1 = 32），两个 Tile 均以行主序声明：

using TileT = Tile<TileType::Vec, T, kRows, kCols, BLayout::RowMajor, -1, -1>; TileT src0Tile(kRows, kCols); // 64×64，全有效 TileT src1Tile(kValidRows1, kValidCols1); // 32×32，部分有效 TileT dstTile(kRows, kCols); TASSIGN(src0Tile, 0); TASSIGN(src1Tile, kRows * kCols * sizeof(typename TileT::DType)); TASSIGN(dstTile, 2 * kRows * kCols * sizeof(typename TileT::DType)); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); TPARTMAX(dstTile, src0Tile, src1Tile); TSTORE(dstGlobal, dstTile);

A2/A3 版本（tests/npu/a2a3/src/st/testcase/tpartmax/tpartmax_kernel.cpp）展示了更多样的有效区域组合，包括 16×16 对 16×32、行少列等、行等列少、以及某一维有效为 0（16, 0/0, 32）等边界情形，覆盖了TPartInstr的三条分支路径。非自动模式下手写内核还需通过set_flag/wait_flag（如PIPE_MTE2 → PIPE_V、PIPE_V → PIPE_MTE3）显式管理 MTE2 加载、向量计算与 MTE3 存储之间的流水同步（见同一文件的第 50-58 行）。

汇编示例

自动模式

# 自动模式：由编译器/运行时负责资源放置与调度。 %dst = pto.tpartmax %src0, %src1 : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

手动模式

手动模式要求先显式绑定资源（对包含 tile 操作数的指令，tassign为可选项），再发射指令：

# 手动模式：先显式绑定资源，再发射指令。 # 可选（当该指令包含 tile 操作数时）： # pto.tassign %arg0, @tile(0x1000) # pto.tassign %arg1, @tile(0x2000) %dst = pto.tpartmax %src0, %src1 : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

PTO 汇编形式

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

测试验证

仓库为 TPARTMAX 提供了覆盖 CPU、A2/A3、A5、kirin 系列等多平台的 ST 测试，便于读者对照验证行为：

• CPU：tests/cpu/st/testcase/tpartmax/（含main.cpp、tpartmax_kernel.cpp、gen_data.py），覆盖float及int8/uint8/int16/uint16/int32/uint32/int64/uint64共 9 种类型；
• A2/A3：tests/npu/a2a3/src/st/testcase/tpartmax/，覆盖部分有效区域的多种组合与边界（有效行/列为 0）；
• A5：tests/npu/a5/src/st/testcase/tpartmax/；
• kirin：tests/npu/kirin9030/src/st/testcase/tpartmax/、tests/npu/kirinDev0000/src/st/testcase/tpartmax/。

这些测试通过 tests/run_st.sh 统一调度，gen_data.py负责生成输入与 golden 数据，main.cpp负责 host 侧 launch 与结果比对。测试中"src0 全有效 + src1 部分有效"的典型形态，正是 TPARTMAX 区别于普通TMAX（要求两个输入有效区域完全对齐，见 docs/isa/TMAX.md）的核心场景。

小结

TPARTMAX 通过"有效区域驱动的三分支选择"机制，让最大值运算能够安全作用于 shape 不完全对齐的两个 Tile，免去开发者手写掩码与边界处理。理解它需要抓住三条主线：数学语义（双有效取 max、单有效透传）、有效区域约束（至少一个输入与 dst 完全一致、另一输入不得越界）、平台落地（CPU 逐元素模拟、A2/A3 用 vmax + 拷贝组合、A5 用带谓词寄存器的 vmax）。在实际算子开发中，建议先用 CPU 参考实现验证语义（可运行tests/cpu/st/testcase/tpartmax/用例），再按目标平台的类型支持与行主序要求编写内核，最后通过对应平台的 ST 测试框架完成验证。

• 人工智能
• 指令集
• 算子库
• CANN
• Ascend