PyPTO 张量向下取整除法 floor_div 算子实战:API 用法、TileShape 设置与向量内核实现原理

PyPTO 张量向下取整除法 floor_div 算子实战:API 用法、TileShape 设置与向量内核实现原理 PyPTO 张量向下取整除法 floor_div 算子实战API 用法、TileShape 设置与向量内核实现原理【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptopypto.floor_div是 CANN PyPTOParallel Tensor/Tile Operation 编程范式张量 API 中用于逐元素向下取整除法floor division的核心算子它对input的每个元素除以other对应位置的元素后向下取整天然契合归一化、量化、索引映射等以整数除法为核心的算子开发场景。读完本文你将掌握pypto.floor_div的函数原型、参数与数据类型约束、多维广播规则、set_vec_tile_shapes切分配置方法并能从 Python 前端封装 与 向量内核实现 两个层面理解其底层运行机制快速上手编写可运行的下采样/取整算子。产品支持情况根据 pypto-floor_div.md 的官方说明pypto.floor_div在以下昇腾硬件产品上得到支持Ascend 950PR / Ascend 950DT支持。Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持。Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持。需要说明的是不同型号产品在Tensor支持的数据类型上存在差异具体差异见下文约束说明小节。编写算子时建议结合目标产品型号如通过pypto.platform或运行时设备信息确认选择合适的输入数据类型。功能说明floor_div将input的每个元素除以other中对应位置的元素并向下取整向负无穷方向取整数学表达式为$$ res_i floor(\frac{input_{i}}{other_{i}}) $$与常规的截断除法truncated division向零取整不同向下取整除法在遇到负商时结果会向负无穷方向靠拢。例如-7 / 2的截断除法结果是-3而向下取整除法的结果是-4。在 PyTorch 中对应的语义是torch.floor_dividePyPTO 的 ST 测试正是以torch.floor_divide作为参照基准进行数值对齐的见 test_floordiv.py。函数原型floor_div(input: Tensor, other: Union[Tensor, int, float]) - Tensor从 Python 前端封装 可以看到floor_div是op_wrapper装饰的模块级函数other既可以是同维Tensor走pypto_impl.FloorDiv(input, other)二元路径也可以是int/float标量先做类型检查与裁剪再包装为pypto_impl.Element走二元-标量路径op_wrapper def floor_div(input: Tensor, other: Union[Tensor, int, float]) - Tensor: if isinstance(other, pypto_impl.Tensor): return pypto_impl.FloorDiv(input, other) else: _check_scalar_type(floor_div, input.dtype, other) other _clip_scalar_to_dtype(input.dtype, other) return pypto_impl.FloorDiv(input, pypto_impl.Element(input.dtype, other))在 C/pybind11 绑定层operation.cpp中FloorDiv针对Tensor × Tensor与Tensor × Element分别注册了两个重载分别对应上面两条调用路径。标量other的类型检查与裁剪other传标量时前端会执行两道预处理math.py_check_scalar_type若other是float而input.dtype不是DT_FP32/DT_FP16/DT_BF16等浮点类型会抛出PyptoError错误码0xF00001提示改用整型标量或将张量转换为浮点类型。_clip_scalar_to_dtype若other是int且张量为整型DT_INT8/DT_INT16/DT_INT32/DT_INT64/DT_UINT8/DT_UINT16等会把标量按张量 dtype 对应的 numpy 类型做一次范围裁剪防止标量溢出目标位宽。参数说明参数名输入/输出说明input输入源操作数被除数。支持的类型为Tensor。Tensor支持的数据类型不同型号有所差异详细请参见约束说明。不支持空 TensorShape 仅支持 1-4 维支持多维度广播到相同形状Shape Size 不大于 2147483647即 INT32_MAX。other输入源操作数除数。支持的类型为Tensor、int、float。Tensor支持的数据类型不同型号有所差异详细请参见约束说明。不支持空 TensorShape 仅支持 1-4 维支持多维度广播到相同形状Shape Size 不大于 2147483647即 INT32_MAX。返回值说明返回输出Tensor。输出Tensor的数据类型与input、other相同Shape 为input和other广播后的形状。当other为标量时输出 Shape 与input一致。约束说明使用pypto.floor_div前需要满足以下约束input与other的数据类型应相同。支持多维度广播broadcast广播规则与常见深度学习框架一致维度从尾部对齐长度为 1 的维度会向另一个操作数的对应维度广播扩展。Tensor数据类型说明Ascend 950PR / Ascend 950DTDT_FP16、DT_BF16、DT_FP32、DT_INT32、DT_INT8、DT_UINT8、DT_INT64。Atlas A3 训练系列产品 / Atlas A3 推理系列产品DT_FP16、DT_BF16、DT_FP32、DT_INT32、DT_INT8、DT_UINT8。Atlas A2 训练系列产品 / Atlas A2 推理系列产品DT_FP16、DT_BF16、DT_FP32、DT_INT32、DT_INT8、DT_UINT8。可见 Ascend 950 系列额外支持DT_INT64A2/A3 系列则不支持该类型。Tensor类型输入不支持TileOpFormat.TILEOP_NZ格式请使用ND等布局ST 测试用例中均以format: ND构造输入输出见 floordiv_test_case.py。调用示例TileShape 设置示例调用该 operation 接口前应通过set_vec_tile_shapes设置向量切分形状TileShape。TileShape 的维度应与输出一致非广播场景输入inputshape 为[m, n]other为[m, n]输出为[m, n]TileShape 设置为[m1, n1]则m1、n1分别用于切分m、n轴。广播场景输入inputshape 为[m, n]other为[m, 1]输出为[m, n]TileShape 设置为[m1, n1]则m1、n1分别用于切分m、n轴。pypto.set_vec_tile_shapes(4, 16)在 ST 测试中切分配置的典型取值为view_shape(128, 512)、tile_shape(64, 64)见 floordiv_test_case.py即在每次循环中通过pypto.view切出执行视图后再以(64, 64)的 Tile 粒度逐块计算floor_div并pypto.assemble回输出张量见 test_floordiv.py。接口调用示例最简单的 Tensor-Tensor 调用方式如下a pypto.tensor([1, 3], pypto.DT_INT32) b pypto.tensor([1, 3], pypto.DT_INT32) out pypto.floor_div(a, b)结果示例如下输入数据a: [[2 4 6]] 输入数据b: [[4 2 5]] 输出数据out: [[0 2 1]]即2//40、4//22、6//51。other也支持标量形式例如a pypto.tensor([2, 3], pypto.DT_INT32) out pypto.floor_div(a, 3)该路径在前端会先做标量类型检查与位宽裁剪再包装为Element后调用底层FloorDiv。源码实现从 pybind 绑定到向量内核内核入口与分派pypto.floor_div在 C 侧最终落到 tileop 层的内核模板TFloorDiv定义于 framework/src/interface/tileop/vector/binary/floor_div.h。内核按照输出 shape 的前三维dstShape0~dstShape2进行循环在每一轮中通过MakeElementwiseOperandExecTile为src0/src1构造按元素对齐的执行 Tile自动处理广播时的地址映射依据目标 dtype 分派到不同的计算函数half/bfloat16_t→FloorDivFp32TmpCompute先提升到 FP32 计算避免半精度除法精度不足float→FloorDivFloatComputeFP32 直接TDIVCAST_FLOOR整型int8/uint8/int32/int64→ 根据是否__DAV_V220宏选择 V220 或非 V220 的整型专用实现。浮点路径FP16/BF16 先升 FP32 再取整以FloorDivFp32TmpCompute为例floor_div.h半精度数据会被TCVT无损提升为 FP32 Tile执行TDIVDivAlgorithm::HIGH_PRECISION高精度除法后用CAST_FLOOR向下取整最后再CAST_RINT转回原半精度类型。FP32 路径FloorDivFloatCompute则直接对dst执行TDIVCAST_FLOOR无需临时空间。整型路径余数修正的精确取整算法整型除法不能直接借用浮点除法结果大数在 FP32 中会丢失精度。以非 V220 的FloorDivNonV220Int32Compute为例floor_div.h其核心是通用的IntFloorDiv模板floor_div.h算法要点如下除数为 0 的兜底先用TCMPS判断被除数符号再通过TSELS将被除数替换为预设的哨兵常量。整型路径分别定义了INT32_NEGATIVE_DIVIDEND_ZERO_DIVISOR_RESULT 0x7FFF7F7F与INT32_NONNEGATIVE_DIVIDEND_ZERO_DIVISOR_RESULT 0x8000808064 位对应0x7FFFFFFFFFFFFFFF与0x8000000000000000见 floor_div.h保证除零时不产生硬件异常。符号与余数判断计算sign_differ (src0 0) ! (src1 0)随后TDIV求商、TREM求余。向下取整修正当sign_differ rem ! 0时结果再减 1即向负无穷方向修正否则保持商不变最终TSEL选出结果。V220 平台上整型路径FloorDivV220Int32Computefloor_div.h采用“FP32 近似商 精确余数逐步修正”的多阶段流水先由浮点除法得到近似商q再用q反推余数r以floor(r/x2)修正q最后通过余数的符号归一化与±1校正保证余数满足0 r2 * sign(x2) |x2|的 floor-div 语义。int8/uint8 同理有Int8Compute/Uint8Compute的专用实现。标量变体binary_scalar 路径当other为标量时内核走 framework/src/interface/tileop/vector/binary_scalar/floor_div.h 中的FloorDivS*系列函数浮点路径使用TDIVStensor-scalar 除法指令配合CAST_FLOOR整型路径复用binary/floor_div.h中的哨兵常量与余数修正思路将TMUL/TADD替换为TMULS/TADDS等 tensor-scalar 变体其中还特意使用TSEL而非TSELS以避免 A2/A3 上 tensor-scalar 选择路径首 lane 不稳定的问题见 binary_scalar/floor_div.h 对应的注释说明。测试验证如何确认实现正确性仓库中floor_div的 ST 测试覆盖了二元与标量两种输入形态Tensor-Tensortest_floordiv.py 中floordiv_2d_2input_kernel通过双层pypto.loop分块遍历 2D 输入逐块执行pypto.floor_div(input0_view, input1_view)并assemble回输出再与torch.floor_divide的结果逐元素比对。Tensor-Scalartest_floordivs.py 使用pypto.floor_div(input0_view, config.scalar)形式参照同样是torch.floor_divide(inputs_cpu[0], config.scalar)。1D 向量切片场景test_vector_operation_1d.py 中test_vector_operation_1d_floor_div演示了在动态循环 切片模式下使用tile.move(pypto.floor_div(tile_a, tile_b))的写法并同样以torch.floor_divide校验。测试用例配置floordiv_test_case.py默认使用(1024, 128)的 int32 ND 输入、view_shape(128, 512)、tile_shape(64, 64)可作为自定义算子的切分与校验参考基线。常见使用注意事项先设 TileShape 再调用floor_div属于向量Vector运算调用前必须通过pypto.set_vec_tile_shapes(...)设置切分形状且 TileShape 维度必须与输出维度一致否则分块计算无法正确映射。广播维度的 Tile 切分广播场景如other为[m, 1]下TileShape 依然按输出形状[m, n]设置m1/n1分别切分m/n轴底层MakeElementwiseOperandExecTile会自动完成广播维的地址映射。数据类型一致性input与otherTensor 形态必须同 dtype标量形态时浮点标量不能作用于整型张量整型标量会按张量 dtype 做位宽裁剪。布局限制输入不支持TileOpFormat.TILEOP_NZ格式应使用ND布局。向下取整语义对负数结果而言floor_div与截断除法不同如-7 // 2 -4设计取整类算子时需确认业务期望的舍入方向。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考