PyPTO pypto.exp2 逐元素 2 的幂运算：从接口参数、TileShape 切分到源码实现的完整指南

PyPTO pypto.exp2 逐元素 2 的幂运算：从接口参数、TileShape 切分到源码实现的完整指南

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

本篇围绕 PyPTO（Parallel Tensor/Tile Operation）Tensor API 中的pypto.exp2接口展开，系统讲解该逐元素 2 的幂运算函数的产品支持范围、参数约束、返回值类型规则、临时空间计算方式与 TileShape 切分配置，并结合当前仓库的 Python 前端实现、C++ 框架层实现与系统测试用例，说明从接口调用到底层向量切分执行的完整链路。读完本文，你可以直接在 NPU 上编写可运行的 exp2 kernel，并理解整数输入为何会升精度为 FP32、临时空间公式从何而来。

功能说明与产品支持

pypto.exp2计算输入 Tensor 中每个元素的 2 的指数（即out = 2^input），属于逐元素（element-wise）一元运算，返回与输入形状相同的 Tensor。

当前仓库中该接口的产品支持情况如下：

• Ascend 950PR / Ascend 950DT：支持
• Atlas A3 训练系列产品 / Atlas A3 推理系列产品：支持
• Atlas A2 训练系列产品 / Atlas A2 推理系列产品：支持

接口定义与完整文档见 pypto-exp2 API 文档。

函数原型

exp2(input: Tensor) -> Tensor

exp2在 Python 前端定义于 python/pypto/op/math.py，通过@op_wrapper装饰器包装并对外暴露。同时，Tensor 方法形式Tensor.exp2()也在 python/pypto/tensor.py 中提供，内部直接委托给pypto.exp2(self)，两种写法等价：

y = pypto.exp2(x) # 函数形式 y = x.exp2() # 方法形式

参数说明

从源码结构看，这些约束在 C++ 框架层 framework/src/interface/operation/vector/exponential.cpp 的Exp2入口函数中被逐一校验：

• CheckTensorDimRange(..., 1, 4, "EXP2")校验维度范围为 1~4 维；
• CheckTensorShapeSize(...)校验 Shape Size 上限；
• 框架层声明的直接支持类型集合为{DT_FP32, DT_FP16, DT_BF16, DT_INT32, DT_INT16}。

返回值说明

返回输出 Tensor，其数据类型按输入类型区分：

• 当输入是 DT_FP32、DT_FP16、DT_BF16 时，输出 Tensor 的数据类型与 input 相同，Shape 与 input 相同；
• 当输入是 DT_INT32、DT_INT16、DT_INT8、DT_UINT8 时，输出 Tensor 的数据类型为 DT_FP32，Shape 与 input 相同。

整数输入升精度的原因与 C++ 实现一致：在 framework/src/interface/operation/vector/exponential.cpp 的TensorExp2中，当输入为 DT_INT32 或 DT_INT16 时，结果 LogicalTensor 会显式构造为 DT_FP32。而 DT_INT8 / DT_UINT8 输入则由 Python 前端先行两次 cast（INT8/UINT8 → FP16 → FP32）后走浮点路径，见后文“源码实现”一节。

约束说明

1. 值域范围：input 中的值需要在 [-2^24, 2^24] 范围内，以确保在计算过程中能精确转换为 float32。这是因为 2 的幂运算的底数与中间量需要以 float32 表示，超出该范围的整数无法在 float32 中无损表达。
2. 格式限制：Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式。该限制对应源码中 framework/src/interface/operation/vector/exponential.cpp 的CheckTensorFormat(self.GetStorage(), {TileOpFormat::TILEOP_NZ}, "Exp2")调用，即 NZ 格式输入会被拒绝。
3. 临时空间（Workspace）大小：若 TileShape 为[a,b,c,d]，需要的临时空间大小按产品区分：
   • Ascend 950PR / Ascend 950DT：a*b*c*d*sizeof(DT_FP32)
   • Atlas A3 训练系列 / Atlas A3 推理系列：a*b*c*(2*d + 3*d/8)*sizeof(DT_FP32)
   • Atlas A2 训练系列 / Atlas A2 推理系列：a*b*c*(2*d + 3*d/8)*sizeof(DT_FP32)

临时空间公式与框架实现可以相互印证。在 framework/src/interface/operation/vector/exponential.cpp 的TiledExp2中，对于非 FP32 输入会额外申请tmpTensor与tmpTensorNext两个 DT_FP32 临时张量：其中tmpShape对最后两维（1D 输入退化为 1 维）做BLOCK_SIZE / BytesOf(DT_FP32)对齐，tmpShape2对最后 1~2 维做 FP32 对齐。从源码结构看，2*d + 3*d/8这一项即来源于这类“输入拷贝缓冲 + 对齐后的 FP32 中间缓冲”的合计开销，且所有临时缓冲均为 FP32，因此公式中统一使用sizeof(DT_FP32)。

TileShape 设置示例

调用exp2接口前，应通过pypto.set_vec_tile_shapes设置 TileShape，且 TileShape 的维度应与输出一致。

示例 1：输入 input shape 为[m, n]，输出为[m, n]，TileShape 设置为[m1, n1]，则m1、n1分别用于切分m、n轴：

pypto.set_vec_tile_shapes(4, 16)

从源码实现看，TileShape 的切分逻辑正体现在TiledExp2的递归循环中（framework/src/interface/operation/vector/exponential.cpp）：沿当前维度以vecTile[cur]为步长遍历，每步生成一个View片段并递归进入下一维，直到所有维度遍历完毕后才在叶子节点真正插入OP_EXP2算子。这也解释了为何 TileShape 维度必须与输出维度一致——切分器是按维度逐级展开的。

