已合并
[sync] PR-35289: [feat] 去掉codegen随机性,保证生成的文件内容固定,使ccache能命中缓存 #35438
ascend-robot创建于 5月12日
[sync] PR-35289: [feat] 去掉codegen随机性,保证生成的文件内容固定,使ccache能命中缓存 #35438
已合并
共 3 个文件变更+748-437
| @@ -191,6 +191,9 @@ TAGS | |||
| 191 | .clangd/ | 191 | .clangd/ |
| 192 | .cache/ | 192 | .cache/ |
| 193 | 193 | ||
| 194 | +# lintrunner binary cache | ||
| 195 | +.lintbin/ | ||
| 196 | + | ||
| 194 | # bazel symlinks | 197 | # bazel symlinks |
| 195 | bazel-* | 198 | bazel-* |
| 196 | 199 | ||
| @@ -13,43 +13,49 @@ | |||
| 13 | # See the License for the specific language governing permissions and | 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. | 14 | # limitations under the License. |
| 15 | 15 | ||
| 16 | +import itertools | ||
| 16 | import os | 17 | import os |
| 17 | -import sys | ||
| 18 | import stat | 18 | import stat |
| 19 | +import sys | ||
| 19 | import traceback | 20 | import traceback |
| 20 | import warnings | 21 | import warnings |
| 21 | -import itertools | ||
| 22 | -from typing import List, Optional, Set, Dict, Union, Sequence, Iterator, Tuple | ||
| 23 | from collections import defaultdict | 22 | from collections import defaultdict |
| 23 | +from collections.abc import Iterator, Sequence | ||
| 24 | from dataclasses import dataclass | 24 | from dataclasses import dataclass |
| 25 | -import yaml | 25 | +from typing import Optional |
| 26 | +from typing_extensions import assert_never | ||
| 26 | 27 | ||
| 27 | import torch | 28 | import torch |
| 28 | -from torchgen.api.types.signatures import NativeSignature, DispatcherSignature | 29 | +import yaml |
| 29 | -from torchgen.context import native_function_manager | ||
| 30 | -from torchgen.code_template import CodeTemplate | ||
| 31 | -from torchgen.model import ( | ||
| 32 | - Arguments, | ||
| 33 | - BackendIndex, | ||
| 34 | - BackendMetadata, | ||
| 35 | - DispatchKey, | ||
| 36 | - is_cuda_dispatch_key, | ||
| 37 | - NativeFunction, | ||
| 38 | - NativeFunctionsGroup, | ||
| 39 | - FunctionSchema, | ||
| 40 | - OperatorName, | ||
| 41 | - TensorOptionsArguments, | ||
| 42 | - SchemaKind, | ||
| 43 | - DeviceCheckType, | ||
| 44 | - Argument, | ||
| 45 | -) | ||
| 46 | -from torchgen.native_function_generation import pre_group_native_functions | ||
| 47 | -from torchgen.utils import concatMap | ||
| 48 | from torchgen.api import cpp | 30 | from torchgen.api import cpp |
| 49 | from torchgen.api.translate import translate | 31 | from torchgen.api.translate import translate |
| 50 | from torchgen.api.types import Binding, CppSignatureGroup, kernel_signature | 32 | from torchgen.api.types import Binding, CppSignatureGroup, kernel_signature |
| 51 | -from torchgen.utils import Target | 33 | +from torchgen.api.types.signatures import ( |
| 34 | + CppSignature, | ||
| 35 | + DispatcherSignature, | ||
| 36 | + NativeSignature, | ||
| 37 | +) | ||
| 38 | +from torchgen.code_template import CodeTemplate | ||
| 39 | +from torchgen.context import native_function_manager | ||
| 52 | from torchgen.dest.register_dispatch_key import RegisterDispatchKey | 40 | from torchgen.dest.register_dispatch_key import RegisterDispatchKey |
| 41 | +from torchgen.model import ( | ||
| 42 | + Argument, | ||
| 43 | + Arguments, | ||
| 44 | + BackendIndex, | ||
| 45 | + BackendMetadata, | ||
| 46 | + DeviceCheckType, | ||
| 47 | + DispatchKey, | ||
| 48 | + FunctionSchema, | ||
| 49 | + NativeFunction, | ||
| 50 | + NativeFunctionsGroup, | ||
| 51 | + OperatorName, | ||
| 52 | + SchemaKind, | ||
| 53 | + SelfArgument, | ||
| 54 | + TensorOptionsArguments, | ||
| 55 | +) | ||
| 56 | +from torchgen.native_function_generation import pre_group_native_functions | ||
| 57 | +from torchgen.utils import concatMap, Target | ||
| 58 | + | ||
| 53 | 59 | ||
| 54 | GLOBAL_STRUCTURED_OP_INFO_CACHE = defaultdict(str) | 60 | GLOBAL_STRUCTURED_OP_INFO_CACHE = defaultdict(str) |
| 55 | GLOBAL_OPAPI_INFO_CACHE = set() | 61 | GLOBAL_OPAPI_INFO_CACHE = set() |
| @@ -62,7 +68,6 @@ DEVICE_CHECK_NOTSUPPORT_TYPE = {"Tensor[]?"} | |||
| 62 | 68 | ||
| 63 | 69 | ||
| 64 | class PathManager: | 70 | class PathManager: |
| 65 | - | ||
| 66 | 71 | ||
| 67 | def check_path_owner_consistent(cls, path: str): | 72 | def check_path_owner_consistent(cls, path: str): |
| 68 | """ | 73 | """ |
| @@ -106,15 +111,16 @@ class PathManager: | |||
| 106 | os.remove(path) | 111 | os.remove(path) |
| 107 | 112 | ||
| 108 | 113 | ||
| 109 | -def parse_npu_yaml(custom_path: str) -> Dict: | 114 | +def parse_npu_yaml(custom_path: str) -> dict: |
| 110 | if not os.path.exists(custom_path): | 115 | if not os.path.exists(custom_path): |
| 111 | return {} | 116 | return {} |
| 112 | from io import StringIO | 117 | from io import StringIO |
| 118 | + | ||
| 113 | f_str = StringIO() | 119 | f_str = StringIO() |
| 114 | PathManager.check_directory_path_readable(custom_path) | 120 | PathManager.check_directory_path_readable(custom_path) |
| 115 | - with open(custom_path, 'r') as f: | 121 | + with open(custom_path) as f: |
| 116 | for line in f: | 122 | for line in f: |
| 117 | - if ':' not in line: | 123 | + if ":" not in line: |
| 118 | continue | 124 | continue |
| 119 | f_str.write(line) | 125 | f_str.write(line) |
| 120 | 126 | ||
| @@ -132,6 +138,7 @@ def merge_yaml(base_data, additional_data): | |||
| 132 | return map_dict[x] | 138 | return map_dict[x] |
| 133 | else: | 139 | else: |
| 134 | return x | 140 | return x |
| 141 | + | ||
| 135 | if isinstance(base_data, dict): | 142 | if isinstance(base_data, dict): |
| 136 | for key, value in additional_data.items(): | 143 | for key, value in additional_data.items(): |
| 137 | if key_map(key) not in base_data: | 144 | if key_map(key) not in base_data: |
| @@ -147,18 +154,23 @@ def merge_yaml(base_data, additional_data): | |||
| 147 | 154 | ||
| 148 | def merge_custom_yaml(pta_path, op_plugin_path): | 155 | def merge_custom_yaml(pta_path, op_plugin_path): |
| 149 | PathManager.check_directory_path_readable(pta_path) | 156 | PathManager.check_directory_path_readable(pta_path) |
| 150 | - with open(pta_path, 'r') as pta_file: | 157 | + with open(pta_path) as pta_file: |
| 151 | pta_es = yaml.safe_load(pta_file) | 158 | pta_es = yaml.safe_load(pta_file) |
| 152 | PathManager.check_directory_path_readable(op_plugin_path) | 159 | PathManager.check_directory_path_readable(op_plugin_path) |
| 153 | - with open(op_plugin_path, 'r') as op_plugin_file: | 160 | + with open(op_plugin_path) as op_plugin_file: |
| 154 | op_es = yaml.safe_load(op_plugin_file) | 161 | op_es = yaml.safe_load(op_plugin_file) |
| 155 | 162 | ||
| 156 | merged_yaml = merge_yaml(pta_es, op_es) | 163 | merged_yaml = merge_yaml(pta_es, op_es) |
| 157 | merged_yaml_path = gen_custom_yaml_path(pta_path) | 164 | merged_yaml_path = gen_custom_yaml_path(pta_path) |
| 158 | PathManager.remove_path_safety(merged_yaml_path) | 165 | PathManager.remove_path_safety(merged_yaml_path) |
| 159 | - with os.fdopen(os.open(merged_yaml_path, os.O_RDWR | os.O_CREAT, stat.S_IWUSR | stat.S_IRUSR), "w") as outfile: | 166 | + with os.fdopen( |
| 167 | + os.open(merged_yaml_path, os.O_RDWR | os.O_CREAT, stat.S_IWUSR | stat.S_IRUSR), | ||
| 168 | + "w", | ||
| 169 | + ) as outfile: | ||
| 160 | yaml.dump(merged_yaml, outfile, default_flow_style=False, width=float("inf")) | 170 | yaml.dump(merged_yaml, outfile, default_flow_style=False, width=float("inf")) |
| 161 | - os.chmod(merged_yaml_path, stat.S_IRUSR | stat.S_IEXEC | stat.S_IRGRP | stat.S_IXGRP) | 171 | + os.chmod( |
| 172 | + merged_yaml_path, stat.S_IRUSR | stat.S_IEXEC | stat.S_IRGRP | stat.S_IXGRP | ||
| 173 | + ) | ||
| 162 | return merged_yaml | 174 | return merged_yaml |
| 163 | 175 | ||
| 164 | 176 | ||
| @@ -166,47 +178,55 @@ def field_tag(custom_es): | |||
| 166 | for i, es in enumerate(custom_es): | 178 | for i, es in enumerate(custom_es): |
| 167 | if not isinstance(es, dict): | 179 | if not isinstance(es, dict): |
| 168 | continue | 180 | continue |
| 169 | - custom_es[i] = {key: custom_es[i][key] for key in FIELDS_TO_USE if key in custom_es[i]} | 181 | + custom_es[i] = { |
| 182 | + key: custom_es[i][key] for key in FIELDS_TO_USE if key in custom_es[i] | ||
| 183 | + } | ||
| 170 | return custom_es | 184 | return custom_es |
| 171 | 185 | ||
| 172 | 186 | ||
| 173 | def filt_exposed_api(custom_path: str): | 187 | def filt_exposed_api(custom_path: str): |
| 174 | source_es = parse_npu_yaml(custom_path) | 188 | source_es = parse_npu_yaml(custom_path) |
| 175 | - custom_es = source_es.get('custom', []) + source_es.get('custom_autograd', []) | 189 | + custom_es = source_es.get("custom", []) + source_es.get("custom_autograd", []) |
| 176 | exposed_set = set() | 190 | exposed_set = set() |
| 177 | for es in custom_es: | 191 | for es in custom_es: |
| 178 | - if es.get('exposed', False): | 192 | + if es.get("exposed", False): |
| 179 | - exposed_set.add(es.get('func').split('(')[0].split('.')[0]) | 193 | + exposed_set.add(es.get("func").split("(")[0].split(".")[0]) |
| 180 | return list(exposed_set) | 194 | return list(exposed_set) |
| 181 | 195 | ||
| 182 | 196 | ||
| 183 | -# Different implements of ops from origin torch. | 197 | +# Different implements of ops from origin torch. |
| 184 | # Native ops with dispatchkey CompositeImplicitAutograd but implemented as a kernel op in pta | 198 | # Native ops with dispatchkey CompositeImplicitAutograd but implemented as a kernel op in pta |
| 185 | COMPOSITEIMPLICITAUTOGRAD_EXCEPT_LIST = [ | 199 | COMPOSITEIMPLICITAUTOGRAD_EXCEPT_LIST = [ |
| 186 | - 'isclose', | 200 | + "isclose", |
| 187 | - 'isfinite', | 201 | + "isfinite", |
| 188 | ] | 202 | ] |
| 189 | 203 | ||
| 190 | 204 | ||
| 191 | -def filt_dispath_key(api_name: str) -> List: | 205 | +def filt_dispath_key(api_name: str) -> list: |
| 192 | dispatch_dump = torch._C._dispatch_dump(f"aten::{api_name}") | 206 | dispatch_dump = torch._C._dispatch_dump(f"aten::{api_name}") |
| 193 | - return [dump.split(":")[0] for dump in dispatch_dump.split('\n')] | 207 | + return [dump.split(":")[0] for dump in dispatch_dump.split("\n")] |
| 194 | 208 | ||
| 195 | 209 | ||
| 196 | def filt_compositeimplicitautograd_api(native_yaml_path, npu_supported): | 210 | def filt_compositeimplicitautograd_api(native_yaml_path, npu_supported): |
| 197 | PathManager.check_directory_path_readable(native_yaml_path) | 211 | PathManager.check_directory_path_readable(native_yaml_path) |
| 198 | - with open(native_yaml_path, 'r') as f: | 212 | + with open(native_yaml_path) as f: |
| 199 | es = yaml.safe_load(f) | 213 | es = yaml.safe_load(f) |
| 200 | 214 | ||
| 201 | from torchnpugen.autograd.utils import TORCH_AUTOGRAD_FUNCTION | 215 | from torchnpugen.autograd.utils import TORCH_AUTOGRAD_FUNCTION |
| 216 | + | ||
| 202 | supported_autograd = [] | 217 | supported_autograd = [] |
| 203 | for e in es: | 218 | for e in es: |
| 204 | - api_name = e['func'].split('(')[0] | 219 | + api_name = e["func"].split("(")[0] |
| 205 | dispatch_keys = filt_dispath_key(api_name) | 220 | dispatch_keys = filt_dispath_key(api_name) |
| 206 | - is_compositekey = "CompositeImplicitAutograd[alias]" in dispatch_keys and \ | 221 | + is_compositekey = ( |
| 207 | - "Autograd[alias]" not in dispatch_keys and \ | 222 | + "CompositeImplicitAutograd[alias]" in dispatch_keys |
| 208 | - api_name not in TORCH_AUTOGRAD_FUNCTION | 223 | + and "Autograd[alias]" not in dispatch_keys |
| 209 | - is_npu_api = api_name in npu_supported and api_name not in COMPOSITEIMPLICITAUTOGRAD_EXCEPT_LIST | 224 | + and api_name not in TORCH_AUTOGRAD_FUNCTION |
| 225 | + ) | ||
| 226 | + is_npu_api = ( | ||
| 227 | + api_name in npu_supported | ||
| 228 | + and api_name not in COMPOSITEIMPLICITAUTOGRAD_EXCEPT_LIST | ||
| 229 | + ) | ||
| 210 | if is_npu_api and is_compositekey: | 230 | if is_npu_api and is_compositekey: |
| 211 | supported_autograd.append(api_name) | 231 | supported_autograd.append(api_name) |
| 212 | return supported_autograd | 232 | return supported_autograd |
| @@ -234,6 +254,7 @@ def get_torchgen_dir(): | |||
| 234 | # get path of torchgen, then get tags.yaml and native_functions.yaml | 254 | # get path of torchgen, then get tags.yaml and native_functions.yaml |
| 235 | try: | 255 | try: |
| 236 | import torchgen | 256 | import torchgen |
| 257 | + | ||
| 237 | return os.path.dirname(os.path.realpath(torchgen.__file__)) | 258 | return os.path.dirname(os.path.realpath(torchgen.__file__)) |
| 238 | except Exception: | 259 | except Exception: |
| 239 | _, _, exc_traceback = sys.exc_info() | 260 | _, _, exc_traceback = sys.exc_info() |
| @@ -241,34 +262,36 @@ def get_torchgen_dir(): | |||
| 241 | return os.path.dirname(frame_summary.filename) | 262 | return os.path.dirname(frame_summary.filename) |
| 242 | 263 | ||
| 243 | 264 | ||
| 244 | -def gen_op_hook_post_code(sig: Union[NativeSignature, DispatcherSignature]) -> Tuple[str, str]: | 265 | +def gen_op_hook_post_code( |
| 266 | + sig: NativeSignature | DispatcherSignature, | ||
| 267 | +) -> tuple[str, str]: | ||
| 245 | res_code: str = None | 268 | res_code: str = None |
| 246 | return_code: str = None | 269 | return_code: str = None |
| 247 | 270 | ||
| 248 | if sig.returns_type().cpp_type() == "void": | 271 | if sig.returns_type().cpp_type() == "void": |
| 249 | res_code = "" | 272 | res_code = "" |
| 250 | - return_code = f"""at_npu::native::OpHook::GetInstance().PostHook(); | 273 | + return_code = """at_npu::native::OpHook::GetInstance().PostHook(); |
| 251 | return;""" | 274 | return;""" |
| 252 | else: | 275 | else: |
| 253 | res_code = f"""{sig.returns_type().cpp_type()} res = """ | 276 | res_code = f"""{sig.returns_type().cpp_type()} res = """ |
| 254 | - return_code = f"""at_npu::native::OpHook::GetInstance().PostHook(res); | 277 | + return_code = """at_npu::native::OpHook::GetInstance().PostHook(res); |
| 255 | return res;""" | 278 | return res;""" |
| 256 | 279 | ||
| 257 | return res_code, return_code | 280 | return res_code, return_code |
| 258 | 281 | ||
| 259 | 282 | ||
| 260 | OVERWRITE_API_LIST = [ | 283 | OVERWRITE_API_LIST = [ |
| 261 | - 'matmul', | 284 | + "matmul", |
| 262 | - 'matmul.out', | 285 | + "matmul.out", |
| 263 | - 'matmul_backward', | 286 | + "matmul_backward", |
| 264 | - 'matmul_double_backward', | 287 | + "matmul_double_backward", |
| 265 | ] | 288 | ] |
| 266 | 289 | ||
| 267 | 290 | ||
| 268 | # This function is to add profiler information for each operator, which is later extended in the official | 291 | # This function is to add profiler information for each operator, which is later extended in the official |
| 269 | def gen_unstructured( | 292 | def gen_unstructured( |
| 270 | - self, f: NativeFunction, g: Optional[NativeFunctionsGroup] = None | 293 | + self, f: NativeFunction, g: NativeFunctionsGroup | None = None |
| 271 | -) -> Optional[str]: | 294 | +) -> str | None: |
| 272 | with native_function_manager(f): | 295 | with native_function_manager(f): |
| 273 | inplace_meta = False | 296 | inplace_meta = False |
| 274 | gets_out_inplace_wrapper = False | 297 | gets_out_inplace_wrapper = False |
| @@ -291,7 +314,9 @@ def gen_unstructured( | |||
| 291 | args_str = ", ".join(a.defn() for a in args) | 314 | args_str = ", ".join(a.defn() for a in args) |
| 292 | 315 | ||
| 293 | op_name = str(f.func.name.name) | 316 | op_name = str(f.func.name.name) |
| 294 | - force_aclnn = f"at_npu::native::ForceAclnn::GetInstance().IsForceAclnnOp(\"{op_name}\")" | 317 | + force_aclnn = ( |
| 318 | + f'at_npu::native::ForceAclnn::GetInstance().IsForceAclnnOp("{op_name}")' | ||
| 319 | + ) | ||
| 295 | # See Note [Direct dispatch bindings] | 320 | # See Note [Direct dispatch bindings] |
| 296 | cpp_sig_group = CppSignatureGroup.from_native_function( | 321 | cpp_sig_group = CppSignatureGroup.from_native_function( |
| 297 | f, method=False, fallback_binding=False | 322 | f, method=False, fallback_binding=False |
| @@ -307,7 +332,7 @@ def gen_unstructured( | |||
| 307 | def generate_defn(cpp_sig: CppSignature) -> str: | 332 | def generate_defn(cpp_sig: CppSignature) -> str: |
| 308 | return f""" | 333 | return f""" |
| 309 | {cpp_sig.defn()} {{ | 334 | {cpp_sig.defn()} {{ |
| 310 | -return {sig.name()}({', '.join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))}); | 335 | +return {sig.name()}({", ".join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))}); |
| 311 | }} | 336 | }} |
| 312 | """ | 337 | """ |
| 313 | 338 | ||
| @@ -352,7 +377,10 @@ return {self_arg_name}; | |||
| 352 | 377 | ||
| 353 | device_check = " // No device check\n" | 378 | device_check = " // No device check\n" |
| 354 | # Backends that require device guards presumably also require device checks. | 379 | # Backends that require device guards presumably also require device checks. |
| 355 | - if self.backend_index.device_guard and str(f.func.name) not in DEVICE_NOCHECK_SET: | 380 | + if ( |
| 381 | + self.backend_index.device_guard | ||
| 382 | + and str(f.func.name) not in DEVICE_NOCHECK_SET | ||
| 383 | + ): | ||
| 356 | device_check_args = itertools.chain( | 384 | device_check_args = itertools.chain( |
| 357 | f.func.arguments.out, f.func.arguments.flat_positional | 385 | f.func.arguments.out, f.func.arguments.flat_positional |
| 358 | ) | 386 | ) |
| @@ -390,27 +418,37 @@ torch_npu::profiler::NPURecordFunction guard; | |||
| 390 | for a in candidate_args: | 418 | for a in candidate_args: |
| 391 | if a.type.is_tensor_like(): | 419 | if a.type.is_tensor_like(): |
| 392 | candidate_tensor_args.append(f"{a.name}") | 420 | candidate_tensor_args.append(f"{a.name}") |
| 393 | - | 421 | + # Sort for deterministic output |
| 394 | - candidate_tensor_args = list(set(candidate_tensor_args)) | 422 | + candidate_tensor_args = sorted(set(candidate_tensor_args)) |
| 395 | unsafe_tensor_check = """// No data check.""" | 423 | unsafe_tensor_check = """// No data check.""" |
| 396 | if len(candidate_tensor_args) > 0: | 424 | if len(candidate_tensor_args) > 0: |
| 397 | unsafe_tensor_check = """ | 425 | unsafe_tensor_check = """ |
| 398 | if (c10_npu::get_npu_data_unsafe_flag()) {""" | 426 | if (c10_npu::get_npu_data_unsafe_flag()) {""" |
| 399 | 427 | ||
| 400 | if name in ["wrapper_NPU__copy_", "wrapper_NPU___foreach_copy_"]: | 428 | if name in ["wrapper_NPU__copy_", "wrapper_NPU___foreach_copy_"]: |
| 401 | - tensor_arg = candidate_tensor_args[0] + ", " +\ | 429 | + tensor_arg = ( |
| 402 | - candidate_tensor_args[1] | 430 | + candidate_tensor_args[0] + ", " + candidate_tensor_args[1] |
| 403 | - unsafe_tensor_check = unsafe_tensor_check + f""" | 431 | + ) |
| 432 | + unsafe_tensor_check = ( | ||
| 433 | + unsafe_tensor_check | ||
| 434 | + + """ | ||
| 404 | c10_npu::check_and_update_npu_tensor_for_copy(self, src);""" | 435 | c10_npu::check_and_update_npu_tensor_for_copy(self, src);""" |
| 436 | + ) | ||
| 405 | else: | 437 | else: |
| 406 | for tensor_arg in candidate_tensor_args: | 438 | for tensor_arg in candidate_tensor_args: |
| 407 | - unsafe_tensor_check = unsafe_tensor_check + f""" | 439 | + unsafe_tensor_check = ( |
| 440 | + unsafe_tensor_check | ||
| 441 | + + f""" | ||
| 408 | c10_npu::check_npu_tensor_is_safe({tensor_arg});""" | 442 | c10_npu::check_npu_tensor_is_safe({tensor_arg});""" |
| 409 | - | 443 | + ) |
| 444 | + | ||
| 410 | if len(candidate_tensor_args) > 0: | 445 | if len(candidate_tensor_args) > 0: |
| 411 | - unsafe_tensor_check = unsafe_tensor_check + """ | 446 | + unsafe_tensor_check = ( |
| 447 | + unsafe_tensor_check | ||
| 448 | + + """ | ||
| 412 | } | 449 | } |
| 413 | """ | 450 | """ |
| 451 | + ) | ||
| 414 | candidate_args = itertools.chain( | 452 | candidate_args = itertools.chain( |
| 415 | self_arg, | 453 | self_arg, |
| 416 | f.func.arguments.out, | 454 | f.func.arguments.out, |
| @@ -419,11 +457,7 @@ if (c10_npu::get_npu_data_unsafe_flag()) {""" | |||
| 419 | 457 | ||
| 420 | # Only tensor like arguments are eligible | 458 | # Only tensor like arguments are eligible |
| 421 | device_of = next( | 459 | device_of = next( |
| 422 | - ( | 460 | + (f"{a.name}" for a in candidate_args if a.type.is_tensor_like()), |
| 423 | - f"{a.name}" | ||
| 424 | - for a in candidate_args | ||
| 425 | - if a.type.is_tensor_like() | ||
| 426 | - ), | ||
| 427 | None, | 461 | None, |
| 428 | ) | 462 | ) |
| 429 | if has_tensor_options and device_of is not None: | 463 | if has_tensor_options and device_of is not None: |
| @@ -448,7 +482,9 @@ const DeviceGuard device_guard(device_or_default(device));""" | |||
| 448 | impl_name = f"op_plugin::{GLOBAL_STRUCTURED_OP_INFO_CACHE[op_key]}" | 482 | impl_name = f"op_plugin::{GLOBAL_STRUCTURED_OP_INFO_CACHE[op_key]}" |
| 449 | 483 | ||
| 450 | # for op_hook_check | 484 | # for op_hook_check |
| 451 | - res_of_op_hook_post_code, return_of_op_hook_post_code = gen_op_hook_post_code(sig) | 485 | + res_of_op_hook_post_code, return_of_op_hook_post_code = ( |
| 486 | + gen_op_hook_post_code(sig) | ||
| 487 | + ) | ||
| 452 | 488 | ||
| 453 | args_exprs_str_list = [ | 489 | args_exprs_str_list = [ |
| 454 | e.expr | 490 | e.expr |
| @@ -462,7 +498,9 @@ const DeviceGuard device_guard(device_or_default(device));""" | |||
| 462 | if any(lvalue_ref in args_expr for lvalue_ref in lvalue_ref_list): | 498 | if any(lvalue_ref in args_expr for lvalue_ref in lvalue_ref_list): |
| 463 | auto_lvalue += f""" | 499 | auto_lvalue += f""" |
| 464 | auto {kernel_sig.arguments()[idx].name}_tmp = {args_expr};""" | 500 | auto {kernel_sig.arguments()[idx].name}_tmp = {args_expr};""" |
| 465 | - args_exprs_str_list[idx] = f"""{kernel_sig.arguments()[idx].name}_tmp""" | 501 | + args_exprs_str_list[idx] = ( |
| 502 | + f"""{kernel_sig.arguments()[idx].name}_tmp""" | ||
| 503 | + ) | ||
| 466 | 504 | ||
| 467 | args_exprs_str_for_op_hook = ", ".join(e for e in args_exprs_str_list) | 505 | args_exprs_str_for_op_hook = ", ".join(e for e in args_exprs_str_list) |
| 468 | 506 | ||
| @@ -475,7 +513,9 @@ const DeviceGuard device_guard(device_or_default(device));""" | |||
| 475 | tensor_check_list = [] | 513 | tensor_check_list = [] |
| 476 | for a in args: | 514 | for a in args: |
| 477 | if a.argument.type.is_tensor_like(): | 515 | if a.argument.type.is_tensor_like(): |
| 478 | - tensor_check_list.append(f"at_npu::native::FormatHelper::IsOpInputBaseFormat({a.name})") | 516 | + tensor_check_list.append( |
| 517 | + f"at_npu::native::FormatHelper::IsOpInputBaseFormat({a.name})" | ||
| 518 | + ) | ||
| 479 | if tensor_check_list: | 519 | if tensor_check_list: |
| 480 | tensor_check_str = f" && {' && '.join(tensor_check_list)}" | 520 | tensor_check_str = f" && {' && '.join(tensor_check_list)}" |
| 481 | 521 | ||
| @@ -567,16 +607,19 @@ namespace {{ | |||
| 567 | else: | 607 | else: |
| 568 | payload = f"TORCH_FN({name})" | 608 | payload = f"TORCH_FN({name})" |
| 569 | if f"{f.func.name}" in OVERWRITE_API_LIST: | 609 | if f"{f.func.name}" in OVERWRITE_API_LIST: |
| 570 | - return f"""if (std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") == nullptr || std::string(std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL")) != "1") {{ | 610 | + return ( |
| 571 | - m.impl("{f.func.name}", {payload}); | 611 | + f'if (std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") == nullptr || ' |
| 572 | -}}""" | 612 | + f'std::string(std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL")) != "1") {{\n' |
| 613 | + f' m.impl("{f.func.name}", {payload});\n' | ||
| 614 | + f"}}" | ||
| 615 | + ) | ||
| 573 | return f'm.impl("{f.func.name}",\n{payload});\n' | 616 | return f'm.impl("{f.func.name}",\n{payload});\n' |
| 574 | else: | 617 | else: |
| 575 | assert_never(self.target) | 618 | assert_never(self.target) |
| 576 | 619 | ||
| 577 | 620 | ||
| 578 | def gen_device_check( | 621 | def gen_device_check( |
| 579 | - type: DeviceCheckType, args: List[Argument], method_name: str | 622 | + type: DeviceCheckType, args: list[Argument], method_name: str |
| 580 | ) -> str: | 623 | ) -> str: |
| 581 | if type == DeviceCheckType.NoCheck: | 624 | if type == DeviceCheckType.NoCheck: |
| 582 | return " // No device check\n" | 625 | return " // No device check\n" |
| @@ -585,9 +628,14 @@ def gen_device_check( | |||
| 585 | device_check += "(void)common_device; // Suppress unused variable warning\n" | 628 | device_check += "(void)common_device; // Suppress unused variable warning\n" |
| 586 | for arg in args: | 629 | for arg in args: |
| 587 | # Only tensor like arguments are eligible | 630 | # Only tensor like arguments are eligible |
| 588 | - if arg.type.is_tensor_like() and str(arg.type) not in DEVICE_CHECK_NOTSUPPORT_TYPE: | 631 | + if ( |
| 589 | - device_check += \ | 632 | + arg.type.is_tensor_like() |
| 590 | -f"""c10::impl::check_and_update_common_device(common_device, {arg.name}, "{method_name}", "{arg.name}");\n""" | 633 | + and str(arg.type) not in DEVICE_CHECK_NOTSUPPORT_TYPE |
| 634 | + ): | ||
| 635 | + device_check += ( | ||
| 636 | + f"c10::impl::check_and_update_common_device(" | ||
| 637 | + f'common_device, {arg.name}, "{method_name}", "{arg.name}");\n' | ||
| 638 | + ) | ||
| 591 | return device_check | 639 | return device_check |
| 592 | 640 | ||
| 593 | 641 | ||
| @@ -597,9 +645,9 @@ def arguments( | |||
| 597 | faithful: bool, | 645 | faithful: bool, |
| 598 | symint: bool = False, | 646 | symint: bool = False, |
| 599 | method: bool, | 647 | method: bool, |
| 600 | - cpp_no_default_args: Set[str], | 648 | + cpp_no_default_args: set[str], |
| 601 | -) -> List[Binding]: | 649 | +) -> list[Binding]: |
| 602 | - args: List[Union[Argument, TensorOptionsArguments, SelfArgument]] = [] | 650 | + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] |
| 603 | args.extend(arguments.non_out) | 651 | args.extend(arguments.non_out) |
| 604 | args.extend(arguments.out) | 652 | args.extend(arguments.out) |
| 605 | result = [] | 653 | result = [] |
| @@ -621,35 +669,41 @@ def arguments( | |||
| 621 | 669 | ||
| 622 | def add_header_to_template_file(): | 670 | def add_header_to_template_file(): |
| 623 | torchgen_path = get_torchgen_dir() | 671 | torchgen_path = get_torchgen_dir() |
| 624 | - template_dir = os.path.join(torchgen_path, "packaged/ATen/templates/DispatchKeyNativeFunctions.h") | 672 | + template_dir = os.path.join( |
| 673 | + torchgen_path, "packaged/ATen/templates/DispatchKeyNativeFunctions.h" | ||
| 674 | + ) | ||
| 625 | PathManager.check_directory_path_readable(template_dir) | 675 | PathManager.check_directory_path_readable(template_dir) |
| 626 | - with open(template_dir, "r") as file: | 676 | + with open(template_dir) as file: |
| 627 | template_content = file.read() | 677 | template_content = file.read() |
| 628 | if "#include <ATen/ATen.h>" not in template_content: | 678 | if "#include <ATen/ATen.h>" not in template_content: |
| 629 | - template_content = template_content.replace("#include <ATen/Tensor.h>", | 679 | + template_content = template_content.replace( |
| 630 | - "#include <ATen/Tensor.h>\n#include <ATen/ATen.h>") | 680 | + "#include <ATen/Tensor.h>", |
| 631 | - with os.fdopen(os.open(template_dir, os.O_WRONLY, stat.S_IWUSR | stat.S_IRUSR), "w") as file: | 681 | + "#include <ATen/Tensor.h>\n#include <ATen/ATen.h>", |
| 682 | + ) | ||
| 683 | + with os.fdopen( | ||
| 684 | + os.open(template_dir, os.O_WRONLY, stat.S_IWUSR | stat.S_IRUSR), "w" | ||
| 685 | + ) as file: | ||
| 632 | file.write(template_content) | 686 | file.write(template_content) |
| 633 | 687 | ||
| 634 | 688 | ||
| 635 | def enable_opplugin() -> bool: | 689 | def enable_opplugin() -> bool: |
| 636 | # enable op_plugin, if path of third_party/op-plugin is valid. | 690 | # enable op_plugin, if path of third_party/op-plugin is valid. |
| 637 | - | 691 | + |
| 638 | - env_aclnn_extension_switch = os.getenv('ACLNN_EXTENSION_SWITCH') | 692 | + env_aclnn_extension_switch = os.getenv("ACLNN_EXTENSION_SWITCH") |
| 639 | - env_aclnn_extension_path = os.getenv('ACLNN_EXTENSION_PATH') | 693 | + env_aclnn_extension_path = os.getenv("ACLNN_EXTENSION_PATH") |
| 640 | # if apply aclnn extension | 694 | # if apply aclnn extension |
| 641 | if env_aclnn_extension_switch and os.path.exists(env_aclnn_extension_path): | 695 | if env_aclnn_extension_switch and os.path.exists(env_aclnn_extension_path): |
| 642 | - op_plugin_path = os.path.join(env_aclnn_extension_path, 'op_plugin') | 696 | + op_plugin_path = os.path.join(env_aclnn_extension_path, "op_plugin") |
| 643 | # original code logic | 697 | # original code logic |
| 644 | else: | 698 | else: |
| 645 | base_dir = os.path.dirname(os.path.realpath(__file__)) | 699 | base_dir = os.path.dirname(os.path.realpath(__file__)) |
| 646 | - op_plugin_path = os.path.join(base_dir, '../third_party/op-plugin/op_plugin') | 700 | + op_plugin_path = os.path.join(base_dir, "../third_party/op-plugin/op_plugin") |
| 647 | - | 701 | + |
| 648 | return os.path.exists(op_plugin_path) | 702 | return os.path.exists(op_plugin_path) |
| 649 | 703 | ||
| 650 | 704 | ||
| 651 | def is_op_valid(op_key: str) -> bool: | 705 | def is_op_valid(op_key: str) -> bool: |
| 652 | - return True if op_key in GLOBAL_STRUCTURED_OP_INFO_CACHE else False | 706 | + return op_key in GLOBAL_STRUCTURED_OP_INFO_CACHE |
| 653 | 707 | ||
| 654 | 708 | ||
| 655 | def get_opplugin_wrap_name(func) -> str: | 709 | def get_opplugin_wrap_name(func) -> str: |
| @@ -670,7 +724,9 @@ def update_opapi_info(op_info): | |||
| 670 | if op_info.get("op_api", False): | 724 | if op_info.get("op_api", False): |
| 671 | GLOBAL_OPAPI_INFO_CACHE.add(op_info.get("func").split("(")[0]) | 725 | GLOBAL_OPAPI_INFO_CACHE.add(op_info.get("func").split("(")[0]) |
| 672 | else: | 726 | else: |
| 673 | - print(f"Warning: Unsupported parameter types, only str and dict is supported, but input is {type(op_info)}") | 727 | + print( |
| 728 | + f"Warning: Unsupported parameter types, only str and dict is supported, but input is {type(op_info)}" | ||
| 729 | + ) | ||
| 674 | 730 | ||
| 675 | 731 | ||
| 676 | def is_opapi(op_key): | 732 | def is_opapi(op_key): |
| @@ -684,9 +740,13 @@ def update_internal_format_opapi_info(op_info): | |||
| 684 | return | 740 | return |
| 685 | elif isinstance(op_info, dict): | 741 | elif isinstance(op_info, dict): |
| 686 | if op_info.get("internal_format_opapi", False): | 742 | if op_info.get("internal_format_opapi", False): |
| 687 | - GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE.add(op_info.get("func").split("(")[0]) | 743 | + GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE.add( |
| 744 | + op_info.get("func").split("(")[0] | ||
| 745 | + ) | ||
| 688 | else: | 746 | else: |
| 689 | - print(f"Warning: Unsupported parameter types, only str and dict is supported, but input is {type(op_info)}") | 747 | + print( |
| 748 | + f"Warning: Unsupported parameter types, only str and dict is supported, but input is {type(op_info)}" | ||
| 749 | + ) | ||
| 690 | 750 | ||
| 691 | 751 | ||
| 692 | def is_opapi_support_internal_format(op_key): | 752 | def is_opapi_support_internal_format(op_key): |
| @@ -694,13 +754,13 @@ def is_opapi_support_internal_format(op_key): | |||
| 694 | return op_key in GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE | 754 | return op_key in GLOBAL_INTERNAL_FORMAT_OPAPI_INFO_CACHE |
| 695 | 755 | ||
| 696 | 756 | ||
| 697 | -def get_target_functions(yaml_path: str, target_op_type: str = None) -> List: | 757 | +def get_target_functions(yaml_path: str, target_op_type: str | None = None) -> list: |
| 698 | source_es = parse_npu_yaml(yaml_path) | 758 | source_es = parse_npu_yaml(yaml_path) |
| 699 | 759 | ||
| 700 | - custom = source_es.pop('custom', []) | 760 | + custom = source_es.pop("custom", []) |
| 701 | if custom is None: | 761 | if custom is None: |
| 702 | custom = [] # Allow an empty list of supported ops | 762 | custom = [] # Allow an empty list of supported ops |
| 703 | - official = source_es.pop('official', []) | 763 | + official = source_es.pop("official", []) |
| 704 | if official is None: | 764 | if official is None: |
| 705 | official = [] # Allow an empty list of supported ops | 765 | official = [] # Allow an empty list of supported ops |
| 706 | 766 | ||
| @@ -709,8 +769,8 @@ def get_target_functions(yaml_path: str, target_op_type: str = None) -> List: | |||
| 709 | symint = source_es.pop("symint", []) | 769 | symint = source_es.pop("symint", []) |
| 710 | if symint is None: | 770 | if symint is None: |
| 711 | symint = [] | 771 | symint = [] |
| 712 | - symint = [op['func'] if isinstance(op, Dict) else op for op in symint] | 772 | + symint = [op["func"] if isinstance(op, dict) else op for op in symint] |
| 713 | - symint_set = set([str(FunctionSchema.parse(op).name) for op in symint]) | 773 | + symint_set = {str(FunctionSchema.parse(op).name) for op in symint} |
| 714 | 774 | ||
| 715 | global GLOBAL_STRUCTURED_OP_INFO_CACHE | 775 | global GLOBAL_STRUCTURED_OP_INFO_CACHE |
| 716 | GLOBAL_STRUCTURED_OP_INFO_CACHE.clear() | 776 | GLOBAL_STRUCTURED_OP_INFO_CACHE.clear() |
| @@ -718,20 +778,22 @@ def get_target_functions(yaml_path: str, target_op_type: str = None) -> List: | |||
| 718 | for op in support_ops: | 778 | for op in support_ops: |
| 719 | funcs = op.get("func", None) | 779 | funcs = op.get("func", None) |
| 720 | if not isinstance(funcs, str): | 780 | if not isinstance(funcs, str): |
| 721 | - raise TypeError(f'not a str : {funcs}') | 781 | + raise TypeError(f"not a str : {funcs}") |
| 722 | func = FunctionSchema.parse(funcs) | 782 | func = FunctionSchema.parse(funcs) |
| 723 | wrap_name = cpp.name(func) | 783 | wrap_name = cpp.name(func) |
| 724 | op_key = str(func.name) | 784 | op_key = str(func.name) |
| 725 | if op_key in symint_set: | 785 | if op_key in symint_set: |
| 726 | wrap_name += "_symint" | 786 | wrap_name += "_symint" |
| 727 | if target_op_type is not None: | 787 | if target_op_type is not None: |
| 728 | - if target_op_type not in op.keys(): | 788 | + if target_op_type not in op: |
| 729 | continue | 789 | continue |
| 730 | wrap_name += "_" + target_op_type | 790 | wrap_name += "_" + target_op_type |
| 731 | cur_wrap_name = GLOBAL_STRUCTURED_OP_INFO_CACHE.get(op_key, "") | 791 | cur_wrap_name = GLOBAL_STRUCTURED_OP_INFO_CACHE.get(op_key, "") |
| 732 | if cur_wrap_name and cur_wrap_name != wrap_name: | 792 | if cur_wrap_name and cur_wrap_name != wrap_name: |
| 733 | - print(f"Find different wrap_name for {cur_wrap_name} and {wrap_name} between pta and opplugin, ", | 793 | + print( |
| 734 | - f"with {wrap_name} being used as the actual wrap_name") | 794 | + f"Find different wrap_name for {cur_wrap_name} and {wrap_name} between pta and opplugin, ", |
| 795 | + f"with {wrap_name} being used as the actual wrap_name", | ||
| 796 | + ) | ||
| 735 | GLOBAL_STRUCTURED_OP_INFO_CACHE[op_key] = wrap_name | 797 | GLOBAL_STRUCTURED_OP_INFO_CACHE[op_key] = wrap_name |
| 736 | target_funcs.append(func) | 798 | target_funcs.append(func) |
| 737 | 799 | ||
| @@ -740,27 +802,36 @@ def get_target_functions(yaml_path: str, target_op_type: str = None) -> List: | |||
| 740 | 802 | ||
| 741 | def get_target_native_registration( | 803 | def get_target_native_registration( |
| 742 | dispatch_key: DispatchKey, | 804 | dispatch_key: DispatchKey, |
| 743 | - backend_indices: Dict[DispatchKey, BackendIndex], | 805 | + backend_indices: dict[DispatchKey, BackendIndex], |
| 744 | - metadata: Dict[OperatorName, BackendMetadata], | 806 | + metadata: dict[OperatorName, BackendMetadata], |
| 745 | - native_functions: List[NativeFunction], | 807 | + native_functions: list[NativeFunction], |
| 746 | ): | 808 | ): |
| 747 | if native_functions is None: | 809 | if native_functions is None: |
| 748 | return "" | 810 | return "" |
| 749 | - cpu_dispatch_key = DispatchKey.parse(dispatch_key.name.replace("PrivateUse1", "CPU")) | 811 | + cpu_dispatch_key = DispatchKey.parse( |
| 812 | + dispatch_key.name.replace("PrivateUse1", "CPU") | ||
| 813 | + ) | ||
| 750 | cpu_backend_indices = backend_indices[cpu_dispatch_key] | 814 | cpu_backend_indices = backend_indices[cpu_dispatch_key] |
| 751 | - cuda_dispatch_key = DispatchKey.parse(dispatch_key.name.replace("PrivateUse1", "CUDA")) | 815 | + cuda_dispatch_key = DispatchKey.parse( |
| 816 | + dispatch_key.name.replace("PrivateUse1", "CUDA") | ||
| 817 | + ) | ||
| 752 | cuda_backend_indices = backend_indices[cuda_dispatch_key] | 818 | cuda_backend_indices = backend_indices[cuda_dispatch_key] |
| 753 | 819 | ||
| 754 | target_native_functions_kernels = {} | 820 | target_native_functions_kernels = {} |
| 755 | - for op_name in cpu_backend_indices.index.keys(): | 821 | + for op_name in cpu_backend_indices.index: |
| 756 | - if op_name not in cuda_backend_indices.index.keys(): | 822 | + if op_name not in cuda_backend_indices.index: |
| 757 | continue | 823 | continue |
| 758 | - if cpu_backend_indices.index[op_name].kernel == cuda_backend_indices.index[op_name].kernel \ | 824 | + if ( |
| 759 | - and op_name not in metadata.keys(): | 825 | + cpu_backend_indices.index[op_name].kernel |
| 760 | - target_native_functions_kernels[op_name] = cpu_backend_indices.index[op_name].kernel | 826 | + == cuda_backend_indices.index[op_name].kernel |
| 827 | + and op_name not in metadata | ||
| 828 | + ): | ||
| 829 | + target_native_functions_kernels[op_name] = cpu_backend_indices.index[ | ||
| 830 | + op_name | ||
| 831 | + ].kernel | ||
| 761 | target_native_functions = [] | 832 | target_native_functions = [] |
| 762 | for f in native_functions: | 833 | for f in native_functions: |
| 763 | - if f.func.name in target_native_functions_kernels.keys(): | 834 | + if f.func.name in target_native_functions_kernels: |
| 764 | target_native_functions.append(f) | 835 | target_native_functions.append(f) |
| 765 | 836 | ||
| 766 | native_functions_registration_template = CodeTemplate( | 837 | native_functions_registration_template = CodeTemplate( |
| @@ -773,27 +844,36 @@ TORCH_LIBRARY_IMPL(aten, ${dispatch_key}, m) { | |||
| 773 | ${native_kernels} | 844 | ${native_kernels} |
| 774 | } | 845 | } |
| 775 | } | 846 | } |
| 776 | -""") | 847 | +""" |
| 848 | + ) | ||
| 777 | 849 | ||
| 778 | def wrap_native_function(f): | 850 | def wrap_native_function(f): |
| 779 | kernel_name = target_native_functions_kernels[f.func.name] | 851 | kernel_name = target_native_functions_kernels[f.func.name] |
| 780 | - wrap_func_name = f"wrap_{dispatch_key}_{str(f.func.name.name)}_{f.func.name.overload_name}" | 852 | + wrap_func_name = ( |
| 853 | + f"wrap_{dispatch_key}_{str(f.func.name.name)}_{f.func.name.overload_name}" | ||
| 854 | + ) | ||
| 781 | with native_function_manager(f): | 855 | with native_function_manager(f): |
| 782 | - sig = NativeSignature(f.func, prefix='', symint=kernel_name.endswith('symint')) | 856 | + sig = NativeSignature( |
| 783 | - args_exprs_str = ', '.join(a.name for a in sig.arguments()) | 857 | + f.func, prefix="", symint=kernel_name.endswith("symint") |
| 858 | + ) | ||
| 859 | + args_exprs_str = ", ".join(a.name for a in sig.arguments()) | ||
| 784 | return f"""{sig.decl(name=wrap_func_name)} {{ | 860 | return f"""{sig.decl(name=wrap_func_name)} {{ |
| 785 | return at::native::{kernel_name}({args_exprs_str}); | 861 | return at::native::{kernel_name}({args_exprs_str}); |
| 786 | }} | 862 | }} |
| 787 | """ | 863 | """ |
| 788 | 864 | ||
| 789 | def register_wrap_native_function(f): | 865 | def register_wrap_native_function(f): |
| 790 | - wrap_func_name = f"wrap_{dispatch_key}_{str(f.func.name.name)}_{f.func.name.overload_name}" | 866 | + wrap_func_name = ( |
| 791 | - return f"m.impl(\"{f.func.name}\", TORCH_FN({wrap_func_name}));" | 867 | + f"wrap_{dispatch_key}_{str(f.func.name.name)}_{f.func.name.overload_name}" |
| 868 | + ) | ||
| 869 | + return f'm.impl("{f.func.name}", TORCH_FN({wrap_func_name}));' | ||
| 792 | 870 | ||
| 793 | return native_functions_registration_template.substitute( | 871 | return native_functions_registration_template.substitute( |
| 794 | dispatch_helpers=[wrap_native_function(f) for f in target_native_functions], | 872 | dispatch_helpers=[wrap_native_function(f) for f in target_native_functions], |
| 795 | dispatch_key=dispatch_key.name, | 873 | dispatch_key=dispatch_key.name, |
| 796 | - native_kernels=[register_wrap_native_function(f) for f in target_native_functions] | 874 | + native_kernels=[ |
| 875 | + register_wrap_native_function(f) for f in target_native_functions | ||
| 876 | + ], | ||
| 797 | ) | 877 | ) |
| 798 | 878 | ||
| 799 | 879 | ||
| @@ -802,9 +882,9 @@ ${native_kernels} | |||
| 802 | 882 | ||
| 803 | class NativeFunctionsGroupOptionalOut: | 883 | class NativeFunctionsGroupOptionalOut: |
| 804 | functional: NativeFunction | 884 | functional: NativeFunction |
| 805 | - inplace: Optional[NativeFunction] | 885 | + inplace: NativeFunction | None |
| 806 | - mutable: Optional[NativeFunction] | 886 | + mutable: NativeFunction | None |
| 807 | - out: Optional[NativeFunction] | 887 | + out: NativeFunction | None |
| 808 | 888 | ||
| 809 | 889 | ||
| 810 | def root_name(self) -> str: | 890 | def root_name(self) -> str: |
| @@ -821,10 +901,10 @@ class NativeFunctionsGroupOptionalOut: | |||
| 821 | 901 | ||
| 822 | 902 | ||
| 823 | def from_dict( | 903 | def from_dict( |
| 824 | - d: Dict[SchemaKind, NativeFunction] | 904 | + d: dict[SchemaKind, NativeFunction], |
| 825 | ) -> Optional["NativeFunctionsGroupOptionalOut"]: | 905 | ) -> Optional["NativeFunctionsGroupOptionalOut"]: |
| 826 | if len(d) == 0: | 906 | if len(d) == 0: |
| 827 | - raise RuntimeError('The variable d is empty') | 907 | + raise RuntimeError("The variable d is empty") |
| 828 | if len(d) == 1: | 908 | if len(d) == 1: |
| 829 | return None | 909 | return None |
| 830 | d = dict(d) # non-destructive updates please | 910 | d = dict(d) # non-destructive updates please |
| @@ -833,10 +913,10 @@ class NativeFunctionsGroupOptionalOut: | |||
| 833 | mutable = d.pop(SchemaKind.mutable, None) | 913 | mutable = d.pop(SchemaKind.mutable, None) |
| 834 | out = d.pop(SchemaKind.out, None) | 914 | out = d.pop(SchemaKind.out, None) |
| 835 | if len(d) != 0: | 915 | if len(d) != 0: |
| 836 | - raise RuntimeError('The variable d is not empty after popping keys') | 916 | + raise RuntimeError("The variable d is not empty after popping keys") |
| 837 | 917 | ||
| 838 | if functional is None: | 918 | if functional is None: |
| 839 | - raise RuntimeError('The variable functional is None') | 919 | + raise RuntimeError("The variable functional is None") |
| 840 | 920 | ||
| 841 | return NativeFunctionsGroupOptionalOut( | 921 | return NativeFunctionsGroupOptionalOut( |
| 842 | functional=functional, | 922 | functional=functional, |
| @@ -848,17 +928,18 @@ class NativeFunctionsGroupOptionalOut: | |||
| 848 | 928 | ||
| 849 | def get_grouped_native_functions_optional_out( | 929 | def get_grouped_native_functions_optional_out( |
| 850 | native_functions: Sequence[NativeFunction], | 930 | native_functions: Sequence[NativeFunction], |
| 851 | -) -> Sequence[Union[NativeFunction, NativeFunctionsGroupOptionalOut]]: | 931 | +) -> Sequence[NativeFunction | NativeFunctionsGroupOptionalOut]: |
| 852 | - | ||
| 853 | def flatten_pre_group( | 932 | def flatten_pre_group( |
| 854 | - d: Dict[SchemaKind, NativeFunction] | 933 | + d: dict[SchemaKind, NativeFunction], |
| 855 | - ) -> Sequence[Union[NativeFunction, NativeFunctionsGroupOptionalOut]]: | 934 | + ) -> Sequence[NativeFunction | NativeFunctionsGroupOptionalOut]: |
| 856 | r = NativeFunctionsGroupOptionalOut.from_dict(d) | 935 | r = NativeFunctionsGroupOptionalOut.from_dict(d) |
| 857 | if r is None: | 936 | if r is None: |
| 858 | # Invariant: any NativeFunctions that are code-generated | 937 | # Invariant: any NativeFunctions that are code-generated |
| 859 | # should have been grouped into NativeFunctionsGroupOptionalOut objects | 938 | # should have been grouped into NativeFunctionsGroupOptionalOut objects |
| 860 | if any("generated" in f.tags for f in d.values()): | 939 | if any("generated" in f.tags for f in d.values()): |
| 861 | - raise RuntimeError("The variable d contains 'generated' in function tags") | 940 | + raise RuntimeError( |
| 941 | + "The variable d contains 'generated' in function tags" | ||
| 942 | + ) | ||
| 862 | return list(d.values()) | 943 | return list(d.values()) |
| 863 | else: | 944 | else: |
| 864 | return [r] | 945 | return [r] |