PyPTO pypto.exp2 逐元素 2 的幂运算从接口参数、TileShape 切分到源码实现的完整指南【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本篇围绕 PyPTOParallel Tensor/Tile OperationTensor 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) - Tensorexp2在 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() # 方法形式参数说明参数名输入/输出说明input输入源操作数。支持的类型为 Tensor。支持的数据类型为DT_FP32、DT_FP16、DT_BF16、DT_INT32、DT_INT16、DT_INT8、DT_UINT8。不支持空 TensorShape 仅支持 1~4 维Shape Size 不大于 2147483647即 INT32_MAX。从源码结构看这些约束在 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_FP32Shape 与 input 相同。整数输入升精度的原因与 C 实现一致在 framework/src/interface/operation/vector/exponential.cpp 的TensorExp2中当输入为 DT_INT32 或 DT_INT16 时结果 LogicalTensor 会显式构造为 DT_FP32。而 DT_INT8 / DT_UINT8 输入则由 Python 前端先行两次 castINT8/UINT8 → FP16 → FP32后走浮点路径见后文“源码实现”一节。约束说明值域范围input 中的值需要在 [-2^24, 2^24] 范围内以确保在计算过程中能精确转换为 float32。这是因为 2 的幂运算的底数与中间量需要以 float32 表示超出该范围的整数无法在 float32 中无损表达。格式限制Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式。该限制对应源码中 framework/src/interface/operation/vector/exponential.cpp 的CheckTensorFormat(self.GetStorage(), {TileOpFormat::TILEOP_NZ}, Exp2)调用即 NZ 格式输入会被拒绝。临时空间Workspace大小若 TileShape 为[a,b,c,d]需要的临时空间大小按产品区分Ascend 950PR / Ascend 950DTa*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_shapesloopviewexp2assemble的组合模式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]”的约束浮点输入则保持原 dtypeexp2最终以pow(2, input)底数 2 的 Full 张量与输入做幂运算的方式组合实现。从源码结构看这是把“2 的幂”归约到通用pow算子路径的写法。框架层C则在 framework/src/interface/operation/vector/exponential.cpp 中完成真正的图构建Exp2入口做格式、数据类型、维度、Shape Size 校验TensorExp2创建输出 LogicalTensorINT32/INT16 时输出 DT_FP32并登记Opcode::OP_EXP2操作算子名 TExp2见 framework/src/interface/operation/opcode.cppTiledExp2按 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 端算子级 STframework/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.pypypto.exp/pypto.expm1同族指数类逐元素算子支持 PrecisionType 参数定义于 python/pypto/op/math.pypl.simt.exp2SIMT标量路径下的 2 的幂运算用于 Pro 端 SIMT 编程见 python/pypto_pro/ir/op/simt_ops.py 与 python/pypto_pro/language/_simt_api.py。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考