已合并
fix(op-plugin): identity round for integers; delegate round_ to round.out for op_api. #4917
Margaret_wangrui创建于 5月9日
fix(op-plugin): identity round for integers; delegate round_ to round.out for op_api. #4917
已合并
共 4 个文件变更+49-1
| @@ -4316,6 +4316,7 @@ official: | |||
| 4316 | acl_op: all_version | 4316 | acl_op: all_version |
| 4317 | op_api: all_version | 4317 | op_api: all_version |
| 4318 | gen_opapi: | 4318 | gen_opapi: |
| 4319 | + integral_identity_tensor: self | ||
| 4319 | out: | 4320 | out: |
| 4320 | size: self | 4321 | size: self |
| 4321 | dtype: self | 4322 | dtype: self |
| @@ -4326,7 +4327,7 @@ official: | |||
| 4326 | acl_op: all_version | 4327 | acl_op: all_version |
| 4327 | op_api: all_version | 4328 | op_api: all_version |
| 4328 | gen_opapi: | 4329 | gen_opapi: |
| 4329 | - exec: aclnnInplaceRound | 4330 | + structured_inherit: round.out |
| 4330 | 4331 | ||
| 4331 | - func: round_.decimals(Tensor(a!) self, *, int decimals) -> Tensor(a!) | 4332 | - func: round_.decimals(Tensor(a!) self, *, int decimals) -> Tensor(a!) |
| 4332 | acl_op: all_version | 4333 | acl_op: all_version |
| @@ -20,6 +20,12 @@ namespace acl_op { | |||
| 20 | using npu_preparation = at_npu::native::OpPreparation; | 20 | using npu_preparation = at_npu::native::OpPreparation; |
| 21 | using npu_utils = at_npu::native::NpuUtils; | 21 | using npu_utils = at_npu::native::NpuUtils; |
| 22 | 22 | ||
| 23 | +bool round_integral_identity(const at::Tensor& self) { | ||
| 24 | + // Match PyTorch: round on integer tensors is identity. CANN aclnn round ops | ||
| 25 | + // do not support int8 / int16 / uint8; use the same path for all integral dtypes. | ||
| 26 | + return at::isIntegralType(self.scalar_type(), /*includeBool=*/false); | ||
| 27 | +} | ||
| 28 | + | ||
| 23 | namespace { | 29 | namespace { |
| 24 | 30 | ||
| 25 | at::Tensor& round_out_npu_nocheck(at::Tensor& result, const at::Tensor& self) | 31 | at::Tensor& round_out_npu_nocheck(at::Tensor& result, const at::Tensor& self) |
| @@ -38,6 +44,17 @@ at::Tensor& round_out(const at::Tensor& self, at::Tensor& out) | |||
| 38 | { | 44 | { |
| 39 | npu_preparation::CheckOut({self}, out, self); | 45 | npu_preparation::CheckOut({self}, out, self); |
| 40 | 46 | ||
| 47 | + if (round_integral_identity(self)) { | ||
| 48 | + if (!npu_utils::check_match(&out)) { | ||
| 49 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); | ||
| 50 | + contiguous_result.copy_(self); | ||
| 51 | + npu_utils::format_fresh_view(out, contiguous_result); | ||
| 52 | + } else { | ||
| 53 | + out.copy_(self); | ||
| 54 | + } | ||
| 55 | + return out; | ||
| 56 | + } | ||
| 57 | + | ||
| 41 | if (!npu_utils::check_match(&out)) { | 58 | if (!npu_utils::check_match(&out)) { |
| 42 | at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 59 | at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 43 | round_out_npu_nocheck(contiguous_result, self); | 60 | round_out_npu_nocheck(contiguous_result, self); |
| @@ -50,6 +67,9 @@ at::Tensor& round_out(const at::Tensor& self, at::Tensor& out) | |||
| 50 | 67 | ||
| 51 | at::Tensor round(const at::Tensor& self) | 68 | at::Tensor round(const at::Tensor& self) |
| 52 | { | 69 | { |
| 70 | + if (round_integral_identity(self)) { | ||
| 71 | + return self.clone(); | ||
| 72 | + } | ||
| 53 | at::Tensor result = npu_preparation::apply_tensor(self); | 73 | at::Tensor result = npu_preparation::apply_tensor(self); |
| 54 | round_out_npu_nocheck(result, self); | 74 | round_out_npu_nocheck(result, self); |
| 55 | 75 | ||
| @@ -92,6 +92,26 @@ class TestRound(TestCase): | |||
| 92 | self.assertRtolEqual(cpu_output, npu_output1) | 92 | self.assertRtolEqual(cpu_output, npu_output1) |
| 93 | self.assertRtolEqual(cpu_output, npu_output2) | 93 | self.assertRtolEqual(cpu_output, npu_output2) |
| 94 | 94 | ||
| 95 | + def test_round_integer_identity_npu(self): | ||
| 96 | + """Integer round/round_ is identity on NPU (int8/int16/uint8 unsupported by aclnn round; see op_plugin yaml).""" | ||
| 97 | + dtypes = [ | ||
| 98 | + torch.int8, | ||
| 99 | + torch.uint8, | ||
| 100 | + torch.int16, | ||
| 101 | + torch.int32, | ||
| 102 | + torch.int64, | ||
| 103 | + ] | ||
| 104 | + for dt in dtypes: | ||
| 105 | + cpu_x = torch.tensor([[1, -2, 7], [-3, 0, 42]], dtype=dt) | ||
| 106 | + npu_x = cpu_x.npu() | ||
| 107 | + | ||
| 108 | + self.assertEqual(torch.round(cpu_x), cpu_x) | ||
| 109 | + self.assertEqual(torch.round(npu_x).cpu(), cpu_x) | ||
| 110 | + | ||
| 111 | + npu_inplace = cpu_x.clone().npu() | ||
| 112 | + npu_inplace.round_() | ||
| 113 | + self.assertEqual(npu_inplace.cpu(), cpu_x) | ||
| 114 | + | ||
| 95 | 115 | ||
| 96 | if __name__ == "__main__": | 116 | if __name__ == "__main__": |
| 97 | run_tests() | 117 | run_tests() |
| @@ -206,6 +206,13 @@ def compute_op_api_definition(struct: StructInfo, env_aclnn_extension_switch: bo | |||
| 206 | tensor=struct.integral_identity_tensor, func_name=name, args_exprs_str=args_exprs_str | 206 | tensor=struct.integral_identity_tensor, func_name=name, args_exprs_str=args_exprs_str |
| 207 | ) | 207 | ) |
| 208 | 208 | ||
| 209 | + integral_identity_guard = "" | ||
| 210 | + if struct.integral_identity_tensor is not None: | ||
| 211 | + integral_identity_guard = INTEGRAL_IDENTITY_GUARD.substitute( | ||
| 212 | + tensor=struct.integral_identity_tensor, | ||
| 213 | + func_name=name, | ||
| 214 | + args_exprs_str=args_exprs_str) | ||
| 215 | + | ||
| 209 | tensor_arguments = ", ".join(filt_input_tensor(f.func.arguments.flat_non_out)) | 216 | tensor_arguments = ", ".join(filt_input_tensor(f.func.arguments.flat_non_out)) |
| 210 | 217 | ||
| 211 | new_params_def = "".join( | 218 | new_params_def = "".join( |