接口调用示例

最简调用如下：

x = pypto.tensor([3], pypto.DT_FP32) y = pypto.exp2(x)

结果示例：

输入数据x: [0.0 1.0 2.0] 输出数据y: [1.0 2.0 4.0]

下面给出一个完整的逐元素 kernel 写法（取自仓库系统测试 python/tests/st/operation/vector/test_exp2.py），展示set_vec_tile_shapes+loop+view+exp2+assemble的组合模式：

import pypto @pypto.frontend.jit(debug_options={"runtime_debug_mode": 0, "compile_debug_mode": 0}) def exp2_2d_1input_kernel(input0: pypto.Tensor(), output: pypto.Tensor(), config: Exp2Config): pypto.set_vec_tile_shapes(*config.tile_shape) for index_0 in pypto.loop(config.loop_ranges[0]): for index_1 in pypto.loop(config.loop_ranges[1]): offsets = [index_0 * config.execution_view_shape[0], index_1 * config.execution_view_shape[1]] input0_offset = [0 if config.input_shapes[0][axis] == 1 else offsets[axis] for axis in range(2)] input0_view = pypto.view(input0, config.input_view_shapes[0], input0_offset) result = pypto.exp2(input0_view) output_offset = [0 if config.output_offset_map[axis] < 0 else offsets[config.output_offset_map[axis]] for axis in range(len(config.execution_view_shape))] pypto.assemble(result, output_offset, output)

该 kernel 以 TileShape 为单位循环遍历输出张量：每次循环用pypto.view取出一个 tile 片段，对其调用pypto.exp2，再用pypto.assemble把结果片段拼回输出张量，是逐元素算子在 PyPTO 中最典型的调用范式。

源码实现：整数输入为何升精度、exp2 如何落地

Python 前端的exp2实现位于 python/pypto/op/math.py，核心逻辑可以拆成三步：

valid_shape = input.GetValidShape() if input.dtype in (pypto.DT_INT8, pypto.DT_UINT8): input = pypto.cast(input, pypto.DT_FP16) input = pypto.cast(input, pypto.DT_FP32) elif input.dtype in (pypto.DT_INT16, pypto.DT_INT32): input = pypto.cast(input, pypto.DT_FP32) two_element = pypto_impl.Element(input.dtype, 2) x = pypto_impl.Full(two_element, input.dtype, input.shape, valid_shape) return pow(x, input)

• 对 DT_INT8 / DT_UINT8 输入，前端先 cast 到 DT_FP16，再 cast 到 DT_FP32，保证整数先被无损地提升到浮点表示，最终输出为 FP32，与 API 文档“整数输入输出为 DT_FP32”的约定一致；
• 对 DT_INT16 / DT_INT32 输入，直接 cast 到 DT_FP32（该范围内整数值可被 float32 精确表示，呼应“值域 [-2^24, 2^24]”的约束）；
• 浮点输入则保持原 dtype，exp2最终以pow(2, input)（底数 2 的 Full 张量与输入做幂运算）的方式组合实现。从源码结构看，这是把“2 的幂”归约到通用pow算子路径的写法。

框架层（C++）则在 framework/src/interface/operation/vector/exponential.cpp 中完成真正的图构建：

1. Exp2入口做格式、数据类型、维度、Shape Size 校验；
2. TensorExp2创建输出 LogicalTensor（INT32/INT16 时输出 DT_FP32）并登记Opcode::OP_EXP2操作（算子名 "TExp2"，见 framework/src/interface/operation/opcode.cpp）；
3. TiledExp2按 TileShape 递归切分输入，在非 FP32 输入时构造按BLOCK_SIZE对齐的 FP32 临时张量，最终调用function.AddOperation(Opcode::OP_EXP2, {tile}, {resultTile, tmpTensor, tmpTensorNext})生成设备端指令。

此外，该算子也接入了解释器路径，REGISTER_CALC_OP(OP_EXP2, ...)位于 framework/src/interface/interpreter/calc_vector.cpp，可用于本地数值核算。

系统测试验证

仓库为exp2提供了多层级测试，可直接作为行为参考：

• Python 端系统测试：python/tests/st/operation/vector/test_exp2.py，通过@pytest.mark.parametrize驱动多个用例，期望值用torch.exp2生成后对比；用例配置（如 (96, 8192) int16 输入 / fp32 输出、(512, 256, 256) fp32 输入等）见 python/tests/st/operation/vector/vector_testcase/exp2_test_case.py，其中view_shape与tile_shape的组合覆盖了 TileShape 与张量形状不对齐的场景；
• C++ 端算子级 ST：framework/tests/st/operation/src/test_exp2_operation.cpp，覆盖 2/3/4 维切分执行；
• 代码生成单测：framework/tests/ut/codegen/src/test_dynamic/test_codegen_dyn_vector/test_codegen_dyn_exp2.cpp。

相关接口

• Tensor.exp2()：方法形式的同一接口，见 python/pypto/tensor.py；
• pypto.exp/pypto.expm1：同族指数类逐元素算子（支持 PrecisionType 参数），定义于 python/pypto/op/math.py；
• pl.simt.exp2：SIMT（标量）路径下的 2 的幂运算，用于 Pro 端 SIMT 编程，见 python/pypto_pro/ir/op/simt_ops.py 与 python/pypto_pro/language/_simt_api.py。

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto