PyPTO pypto.asin 接口详解:逐元素反正弦运算的参数约束、TileShape 设置与源码实现 PyPTO pypto.asin 接口详解逐元素反正弦运算的参数约束、TileShape 设置与源码实现【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本文围绕 PyPTO 的pypto.asin接口展开完整覆盖其功能定义、参数与返回值约束、TileShape 设置方法及 UB 临时空间占用公式并结合仓库源码剖析反正弦运算在 NPU 向量单元上的实际实现路径类型提升、区间缩减与 Horner 多项式求值帮助读者在 Ascend 系列设备上正确调用该逐元素三角运算接口并理解其底层原理。功能说明与数学定义pypto.asin用于计算输入 Tensor 中每个元素的反正弦值属于逐元素element-wise运算计算公式为$$ y_i \arcsin(x_i) $$关键数值行为输出结果落在 $[-\pi/2, \pi/2]$ 区间内当输入元素的绝对值大于 1 时对应输出为 NaN。因此调用前应确保输入数据基本落在 $[-1, 1]$ 范围内例如归一化后的张量、注意力中的概率/相似度投影结果等场景。产品支持情况来自 接口文档产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持函数原型与参数说明函数原型如下asin(input: Tensor) - Tensor参数说明参数名输入/输出说明input输入源操作数。支持的类型为Tensor。Tensor 支持的数据类型为DT_FP32DT_FP16DT_BF16。不支持空 TensorShape 仅支持 1-4 维Shape Size 不大于 2147483647即 INT32_MAX。返回值为 Tensor 类型其 Shape 与输入 Tensor 一致数据类型与输入 Tensor 一致元素为输入 Tensor 对应元素的反正弦值。源码中的参数校验实现上述约束并非仅停留在文档层面。在 C 侧实现 framework/src/interface/operation/vector/trigonometric.cpp#L106-L121 中Asin函数逐项执行了校验Tensor Asin(const Tensor self) { DECLARE_TRACER(); CheckTensorFormat(self.GetStorage(), {TileOpFormat::TILEOP_NZ}, Asin); std::unordered_setDataType supportedTypes {DT_FP16, DT_FP32, DT_BF16}; CheckTensorDataType(self.GetStorage(), supportedTypes, ASIN); CheckTensorDimRange(self.GetStorage(), 1, NUM_VALUE_4, ASIN); CheckTensorShapeSize(self.GetStorage(), ASIN); auto castSelf Cast(self, DataType::DT_FP32); auto result CALL(UnaryOperationUnaryOpType::ASIN, *Program::GetInstance().GetCurrentFunction(), castSelf.GetStorage()); auto castResult Cast(result, self.GetDataType()); return castResult; }从源码结构看可以确认三个实现事实格式限制CheckTensorFormat显式拒绝TileOpFormat::TILEOP_NZ格式与文档中Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式的约束一一对应类型与形状校验允许的数据类型集合为{DT_FP16, DT_FP32, DT_BF16}维度范围通过CheckTensorDimRange(1, 4)限制为 1~4 维总元素数限制由CheckTensorShapeSize保证不超过 INT32_MAX计算精度路径无论输入是 FP16 还是 BF16都会先Cast到DT_FP32执行一元运算计算完成后再Cast回输入类型。这意味着低精度类型下反正弦的数值计算实际在 FP32 精度上完成最后一步转换回原类型。Python 侧的入口位于 python/pypto/op/math.py#L1110-L1141通过op_wrapper装饰器包装后转发到 C 绑定op_wrapper def asin(self: Tensor) - Tensor: Computes the element-wise arcsine of self. ... return pypto_impl.Asin(self)同时asin也以Tensor方法的形式暴露即可以调用x.asin()等价于pypto.asin(x)。约束说明TileShape 与 UB 空间占用在综合考虑输入、输出及临时存储空间占用的前提下TileShape 大小需满足额外约束条件。假设 TileShape 为[a, b, c, d]记 $d_{align} CeilAlign(d, 64)$则总的 UB 空间占用为$$ 6 \cdot a \cdot b \cdot c \cdot d_{align} \cdot sizeof(DT_FP32) \le UB $$其中系数 6 表示该运算在 UB 上需要 6 个以 FP32 计量的 tile 空间。对照 tile 级实现 framework/src/interface/tileop/vector/unary/inverse.h#L151-L194 中的TAsinAcosImpl可以推断这 6 个空间单元对应槽位用途src输入数据 tile已提升为 FP32dst输出数据 tiletmp0先存放 $x$后复用于大区间分支结果 / 取负结果tmp1大区间分支参数 $\sqrt{1-x^2}$后复用为 Horner 求值 /TSEL暂存tmp2区间缩减后的公共参数mask比较掩码 tile定义为uint8_t且宽度为4 * tileW恰好等价于一个 FP32 tile 的字节数末维对齐到 64$d_{align}$则来自向量 tile 硬件按 64 元素对齐访问的要求。因此在配置 TileShape 时末维应适当选取以减小对齐浪费当总占用超出 UB 容量时需减小 TileShape 或拆分循环。此外Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式。底层实现原理区间缩减 Horner 多项式pypto.asin在向量 tile 层面并非调用库函数而是基于源码 framework/src/interface/tileop/vector/unary/inverse.h 中的手工组合指令实现TAsin与TAcos共用TAsinAcosImpl模板体仅通过IsAsin布尔参数区分。核心数值算法分为四步1. 取绝对值tmp0 |x|先把奇函数问题转化为非负区间问题。2. 区间缩减到 $[0, 1/\sqrt{2}]$计算tmp1 sqrt(1 - x^2)用比较指令生成掩码|x| 1/sqrt(2)常量ASIN_THRESHOLD 0.70710678f通过TSEL选择小区间分支取 $t |x|$大区间分支取 $t \sqrt{1 - x^2}$利用了恒等式 $\arcsin(x) \pi/2 - \arcsin(\sqrt{1-x^2})$。3. 一次性求值 8 项 Taylor 多项式ArcsinPolyHorner// arcsin(t) t * (c0 c1*s c2*s^2 ... c7*s^7), s t^2 constexpr float ASIN_C0 1.0f; // 1 constexpr float ASIN_C1 0.16666667f; // 1/6 constexpr float ASIN_C2 0.075f; // 3/40 constexpr float ASIN_C3 0.04464286f; // 5/112 ...系数正是反正弦级数 $\arcsin(t) \sum_{k0}^{\infty} \frac{(2k)!}{4^k (k!)^2 (2k1)} t^{2k1}$ 的前 8 项用 Horner 形式$acc acc \cdot s c_k$以乘加指令序列高效求值最后乘以 $t$ 得到 $\arcsin(|x|)$大区间分支再计算 $\pi/2 - poly(t)$ 并回填。4. 符号恢复利用反正弦的奇函数性质src 0时取dst否则取-dst通过比较 TSEL完成。这套实现只依赖TABS、TMUL、TSQRT、TCMPS、TSEL等基础向量指令配合SyncV()同步点组织数据依赖也解释了为何该接口在 FP16/BF16 输入下仍能以 FP32 精度执行全部中间计算。调用示例TileShape 设置示例调用该 operation 接口前应通过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)注意末维 16 会按 64 对齐计入 UB 占用$d_{align} CeilAlign(16, 64) 64$选型时可结合前述 UB 公式自行核算空间是否充足。接口调用示例x pypto.tensor([3], pypto.DT_FP32) y pypto.asin(x)结果示例如下输入数据x: [-1.0000, 0.0000, 1.0000] 输出数据y: [-1.5708, 0.0000, 1.5708]真实工程化调用从 ST 测试看完整 kernel 写法仓库中的系统测试 python/tests/st/operation/vector/test_asin.py 展示了一个与文档示例同构但面向真实大张量场景的 2D kernel 写法先设置 TileShape再用双重循环逐 tile 处理通过pypto.view取出输入分块、调用pypto.asin、最后用pypto.assemble写回输出张量pypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def asin_2d_1input_kernel(input0: pypto.Tensor(), output: pypto.Tensor(), config: AsinConfig): 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.asin(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)该测试以torch.asin的 CPU 结果为 golden在 NPU 上执行后逐元素比对验证了接口结果与标准数学库行为一致同时测试开启pypto.options(pass_options{enable_slice: True})说明asin与编译器切片 pass 兼容。相关算子注册可在 framework/src/interface/operation/vector/trigonometric.cpp#L274-L275 中确认OP_ASIN通过AsinAcosOperationTileFuncUnaryOpType::ASIN统一注册到 tile 执行框架。小结pypto.asin是 1~4 维、Shape Size 不超过 INT32_MAX 的逐元素反正弦接口支持 DT_FP32 / DT_FP16 / DT_BF16不支持空 Tensor 与TILEOP_NZ格式调用前必须用pypto.set_vec_tile_shapes设置与输出一致的 TileShape并按 $6 \cdot a \cdot b \cdot c \cdot d_{align} \cdot 4 \le UB$ 的公式核算临时空间实现上先提升到 FP32再以绝对值 区间缩减 8 项 Horner Taylor 符号恢复的纯向量指令序列完成计算输入绝对值大于 1 时输出 NaN。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考