多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

PyPTO Tensor 操作全指南:数学运算与逻辑结构变换实战

PyPTO Tensor 操作全指南:数学运算与逻辑结构变换实战 PyPTO Tensor 操作全指南数学运算与逻辑结构变换实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoPyPTO 作为 Parallel Tensor/Tile Operation 编程范式为昇腾 NPU 上的张量编程提供了一套覆盖向量、矩阵与结构变换的完整操作体系。本文以官方编程指南中《Tensor的操作》章节为骨架结合仓库源码逐项剖析每个操作的函数签名、参数语义与底层实现帮助读者在编写算子时准确选用 API并能基于这些基础操作组合出 Softmax 等复杂计算逻辑。支持数学运算PyPTO 为张量计算提供了一套全面的操作旨在为用户提供高效、灵活的计算能力。这些操作分为向量操作、矩阵操作并可进一步将相关操作进行组合以实现更复杂的计算逻辑。从源码结构看所有数学运算集中在 python/pypto/op/math.py元素级数学函数、python/pypto/op/reduction.py归约操作与 python/pypto/op/matmul.py矩阵乘法中通过op_wrapper装饰器统一暴露为pypto.add、pypto.sum这类 Python 层接口。向量运算向量运算是在张量上执行元素级别的计算适用于各种基本数学操作包括算术运算、数学函数、激活函数、比较、归约等同时也支持原地修改与广播模式。算术运算# 加法 result pypto.add(a, b) result pypto.add(a, scalar) # 将标量添加到张量 result pypto.add(a, b, alpha2.0) # a 2.0 * b result a b # 减法 result pypto.sub(a, b) result pypto.sub(a, scalar) result pypto.sub(a, b, alpha2.0) # a - 2.0 * b result a - b # 乘法 result pypto.mul(a, b) result pypto.mul(a, scalar) result a * b # 除法 result pypto.div(a, b) result pypto.div(a, scalar) result a / b # 指数 result pypto.pow(a, scalar) # a ** scalar从源码看这些算术运算内部将标量操作数统一包装为Element后再参与计算。以add为例python/pypto/op/math.py 第 60104 行当other是 Tensor 时直接调用pypto_impl.Add(input_tensor, other)当other是标量时会先执行_check_scalar_type做类型校验例如浮点标量不能与整型张量直接相加再通过_clip_scalar_to_dtype将标量裁剪到张量数据类型可表示的范围内最后包装为pypto_impl.Element(input_tensor.dtype, other)参与计算。值得留意的是div和pow都提供了precision_type参数python/pypto/op/math.py 第 256311 行、第 693740 行PrecisionType.HIGH_PRECISION默认使用更高精度的计算方式降低精度损失适合对精度敏感的数值场景PrecisionType.INTRINSIC直接使用芯片指令计算更快适合对速度敏感、对精度要求不高的场景。此外pow在源码中对不同类型的标量指数做了区分整数指数包装为DT_INT32的Element浮点指数则包装为DT_DOUBLE的Element从而保证底数与指数类型匹配。数学函数# 指数和对数 result pypto.exp(x) # e ** x result pypto.log(x) # ln(x) result x.exp() result x.log() # 开根号 result pypto.sqrt(x) # √x result pypto.rsqrt(x) # 1/√x result x.sqrt() result x.rsqrt() # 三角函数 result pypto.sin(x) result pypto.cos(x) result x.sin() result x.cos() # 绝对值 result pypto.abs(x) result x.abs() # 相反数 result pypto.neg(x) # -x result x.neg()源码中数学函数家族相当完整除exp/log/sqrt/rsqrt/sin/cos/abs/neg外还包括exp2/expm1、log2/log10/log1p、tan/atan/atan2、sinh/cosh/tanh/asinh/acosh/atanh、asin/acos、erf/erfc、ceil/floor/trunc/round、sign/signbit、reciprocal、cbrt等见 python/pypto/op/math.py。exp、sqrt、rsqrt、log、log2、log10均支持precision_type参数。其中exp、sqrt、rsqrt、log系列默认使用INTRINSIC芯片指令速度优先div、pow、reciprocal默认使用HIGH_PRECISION精度优先。对数函数对输入域有要求log要求input 0log2/log10要求输入为正数log1p要求input -1。反三角与反双曲函数同样有输入域约束asin/acos的输入期望落在[-1, 1]超出范围的输入会产生 NaNatanh需要(-1, 1)区间内才有实数输出acosh要求input 1。exp2在源码中通过构造全 2 张量再调用pow(x, 2)实现而cbrt则由copysign(pow(abs(self), 1.0/3.0), self)组合实现python/pypto/op/math.py 第 22952321 行体现了基础操作自由组合的设计理念。激活函数# Sigmoid result pypto.sigmoid(x) # 1 / (1 exp(-x)) result x.sigmoid() # ReLU result pypto.relu(x) # 最大值(0, x) result x.relu() # Softmax result pypto.softmax(x, dim-1) # 沿着dim维度做Softmax result x.softmax(dim-1)激活函数并非全部是原子算子。以sigmoid为例python/pypto/operator.py 第 2581 行其在 Python 层由cast → mul → exp → add → full → div组合而成先将输入统一 cast 到 FP32乘以 -1 后取指数加 1 得到分母再用全 1 张量做INTRINSIC精度的除法最后按需 cast 回原类型。在 Lite NPU 上走sigmoid_no_cast分支跳过 cast普通 NPU 走sigmoid_fp32_cast分支。这印证了文档中组合相关操作以实现更复杂的计算逻辑的架构思想。relu则是对应底层算子pypto_impl.Relupython/pypto/op/math.py 第 14151441 行。比较操作# 最大值和最小值 result pypto.maximum(a, b) # 逐元素最大值 result pypto.minimum(a, b) # 逐元素最小值 # 截断 result pypto.clip(x, min_val, max_val) # 将x的取值范围进行截断源码细节maximum/minimumpython/pypto/op/reduction.py 第 90161 行支持 Tensor 与标量混合输入若只有一个操作数是 Tensor会自动交换顺序标量操作数会被包装成与张量同类型的Element。两者都要求至少一方是 Tensor否则抛出类型错误。clippython/pypto/op/math.py 第 19291979 行的min和max均可省略传入None表示单边截断省略一侧时会使用DT_BOTTOM类型的默认Element占位标量上下界会统一转换为输入张量的数据类型。归约操作# 求和 result pypto.sum(x, dim-1, keepdimFalse) result x.sum(dim-1, keepdimFalse) # 最大值 result pypto.amax(x, dim-1, keepdimFalse) result x.amax(dim-1, keepdimFalse) # 最小值 result pypto.amin(x, dim-1, keepdimFalse) result x.amin(dim-1, keepdimFalse)归约操作集中在 python/pypto/op/reduction.py参数语义统一dim沿哪个维度归约支持负索引-1表示最后一个维度不可省略keepdim为True时输出在该维度保留长度为 1 的轴为False默认时该维度被挤压输出少一维。除sum/amax/amin外该文件还提供prod连乘、argmax/argmin返回极值下标dim默认-1。配合where、pad见 python/pypto/op/other.py等算子可以进一步构建归一化、掩码、填充等逻辑。原地修改inplace# 使用move()进行就地操作 output.move(pypto.add(a, b)) # 高效无拷贝 # 或者使用赋值 output[:] pypto.add(a, b) # 也是高效的Tensor.move()是 PyPTO 中实现原地写入的核心方法python/pypto/tensor.py 第 640 行附近它将计算结果直接搬移到目标张量的存储中避免额外分配与拷贝。同理整张切片赋值output[:] ...在框架内部同样走高效的存储搬运路径。在分块计算tiling场景中使用move复用输出缓冲区是减少内存搬运、提升性能的常用手段。广播模式broadcast许多操作支持广播这使得操作更加灵活和高效。# 张量标量张量为标量 result pypto.add(tensor, 2.0) # 张量一维张量一维张量 bias pypto.tensor([features], pypto.DT_BF16, bias) result pypto.add(tensor, bias)广播的语义是两个操作数从最后一个维度开始逐维对齐维度为 1 的轴自动扩展直至形状一致。上例中[..., features]形状的张量与[features]形状的 bias 相加时bias 会沿前导维自动广播。广播不仅适用于addsub/mul/div/maximum/minimum/pow以及逻辑、位运算等绝大多数元素级操作均支持相关实现可见 python/pypto/op/math.py 中各算子 docstring 对 broadcasting 的说明。矩阵运算矩阵运算针对 NPU 的 Cube 核进行了优化适用于大规模矩阵计算。与向量运算走 Vector/AICore 路径不同矩阵乘法在编译期会被调度到 Cube 单元执行因此在算子性能敏感路径上应优先使用matmul而非逐元素展开。# 基本矩阵乘法 # C A B其中A: [M, K], B: [K, N], C: [M, N] result pypto.matmul(A, B, out_dtypepypto.DT_BF16) # 用转置其中A: [M, K], B: [N, K] result pypto.matmul(A, B, out_dtypepypto.DT_BF16, a_transFalse, b_transTrue) # 批量矩阵乘法 # A: [B, M, K], B: [B, K, N]结果[B, M, N] result pypto.matmul(A, B, out_dtypepypto.DT_BF16) # 带有偏置 bias pypto.tensor([1, N], pypto.DT_BF16, bias) result pypto.matmul(A, B, out_dtypepypto.DT_BF16, extend_params{bias_tensor: bias}) # 用转置其中A: [M, K], B: [N, K],输出为NZ格式 result pypto.matmul(A, B, out_dtypepypto.DT_BF16, a_transFalse, b_transTrue, c_matrix_nzTrue)矩阵乘法参数对应 python/pypto/op/matmul.py 第 22143 行的matmul函数签名参数含义默认值input左矩阵[M, K]或[B, M, K]必填mat2右矩阵[K, N]或[B, K, N]必填out_dtype输出数据类型必填a_trans是否转置左矩阵Falseb_trans是否转置右矩阵Falsec_matrix_nz是否以 NZ 格式输出Falseextend_params扩展特征偏置、去量化等None补充源码级细节维度与广播input与mat2支持 2 维、3 维、4 维张量且两者维度必须一致。3/4 维输入按批量矩阵乘法处理内部走BatchMatmul批量维支持广播例如input为[1, M, K]、mat2为[B, K, N]时输出为[B, M, N]。数据类型组合源码中的INPUT_COMBOS规定了合法输入组合包括 FP16×FP16、BF16×BF16、FP32×FP32、INT8×INT8以及多种 FP8FP8E5M2/FP8E4M3与 HF8 组合python/pypto/op/matmul.py 第 312322 行。out_dtype也需与输入类型匹配见BASIC_OUT_DTYPES表如 INT8 输入对应 INT32 输出。NZ 格式限制c_matrix_nzTrue用于以昇腾 NZNon-Zero分形数据布局格式输出便于后续 Cube 算子直接消费但源码NZ_UNSUPPORTED_INPUT_DTYPES表明 FP32、FP8E5M2、HF8 输入不支持 NZ 输出。extend_params 扩展能力extend_params字典支持以下键源码 docstring 与MatmulExtendParam转换逻辑bias_tensor偏置张量实现C A B biasscale浮点缩放系数用于反量化C DEQF16(ReLU(A B)) * scalescale_tensor逐通道缩放张量实现 per-channel 反量化relu_type反量化前应用的 ReLU 类型如pypto.ReLuType.RELUtrans_modeFP32 转 TF32 的舍入模式CAST_NONE关闭转换、CAST_RINT四舍五入到最近偶数、CAST_ROUND四舍五入远离零。另外 python/pypto/op/matmul.py 第 146 行起还提供了scaled_mm面向 FP8 微缩放MX格式的矩阵乘支持scale_a/scale_b缩放张量及各自的转置开关可满足 FP8 训练推理场景。组合操作可以组合上述基本操作以实现更复杂的计算逻辑。文档给出了一个经典的行 Softmax 组合示例def softmax_core(x: pypto.Tensor) - pypto.Tensor: row_max pypto.amax(x, dim-1, keepdimTrue) # 计算行最大值 sub x - row_max # 值归一化 exp pypto.exp(sub) # 指数运算 esum pypto.sum(exp, dim-1, keepdimTrue) # 求和 return exp / esum # 概率归一化 def softmax_kernel(x: pypto.Tensor, y: pypto.Tensor) - None: ... for idx in pypto.loop(x.shape[0]): ... softmax_out softmax_core(x) ...该实现分为三步先用amax求出每行的最大值keepdimTrue保留维度以便后续广播再让原张量减去行最大值做数值稳定性归一化防止exp溢出最后exp后求和并相除得到概率分布。keepdim在这里至关重要它保证了x - row_max可以正确广播。这一组合模式在仓库中已有真实落地官方pypto.softmax的底层实现python/pypto/operator.py 第 84 行起以及 examples/02_intermediate/operators/softmax/softmax.py 示例均体现了归约取极值 → 减均值/极值 → 指数 → 归约求和 → 归一化的经典流水线examples/01_beginner/compute/reduce_ops.py 也展示了归约类操作的组合用法。这类由基础算子自由拼装的能力正是 PyPTO组合成更复杂计算逻辑设计目标的体现。支持逻辑结构变换逻辑结构变换操作允许用户在张量上进行形状、维度和类型等变换以适应不同的计算需求。这些操作本身不引入数学语义只改变张量的视图、存储布局或数据类型是实现分块tiling、拼接与数据排布转换的基础工具对应实现集中在 python/pypto/operation.py 与 python/pypto/op/mutating.py。视图和组装视图视图操作创建对相同基础数据的新张量引用适用于分块处理和局部计算。# 视图 view pypto.view(tensor, view_shape, offset, valid_shape)示例# 创建具有特定形状和偏移量的视图 view pypto.view( tensor, view_shape[32, 32], offset[10, 20], valid_shape[actual_h, actual_w] # 可选 )参数说明tensor源张量view_shape视图的形状offset源张量的起始位置valid_shape实际有效尺寸用于边界处理可选view的典型用途是在大张量上切出一块参与计算的窗口例如滑窗卷积、分块累加创建视图零拷贝、只维护元信息因而成本极低。索引与切片Tensor 支持 Python 风格的索引和切片适用于灵活的数据访问。tensor pypto.tensor([10, 20], pypto.DT_FP16, tensor) # 单个元素创建视图 element tensor[0, 0] # 只支持INT32 # 切片创建视图 slice_tensor tensor[0:5, 10:20] # 省略号 ellipsis_slice tensor[..., 0:10]注意三点单个元素索引tensor[0, 0]返回的是视图而非拷贝其数据类型限定为 INT32切片与省略号同样创建视图后续修改视图会影响底层数据更精细的按下标取值/赋值能力可借助gather、index_select、scatter_update、index_put_、index_add_等算子见 python/pypto/op/indexing.py它们支持任意位置的下标收集与散布是实现数据重排的进阶手段。组装assemble函数在指定的偏移处将较小的张量放入较大的张量适用于分块处理后的结果合并。# 将一个小的张量组装成一个大的张量 pypto.assemble( small_tensor, # 源张量 [10, 20], # 目标位置 large_tensor # 目标张量 )示例# 小张量结果 tile_result pypto.tensor([32, 32], pypto.DT_FP16, tile) # 大输出张量 output pypto.tensor([100, 200], pypto.DT_FP16, output) # 在 [10, 20] 位置组装输出张量 pypto.assemble(tile_result, [10, 20], output)assemble与view是一对互逆操作view从大张量切出小视图assemble把小张量写回大张量的指定偏移python/pypto/operation.py 第 30 行起提供多组重载支持parallel并行写入开关Tensor 上也有对应的成员方法tensor.assemble(input, offsets)见 python/pypto/tensor.py 第 731 行附近。两者配合构成了典型的分块计算闭环切块计算 → 结果组装回大张量。仓库中 examples/01_beginner/transform/add_scalar_loop_view_assemble.py 与 python/tests/st/test_tensor_slice.py、python/tests/st/test_tensor_slice_assemble.py 均围绕该模式给出了可直接运行的示例与测试。重塑形状/维度# 重塑张量 reshaped pypto.reshape(tensor, [new_shape]) # 转置 transposed pypto.transpose(tensor, dim00, dim11)重塑不会改变数据只会改变张量的视图。源码层面的更多细节reshapepython/pypto/operation.py 第 183258 行除shape外还支持两个可选参数valid_shape指定新形状下的有效尺寸用于局部重塑或填充场景例如把[2, 2]重塑为[4, 1]并声明有效部分[2, 1]inplace为True时与输入共享存储零拷贝为False时创建新张量动态形状shape 含SymbolicScalar必须使用inplaceTrue。shape[-1]表示将张量展平为一维。transposepython/pypto/op/mutating.py 第 5892 行仅交换dim0与dim1两个维度若需任意维排列可使用同文件中的permute对应pypto_impl.Permute。相关能力还包括unsqueeze插入长度为 1 的新维度、squeeze去除长度为 1 的维度与clone深拷贝张量。类型转换# 转换为不同的数据类型 result pypto.cast(tensor, pypto.DT_FP32, modepypto.CastMode.CAST_NONE)castpython/pypto/op/mutating.py 第 151 行起的完整签名为def cast(input: Tensor, dtype: DataType, mode: CastMode CastMode.CAST_NONE, satmode: SaturationMode SaturationMode.OFF) - Tensordtype目标数据类型如pypto.DT_FP32、pypto.DT_FP16、pypto.DT_BF16、pypto.DT_INT32等mode舍入模式CastMode.CAST_NONE为默认直接截断/按硬件默认行为转换另有CAST_RINT就近舍入到偶数等模式satmode浮点转整数时的饱和处理开关默认OFF截断行为置为ON时对超出目标类型范围的数值做钳位saturation/clamping避免回绕。若目标类型与原类型一致cast不执行拷贝直接返回原对象。类型转换常与精度控制配合使用例如在 FP16/BF16 计算前先 cast 到 FP32 提升中间精度或在axpy_等混合精度算子中显式控制转换模式见 python/pypto/op/math.py 第 107163 行的axpy_实现其内部即用CastMode.CAST_RINT完成 FP32 结果回写 FP16 的舍入。小结本文完整梳理了 PyPTO Tensor 操作的两大支柱数学运算以add/sub/mul/div/pow为代表的向量算术、以exp/log/sqrt/sin/cos/abs/neg为代表的数学函数、sigmoid/relu/softmax激活、maximum/minimum/clip比较、sum/amax/amin归约以及move原地写入与广播模式矩阵侧则聚焦 Cube 核的matmul覆盖转置、批量、NZ 输出与extend_params偏置、反量化、TF32 转换等扩展能力。逻辑结构变换view/索引切片/assemble构成分块计算闭环reshape/transpose/permute改变视图与维度cast在类型间自由转换并支持舍入与饱和模式控制。这些操作的源码实现集中在 python/pypto/op/math.py、python/pypto/op/reduction.py、python/pypto/op/matmul.py、python/pypto/op/mutating.py 与 python/pypto/operation.py配套示例位于 examples/01_beginner/compute 与 examples/02_intermediate/operators。读者在编写算子时可以遵循基础算子组合 → 视图/组装组织数据 → cast 控制精度 → move 复用缓冲区的通用思路在保持代码可读性的同时兼顾 NPU 上的执行效率。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表