已合并
support strutured meta #36643
support strutured meta #36643
已合并
maoyuanpeng1创建于 5月25日
5 个文件变更+597-23
@@ -22,6 +22,7 @@ import stat
22from collections import Counter, defaultdict, namedtuple22from collections import Counter, defaultdict, namedtuple
23from collections.abc import Callable, Sequence23from collections.abc import Callable, Sequence
24from dataclasses import dataclass24from dataclasses import dataclass
25+from typing import Optional, List, Set, Union
25 26 
26import torchgen27import torchgen
27import torchgen.api.dispatcher as dispatcher28import torchgen.api.dispatcher as dispatcher
@@ -30,6 +31,11 @@ import torchgen.dest as dest
30import yaml31import yaml
31from torchgen.api.cpp import JIT_TO_CPP_DEFAULT32from torchgen.api.cpp import JIT_TO_CPP_DEFAULT
32from torchgen.code_template import CodeTemplate33from torchgen.code_template import CodeTemplate
34+from torchgen.model import (
35+ BackendIndex,
36+ SchemaKind,
37+ TensorOptionsArguments,
38+)
33from torchgen.gen import (39from 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 (
49from torchgen.native_function_generation import add_generated_native_functions55from torchgen.native_function_generation import add_generated_native_functions
50from torchgen.selective_build.selector import SelectiveBuilder56from torchgen.selective_build.selector import SelectiveBuilder
51from torchgen.utils import concatMap, context, NamespaceHelper, Target57from 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 
53from torchnpugen.custom_functions import (75from torchnpugen.custom_functions import (
54 gen_custom_functions_dispatch,76 gen_custom_functions_dispatch,
@@ -88,6 +110,9 @@ from torchnpugen.utils import (
88 110 
89torchgen.model.dispatch_keys.append(torchgen.model.DispatchKey.AutogradPrivateUse1)111torchgen.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.
93def create_backend_index(118def 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] = m142 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+ 
262def parse_backend_yaml(366def 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] = f376 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] = func379 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_list497 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_namespace554+ 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+ 
561def gen_dispatcher_registrations(748def 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.h781 # 根据环境变量决定是否包含FormatHelper.h
@@ -636,7 +824,7 @@ $dispatch_registrations_body
636 backend_index, per_operator_headers=False, rocm=False824 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
Hhuangyunlong5月28日

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

likedislike
maoyuanpeng1
5月28日 评论:
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+ 
689def get_supported_grouped_native_functions(1180def 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(
1289def apply_torchgen_patch():1786def apply_torchgen_patch():
1290 dest.RegisterDispatchKey.gen_unstructured = gen_unstructured1787 dest.RegisterDispatchKey.gen_unstructured = gen_unstructured
1291 dest.RegisterDispatchKey.gen_device_check = gen_device_check1788 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 arguments1834 # 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.arguments1836 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
Hhuangyunlong5月27日

这里为啥要删除

likedislike
maoyuanpeng1
5月27日 评论:
huangyunlong
5月28日 评论:
maoyuanpeng1
5月28日 评论:
huangyunlong
5月28日 评论:
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
Hhuangyunlong5月27日

这些模板不能复用的原因是什么

likedislike
maoyuanpeng1
5月27日 评论:
2+ 
3+#include "torch_npu/csrc/aten/NPUStructuredNativeFunctions.h"
4+ 
5+#include "op_plugin/OpInterface.h"
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+#pragma once
2+ 
3+// ${generated_comment}
4+ 
5+#include <ATen/Tensor.h>
6+#include <ATen/TensorIterator.h>
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()
63GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE = set()63GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE = set()
64 64 
65CUSTOM_YAML_NAME = "npu_native_functions_by_codegen.yaml"65CUSTOM_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+]
67DEVICE_NOCHECK_SET = set()76DEVICE_NOCHECK_SET = set()
68DEVICE_CHECK_NOTSUPPORT_TYPE = {"Tensor[]?"}77DEVICE_CHECK_NOTSUPPORT_TYPE = {"Tensor[]?"}
69 78