Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/proto.py: 96%
293 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 20:49 +0800
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# -----------------------------------------------------------------------------------------------------------
4# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6# CANN Open Software License Agreement Version 2.0 (the "License").
7# Please refer to the License for details. You may not use this file except in compliance with the License.
8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10# See LICENSE in the root of the software repository for the full text of the License.
11# -----------------------------------------------------------------------------------------------------------
13"""Python custom op prototype models, parser, and internal registry."""
15import inspect
16import math
17import threading
18from dataclasses import dataclass, field
19from typing import (
20 Callable,
21 Dict,
22 List,
23 Optional,
24 Tuple,
25 Union,
26 get_type_hints,
27)
29from ge.graph import DataType
30from ge.runtime import Tensor, TensorDesc
32from ._ir_types import AttrType, InputType, OutputType
35def _get_origin(annotation):
36 return getattr(annotation, "__origin__", None)
39def _get_args(annotation):
40 return getattr(annotation, "__args__", ())
43def _freeze_default(value):
44 if type(value) is list:
45 return tuple(_freeze_default(item) for item in value)
46 return value
49@dataclass(frozen=True)
50class OpInput:
51 name: str
52 index: int
53 kind: InputType
56@dataclass(frozen=True)
57class OpAttr:
58 name: str
59 index: int
60 type: str
61 is_required: bool
62 default: object = None
64 def __post_init__(self) -> None:
65 object.__setattr__(self, "default", _freeze_default(self.default))
68@dataclass(frozen=True)
69class OpOutput:
70 name: str
71 index: int
72 kind: OutputType
75@dataclass(frozen=True)
76class OpProtoDescriptor:
77 descriptor_key: str
78 op_type: str
79 module_name: str
80 func_name: str
81 inputs: Tuple[OpInput, ...]
82 attrs: Tuple[OpAttr, ...]
83 outputs: Tuple[OpOutput, ...]
84 infer_func: Callable[..., object] = field(compare=False, repr=False)
86 def __post_init__(self) -> None:
87 object.__setattr__(self, "inputs", tuple(self.inputs))
88 object.__setattr__(self, "attrs", tuple(self.attrs))
89 object.__setattr__(self, "outputs", tuple(self.outputs))
91 def to_bridge_dict(self) -> dict:
92 return {
93 "descriptor_key": self.descriptor_key,
94 "op_type": self.op_type,
95 "module_name": self.module_name,
96 "func_name": self.func_name,
97 "inputs": [
98 {"name": item.name, "kind": int(item.kind)} for item in self.inputs
99 ],
100 "attrs": [
101 {
102 "name": item.name,
103 "type": item.type,
104 "is_required": item.is_required,
105 "default": _thaw_default(item.default),
106 }
107 for item in self.attrs
108 ],
109 "outputs": [
110 {"name": item.name, "kind": int(item.kind)} for item in self.outputs
111 ],
112 }
115def _definition_values_equal(existing, current) -> bool:
116 if type(existing) is not type(current):
117 return False
118 if type(existing) is tuple:
119 return len(existing) == len(current) and all(
120 _definition_values_equal(existing_item, current_item)
121 for existing_item, current_item in zip(existing, current)
122 )
123 if type(existing) is float and math.isnan(existing) and math.isnan(current):
124 return True
125 return existing == current
128def _descriptor_definition(descriptor: OpProtoDescriptor) -> tuple:
129 return (
130 descriptor.op_type,
131 tuple((item.name, item.index, item.kind) for item in descriptor.inputs),
132 tuple(
133 (
134 item.name,
135 item.index,
136 item.type,
137 item.is_required,
138 item.default,
139 )
140 for item in descriptor.attrs
141 ),
142 tuple((item.name, item.index, item.kind) for item in descriptor.outputs),
143 )
146def _descriptor_definitions_equal(
147 existing: OpProtoDescriptor, current: OpProtoDescriptor
148) -> bool:
149 return _definition_values_equal(
150 _descriptor_definition(existing), _descriptor_definition(current)
151 )
154def _format_descriptor_source(label: str, descriptor: OpProtoDescriptor) -> str:
155 return (
156 f"{label} source: module_name.func_name "
157 f"'{descriptor.module_name}.{descriptor.func_name}', "
158 f"descriptor_key '{descriptor.descriptor_key}'"
159 )
162class _OpProtoRegistry:
163 def __init__(self) -> None:
164 self._lock = threading.RLock()
165 self._descriptor_key_to_desc: Dict[str, OpProtoDescriptor] = {}
166 self._op_type_to_desc: Dict[str, OpProtoDescriptor] = {}
168 def clear(self) -> None:
169 with self._lock:
170 self._descriptor_key_to_desc.clear()
171 self._op_type_to_desc.clear()
173 def register(self, descriptor: OpProtoDescriptor) -> OpProtoDescriptor:
174 with self._lock:
175 existing = self._descriptor_key_to_desc.get(descriptor.descriptor_key)
176 if existing is not None:
177 if not _descriptor_definitions_equal(existing, descriptor):
178 raise ValueError(
179 "python op proto descriptor content changed for "
180 f"descriptor_key '{descriptor.descriptor_key}'; "
181 f"{_format_descriptor_source('existing', existing)}; "
182 f"{_format_descriptor_source('current', descriptor)}"
183 )
184 return existing
186 existing = self._op_type_to_desc.get(descriptor.op_type)
187 if existing is not None:
188 raise ValueError(
189 f"python op proto op type '{descriptor.op_type}' already registered "
190 f"by '{existing.descriptor_key}'; "
191 f"{_format_descriptor_source('existing', existing)}; "
192 f"{_format_descriptor_source('current', descriptor)}"
193 )
195 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor
196 self._op_type_to_desc[descriptor.op_type] = descriptor
197 return descriptor
199 def get_by_descriptor_key(self, descriptor_key: str) -> Optional[OpProtoDescriptor]:
200 with self._lock:
201 return self._descriptor_key_to_desc.get(descriptor_key)
203 def get_all(self) -> List[OpProtoDescriptor]:
204 with self._lock:
205 return sorted(self._op_type_to_desc.values(), key=lambda item: item.op_type)
208_OP_PROTO_REGISTRY = _OpProtoRegistry()
209_EMPTY = inspect.Signature.empty
210_NONE_TYPE = type(None)
212_SCALAR_ATTR_TYPES = {
213 int: AttrType.INT,
214 float: AttrType.FLOAT,
215 bool: AttrType.BOOL,
216 str: AttrType.STRING,
217 DataType: AttrType.DATA_TYPE,
218 Tensor: AttrType.TENSOR,
219}
221_LIST_ATTR_TYPES = {
222 int: AttrType.LIST_INT,
223 float: AttrType.LIST_FLOAT,
224 bool: AttrType.LIST_BOOL,
225 str: AttrType.LIST_STRING,
226 DataType: AttrType.LIST_DATA_TYPE,
227}
230def _build_descriptor_key(module_name: str, func_name: str, op_type: str) -> str:
231 return f"{module_name}:{func_name}:{op_type}"
234def _normalize_op_type(op_type: str) -> str:
235 if not isinstance(op_type, str) or not op_type:
236 raise TypeError(
237 f"register_op op_type must be a non-empty string, got {op_type!r}"
238 )
239 return op_type
242def _get_optional_value_type(annotation):
243 origin = _get_origin(annotation)
244 args = _get_args(annotation)
245 if origin is Union and len(args) == 2 and _NONE_TYPE in args:
246 return args[0] if args[1] is _NONE_TYPE else args[1]
247 return _EMPTY
250def _is_list_annotation(annotation) -> bool:
251 return _get_origin(annotation) is list
254def _parse_input_annotation(name: str, annotation) -> InputType:
255 if annotation is TensorDesc:
256 return InputType.REQUIRED
258 args = _get_args(annotation)
259 if _get_optional_value_type(annotation) is TensorDesc:
260 return InputType.OPTIONAL
261 if _is_list_annotation(annotation) and args == (TensorDesc,):
262 return InputType.DYNAMIC
263 raise TypeError(f"unsupported input annotation for '{name}': {annotation!r}")
266def _parse_attr_annotation(name: str, annotation) -> str:
267 ir_type = _SCALAR_ATTR_TYPES.get(annotation)
268 if ir_type is not None:
269 return ir_type
271 args = _get_args(annotation)
272 if _is_list_annotation(annotation) and len(args) == 1:
273 element_type = args[0]
274 ir_type = _LIST_ATTR_TYPES.get(element_type)
275 if ir_type is not None:
276 return ir_type
277 if _is_list_annotation(element_type) and _get_args(element_type) == (int,):
278 return AttrType.LIST_LIST_INT
279 raise TypeError(f"unsupported attr annotation for '{name}': {annotation!r}")
282def _parse_output_kinds(annotation) -> Tuple[OutputType, ...]:
283 if annotation is _NONE_TYPE:
284 return ()
285 if annotation is TensorDesc:
286 return (OutputType.REQUIRED,)
287 if _is_list_annotation(annotation) and _get_args(annotation) == (TensorDesc,):
288 return (OutputType.DYNAMIC,)
290 if _get_origin(annotation) is tuple:
291 output_annotations = _get_args(annotation)
292 if not output_annotations or Ellipsis in output_annotations:
293 raise TypeError(f"unsupported return annotation: {annotation!r}")
294 output_kinds = []
295 for output_index, output_annotation in enumerate(output_annotations):
296 if output_annotation is TensorDesc:
297 output_kinds.append(OutputType.REQUIRED)
298 elif _is_list_annotation(output_annotation) and _get_args(
299 output_annotation
300 ) == (TensorDesc,):
301 output_kinds.append(OutputType.DYNAMIC)
302 else:
303 raise TypeError(
304 f"unsupported return annotation at output index {output_index}: "
305 f"{output_annotation!r}"
306 )
307 return tuple(output_kinds)
308 raise TypeError(f"unsupported return annotation: {annotation!r}")
311def _validate_scalar_default(name: str, value, expected_type) -> None:
312 if type(value) is not expected_type:
313 raise TypeError(
314 f"default value for attr '{name}' must be {expected_type.__name__}, "
315 f"got {type(value).__name__}"
316 )
317 if expected_type is DataType and value is DataType.DT_MAX:
318 raise TypeError(f"default value for attr '{name}' must be a valid DataType")
321def _validate_list_default(name: str, value, element_type) -> None:
322 if type(value) is not list:
323 raise TypeError(
324 f"default value for attr '{name}' must be list, got {type(value).__name__}"
325 )
326 for element in value:
327 if _is_list_annotation(element_type):
328 _validate_list_default(name, element, _get_args(element_type)[0])
329 elif type(element) is not element_type:
330 raise TypeError(
331 f"default value for attr '{name}' contains {type(element).__name__}, "
332 f"expected {element_type.__name__}"
333 )
334 elif element_type is DataType and element is DataType.DT_MAX:
335 raise TypeError(
336 f"default value for attr '{name}' must contain valid DataType values"
337 )
340def _validate_default(name: str, annotation, value) -> None:
341 if annotation is Tensor:
342 raise TypeError(f"Tensor attr '{name}' does not support a default value")
343 if _is_list_annotation(annotation):
344 _validate_list_default(name, value, _get_args(annotation)[0])
345 return
346 _validate_scalar_default(name, value, annotation)
349def _thaw_default(value):
350 if type(value) is tuple:
351 return [_thaw_default(item) for item in value]
352 return value
355def _parse_mutates_args(mutates_args, inputs, output_count: int) -> Dict[int, str]:
356 if isinstance(mutates_args, str) or not isinstance(mutates_args, (list, tuple)):
357 raise TypeError("mutates_args must be a list or tuple")
358 if not mutates_args:
359 return {}
361 has_name = any(isinstance(item, str) for item in mutates_args)
362 has_explicit = any(not isinstance(item, str) for item in mutates_args)
363 if has_name and has_explicit:
364 raise TypeError("mutates_args sequential and explicit forms cannot be mixed")
366 if has_name:
367 if len(mutates_args) > output_count:
368 raise ValueError("mutates_args has more entries than outputs")
369 bindings = {index: name for index, name in enumerate(mutates_args)}
370 else:
371 bindings = {}
372 for item in mutates_args:
373 if not isinstance(item, (list, tuple)) or len(item) != 2:
374 raise TypeError(
375 "mutates_args explicit entries must be (input_name, output_index)"
376 )
377 name, output_index = item
378 if not isinstance(name, str) or type(output_index) is not int:
379 raise TypeError("mutates_args explicit entries must be (str, int)")
380 if output_index < 0 or output_index >= output_count:
381 raise ValueError(
382 f"mutates_args output index out of range: {output_index}"
383 )
384 if output_index in bindings:
385 raise ValueError(
386 f"mutates_args output index is duplicated: {output_index}"
387 )
388 bindings[output_index] = name
390 input_names = {item.name for item in inputs}
391 bound_names = set()
392 for name in bindings.values():
393 if name not in input_names:
394 raise ValueError(f"mutates_args input does not exist: '{name}'")
395 if name in bound_names:
396 raise ValueError(f"mutates_args input is duplicated: '{name}'")
397 bound_names.add(name)
398 return bindings
401def _build_outputs(
402 output_kinds: Tuple[OutputType, ...], inputs: Tuple[OpInput, ...], mutates_args
403) -> Tuple[OpOutput, ...]:
404 mutations = _parse_mutates_args(mutates_args, inputs, len(output_kinds))
405 used_names = {item.name for item in inputs}
406 output_names = set()
407 outputs = []
408 for index, kind in enumerate(output_kinds):
409 name = mutations.get(index)
410 if name is None:
411 base_name = f"output{index}"
412 name = base_name
413 suffix = 1
414 while name in used_names or name in output_names:
415 name = f"{base_name}_{suffix}"
416 suffix += 1
417 output_names.add(name)
418 outputs.append(OpOutput(name=name, index=index, kind=kind))
419 return tuple(outputs)
422def _build_descriptor(
423 fn: Callable[..., object], op_type: str, mutates_args
424) -> OpProtoDescriptor:
425 signature = inspect.signature(fn)
426 try:
427 type_hints = get_type_hints(fn)
428 except (NameError, TypeError) as exc:
429 raise TypeError(
430 f"failed to resolve annotations for '{fn.__qualname__}': {exc}"
431 ) from exc
433 inputs = []
434 attrs = []
435 for parameter in signature.parameters.values():
436 if parameter.kind in (
437 inspect.Parameter.VAR_POSITIONAL,
438 inspect.Parameter.VAR_KEYWORD,
439 ):
440 raise TypeError(
441 f"register_op does not support variadic parameter '{parameter.name}'"
442 )
443 annotation = type_hints.get(parameter.name, _EMPTY)
444 if parameter.kind in (
445 inspect.Parameter.POSITIONAL_ONLY,
446 inspect.Parameter.POSITIONAL_OR_KEYWORD,
447 ):
448 if annotation is _EMPTY:
449 raise TypeError(f"input '{parameter.name}' must have a type annotation")
450 if parameter.default is not _EMPTY:
451 raise TypeError(
452 f"input '{parameter.name}' must not have a default value"
453 )
454 inputs.append(
455 OpInput(
456 name=parameter.name,
457 index=len(inputs),
458 kind=_parse_input_annotation(parameter.name, annotation),
459 )
460 )
461 continue
462 if parameter.kind is inspect.Parameter.KEYWORD_ONLY:
463 if annotation is _EMPTY:
464 raise TypeError(f"attr '{parameter.name}' must have a type annotation")
465 # Python 3.7 wraps T in Optional[T] when its default value is None.
466 if parameter.default is None:
467 value_type = _get_optional_value_type(annotation)
468 if value_type is not _EMPTY:
469 annotation = value_type
470 ir_type = _parse_attr_annotation(parameter.name, annotation)
471 is_required = parameter.default is _EMPTY
472 default = None
473 if not is_required:
474 _validate_default(parameter.name, annotation, parameter.default)
475 default = parameter.default
476 attrs.append(
477 OpAttr(
478 name=parameter.name,
479 index=len(attrs),
480 type=ir_type,
481 is_required=is_required,
482 default=default,
483 )
484 )
485 continue
486 raise TypeError(f"unsupported parameter kind for '{parameter.name}'")
488 if "return" not in type_hints:
489 raise TypeError("register_op return type annotation is required")
490 input_tuple = tuple(inputs)
491 outputs = _build_outputs(
492 _parse_output_kinds(type_hints["return"]), input_tuple, mutates_args
493 )
494 module_name = fn.__module__
495 func_name = fn.__qualname__
496 return OpProtoDescriptor(
497 descriptor_key=_build_descriptor_key(module_name, func_name, op_type),
498 op_type=op_type,
499 module_name=module_name,
500 func_name=func_name,
501 inputs=input_tuple,
502 attrs=tuple(attrs),
503 outputs=outputs,
504 infer_func=fn,
505 )
508def register_op(*, op_type: str, mutates_args=()) -> callable:
509 """Collect a Python custom op prototype without registering it in C++."""
511 normalized_op_type = _normalize_op_type(op_type)
513 def decorator(fn: Callable[..., object]) -> Callable[..., object]:
514 try:
515 if not inspect.isfunction(fn):
516 raise TypeError("register_op expects a Python function")
517 descriptor = _OP_PROTO_REGISTRY.register(
518 _build_descriptor(fn, normalized_op_type, mutates_args)
519 )
520 except (TypeError, ValueError) as exc:
521 raise type(exc)(
522 f"register_op op_type '{normalized_op_type}' failed: {exc}"
523 ) from exc
524 setattr(fn, "__ge_op_proto_descriptor__", descriptor)
525 return fn
527 return decorator
530def clear_registered_op_protos() -> None:
531 _OP_PROTO_REGISTRY.clear()
534def get_registered_op_protos() -> List[OpProtoDescriptor]:
535 return _OP_PROTO_REGISTRY.get_all()
538def get_registered_op_proto_dicts() -> List[dict]:
539 return [item.to_bridge_dict() for item in get_registered_op_protos()]
542def get_registered_op_proto_by_descriptor_key(
543 descriptor_key: str,
544) -> Optional[OpProtoDescriptor]:
545 return _OP_PROTO_REGISTRY.get_by_descriptor_key(descriptor_key)