已合并
support strutured meta #36643
maoyuanpeng1创建于 5月25日
support strutured meta #36643
已合并
共 5 个文件变更+597-23
| @@ -22,6 +22,7 @@ import stat | |||
| 22 | from collections import Counter, defaultdict, namedtuple | 22 | from collections import Counter, defaultdict, namedtuple |
| 23 | from collections.abc import Callable, Sequence | 23 | from collections.abc import Callable, Sequence |
| 24 | from dataclasses import dataclass | 24 | from dataclasses import dataclass |
| 25 | +from typing import Optional, List, Set, Union | ||
| 25 | 26 | ||
| 26 | import torchgen | 27 | import torchgen |
| 27 | import torchgen.api.dispatcher as dispatcher | 28 | import torchgen.api.dispatcher as dispatcher |
| @@ -30,6 +31,11 @@ import torchgen.dest as dest | |||
| 30 | import yaml | 31 | import yaml |
| 31 | from torchgen.api.cpp import JIT_TO_CPP_DEFAULT | 32 | from torchgen.api.cpp import JIT_TO_CPP_DEFAULT |
| 32 | from torchgen.code_template import CodeTemplate | 33 | from torchgen.code_template import CodeTemplate |
| 34 | +from torchgen.model import ( | ||
| 35 | + BackendIndex, | ||
| 36 | + SchemaKind, | ||
| 37 | + TensorOptionsArguments, | ||
| 38 | +) | ||
| 33 | from torchgen.gen import ( | 39 | from torchgen.gen import ( |
| 34 | error_check_native_functions, | 40 | error_check_native_functions, |
| 35 | FileManager, | 41 | FileManager, |
| @@ -49,6 +55,22 @@ from torchgen.model import ( | |||
| 49 | from torchgen.native_function_generation import add_generated_native_functions | 55 | from torchgen.native_function_generation import add_generated_native_functions |
| 50 | from torchgen.selective_build.selector import SelectiveBuilder | 56 | from torchgen.selective_build.selector import SelectiveBuilder |
| 51 | from torchgen.utils import concatMap, context, NamespaceHelper, Target | 57 | from torchgen.utils import concatMap, context, NamespaceHelper, Target |
| 58 | +from torchgen.context import native_function_manager | ||
| 59 | +from torchgen.api.translate import translate | ||
| 60 | +from torchgen.api.types import ( | ||
| 61 | + BaseCType, | ||
| 62 | + ConstRefCType, | ||
| 63 | + CppSignatureGroup, | ||
| 64 | + Expr, | ||
| 65 | + MutRefCType, | ||
| 66 | + NamedCType, | ||
| 67 | + NativeSignature, | ||
| 68 | + tensorT | ||
| 69 | +) | ||
| 70 | +from torchgen.dest.register_dispatch_key import StructuredRegisterDispatchKey | ||
| 71 | +from torchgen.gen_backend_stubs import gen_dispatchkey_nativefunc_headers | ||
| 72 | +import torchgen.api.meta as meta | ||
| 73 | +import torchgen.api.structured as structured | ||
| 52 | 74 | ||
| 53 | from torchnpugen.custom_functions import ( | 75 | from torchnpugen.custom_functions import ( |
| 54 | gen_custom_functions_dispatch, | 76 | gen_custom_functions_dispatch, |
| @@ -88,6 +110,9 @@ from torchnpugen.utils import ( | |||
| 88 | 110 | ||
| 89 | torchgen.model.dispatch_keys.append(torchgen.model.DispatchKey.AutogradPrivateUse1) | 111 | torchgen.model.dispatch_keys.append(torchgen.model.DispatchKey.AutogradPrivateUse1) |
| 90 | 112 | ||
| 113 | +NPU_STRUCTURED_INPLACE_OPS: set[str] = set() | ||
| 114 | +NPU_STRUCTURED_PRECOMPUTED_OPS: set[str] = set() | ||
| 115 | + | ||
| 91 | 116 | ||
| 92 | # Create backend_indices map for func retrieval with the key of each func we supported. | 117 | # Create backend_indices map for func retrieval with the key of each func we supported. |
| 93 | def create_backend_index( | 118 | def create_backend_index( |
| @@ -96,23 +121,28 @@ def create_backend_index( | |||
| 96 | dispatch_key: DispatchKey, | 121 | dispatch_key: DispatchKey, |
| 97 | native_funcs_map: dict[OperatorName, NativeFunction], | 122 | native_funcs_map: dict[OperatorName, NativeFunction], |
| 98 | cpp_namespace: str, | 123 | cpp_namespace: str, |
| 124 | + structured_ops: Optional[set[str]] = None, | ||
| 99 | ) -> BackendIndex: | 125 | ) -> BackendIndex: |
| 100 | metadata: dict[OperatorName, BackendMetadata] = {} | 126 | metadata: dict[OperatorName, BackendMetadata] = {} |
| 101 | for op in backend_ops: | 127 | for op in backend_ops: |
| 102 | op_name = OperatorName.parse(op) | 128 | op_name = OperatorName.parse(op) |
| 103 | if op_name not in native_funcs_map: | 129 | if op_name not in native_funcs_map: |
| 104 | raise KeyError(f"Found an invalid operator name: {op_name}") | 130 | raise KeyError(f"Found an invalid operator name: {op_name}") |
| 131 | + is_structured = structured_ops is not None and op in structured_ops | ||
| 132 | + if is_structured and not native_funcs_map[op_name].structured: | ||
| 133 | + raise ValueError( | ||
| 134 | + f"{op} is marked structured in op_plugin yaml, but the upstream " | ||
| 135 | + "native function is not a structured out variant." | ||
| 136 | + ) | ||
| 105 | # See Note [External Backends Follow Dispatcher API] | 137 | # See Note [External Backends Follow Dispatcher API] |
| 106 | kernel_name = dispatcher.name(native_funcs_map[op_name].func) | 138 | kernel_name = dispatcher.name(native_funcs_map[op_name].func) |
| 107 | - if op in symint_ops: | 139 | + if op in symint_ops and not is_structured: |
| 108 | kernel_name += "_symint" | 140 | kernel_name += "_symint" |
| 109 | - m = BackendMetadata( | 141 | + m = BackendMetadata(kernel=kernel_name, structured=is_structured, cpp_namespace=cpp_namespace) |
| 110 | - kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace | ||
| 111 | - ) | ||
| 112 | metadata[op_name] = m | 142 | metadata[op_name] = m |
| 113 | return BackendIndex( | 143 | return BackendIndex( |
| 114 | dispatch_key=dispatch_key, | 144 | dispatch_key=dispatch_key, |
| 115 | - use_out_as_primary=False, | 145 | + use_out_as_primary=True, |
| 116 | external=True, | 146 | external=True, |
| 117 | device_guard=True, | 147 | device_guard=True, |
| 118 | index=metadata, | 148 | index=metadata, |
| @@ -259,6 +289,80 @@ ParsedExternalYaml = namedtuple( | |||
| 259 | ) | 289 | ) |
| 260 | 290 | ||
| 261 | 291 | ||
| 292 | +def _get_backend_yaml_op_name(op: object) -> str: | ||
| 293 | + if isinstance(op, dict): | ||
| 294 | + func = op.get("func") | ||
| 295 | + if not isinstance(func, str): | ||
| 296 | + raise TypeError(f'expected "func" to be a string, but got: {func}') | ||
| 297 | + return func.split("(", 1)[0].strip() | ||
| 298 | + if isinstance(op, str): | ||
| 299 | + return op | ||
| 300 | + raise TypeError(f"expected an operator name or a dict, but got: {op}") | ||
| 301 | + | ||
| 302 | + | ||
| 303 | +def _check_structured_delegate( | ||
| 304 | + op_name: OperatorName, | ||
| 305 | + op: dict[str, object], | ||
| 306 | + group: NativeFunctionsGroup, | ||
| 307 | +) -> None: | ||
| 308 | + expected_delegate = group.out.func.name | ||
| 309 | + delegate = op.get("structured_delegate") | ||
| 310 | + if delegate is None: | ||
| 311 | + raise ValueError( | ||
| 312 | + f"{op_name} belongs to a structured group whose out variant is " | ||
| 313 | + f"{expected_delegate}, but it does not specify structured_delegate: " | ||
| 314 | + f"{expected_delegate}." | ||
| 315 | + ) | ||
| 316 | + if not isinstance(delegate, str): | ||
| 317 | + raise TypeError( | ||
| 318 | + f'expected "structured_delegate" for {op_name} to be a string, ' | ||
| 319 | + f"but got: {delegate}" | ||
| 320 | + ) | ||
| 321 | + delegate_name = OperatorName.parse(delegate) | ||
| 322 | + if delegate_name != expected_delegate: | ||
| 323 | + raise ValueError( | ||
| 324 | + f"{op_name} specifies structured_delegate: {delegate_name}, " | ||
| 325 | + f"but its structured out variant is {expected_delegate}." | ||
| 326 | + ) | ||
| 327 | + | ||
| 328 | + | ||
| 329 | +def check_structured_group_consistency( | ||
| 330 | + supported_by_name: dict[OperatorName, dict[str, object]], | ||
| 331 | + structured_groups: dict[OperatorName, NativeFunctionsGroup], | ||
| 332 | + structured_group_members: dict[OperatorName, NativeFunctionsGroup], | ||
| 333 | +) -> None: | ||
| 334 | + for out_name, group in sorted(structured_groups.items(), key=lambda item: str(item[0])): | ||
| 335 | + functional_name = group.functional.func.name | ||
| 336 | + functional_op = supported_by_name.get(functional_name) | ||
| 337 | + if functional_op is None: | ||
| 338 | + raise ValueError( | ||
| 339 | + f"{out_name} is marked structured in op_plugin yaml, " | ||
| 340 | + f"but its functional variant {functional_name} is not listed in supported." | ||
| 341 | + ) | ||
| 342 | + _check_structured_delegate(functional_name, functional_op, group) | ||
| 343 | + | ||
| 344 | + if group.inplace is None: | ||
| 345 | + continue | ||
| 346 | + inplace_name = group.inplace.func.name | ||
| 347 | + inplace_op = supported_by_name.get(inplace_name) | ||
| 348 | + if inplace_op is None: | ||
| 349 | + continue | ||
| 350 | + if bool(inplace_op.get("structured", False)): | ||
| 351 | + continue | ||
| 352 | + _check_structured_delegate(inplace_name, inplace_op, group) | ||
| 353 | + | ||
| 354 | + for op_name, group in sorted(structured_group_members.items(), key=lambda item: str(item[0])): | ||
| 355 | + if group.out.func.name in structured_groups: | ||
| 356 | + continue | ||
| 357 | + out_name = group.out.func.name | ||
| 358 | + out_op = supported_by_name.get(out_name) | ||
| 359 | + if out_op is None or not bool(out_op.get("structured", False)): | ||
| 360 | + raise ValueError( | ||
| 361 | + f"{op_name} is annotated as part of a structured group, " | ||
| 362 | + f"but its out variant {out_name} is not marked structured in op_plugin yaml." | ||
| 363 | + ) | ||
| 364 | + | ||
| 365 | + | ||
| 262 | def parse_backend_yaml( | 366 | def parse_backend_yaml( |
| 263 | native_yaml_path: str, | 367 | native_yaml_path: str, |
| 264 | backend_yaml_path: str, | 368 | backend_yaml_path: str, |
| @@ -266,12 +370,14 @@ def parse_backend_yaml( | |||
| 266 | backend_indices: dict[DispatchKey, BackendIndex], | 370 | backend_indices: dict[DispatchKey, BackendIndex], |
| 267 | ) -> ParsedExternalYaml: | 371 | ) -> ParsedExternalYaml: |
| 268 | native_functions_map = {} | 372 | native_functions_map = {} |
| 373 | + native_function_groups = {} | ||
| 269 | for f in grouped_native_functions: | 374 | for f in grouped_native_functions: |
| 270 | if isinstance(f, NativeFunction): | 375 | if isinstance(f, NativeFunction): |
| 271 | native_functions_map[f.func.name] = f | 376 | native_functions_map[f.func.name] = f |
| 272 | else: | 377 | else: |
| 273 | for func in f.functions(): | 378 | for func in f.functions(): |
| 274 | native_functions_map[func.func.name] = func | 379 | native_functions_map[func.func.name] = func |
| 380 | + native_function_groups[func.func.name] = f | ||
| 275 | 381 | ||
| 276 | PathManager.check_directory_path_readable(backend_yaml_path) | 382 | PathManager.check_directory_path_readable(backend_yaml_path) |
| 277 | with open(backend_yaml_path) as f: | 383 | with open(backend_yaml_path) as f: |
| @@ -327,15 +433,67 @@ def parse_backend_yaml( | |||
| 327 | ) | 433 | ) |
| 328 | 434 | ||
| 329 | supported_list = [] | 435 | supported_list = [] |
| 436 | + structured_ops: Set[str] = set() | ||
| 437 | + structured_groups: dict[OperatorName, NativeFunctionsGroup] = {} | ||
| 438 | + structured_group_members: dict[OperatorName, NativeFunctionsGroup] = {} | ||
| 439 | + supported_by_name: dict[OperatorName, dict[str, object]] = {} | ||
| 330 | for op in supported: | 440 | for op in supported: |
| 331 | - if isinstance(op, dict) and op.get("device_check", None) == "NoCheck": | 441 | + op_name = _get_backend_yaml_op_name(op) |
| 332 | - DEVICE_NOCHECK_SET.add(op["func"].split("(")[0]) | 442 | + op_schema_name = OperatorName.parse(op_name) |
| 333 | - if isinstance(op, dict) and ( | 443 | + op_info = op if isinstance(op, dict) else {} |
| 334 | - {"impl_ns", "op_api", "device_check"} & set(op.keys()) | 444 | + supported_by_name[op_schema_name] = op_info |
| 335 | - ): | 445 | + if isinstance(op, dict): |
| 336 | - supported_list.append(op["func"].split("(")[0]) | 446 | + native_func = native_functions_map.get(op_schema_name) |
| 447 | + is_structured = bool(op.get('structured', False)) | ||
| 448 | + has_structured_delegate = op.get('structured_delegate') is not None | ||
| 449 | + if is_structured or has_structured_delegate: | ||
| 450 | + if native_func is None: | ||
| 451 | + raise KeyError(f"Found an invalid structured operator name: {op_schema_name}") | ||
| 452 | + schema_kind = native_func.func.kind() | ||
| 453 | + group = native_function_groups.get(op_schema_name) | ||
| 454 | + if is_structured: | ||
| 455 | + if group is None: | ||
| 456 | + raise ValueError( | ||
| 457 | + f"{op_name} is marked structured in op_plugin yaml, " | ||
| 458 | + "but it is not part of a structured native function group." | ||
| 459 | + ) | ||
| 460 | + if native_func.func.kind() == SchemaKind.out: | ||
| 461 | + structured_ops.add(op_name) | ||
| 462 | + structured_groups[op_schema_name] = group | ||
| 463 | + if op.get('precomputed') is not None: | ||
| 464 | + NPU_STRUCTURED_PRECOMPUTED_OPS.add(op_name) | ||
| 465 | + elif native_func.func.kind() == SchemaKind.inplace: | ||
| 466 | + structured_group_members[op_schema_name] = group | ||
| 467 | + NPU_STRUCTURED_INPLACE_OPS.add(op_name) | ||
| 468 | + else: | ||
| 469 | + raise ValueError( | ||
| 470 | + f"{op_name} is marked structured in op_plugin yaml, " | ||
| 471 | + "but only out and inplace variants are supported." | ||
| 472 | + ) | ||
| 473 | + if has_structured_delegate: | ||
| 474 | + if schema_kind not in (SchemaKind.functional, SchemaKind.inplace): | ||
| 475 | + raise ValueError( | ||
| 476 | + f"{op_name} specifies structured_delegate in op_plugin yaml, " | ||
| 477 | + "but only functional and inplace variants can delegate to a structured out variant." | ||
| 478 | + ) | ||
| 479 | + if group is None: | ||
| 480 | + raise ValueError( | ||
| 481 | + f"{op_name} specifies structured_delegate in op_plugin yaml, " | ||
| 482 | + "but it is not part of a structured native function group." | ||
| 483 | + ) | ||
| 484 | + _check_structured_delegate(op_schema_name, op, group) | ||
| 485 | + structured_group_members[op_schema_name] = group | ||
| 486 | + if isinstance(op, dict) and op.get('device_check', None) == 'NoCheck': | ||
| 487 | + DEVICE_NOCHECK_SET.add(op['func'].split("(")[0]) | ||
| 488 | + if isinstance(op, dict) and ({"impl_ns", "op_api", "device_check"} & set(op.keys())): | ||
| 489 | + supported_list.append(op['func'].split("(")[0]) | ||
| 337 | elif not isinstance(op, dict): | 490 | elif not isinstance(op, dict): |
| 338 | supported_list.append(op) | 491 | supported_list.append(op) |
| 492 | + check_structured_group_consistency( | ||
| 493 | + supported_by_name, | ||
| 494 | + structured_groups, | ||
| 495 | + structured_group_members, | ||
| 496 | + ) | ||
| 339 | supported = supported_list | 497 | supported = supported_list |
| 340 | 498 | ||
| 341 | supported_autograd = [ | 499 | supported_autograd = [ |
| @@ -393,14 +551,14 @@ def parse_backend_yaml( | |||
| 393 | backend_key = DispatchKey.parse(backend) | 551 | backend_key = DispatchKey.parse(backend) |
| 394 | 552 | ||
| 395 | backend_idx = create_backend_index( | 553 | backend_idx = create_backend_index( |
| 396 | - supported, symint_set, backend_key, native_functions_map, cpp_namespace | 554 | + supported, symint_set, backend_key, native_functions_map, cpp_namespace, structured_ops |
| 397 | ) | 555 | ) |
| 398 | opapi_backend_idx = create_backend_index( | 556 | opapi_backend_idx = create_backend_index( |
| 399 | [op for op in supported if is_opapi(op)], | 557 | [op for op in supported if is_opapi(op)], |
| 400 | - symint_set, | 558 | + symint_set, backend_key, |
| 401 | - backend_key, | ||
| 402 | native_functions_map, | 559 | native_functions_map, |
| 403 | cpp_namespace, | 560 | cpp_namespace, |
| 561 | + structured_ops, | ||
| 404 | ) | 562 | ) |
| 405 | if backend_key in backend_indices: | 563 | if backend_key in backend_indices: |
| 406 | raise KeyError("backend_key should not be in backend_indices.") | 564 | raise KeyError("backend_key should not be in backend_indices.") |
| @@ -558,6 +716,35 @@ def main() -> None: | |||
| 558 | ) | 716 | ) |
| 559 | 717 | ||
| 560 | 718 | ||
| 719 | +def gen_npu_structured_registration_helpers(backend_index: BackendIndex) -> List[str]: | ||
| 720 | + helpers = dest.gen_registration_helpers(backend_index) | ||
| 721 | + if backend_index.dispatch_key.name not in ("NPU", "PrivateUse1"): | ||
| 722 | + return helpers | ||
| 723 | + if not any(metadata.structured for metadata in backend_index.index.values()): | ||
| 724 | + return helpers | ||
| 725 | + if any("create_out(" in helper for helper in helpers): | ||
| 726 | + return helpers | ||
| 727 | + return [ | ||
| 728 | + """ | ||
| 729 | +Tensor create_out(IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) { | ||
| 730 | + if (strides.empty()) { | ||
| 731 | + return at::empty(sizes, options); | ||
| 732 | + } | ||
| 733 | + return at::empty_strided(sizes, strides, options); | ||
| 734 | +} | ||
| 735 | +""", | ||
| 736 | + *helpers, | ||
| 737 | + """ | ||
| 738 | +std::optional<Tensor> maybe_create_proxy(const Tensor &out, IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) { | ||
| 739 | + if (!strides.empty() && out.strides() != strides) { | ||
| 740 | + return at::empty_strided(sizes, strides, options); | ||
| 741 | + } | ||
| 742 | + return std::nullopt; | ||
| 743 | +} | ||
| 744 | +""", | ||
| 745 | + ] | ||
| 746 | + | ||
| 747 | + | ||
| 561 | def gen_dispatcher_registrations( | 748 | def gen_dispatcher_registrations( |
| 562 | fm: FileManager, | 749 | fm: FileManager, |
| 563 | class_name: str, | 750 | class_name: str, |
| @@ -588,6 +775,7 @@ def gen_dispatcher_registrations( | |||
| 588 | #include "torch_npu/csrc/aten/NPUNativeFunctions.h" | 775 | #include "torch_npu/csrc/aten/NPUNativeFunctions.h" |
| 589 | #include "torch_npu/csrc/framework/interface/EnvVariables.h" | 776 | #include "torch_npu/csrc/framework/interface/EnvVariables.h" |
| 590 | #include "torch_npu/csrc/aten/NPUOpApiNativeFunctions.h" | 777 | #include "torch_npu/csrc/aten/NPUOpApiNativeFunctions.h" |
| 778 | +#include "torch_npu/csrc/aten/NPUStructuredNativeFunctions.h" | ||
| 591 | """ | 779 | """ |
| 592 | 780 | ||
| 593 | # 根据环境变量决定是否包含FormatHelper.h | 781 | # 根据环境变量决定是否包含FormatHelper.h |
| @@ -636,7 +824,7 @@ $dispatch_registrations_body | |||
| 636 | backend_index, per_operator_headers=False, rocm=False | 824 | backend_index, per_operator_headers=False, rocm=False |
| 637 | ), | 825 | ), |
| 638 | "ops_headers": "", | 826 | "ops_headers": "", |
| 639 | - "dispatch_helpers": dest.gen_registration_helpers(backend_index), | 827 | + "dispatch_helpers": gen_npu_structured_registration_helpers(backend_index), |
| 640 | "dispatch_definitions": fm.substitute_with_template( | 828 | "dispatch_definitions": fm.substitute_with_template( |
| 641 | "RegisterDispatchDefinitions.ini", | 829 | "RegisterDispatchDefinitions.ini", |
| 642 | lambda: { | 830 | lambda: { |
| @@ -686,6 +874,309 @@ $dispatch_registrations_body | |||
| 686 | ) | 874 | ) |
| 687 | 875 | ||
| 688 | 876 | ||
| 877 | +def get_supported_structured_groups( | ||
H | |||
| 878 | + grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], | ||
| 879 | + backend_index: BackendIndex, | ||
| 880 | +) -> List[NativeFunctionsGroup]: | ||
| 881 | + structured_groups: List[NativeFunctionsGroup] = [] | ||
| 882 | + for funcs in grouped_native_functions: | ||
| 883 | + if not isinstance(funcs, NativeFunctionsGroup): | ||
| 884 | + continue | ||
| 885 | + metadata = backend_index.get_kernel(funcs) | ||
| 886 | + if metadata is not None and metadata.structured: | ||
| 887 | + structured_groups.append(funcs) | ||
| 888 | + return structured_groups | ||
| 889 | + | ||
| 890 | + | ||
| 891 | +def npu_structured_uses_precompute(g: NativeFunctionsGroup) -> bool: | ||
| 892 | + return str(g.out.func.name) in NPU_STRUCTURED_PRECOMPUTED_OPS | ||
| 893 | + | ||
| 894 | + | ||
| 895 | +def npu_structured_impl_arguments(g: NativeFunctionsGroup): | ||
| 896 | + if npu_structured_uses_precompute(g): | ||
| 897 | + return structured.impl_arguments(g) | ||
| 898 | + return [*structured.meta_arguments(g), *structured.out_arguments(g)] | ||
| 899 | + | ||
| 900 | + | ||
| 901 | +def npu_structured_call_args(g: NativeFunctionsGroup) -> str: | ||
| 902 | + impl_arg_names = {arg.name for arg in npu_structured_impl_arguments(g)} | ||
| 903 | + call_args: List[str] = [] | ||
| 904 | + | ||
| 905 | + for arg in g.out.func.arguments.non_out: | ||
| 906 | + if isinstance(arg, TensorOptionsArguments): | ||
| 907 | + raise RuntimeError(f"{g.out.func.name} has TensorOptions, which structured kernels do not support") | ||
| 908 | + if hasattr(arg, "argument"): | ||
| 909 | + arg_name = arg.argument.name | ||
| 910 | + else: | ||
| 911 | + arg_name = arg.name | ||
| 912 | + if arg_name not in impl_arg_names: | ||
| 913 | + raise RuntimeError( | ||
| 914 | + f"Cannot call op_plugin::{dispatcher.name(g.out.func)} from structured impl for " | ||
| 915 | + f"{g.out.func.name}: original argument '{arg_name}' is not available after precompute." | ||
| 916 | + ) | ||
| 917 | + call_args.append(arg_name) | ||
| 918 | + | ||
| 919 | + for arg in g.out.func.arguments.out: | ||
| 920 | + call_args.append(f"const_cast<at::Tensor&>({arg.name})") | ||
| 921 | + | ||
| 922 | + return ", ".join(call_args) | ||
| 923 | + | ||
| 924 | + | ||
| 925 | +def npu_structured_inplace_arguments(f: NativeFunction): | ||
| 926 | + with native_function_manager(f): | ||
| 927 | + return dispatcher.arguments(f.func, symint=True) | ||
| 928 | + | ||
| 929 | + | ||
| 930 | +def npu_structured_inplace_call_args(f: NativeFunction) -> str: | ||
| 931 | + return ", ".join(arg.name for arg in npu_structured_inplace_arguments(f)) | ||
| 932 | + | ||
| 933 | + | ||
| 934 | +def npu_structured_return_expr(f: NativeFunction, k: SchemaKind) -> str: | ||
| 935 | + if k is SchemaKind.functional: | ||
| 936 | + if len(f.func.returns) == 1: | ||
| 937 | + return "std::move(op.outputs_[0])" | ||
| 938 | + moved = ", ".join(f"std::move(op.outputs_[{i}])" for i in range(len(f.func.returns))) | ||
| 939 | + return f"std::make_tuple({moved})" | ||
| 940 | + if k is SchemaKind.inplace: | ||
| 941 | + return "self" | ||
| 942 | + if k is SchemaKind.out: | ||
| 943 | + if len(f.func.returns) == 1: | ||
| 944 | + return f.func.arguments.out[0].name | ||
| 945 | + refs = ", ".join(a.name for a in f.func.arguments.out) | ||
| 946 | + return f"std::forward_as_tuple({refs})" | ||
| 947 | + raise AssertionError(f"{k} structured operators are currently not supported") | ||
| 948 | + | ||
| 949 | + | ||
| 950 | +def gen_npu_structured_return_with_op_hook(f: NativeFunction, sig: NativeSignature, ret_expr: str) -> str: | ||
| 951 | + returns_type = sig.returns_type() | ||
| 952 | + if isinstance(returns_type, MutRefCType): | ||
| 953 | + return f"""\ | ||
| 954 | +auto& op_hook_result = {ret_expr}; | ||
| 955 | +if (op_hook_enabled) {{ | ||
| 956 | + at_npu::native::OpHook::GetInstance().PostHook(op_hook_result); | ||
| 957 | +}} | ||
| 958 | +return op_hook_result;""" | ||
| 959 | + return f"""\ | ||
| 960 | +auto op_hook_result = {ret_expr}; | ||
| 961 | +if (op_hook_enabled) {{ | ||
| 962 | + at_npu::native::OpHook::GetInstance().PostHook(op_hook_result); | ||
| 963 | +}} | ||
| 964 | +return op_hook_result;""" | ||
| 965 | + | ||
| 966 | + | ||
| 967 | +def gen_npu_structured_one(self, f: NativeFunction) -> Optional[str]: | ||
| 968 | + if f.manual_kernel_registration: | ||
| 969 | + raise AssertionError(f"Function {f.func.name} has manual_kernel_registration=True") | ||
| 970 | + | ||
| 971 | + if self.target is Target.REGISTRATION and not self.selector.is_native_function_selected(f): | ||
| 972 | + return None | ||
| 973 | + | ||
| 974 | + if not self.backend_index.has_kernel(f): | ||
| 975 | + return None | ||
| 976 | + | ||
| 977 | + metadata = self.backend_index.get_kernel(self.g) | ||
| 978 | + if metadata is None: | ||
| 979 | + raise AssertionError(f"No kernel metadata found for {self.g.functional.func.name}") | ||
| 980 | + | ||
| 981 | + kern = self.backend_index.get_kernel(f) | ||
| 982 | + sig = NativeSignature( | ||
| 983 | + f.func, | ||
| 984 | + prefix=f"wrapper_{self.backend_index.dispatch_key}_", | ||
| 985 | + symint=kern is not None and kern.supports_symint(), | ||
| 986 | + ) | ||
| 987 | + | ||
| 988 | + cpp_sig_group = CppSignatureGroup.from_native_function( | ||
| 989 | + f, method=False, fallback_binding=False | ||
| 990 | + ) | ||
| 991 | + | ||
| 992 | + if self.target is Target.NAMESPACED_DECLARATION: | ||
| 993 | + result = "" | ||
| 994 | + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): | ||
| 995 | + result += f"TORCH_API {cpp_sig.decl()};\n" | ||
| 996 | + return result | ||
| 997 | + | ||
| 998 | + if self.target is Target.NAMESPACED_DEFINITION: | ||
| 999 | + def generate_defn(cpp_sig) -> str: | ||
| 1000 | + return f""" | ||
| 1001 | +{cpp_sig.defn()} {{ | ||
| 1002 | +return {sig.name()}({", ".join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))}); | ||
| 1003 | +}} | ||
| 1004 | +""" | ||
| 1005 | + | ||
| 1006 | + result = "" | ||
| 1007 | + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): | ||
| 1008 | + result += generate_defn(cpp_sig) | ||
| 1009 | + return result | ||
| 1010 | + | ||
| 1011 | + if self.target is Target.REGISTRATION: | ||
| 1012 | + return f'm.impl("{f.func.name}", TORCH_FN({sig.name()}));' | ||
| 1013 | + if self.target is not Target.ANONYMOUS_DEFINITION: | ||
| 1014 | + return None | ||
| 1015 | + | ||
| 1016 | + k = f.func.kind() | ||
| 1017 | + class_name = f"structured_{metadata.kernel}_{k.name}" | ||
| 1018 | + parent_kernel = metadata.kernel | ||
| 1019 | + if k is SchemaKind.inplace and str(f.func.name) in NPU_STRUCTURED_INPLACE_OPS: | ||
| 1020 | + parent_kernel = dispatcher.name(f.func) | ||
| 1021 | + parent_class = f"{metadata.cpp_namespace}::structured_{parent_kernel}" | ||
| 1022 | + | ||
| 1023 | + sig_body: List[str] = [] | ||
| 1024 | + context_args = list(sig.arguments()) | ||
| 1025 | + context_exprs = list(context_args) | ||
| 1026 | + | ||
| 1027 | + if self.backend_index.device_guard: | ||
| 1028 | + device_check_args = [*f.func.arguments.out, *f.func.arguments.flat_positional] | ||
| 1029 | + sig_body.append( | ||
| 1030 | + dest.RegisterDispatchKey.gen_device_check( | ||
| 1031 | + f.device_check, list(device_check_args), sig.name() | ||
| 1032 | + ) | ||
| 1033 | + ) | ||
| 1034 | + | ||
| 1035 | + pre_hook_args = ", ".join(arg.name for arg in context_args) | ||
| 1036 | + pre_hook_args = f", {pre_hook_args}" if pre_hook_args else "" | ||
| 1037 | + sig_body.append("const bool op_hook_enabled = C10_UNLIKELY(at_npu::native::env::CheckOpHookEnable());") | ||
| 1038 | + sig_body.append("if (op_hook_enabled) {") | ||
| 1039 | + sig_body.append(f' at_npu::native::OpHook::GetInstance().PreHook("{f.func.name}"{pre_hook_args});') | ||
| 1040 | + sig_body.append("}") | ||
| 1041 | + | ||
| 1042 | + if k is SchemaKind.functional: | ||
| 1043 | + sig_body.append(f"{class_name} op;") | ||
| 1044 | + elif k is SchemaKind.inplace: | ||
| 1045 | + sig_body.append(f"{class_name} op(self);") | ||
| 1046 | + elif k is SchemaKind.out: | ||
| 1047 | + out_args_str = ", ".join(a.name for a in f.func.arguments.out) | ||
| 1048 | + sig_body.append(f"{class_name} op({out_args_str});") | ||
| 1049 | + else: | ||
| 1050 | + raise AssertionError(f"{k} structured operators are currently not supported") | ||
| 1051 | + | ||
| 1052 | + meta_exprs = ", ".join( | ||
| 1053 | + e.expr for e in translate(context_exprs, structured.meta_arguments(self.g), method=False) | ||
| 1054 | + ) | ||
| 1055 | + if npu_structured_uses_precompute(self.g) and self.g.out.precomputed: | ||
| 1056 | + sig_body.append(f"auto precompute = op.meta({meta_exprs});") | ||
| 1057 | + precomputed_values = [ | ||
| 1058 | + *self.g.out.precomputed.replace.values(), | ||
| 1059 | + self.g.out.precomputed.add, | ||
| 1060 | + ] | ||
| 1061 | + for precomputed_elems in precomputed_values: | ||
| 1062 | + context_exprs.extend( | ||
| 1063 | + Expr( | ||
| 1064 | + expr=f"precompute.{arg.name}", | ||
| 1065 | + type=structured.argument_type(arg, binds=arg.name), | ||
| 1066 | + ) | ||
| 1067 | + for arg in precomputed_elems | ||
| 1068 | + ) | ||
| 1069 | + sig_body.append("(void)precompute;") | ||
| 1070 | + else: | ||
| 1071 | + sig_body.append(f"op.meta({meta_exprs});") | ||
| 1072 | + | ||
| 1073 | + out_args = structured.out_arguments(self.g) | ||
| 1074 | + for i, out_arg in enumerate(out_args): | ||
| 1075 | + if ConstRefCType(BaseCType(tensorT)) != out_arg.nctype.type: | ||
| 1076 | + raise AssertionError( | ||
| 1077 | + f"Expected out_arg type to be ConstRefCType(BaseCType(tensorT)), got {out_arg.nctype.type}" | ||
| 1078 | + ) | ||
| 1079 | + expr = f"op.maybe_get_output({i})" if k is SchemaKind.out else f"op.outputs_[{i}]" | ||
| 1080 | + context_exprs.append( | ||
| 1081 | + Expr( | ||
| 1082 | + expr=expr, | ||
| 1083 | + type=NamedCType(out_arg.nctype.name, MutRefCType(BaseCType(tensorT))), | ||
| 1084 | + ) | ||
| 1085 | + ) | ||
| 1086 | + | ||
| 1087 | + if k is SchemaKind.inplace and str(f.func.name) in NPU_STRUCTURED_INPLACE_OPS: | ||
| 1088 | + impl_exprs = npu_structured_inplace_call_args(f) | ||
| 1089 | + else: | ||
| 1090 | + impl_exprs = ", ".join( | ||
| 1091 | + e.expr for e in translate(context_exprs, npu_structured_impl_arguments(self.g), method=False) | ||
| 1092 | + ) | ||
| 1093 | + sig_body.append(f"op.impl({impl_exprs});") | ||
| 1094 | + | ||
| 1095 | + if k is SchemaKind.out or k is SchemaKind.inplace: | ||
| 1096 | + for i in range(len(f.func.returns)): | ||
| 1097 | + sig_body.append( | ||
| 1098 | + f"if (op.proxy_outputs_[{i}].has_value()) op.outputs_[{i}].get().copy_(*op.proxy_outputs_[{i}]);" | ||
| 1099 | + ) | ||
| 1100 | + | ||
| 1101 | + sig_body.append(gen_npu_structured_return_with_op_hook(f, sig, npu_structured_return_expr(f, k))) | ||
| 1102 | + sig_body_str = "\n".join(sig_body) | ||
| 1103 | + | ||
| 1104 | + return f"""\ | ||
| 1105 | +{self.gen_class( | ||
| 1106 | + f, | ||
| 1107 | + k, | ||
| 1108 | + class_name=class_name, | ||
| 1109 | + parent_class=parent_class, | ||
| 1110 | + generate_super=self.g.out.structured_inherits is not None, | ||
| 1111 | +)} | ||
| 1112 | + | ||
| 1113 | +{sig.defn()} {{ | ||
| 1114 | +{sig_body_str} | ||
| 1115 | +}} | ||
| 1116 | +""" | ||
| 1117 | + | ||
| 1118 | + | ||
| 1119 | +def gen_npu_structured_native_functions( | ||
| 1120 | + fm: FileManager, | ||
| 1121 | + grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], | ||
| 1122 | + backend_index: BackendIndex, | ||
| 1123 | +) -> None: | ||
| 1124 | + structured_groups = get_supported_structured_groups(grouped_native_functions, backend_index) | ||
| 1125 | + includes = sorted({ | ||
| 1126 | + f"#include <ATen/ops/{meta.name(g)}_meta.h>" | ||
| 1127 | + for g in structured_groups | ||
| 1128 | + }) | ||
| 1129 | + | ||
| 1130 | + declarations: List[str] = [] | ||
| 1131 | + definitions: List[str] = [] | ||
| 1132 | + for g in structured_groups: | ||
| 1133 | + metadata = backend_index.get_kernel(g) | ||
| 1134 | + if metadata is None: | ||
| 1135 | + continue | ||
| 1136 | + impl_args = ", ".join(arg.decl() for arg in npu_structured_impl_arguments(g)) | ||
| 1137 | + meta_name = meta.name(g) | ||
| 1138 | + declarations.append(f"""\ | ||
| 1139 | +struct structured_{metadata.kernel} : public at::meta::structured_{meta_name} {{ | ||
| 1140 | + void impl({impl_args}); | ||
| 1141 | +}}; | ||
| 1142 | +""") | ||
| 1143 | + definitions.append(f"""\ | ||
| 1144 | +void structured_{metadata.kernel}::impl({impl_args}) | ||
| 1145 | +{{ | ||
| 1146 | + op_plugin::{metadata.kernel}({npu_structured_call_args(g)}); | ||
| 1147 | +}} | ||
| 1148 | +""") | ||
| 1149 | + for f in g.functions(): | ||
| 1150 | + if f.func.kind() != SchemaKind.inplace or str(f.func.name) not in NPU_STRUCTURED_INPLACE_OPS: | ||
| 1151 | + continue | ||
| 1152 | + inplace_kernel = dispatcher.name(f.func) | ||
| 1153 | + inplace_impl_args = ", ".join(arg.decl() for arg in npu_structured_inplace_arguments(f)) | ||
| 1154 | + declarations.append(f"""\ | ||
| 1155 | +struct structured_{inplace_kernel} : public at::meta::structured_{meta_name} {{ | ||
| 1156 | + void impl({inplace_impl_args}); | ||
| 1157 | +}}; | ||
| 1158 | +""") | ||
| 1159 | + definitions.append(f"""\ | ||
| 1160 | +void structured_{inplace_kernel}::impl({inplace_impl_args}) | ||
| 1161 | +{{ | ||
| 1162 | + op_plugin::{inplace_kernel}({npu_structured_inplace_call_args(f)}); | ||
| 1163 | +}} | ||
| 1164 | +""") | ||
| 1165 | + | ||
| 1166 | + generated_comment = "Autogenerated file by gen_backend_stubs.py. Do not edit directly!" | ||
| 1167 | + template_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "templates") | ||
| 1168 | + structured_fm = FileManager(install_dir=fm.install_dir, template_dir=template_dir, dry_run=fm.dry_run) | ||
| 1169 | + structured_fm.write_with_template("NPUStructuredNativeFunctions.h", "NPUStructuredNativeFunctions.h", lambda: { | ||
| 1170 | + "generated_comment": generated_comment, | ||
| 1171 | + "meta_includes": includes, | ||
| 1172 | + "structured_declarations": declarations, | ||
| 1173 | + }) | ||
| 1174 | + structured_fm.write_with_template("NPUStructuredNativeFunctions.cpp", "NPUStructuredNativeFunctions.cpp", lambda: { | ||
| 1175 | + "generated_comment": generated_comment, | ||
| 1176 | + "structured_definitions": definitions, | ||
| 1177 | + }) | ||
| 1178 | + | ||
| 1179 | + | ||
| 689 | def get_supported_grouped_native_functions( | 1180 | def get_supported_grouped_native_functions( |
| 690 | grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], | 1181 | grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], |
| 691 | backend_index: BackendIndex, | 1182 | backend_index: BackendIndex, |
| @@ -1155,6 +1646,12 @@ def run( | |||
| 1155 | None, | 1646 | None, |
| 1156 | ) | 1647 | ) |
| 1157 | 1648 | ||
| 1649 | + gen_npu_structured_native_functions( | ||
| 1650 | + fm, | ||
| 1651 | + grouped_native_functions, | ||
| 1652 | + backend_indices[backend_dispatch_key], | ||
| 1653 | + ) | ||
| 1654 | + | ||
| 1158 | gen_per_operator_headers( | 1655 | gen_per_operator_headers( |
| 1159 | fm, | 1656 | fm, |
| 1160 | ops_fm, | 1657 | ops_fm, |
| @@ -1289,6 +1786,51 @@ def run( | |||
| 1289 | def apply_torchgen_patch(): | 1786 | def apply_torchgen_patch(): |
| 1290 | dest.RegisterDispatchKey.gen_unstructured = gen_unstructured | 1787 | dest.RegisterDispatchKey.gen_unstructured = gen_unstructured |
| 1291 | dest.RegisterDispatchKey.gen_device_check = gen_device_check | 1788 | dest.RegisterDispatchKey.gen_device_check = gen_device_check |
| 1789 | + original_compute_native_function_declaration = dest.compute_native_function_declaration | ||
| 1790 | + original_gen_class_set_output_body = StructuredRegisterDispatchKey.gen_class_set_output_body | ||
| 1791 | + original_gen_one = StructuredRegisterDispatchKey.gen_one | ||
| 1792 | + | ||
| 1793 | + def compute_native_function_declaration(g, backend_index): | ||
| 1794 | + metadata = backend_index.get_kernel(g) | ||
| 1795 | + if isinstance(g, NativeFunctionsGroup) and metadata is not None and metadata.structured and backend_index.external: | ||
| 1796 | + return [] | ||
| 1797 | + return original_compute_native_function_declaration(g, backend_index) | ||
| 1798 | + | ||
| 1799 | + def gen_class_set_output_body(self, k, maybe_create_proxy): | ||
| 1800 | + if self.backend_index.dispatch_key.name not in ("NPU", "PrivateUse1"): | ||
| 1801 | + return original_gen_class_set_output_body(self, k, maybe_create_proxy) | ||
| 1802 | + | ||
| 1803 | + if maybe_create_proxy: | ||
| 1804 | + create_proxy = """ | ||
| 1805 | +auto maybe_proxy = maybe_create_proxy(out, sizes, strides, options); | ||
| 1806 | +if (C10_UNLIKELY(maybe_proxy.has_value())) { | ||
| 1807 | + proxy_outputs_[output_idx] = std::move(maybe_proxy).value(); | ||
| 1808 | +} | ||
| 1809 | +""" | ||
| 1810 | + else: | ||
| 1811 | + create_proxy = "" | ||
| 1812 | + | ||
| 1813 | + if k is SchemaKind.functional: | ||
| 1814 | + return "outputs_[output_idx] = create_out(sizes, strides, options);" | ||
| 1815 | + elif k is SchemaKind.inplace: | ||
| 1816 | + return f"""const auto& out = outputs_[output_idx].get(); | ||
| 1817 | +check_inplace(out, sizes, options); | ||
| 1818 | +{create_proxy}""" | ||
| 1819 | + elif k is SchemaKind.out: | ||
| 1820 | + return f"""const auto& out = outputs_[output_idx].get(); | ||
| 1821 | +resize_out(out, sizes, strides, options); | ||
| 1822 | +{create_proxy}""" | ||
| 1823 | + else: | ||
| 1824 | + raise AssertionError(f"{k} structured operators are currently not supported") | ||
| 1825 | + | ||
| 1826 | + def gen_one(self, f): | ||
| 1827 | + if self.backend_index.dispatch_key.name in ("NPU", "PrivateUse1"): | ||
| 1828 | + return gen_npu_structured_one(self, f) | ||
| 1829 | + return original_gen_one(self, f) | ||
| 1830 | + | ||
| 1831 | + dest.compute_native_function_declaration = compute_native_function_declaration | ||
| 1832 | + StructuredRegisterDispatchKey.gen_class_set_output_body = gen_class_set_output_body | ||
| 1833 | + StructuredRegisterDispatchKey.gen_one = gen_one | ||
| 1292 | # generate default arguments | 1834 | # generate default arguments |
| 1293 | JIT_TO_CPP_DEFAULT["contiguous_format"] = "c10::MemoryFormat::Contiguous" | 1835 | JIT_TO_CPP_DEFAULT["contiguous_format"] = "c10::MemoryFormat::Contiguous" |
| 1294 | dispatcher.arguments = native.arguments | 1836 | dispatcher.arguments = native.arguments |
| @@ -55,13 +55,7 @@ def _get_backend_index_for_npu( | |||
| 55 | ) -> BackendIndex | None: | 55 | ) -> BackendIndex | None: |
| 56 | for dk in NPU_DISPATCH_KEYS: | 56 | for dk in NPU_DISPATCH_KEYS: |
| 57 | if dk in backend_indices: | 57 | if dk in backend_indices: |
| 58 | - if backend_indices[dk].has_kernel(func) or ( | 58 | + if backend_indices[dk].has_kernel(func): |
H | |||
| 59 | - func.structured_delegate is not None | ||
| 60 | - and func.structured_delegate in structured_func_group_dict | ||
| 61 | - and backend_indices[dk].has_kernel( | ||
| 62 | - structured_func_group_dict[func.structured_delegate] | ||
| 63 | - ) | ||
| 64 | - ): | ||
| 65 | return backend_indices[dk] | 59 | return backend_indices[dk] |
| 66 | if DispatchKey.CompositeExplicitAutograd in backend_indices: | 60 | if DispatchKey.CompositeExplicitAutograd in backend_indices: |
| 67 | if backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func): | 61 | if backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func): |
| @@ -0,0 +1,13 @@ | |||
| 1 | +// ${generated_comment} | ||
H 这些模板不能复用的原因是什么 ![]() ![]() | |||
| 2 | + | ||
| 3 | + | ||
| 4 | + | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +namespace at_npu { | ||
| 8 | +namespace native { | ||
| 9 | + | ||
| 10 | +${structured_definitions} | ||
| 11 | + | ||
| 12 | +} // namespace native | ||
| 13 | +} // namespace at_npu | ||
| @@ -0,0 +1,16 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | +// ${generated_comment} | ||
| 4 | + | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +${meta_includes} | ||
| 9 | + | ||
| 10 | +namespace at_npu { | ||
| 11 | +namespace native { | ||
| 12 | + | ||
| 13 | +${structured_declarations} | ||
| 14 | + | ||
| 15 | +} // namespace native | ||
| 16 | +} // namespace at_npu | ||
| @@ -63,7 +63,16 @@ GLOBAL_OPAPI_INFO_CACHE = set() | |||
| 63 | GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE = set() | 63 | GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE = set() |
| 64 | 64 | ||
| 65 | CUSTOM_YAML_NAME = "npu_native_functions_by_codegen.yaml" | 65 | CUSTOM_YAML_NAME = "npu_native_functions_by_codegen.yaml" |
| 66 | -FIELDS_TO_USE = ["func", "tags", "dispatch", "device_check"] | 66 | +FIELDS_TO_USE = [ |
| 67 | + "func", | ||
| 68 | + "tags", | ||
| 69 | + "dispatch", | ||
| 70 | + "device_check", | ||
| 71 | + "structured", | ||
| 72 | + "structured_delegate", | ||
| 73 | + "structured_inherits", | ||
| 74 | + "precomputed", | ||
| 75 | +] | ||
| 67 | DEVICE_NOCHECK_SET = set() | 76 | DEVICE_NOCHECK_SET = set() |
| 68 | DEVICE_CHECK_NOTSUPPORT_TYPE = {"Tensor[]?"} | 77 | DEVICE_CHECK_NOTSUPPORT_TYPE = {"Tensor[]?"} |
| 69 | 78 | ||


这里新增的有没有可以复用上游社区的, 如果针对structured的代码比较多的话,是不是可以单独抽出一个文件,那样理解起来是不是简单些