openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=44321&platform=gitcode&codeHostingPlatformFlag=gitcode

变更摘要

本次 PR 修复了 TE 后端 5 个单元测试缺陷，主要涉及 avg_pool2d_backward 的自定义 scatter kernel 实现、round 算子的 lowering 覆盖、min2 中非法参数移除、以及 in_out 类型参数在下播时的数据正确性问题，同时更新了 lowering 覆盖列表以纳入新增算子。

主要改动

• 新增 NPUAvgPoolBackwardKernel scatter kernel（ir.py）：针对 k <= stride、pad=0、无 ceil 模式的 avg_pool2d_backward 场景，新增基于 Pointwise 的自定义 kernel，在超缓冲区 [N, C, H_PAD, W_PAD] 上迭代池化空间，避免上游实现中整数除法索引导致的 DiscreteMemAccess 标量化问题；通过 get_size 与 get_pointwise_size 解耦输出域与迭代域，消除 range-tree tiling 对输出轴的模分片。

• 新增 npu_avg_pool2d_backward lowering（lowering.py）：在满足 k <= stride、无 padding、无 ceil、无 divisor_override 的条件下使用 scatter kernel；其他情况回退到上游 lowering。生成超缓冲区后在需要时通过 npu_slice 切片回真实输出尺寸。

• 新增 npu_round lowering（lowering.py）：对 float64 类型输入回退到 fallback_handler，其余类型走上游 round lowering，避免 TE 后端对 float64 的不支持导致 UT 失败。

• 修复 min2 函数（npu_triton_helpers.py）：移除 tl.min 调用中的 propagate_nan=True 参数，该参数在 Ascend 后端 Triton 中不被支持。

• 修复下播包装器对 in_out 参数的处理（npu_triton_heuristics.py）：在 _wrap_launcher_with_downcast 中，对 out_ptr（write-only）输出保持空缓冲区分配，对 in_out/mutated 参数改用 a.to(dst_dtype) 从原始数据初始化下播缓冲区，避免 read-modify-write 类 kernel（如 buf -= 1）因空缓冲区丢失输入数据。

代码审查

✅ 未发现问题

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.