已合并
fix te ut bug #44321
伦创建于 27 天前
fix te ut bug #44321
已合并
共 2 个文件变更+18-1
| @@ -220,6 +220,20 @@ def npu_embedding(weight, indices, padding_idx=-1, scale_grad_by_freq=False, spa | |||
| 220 | overwrite_lowering(aten.embedding, npu_embedding, type_promotion_kind=None) | 220 | overwrite_lowering(aten.embedding, npu_embedding, type_promotion_kind=None) |
| 221 | 221 | ||
| 222 | 222 | ||
| 223 | +_upstream_round = lowerings.get(aten.round.default) | ||
| 224 | + | ||
| 225 | + | ||
| 226 | +def npu_round(x): | ||
| 227 | + if x.get_dtype() == torch.float64: | ||
| 228 | + return fallback_handler(aten.round.default)(x) | ||
| 229 | + if _upstream_round is not None: | ||
| 230 | + return _upstream_round(x) | ||
| 231 | + return fallback_handler(aten.round.default)(x) | ||
| 232 | + | ||
| 233 | + | ||
| 234 | +overwrite_lowering(aten.round.default, npu_round, type_promotion_kind=None) | ||
| 235 | + | ||
| 236 | + | ||
| 223 | def npu_cat(inputs, dim=0): | 237 | def npu_cat(inputs, dim=0): |
| 224 | return fallback_handler(aten.cat.default)(inputs, dim) | 238 | return fallback_handler(aten.cat.default)(inputs, dim) |
| 225 | 239 | ||
| @@ -54,6 +54,7 @@ GENERATE_LIST = [ | |||
| 54 | aten.tanh, | 54 | aten.tanh, |
| 55 | aten.copy, | 55 | aten.copy, |
| 56 | aten.copy_, | 56 | aten.copy_, |
| 57 | + aten.remainder, | ||
| 57 | 58 | ||
| 58 | # Non-pointwise ops | 59 | # Non-pointwise ops |
| 59 | aten.squeeze, | 60 | aten.squeeze, |
| @@ -73,6 +74,7 @@ GENERATE_LIST = [ | |||
| 73 | aten.amax, | 74 | aten.amax, |
| 74 | aten.cat, | 75 | aten.cat, |
| 75 | aten.slice_scatter, | 76 | aten.slice_scatter, |
| 77 | + aten.round, | ||
| 76 | aten.scalar_tensor, | 78 | aten.scalar_tensor, |
| 77 | aten.unbind, | 79 | aten.unbind, |
| 78 | aten.lift_fresh_copy, | 80 | aten.lift_fresh_copy, |
| @@ -100,7 +102,7 @@ GENERATE_LIST = [ | |||
| 100 | aten.convolution_backward, | 102 | aten.convolution_backward, |
| 101 | aten.bmm, | 103 | aten.bmm, |
| 102 | aten.addmm, | 104 | aten.addmm, |
| 103 | - aten.glu | 105 | + aten.glu, |
| 104 | ] | 106 | ] |
| 105 | # A5 (910_95) ONLY: the CANN indirect-mem extension ops behind these exist only | 107 | # A5 (910_95) ONLY: the CANN indirect-mem extension ops behind these exist only |
| 106 | # on A5 (fail to lower on A2/A3/910B). Mirrors torch_npu's INDIRECT_MEM_OVERRIDE_LIST. | 108 | # on A5 (fail to lower on A2/A3/910B). Mirrors torch_npu's INDIRECT_MEM_OVERRIDE_LIST. |
| @@ -128,6 +130,7 @@ KEEP_UPSTREAM_LOWERING = [control_deps] + [ | |||
| 128 | # self / fp32 values hits aclnnIndexPutImpl EZ1001 (self/values dtype mismatch). | 130 | # self / fp32 values hits aclnnIndexPutImpl EZ1001 (self/values dtype mismatch). |
| 129 | "index_put", | 131 | "index_put", |
| 130 | "index_put_", | 132 | "index_put_", |
| 133 | + "_local_scalar_dense", | ||
| 131 | ) | 134 | ) |
| 132 | if hasattr(aten, _name) | 135 | if hasattr(aten, _name) |
| 133 | ] | 136 | ] |