CANN PyPTO pypto.sqrt 逐元素开方算子详解：精度模式、TileShape 设置与调用实战

CANN PyPTO pypto.sqrt 逐元素开方算子详解精度模式、TileShape 设置与调用实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto.sqrt是 CANN PyPTOParallel Tensor/Tile Operation 编程范式提供的逐元素平方根Square Root算子用于计算输入 Tensor 中每个元素的平方根并支持INTRINSIC指令模式与HIGH_PRECISION高精度模式两种精度控制覆盖 Ascend 950 系列、Atlas A3 系列与 Atlas A2 系列产品。本文以官方 API 文档 pypto-sqrt.md 为主线结合 Python 前端实现、C 算子实现 与 ST 测试用例完整讲解函数原型、参数语义、精度模式底层原理、TileShape 设置规则及可运行示例帮助你直接将该算子集成到自己的 PyPTO 内核代码中。产品支持情况pypto.sqrt在当前仓库中声明的产品支持矩阵如下产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持功能说明pypto.sqrt对输入 Tensor 逐元素计算平方根element-wise square root即对于输入中的每个元素x输出对应的√x。该运算为逐元素element-wise一元运算不涉及跨元素的数据聚合。当输入元素为负数时返回NaN因此在调用前需自行确保输入数据的非负性或对结果中的NaN做后续处理。从源码实现看该算子属于向量Vector一元运算族。C 侧实现位于 framework/src/interface/operation/vector/unary_basic.cpp#L55-L67通过CheckTensorFormat校验输入格式详见下方约束说明通过CheckTensorDataType校验数据类型为DT_FP16、DT_BF16、DT_FP32构造UnaryOpType::SQRT一元运算节点将用户传入的precisionType写入算子的OpAttributeKey::precisionType属性供后续代码生成阶段消费。函数原型pypto.sqrt(input, precision_typepypto.PrecisionType.INTRINSIC) - Tensor该函数在 Python 前端中的定义位于 python/pypto/op/math.py#L1693-L1729使用op_wrapper装饰器封装最终通过pypto_impl.Sqrt(input, precision_type)调用 C 底层实现。此外Tensor对象还提供了同名的成员方法 Tensor.sqrt可写作x.sqrt(...)形式。参数说明参数类型说明inputTensor源操作数。支持的类型为Tensor。Tensor 支持的数据类型为 DT_FP16、DT_BF16、DT_FP32。不支持空 Tensor支持的维度1-4 维Shape Size 不大于 2147483647即 INT32_MAX。precision_typePrecisionType可选平方根操作的精度模式。默认值为PrecisionType.INTRINSIC。INTRINSIC直接使用芯片指令进行计算速度更快。HIGH_PRECISION使用更高精度的计算方式减少精度损失。数据类型与形状约束数据类型仅支持DT_FP16、DT_BF16、DT_FP32三种浮点类型与 C 侧supportedTypes {DT_FP16, DT_BF16, DT_FP32}的校验逻辑一致见 unary_basic.cpp#L60-L61。形状支持 1 至 4 维 TensorShape Size各维度的乘积不得超过INT32_MAX2147483647。空 Tensor不支持空 Tensor调用前需确保输入非空。precision_type 精度模式PrecisionType枚举在 C 侧定义于 framework/include/tilefwk/tilefwk_op.h#L62-L65enum class PrecisionType : uint8_t { INTRINSIC 0, // 指令模式直接使用芯片指令 HIGH_PRECISION 1 // 高精度模式 };两种模式的核心区别INTRINSIC默认直接使用芯片内建指令intrinsic instruction完成平方根计算执行速度快、指令开销低适合对精度要求不苛刻、追求吞吐的场景。HIGH_PRECISION采用更高精度的计算路径以减少精度损失适合科学计算、数值敏感场景。从代码生成角度precisionType最终会映射为 TileOp 的模板参数。在 framework/src/codegen/npu/codegen_vector_unary.cpp#L523-L545 的AddUnaryPrecisionTypeParm中OP_SQRT对应的枚举名为SqrtAlgorithm当precisionType 1时生成pto::SqrtAlgorithm::HIGH_PRECISION否则生成pto::SqrtAlgorithm::DEFAULT并以模板参数形式注入生成的 Tile 算子调用。可见该参数并非 Python 层的摆设而是贯穿前端属性、中间表示属性与后端代码生成的完整链路。返回值说明返回 Tensor 类型其 Shape 与数据类型与输入 Tensor 一致每个元素为输入 Tensor 对应元素的平方根。即输入为[x0, x1, ..., xn]输出为[√x0, √x1, ..., √xn]。约束说明Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式即输入需为 ND 等非 NZ 布局。C 侧通过CheckTensorFormat(self.GetStorage(), {TileOpFormat::TILEOP_NZ}, Sqrt)完成格式校验见 unary_basic.cpp#L58。此外结合前述功能说明还需注意输入元素为负数时返回NaN调用侧应自行保证输入非负或处理NaN结果。调用示例前置TileShape 设置调用该 operation 接口前应通过pypto.set_vec_tile_shapes设置 TileShape。TileShape 的维度应与输出 Tensor 维度一致。TileShape 各维度的含义TileShape 的第 i 个值用于切分输出 Tensor 的第 i 个轴。例如输入 input shape 为[m, n]输出为[m, n]TileShape 设置为[m1, n1]则m1、n1分别用于切分m、n轴。set_vec_tile_shapes的实现位于 python/pypto/_controller.py#L46-L71其将各维度的 tile shape 写入当前 scope 的vec_tile_shapes配置中供向量计算vector computation切分使用同时支持传入SymbolicScalar符号标量便于在动态 shape 场景下使用。例如将 2 维输出切分为4 x 16的 Tilepypto.set_vec_tile_shapes(4, 16)接口调用示例默认指令模式x pypto.tensor([5], pypto.DT_FP32) y pypto.sqrt(x)结果示例如下输入数据x: [1.0, 4.0, 9.0, 16.0, 25.0] 输出数据y: [1.0, 2.0, 3.0, 4.0, 5.0]高精度模式示例对DT_FP16数据使用PrecisionType.HIGH_PRECISION以减少精度损失x pypto.tensor([5], pypto.DT_FP16) y pypto.sqrt(x, pypto.PrecisionType.HIGH_PRECISION)指令模式示例显式指定PrecisionType.INTRINSIC使用芯片指令以获取更快速度x pypto.tensor([5], pypto.DT_FP16) y pypto.sqrt(x, pypto.PrecisionType.INTRINSIC)在真实内核中的组合用法官方 ST 测试 python/tests/st/operation/vector/test_sqrt.py 展示了在 2D 循环 view 切片场景下的完整用法可作为将该算子集成进内核的参考模板pypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def sqrt_2d_1input_kernel(input0: pypto.Tensor(), output: pypto.Tensor(), config: SqrtConfig): 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.sqrt(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)该内核通过pypto.loop双层循环遍历分块对每个分块先用pypto.view切出输入视图再调用pypto.sqrt计算平方根最后用pypto.assemble写回输出。测试用例数据定义在 python/tests/st/operation/vector/vector_testcase/sqrt_test_case.py例如Sqrt_test_1使用(16384, 1)的 FP32 输入、(121, 144)的 tile_shape 与(2352, 2)的 view_shape并以torch.sqrt的结果作为期望输出进行数值比对见 test_sqrt.py#L43可用于验证算子正确性。总结pypto.sqrt是一个参数简单、链路完整的向量一元算子调用侧只需提供 Tensor 输入与可选的精度模式底层会自动完成格式/类型校验、算子节点构建、precisionType属性写入以及后端代码生成pto::SqrtAlgorithm::DEFAULT/HIGH_PRECISION模板参数注入。实际使用时需重点把握三点调用前通过pypto.set_vec_tile_shapes设置与输出维度一致的 TileShape输入数据保证非负否则结果为NaN输入格式避免使用TILEOP_NZ。对于精度敏感场景显式指定PrecisionType.HIGH_PRECISION对于吞吐优先场景保持默认的PrecisionType.INTRINSIC即可。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考