已合并
fix: support quantized NPU flip dispatch #36070
hz893创建于 5月19日
fix: support quantized NPU flip dispatch #36070
已合并
共 4 个文件变更+64-0
| @@ -482,6 +482,18 @@ class TestShapeOps(TestCase): | |||
| 482 | self.assertRaises(IndexError, lambda: data.flip(0, 1, 2, 3)) | 482 | self.assertRaises(IndexError, lambda: data.flip(0, 1, 2, 3)) |
| 483 | self.assertRaises(IndexError, lambda: data.flip(3)) | 483 | self.assertRaises(IndexError, lambda: data.flip(3)) |
| 484 | 484 | ||
| 485 | + def test_flip_per_channel_quantized_error(self, device): | ||
| 486 | + data = torch.randn(2, 3, device=device) | ||
| 487 | + scales = torch.ones(3, dtype=torch.float, device=device) * 0.1 | ||
| 488 | + zero_points = torch.zeros(3, dtype=torch.long, device=device) | ||
| 489 | + qdata = torch.quantize_per_channel( | ||
| 490 | + data, scales, zero_points, axis=1, dtype=torch.qint8 | ||
| 491 | + ) | ||
| 492 | + | ||
| 493 | + error_msg = "Setting strides is possible only on uniformly quantized tensor" | ||
| 494 | + with self.assertRaisesRegex(RuntimeError, error_msg): | ||
| 495 | + qdata.flip((0,)) | ||
| 496 | + | ||
| 485 | 497 | ||
| 486 | def _rand_shape(self, dim, min_size, max_size): | 498 | def _rand_shape(self, dim, min_size, max_size): |
| 487 | return tuple(torch.randint(min_size, max_size + 1, (dim,))) | 499 | return tuple(torch.randint(min_size, max_size + 1, (dim,))) |
| @@ -0,0 +1,39 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | + | ||
| 4 | + | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +namespace at_npu { | ||
| 10 | +namespace native { | ||
| 11 | + | ||
| 12 | +at::Tensor quantized_flip(const at::Tensor& self, at::IntArrayRef dims) { | ||
| 13 | + TORCH_CHECK( | ||
| 14 | + self.scalar_type() != at::kQUInt4x2 && | ||
| 15 | + self.scalar_type() != at::kQUInt2x4, | ||
| 16 | + "flip is not supported for tensor with data type ", | ||
| 17 | + self.scalar_type()); | ||
| 18 | + | ||
| 19 | + // Match the native flip validation order before rejecting per-channel | ||
| 20 | + // quantization. | ||
| 21 | + (void)at::dim_list_to_bitset(dims, self.dim()); | ||
| 22 | + TORCH_CHECK( | ||
| 23 | + self.qscheme() == at::kPerTensorAffine, | ||
| 24 | + "Setting strides is possible only on uniformly quantized tensor"); | ||
| 25 | + | ||
| 26 | + at::Tensor repr = self.int_repr(); | ||
| 27 | + at::Tensor flipped = at::flip(repr, dims); | ||
| 28 | + at::Tensor result = at::_empty_affine_quantized( | ||
| 29 | + self.sizes(), | ||
| 30 | + self.options(), | ||
| 31 | + self.q_scale(), | ||
| 32 | + self.q_zero_point(), | ||
| 33 | + self.suggest_memory_format()); | ||
| 34 | + at_npu::native::NPUNativeFunctions::set_(result, flipped); | ||
| 35 | + return result; | ||
| 36 | +} | ||
| 37 | + | ||
| 38 | +} // namespace native | ||
| 39 | +} // namespace at_npu | ||
| @@ -0,0 +1,11 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +namespace at_npu { | ||
| 6 | +namespace native { | ||
| 7 | + | ||
| 8 | +at::Tensor quantized_flip(const at::Tensor& self, at::IntArrayRef dims); | ||
| 9 | + | ||
| 10 | +} // namespace native | ||
| 11 | +} // namespace at_npu | ||
| @@ -842,9 +842,11 @@ SPECIAL_REGISTERS = { | |||
| 842 | filename="QuantizedRegister", | 842 | filename="QuantizedRegister", |
| 843 | header="""\ | 843 | header="""\ |
| 844 | #include <ATen/ops/quantize_per_tensor.h> | 844 | #include <ATen/ops/quantize_per_tensor.h> |
| 845 | +#include "torch_npu/csrc/aten/ops/QuantizedFlipKernelNpu.h" | ||
| 845 | #include "op_plugin/OpInterface.h" | 846 | #include "op_plugin/OpInterface.h" |
| 846 | """, | 847 | """, |
| 847 | extra_impls=[ | 848 | extra_impls=[ |
| 849 | + 'm.impl("flip", TORCH_FN(at_npu::native::quantized_flip));', | ||
| 848 | 'm.impl("q_scale", TORCH_FN(at::native::q_scale_quant));', | 850 | 'm.impl("q_scale", TORCH_FN(at::native::q_scale_quant));', |
| 849 | 'm.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));', | 851 | 'm.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));', |
| 850 | 'm.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));', | 852 | 'm.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));', |