企业官网建设流程全解析

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

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

导读

pypto.sqrt是 CANN PyPTO（Parallel 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在当前仓库中声明的产品支持矩阵如下：

功能说明

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_type=pypto.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(...)形式。

参数说明

数据类型与形状约束

• 数据类型：仅支持DT_FP16、DT_BF16、DT_FP32三种浮点类型，与 C++ 侧supportedTypes = {DT_FP16, DT_BF16, DT_FP32}的校验逻辑一致（见 unary_basic.cpp#L60-L61）。
• 形状：支持 1 至 4 维 Tensor；Shape Size（各维度的乘积）不得超过INT32_MAX（2147483647）。
• 空 Tensor：不支持空 Tensor，调用前需确保输入非空。

precision_type 精度模式

PrecisionType枚举在 C++ 侧定义于 framework/include/tilefwk/tilefwk_op.h#L62-L65：

enum 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]。

约束说明

1. 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的 Tile：

pypto.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-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto