PyPTO 逐元素大于比较算子 pypto.gt函数原型、广播约束与 TileShape 切分实践【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto.gt是 CANN PyPTOParallel Tensor/Tile Operation 编程范式提供的逐元素大于比较算子用于判断输入 Tensor 各位置元素是否严格大于另一个 Tensor 或标量返回同 Shape 的 DT_BOOL 结果。本文以仓库文档 docs/zh/api/tensor_api/operation/pypto-gt.md 为核心结合源码实现与测试用例完整讲解pypto.gt的函数原型、参数与返回语义、各型号产品的数据类型约束、广播与 TileShape 设置要求并给出可直接复用的调用示例与实战验证方法。一、功能定位与适用场景pypto.gt实现的是逐元素大于比较element-wise greater-than运算即对输入 Tensor 中每个位置逐一执行input[i] other[i]的比较比较结果以布尔值True/False逐位写出到输出 Tensor。它是 PyPTO 张量运算体系中比较类算子的核心成员与pypto.ge大于等于、pypto.eq等于、pypto.ne不等于、pypto.lt小于、pypto.le小于等于共同构成完整的六元比较算子族。该算子常见的应用场景包括掩码mask生成如激活函数、归一化、注意力等算子中依据数据是否超过阈值生成布尔掩码数值截断与条件选择与pypto.where等算子配合实现基于比较结果的条件分支数据流统计与过滤统计满足某一阈值的元素个数、位置等。从实现上看pypto.gt是pypto.greater的别名二者在 python/pypto/op/comparison.py 中为完全等价的实现内部统一调用底层 C 算子npu::tile_fwk::Compare比较模式取OpType::GT、输出类型取OutType::BOOL。二、产品支持情况pypto.gt在以下产品系列上均得到支持依据 pypto-gt.mdAscend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持。不同产品系列对input/other所允许的数据类型存在差异详见下文约束说明一节的逐型号数据类型表。三、函数原型gt(input: Tensor, other: Union[Tensor, float, Element]) - Tensor其中input第一个源操作数类型为 Tensorother第二个源操作数支持 Tensor、float 或 Element 三种类型返回值为 Tensor 类型DT_BOOL。与原型完全对应的实现位于 python/pypto/op/comparison.pyop_wrapper def gt(input: Tensor, other: Union[Tensor, float, Element]) - Tensor: if isinstance(other, float): # Tensor vs Scalar comparison return pypto_impl.Compare(input, pypto_impl.Element(input.dtype, other), OpType.GT, OutType.BOOL) return pypto_impl.Compare(input, other, OpType.GT, OutType.BOOL)从这里可以读出两条关键实现细节float 自动转 Element当other传入 Python 的float时实现会以input的 dtype 构造pypto_impl.Element(input.dtype, other)即该浮点标量默认按DT_FP32语义参与比较如需使用其他数据类型如 DT_FP16、DT_INT16 等则应显式通过Element构建见 python/pypto/_element.pyElement(dtype, data)支持 int/float 标量。底层统一走 Compare 算子不论other是 Tensor 还是标量最终都落入 C 层的Compare。对应声明见 framework/include/tilefwk/tilefwk_op.hpybind 绑定见 python/src/bindings/operation.cpp其中比较模式枚举OpType包含EQ / NE / LT / LE / GT / GE输出类型枚举OutType包含BOOL / BIT见 tilefwk_op.hpypto.gt固定使用GT BOOL组合。四、参数说明参数名输入/输出说明input输入源操作数。支持的类型为 Tensor。不同型号支持的 Tensor 数据类型有所差异详细请参见约束说明。不支持空 TensorShape 仅支持 1-4 维Shape Size 不大于 2147483647即 INT32_MAX。other输入源操作数。支持的类型为 Tensor、float、Element。当为 float 类型时会自动转换为 Element 类型float 对应 DT_FP32当需要使用其他数据类型时可以通过 Element 构建。不同型号支持的 Tensor 和 Element 的数据类型有所差异详细请参见约束说明。不支持空 TensorShape 仅支持 1-4 维Shape Size 不大于 2147483647即 INT32_MAX。补充说明input与other的数据类型须保持一致见下文约束第 1 条other为 Tensor 时其 Shape 支持与input相同或可通过广播对齐见下文约束第 2 条维度上限为 4 维且元素总数不能超过 INT32_MAX这是由 Tile 框架内部索引与缓冲区寻址的 32 位粒度决定的在构造大规模输入时需要提前核算。五、返回值说明返回一个Shape 与输入 Tensor 一致、数据类型为 DT_BOOL 的 Tensor若input对应位置的元素值严格大于other对应位置的元素值则返回 True其余位置返回 False。注意严格大于的语义相等或小于均不满足条件这正是gt与gegreater-or-equal的关键区别。比较输出统一为布尔张量便于后续与逻辑算子、选择算子衔接。六、约束说明类型一致性input和other类型须保持一致other为 float 时按其自动转换为的 Element 语义理解。广播支持支持多维度广播到相同形状即other的 Shape 可以比input更小按 PyPTO 的广播规则扩展后再逐元素比较。Tensor 和 Element 数据类型按产品型号区分产品系列支持的数据类型Ascend 950PR / Ascend 950DTDT_FP16、DT_FP32、DT_INT16、DT_INT64、DT_UINT64Atlas A3 训练系列产品 / Atlas A3 推理系列产品DT_FP16、DT_FP32Atlas A2 训练系列产品 / Atlas A2 推理系列产品DT_FP16、DT_FP32可以看到Ascend 950 系列额外支持整型比较DT_INT16 / DT_INT64 / DT_UINT64而 A2 / A3 系列当前仅支持浮点类型 DT_FP16 / DT_FP32。据此若在 A2/A3 产品上对整型数据执行pypto.gt需要先进行类型转换。 4.格式限制Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式即参与比较的 Tensor 需使用 PyPTO 默认的矢量数据排布格式。七、调用示例7.1 TileShape 设置示例PyPTO 的矢量vector类算子在执行前需要通过pypto.set_vec_tile_shapes设置 TileShape分块形状该接口的定义见 docs/zh/api/tensor_api/config/pypto-set_vec_tile_shapes.md原型为set_vec_tile_shapes(*args: int) - None最多可传 4 个维度且每个维度必须大于 0。调用pypto.gt前TileShape 维度应和输出一致并按输出张量的各轴进行切分。示例 1非广播场景输入input的 Shape 为[m, n]other的 Shape 为[m, n]输出 Shape 为[m, n]。此时 TileShape 设置为[m1, n1]m1、n1分别用于切分 m 轴、n 轴pypto.set_vec_tile_shapes(4, 16)示例 2广播场景输入input的 Shape 为[m, n]other的 Shape 为[m, 1]输出 Shape 为[m, n]。TileShape 仍设置为[m1, n1]m1、n1分别用于切分 m 轴、n 轴pypto.set_vec_tile_shapes(4, 16)即广播发生在比较语义层面[m, 1]沿 n 轴扩展为[m, n]TileShape 依然按最终输出 Shape 的两个维度来设置。7.2 接口调用示例最简单的 Tensor-Tensor 比较a pypto.tensor([3], pypto.DT_FP32) b pypto.tensor([3], pypto.DT_FP32) out pypto.gt(a, b)结果示例如下输入数据a: [1.0 2.0 3.0] 输入数据b: [2.0 2.0 2.0] 输出数据out: [False, False, True]即仅当a的元素严格大于b的对应元素3.0 2.0时为 True1.0 2.0、2.0 2.0均为 False。标量比较other传 float等价写法a pypto.tensor([3], pypto.DT_FP32) out pypto.gt(a, 2.0) # 等价于 pypto.gt(a, pypto.Element(pypto.DT_FP32, 2.0))7.3 完整 Kernel 级实战分块循环 TileShape 结果回写仅调用pypto.gt会生成一个比较算子但在 PyPTO 中通常需要配合pypto.function、pypto.loop、pypto.view、pypto.assemble构成完整的 Kernel。仓库测试 python/tests/st/operation/vector/test_greater.py 提供了可复用的完整范式下面以test_vector_operation_greater为例说明关键环节import numpy as np import torch import pypto shape (32, 32) view_shape (16, 16) tile_shape (8, 8) a pypto.tensor(shape, pypto.DT_FP32, Greater_TENSOR_a) b pypto.tensor(shape, pypto.DT_FP32, Greater_TENSOR_b) c pypto.tensor(shape, pypto.DT_BOOL, Greater_TENSOR_c) with pypto.function(Greater, a, b, c): for b_idx in pypto.loop(2, nameLOOP_GREATER_L0, idx_nameb_idx): for s_idx in pypto.loop(2, nameLOOP_GREATER_L1, idx_names_idx): tile_a pypto.view(a, view_shape, [b_idx * view_shape[0], s_idx * view_shape[1]]) tile_b pypto.view(b, view_shape, [b_idx * view_shape[0], s_idx * view_shape[1]]) pypto.set_vec_tile_shapes(tile_shape[0], tile_shape[1]) tile_a.move(pypto.greater(tile_a, tile_b)) # gt 与 greater 等价 pypto.assemble(tile_a, [b_idx * view_shape[0], s_idx * view_shape[1]], c)该示例的核心链路是用pypto.view从全局张量上切出view_shape大小的分块tile调用pypto.set_vec_tile_shapes(tile_shape[0], tile_shape[1])设置矢量算子的内部 TileShape在分块上执行pypto.greater(tile_a, tile_b)与pypto.gt完全等价结果move回分块用pypto.assemble将比较结果写回全局输出张量c。同文件的test_greater_scalartest_greater.py则演示了other传标量50.0的写法即view_block.move(pypto.greater(view_block, scalar_value))。验证方式测试中使用torch.npu.set_device(device_id)指定设备运行后与torch.greater(a_tensor, b_tensor)标量场景为torch.greater(input_data, scalar_value)的结果做assert_allclose对比用于校验设备侧比较结果与参考实现一致。八、与比较算子族的横向对比pypto.gt位于 PyPTO 比较算子族python/pypto/op/comparison.py中与同族算子的语义差异如下均为逐元素、输出 DT_BOOL算子语义示例a[1,2,3], b[2,2,2]pypto.gt(a, b)严格大于[False, False, True]pypto.ge(a, b)大于等于[False, True, True]pypto.eq(a, b)等于[False, True, False]pypto.ne(a, b)不等于[True, False, True]pypto.lt(a, b)严格小于[True, False, False]pypto.le(a, b)小于等于[True, True, False]从源码看ge/eq/ne/lt/le与gt的实现结构完全一致仅在OpType枚举取值上不同见 comparison.py因此本文关于广播、TileShape、数据类型约束的结论对整族算子同样适用。选择哪个算子只需根据比较语义严格 / 非严格确定。九、常见问题与使用建议数据类型不匹配input与other数据类型须一致标量传入 float 时自动按 DT_FP32 处理。若要与其他类型比较应显式Element(dtype, value)构造避免隐式类型歧义。型号不支持整型比较A2 / A3 系列仅支持 DT_FP16 / DT_FP32若需要对整型数据比较可先通过pypto.cast转换类型后再调用pypto.gt详见 docs/zh/api/tensor_api/operation/pypto-cast.md。忘记设置 TileShape矢量算子执行前须调用pypto.set_vec_tile_shapes且 TileShape 各维度需大于 0、维度数与输出一致否则算子无法正确切分数据。NZ 格式限制参与比较的 Tensor 不要使用TileOpFormat.TILEOP_NZ格式排布。Shape 上限输入 Shape 仅支持 1-4 维且 Shape Size 不大于 INT32_MAX超大张量需先拆分处理。十、总结pypto.gt是 PyPTO 中实现逐元素严格大于比较的标准算子功能定义在 docs/zh/api/tensor_api/operation/pypto-gt.mdPython 层入口为 python/pypto/op/comparison.py底层由 CCompare(GT, BOOL)算子承载framework/include/tilefwk/tilefwk_op.h、python/src/bindings/operation.cpp并有完整的 Kernel 级测试用例 python/tests/st/operation/vector/test_greater.py 佐证。掌握其严格大于、输出 BOOL、支持广播、需设 TileShape四大要点即可在算子开发中正确、高效地使用比较运算并可类推至同族六个比较算子。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考