已合并
fix: avoid quantized flip registration for aclnn extensions #38786
hz893创建于 6月17日
fix: avoid quantized flip registration for aclnn extensions #38786
已合并
共 1 个文件变更+33-17
| @@ -33,7 +33,6 @@ import yaml | |||
| 33 | from torchgen.api.cpp import JIT_TO_CPP_DEFAULT | 33 | from torchgen.api.cpp import JIT_TO_CPP_DEFAULT |
| 34 | from torchgen.code_template import CodeTemplate | 34 | from torchgen.code_template import CodeTemplate |
| 35 | from torchgen.model import ( | 35 | from torchgen.model import ( |
| 36 | - BackendIndex, | ||
| 37 | SchemaKind, | 36 | SchemaKind, |
| 38 | TensorOptionsArguments, | 37 | TensorOptionsArguments, |
| 39 | ) | 38 | ) |
| @@ -69,7 +68,6 @@ from torchgen.api.types import ( | |||
| 69 | tensorT | 68 | tensorT |
| 70 | ) | 69 | ) |
| 71 | from torchgen.dest.register_dispatch_key import StructuredRegisterDispatchKey | 70 | from torchgen.dest.register_dispatch_key import StructuredRegisterDispatchKey |
| 72 | -from torchgen.gen_backend_stubs import gen_dispatchkey_nativefunc_headers | ||
| 73 | import torchgen.api.meta as meta | 71 | import torchgen.api.meta as meta |
| 74 | import torchgen.api.structured as structured | 72 | import torchgen.api.structured as structured |
| 75 | 73 | ||
| @@ -851,7 +849,8 @@ $dispatch_registrations_body | |||
| 851 | class_method_name=f'{class_name}', | 849 | class_method_name=f'{class_name}', |
| 852 | skip_dispatcher_op_registration=False, | 850 | skip_dispatcher_op_registration=False, |
| 853 | ), | 851 | ), |
| 854 | - [ g | 852 | + [ |
| 853 | + g | ||
| 855 | for g in grouped_native_functions | 854 | for g in grouped_native_functions |
| 856 | if g.root_name not in custom_op_names | 855 | if g.root_name not in custom_op_names |
| 857 | ] | 856 | ] |
| @@ -1272,6 +1271,35 @@ KERNEL_TEMPLATE = CodeTemplate("""\ | |||
| 1272 | m.impl("${schema}", TORCH_FN(op_plugin::${kernel}));""") | 1271 | m.impl("${schema}", TORCH_FN(op_plugin::${kernel}));""") |
| 1273 | 1272 | ||
| 1274 | 1273 | ||
| 1274 | +def _is_aclnn_extension_codegen() -> bool: | ||
| 1275 | + return bool(os.getenv("ACLNN_EXTENSION_SWITCH")) | ||
| 1276 | + | ||
| 1277 | + | ||
| 1278 | +def _quantized_register_header() -> str: | ||
| 1279 | + headers = [ | ||
| 1280 | + "#include <ATen/ops/quantize_per_tensor.h>", | ||
| 1281 | + ] | ||
| 1282 | + if not _is_aclnn_extension_codegen(): | ||
| 1283 | + headers.append('#include "torch_npu/csrc/aten/ops/QuantizedFlipKernelNpu.h"') | ||
| 1284 | + headers.append('#include "op_plugin/OpInterface.h"') | ||
| 1285 | + return "\n".join(headers) + "\n" | ||
| 1286 | + | ||
| 1287 | + | ||
| 1288 | +def _quantized_extra_impls() -> list[str]: | ||
| 1289 | + extra_impls = [] | ||
| 1290 | + if not _is_aclnn_extension_codegen(): | ||
| 1291 | + extra_impls.append('m.impl("flip", TORCH_FN(at_npu::native::quantized_flip));') | ||
| 1292 | + extra_impls.extend([ | ||
| 1293 | + 'm.impl("q_scale", TORCH_FN(at::native::q_scale_quant));', | ||
| 1294 | + 'm.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));', | ||
| 1295 | + 'm.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));', | ||
| 1296 | + 'm.impl("q_per_channel_zero_points", TORCH_FN(at::native::q_per_channel_zero_points));', | ||
| 1297 | + 'm.impl("q_per_channel_axis", TORCH_FN(at::native::q_per_channel_axis));', | ||
| 1298 | + 'm.impl("qscheme", TORCH_FN(at::native::qscheme_quant));', | ||
| 1299 | + ]) | ||
| 1300 | + return extra_impls | ||
| 1301 | + | ||
| 1302 | + | ||
| 1275 | def _gen_special_registration_body( | 1303 | def _gen_special_registration_body( |
| 1276 | backend_indices: BackendIndex, | 1304 | backend_indices: BackendIndex, |
| 1277 | config: SpecialRegisterConfig, | 1305 | config: SpecialRegisterConfig, |
| @@ -1338,20 +1366,8 @@ SPECIAL_REGISTERS = { | |||
| 1338 | "quantize": SpecialRegisterConfig( | 1366 | "quantize": SpecialRegisterConfig( |
| 1339 | dispatch_key="QuantizedPrivateUse1", | 1367 | dispatch_key="QuantizedPrivateUse1", |
| 1340 | filename="QuantizedRegister", | 1368 | filename="QuantizedRegister", |
| 1341 | - header="""\ | 1369 | + header=_quantized_register_header(), |
| 1342 | -#include <ATen/ops/quantize_per_tensor.h> | 1370 | + extra_impls=_quantized_extra_impls(), |
| 1343 | -#include "torch_npu/csrc/aten/ops/QuantizedFlipKernelNpu.h" | ||
| 1344 | -#include "op_plugin/OpInterface.h" | ||
| 1345 | -""", | ||
| 1346 | - extra_impls=[ | ||
| 1347 | - 'm.impl("flip", TORCH_FN(at_npu::native::quantized_flip));', | ||
| 1348 | - 'm.impl("q_scale", TORCH_FN(at::native::q_scale_quant));', | ||
| 1349 | - 'm.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));', | ||
| 1350 | - 'm.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));', | ||
| 1351 | - 'm.impl("q_per_channel_zero_points", TORCH_FN(at::native::q_per_channel_zero_points));', | ||
| 1352 | - 'm.impl("q_per_channel_axis", TORCH_FN(at::native::q_per_channel_axis));', | ||
| 1353 | - 'm.impl("qscheme", TORCH_FN(at::native::qscheme_quant));', | ||
| 1354 | - ], | ||
| 1355 | ), | 1371 | ), |
| 1356 | "nestedtensor": SpecialRegisterConfig( | 1372 | "nestedtensor": SpecialRegisterConfig( |
| 1357 | dispatch_key="NestedTensorPrivateUse1", | 1373 | dispatch_key="NestedTensorPrivateUse1", |