Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/proto.py: 96%
298 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 10:22 +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_by_op_type(self, op_type: str) -> Optional[OpProtoDescriptor]:
204 with self._lock:
205 return self._op_type_to_desc.get(op_type)
207 def get_all(self) -> List[OpProtoDescriptor]:
208 with self._lock:
209 return sorted(self._op_type_to_desc.values(), key=lambda item: item.op_type)
212_OP_PROTO_REGISTRY = _OpProtoRegistry()
213_EMPTY = inspect.Signature.empty
214_NONE_TYPE = type(None)
216_SCALAR_ATTR_TYPES = {
217 int: AttrType.INT,
218 float: AttrType.FLOAT,
219 bool: AttrType.BOOL,
220 str: AttrType.STRING,
221 DataType: AttrType.DATA_TYPE,
222 Tensor: AttrType.TENSOR,
223}
225_LIST_ATTR_TYPES = {
226 int: AttrType.LIST_INT,
227 float: AttrType.LIST_FLOAT,
228 bool: AttrType.LIST_BOOL,
229 str: AttrType.LIST_STRING,
230 DataType: AttrType.LIST_DATA_TYPE,
231}
234def _build_descriptor_key(module_name: str, func_name: str, op_type: str) -> str:
235 return f"{module_name}:{func_name}:{op_type}"
238def _normalize_op_type(op_type: str) -> str:
239 if not isinstance(op_type, str) or not op_type:
240 raise TypeError(
241 f"register_op op_type must be a non-empty string, got {op_type!r}"
242 )
243 return op_type
246def _get_optional_value_type(annotation):
247 origin = _get_origin(annotation)
248 args = _get_args(annotation)
249 if origin is Union and len(args) == 2 and _NONE_TYPE in args:
250 return args[0] if args[1] is _NONE_TYPE else args[1]
251 return _EMPTY
254def _is_list_annotation(annotation) -> bool:
255 return _get_origin(annotation) is list
258def _parse_input_annotation(name: str, annotation) -> InputType:
259 if annotation is TensorDesc:
260 return InputType.REQUIRED
262 args = _get_args(annotation)
263 if _get_optional_value_type(annotation) is TensorDesc:
264 return InputType.OPTIONAL
265 if _is_list_annotation(annotation) and args == (TensorDesc,):
266 return InputType.DYNAMIC
267 raise TypeError(f"unsupported input annotation for '{name}': {annotation!r}")
270def _parse_attr_annotation(name: str, annotation) -> str:
271 ir_type = _SCALAR_ATTR_TYPES.get(annotation)
272 if ir_type is not None:
273 return ir_type
275 args = _get_args(annotation)
276 if _is_list_annotation(annotation) and len(args) == 1:
277 element_type = args[0]
278 ir_type = _LIST_ATTR_TYPES.get(element_type)
279 if ir_type is not None:
280 return ir_type
281 if _is_list_annotation(element_type) and _get_args(element_type) == (int,):
282 return AttrType.LIST_LIST_INT
283 raise TypeError(f"unsupported attr annotation for '{name}': {annotation!r}")
286def _parse_output_kinds(annotation) -> Tuple[OutputType, ...]:
287 if annotation is _NONE_TYPE:
288 return ()
289 if annotation is TensorDesc:
290 return (OutputType.REQUIRED,)
291 if _is_list_annotation(annotation) and _get_args(annotation) == (TensorDesc,):
292 return (OutputType.DYNAMIC,)
294 if _get_origin(annotation) is tuple:
295 output_annotations = _get_args(annotation)
296 if not output_annotations or Ellipsis in output_annotations:
297 raise TypeError(f"unsupported return annotation: {annotation!r}")
298 output_kinds = []
299 for output_index, output_annotation in enumerate(output_annotations):
300 if output_annotation is TensorDesc:
301 output_kinds.append(OutputType.REQUIRED)
302 elif _is_list_annotation(output_annotation) and _get_args(
303 output_annotation
304 ) == (TensorDesc,):
305 output_kinds.append(OutputType.DYNAMIC)
306 else:
307 raise TypeError(
308 f"unsupported return annotation at output index {output_index}: "
309 f"{output_annotation!r}"
310 )
311 return tuple(output_kinds)
312 raise TypeError(f"unsupported return annotation: {annotation!r}")
315def _validate_scalar_default(name: str, value, expected_type) -> None:
316 if type(value) is not expected_type:
317 raise TypeError(
318 f"default value for attr '{name}' must be {expected_type.__name__}, "
319 f"got {type(value).__name__}"
320 )
321 if expected_type is DataType and value is DataType.DT_MAX:
322 raise TypeError(f"default value for attr '{name}' must be a valid DataType")
325def _validate_list_default(name: str, value, element_type) -> None:
326 if type(value) is not list:
327 raise TypeError(
328 f"default value for attr '{name}' must be list, got {type(value).__name__}"
329 )
330 for element in value:
331 if _is_list_annotation(element_type):
332 _validate_list_default(name, element, _get_args(element_type)[0])
333 elif type(element) is not element_type:
334 raise TypeError(
335 f"default value for attr '{name}' contains {type(element).__name__}, "
336 f"expected {element_type.__name__}"
337 )
338 elif element_type is DataType and element is DataType.DT_MAX:
339 raise TypeError(
340 f"default value for attr '{name}' must contain valid DataType values"
341 )
344def _validate_default(name: str, annotation, value) -> None:
345 if annotation is Tensor:
346 raise TypeError(f"Tensor attr '{name}' does not support a default value")
347 if _is_list_annotation(annotation):
348 _validate_list_default(name, value, _get_args(annotation)[0])
349 return
350 _validate_scalar_default(name, value, annotation)
353def _thaw_default(value):
354 if type(value) is tuple:
355 return [_thaw_default(item) for item in value]
356 return value
359def _parse_mutates_args(mutates_args, inputs, output_count: int) -> Dict[int, str]:
360 if isinstance(mutates_args, str) or not isinstance(mutates_args, (list, tuple)):
361 raise TypeError("mutates_args must be a list or tuple")
362 if not mutates_args:
363 return {}
365 has_name = any(isinstance(item, str) for item in mutates_args)
366 has_explicit = any(not isinstance(item, str) for item in mutates_args)
367 if has_name and has_explicit:
368 raise TypeError("mutates_args sequential and explicit forms cannot be mixed")
370 if has_name:
371 if len(mutates_args) > output_count:
372 raise ValueError("mutates_args has more entries than outputs")
373 bindings = {index: name for index, name in enumerate(mutates_args)}
374 else:
375 bindings = {}
376 for item in mutates_args:
377 if not isinstance(item, (list, tuple)) or len(item) != 2:
378 raise TypeError(
379 "mutates_args explicit entries must be (input_name, output_index)"
380 )
381 name, output_index = item
382 if not isinstance(name, str) or type(output_index) is not int:
383 raise TypeError("mutates_args explicit entries must be (str, int)")
384 if output_index < 0 or output_index >= output_count:
385 raise ValueError(
386 f"mutates_args output index out of range: {output_index}"
387 )
388 if output_index in bindings:
389 raise ValueError(
390 f"mutates_args output index is duplicated: {output_index}"
391 )
392 bindings[output_index] = name
394 input_names = {item.name for item in inputs}
395 bound_names = set()
396 for name in bindings.values():
397 if name not in input_names:
398 raise ValueError(f"mutates_args input does not exist: '{name}'")
399 if name in bound_names:
400 raise ValueError(f"mutates_args input is duplicated: '{name}'")
401 bound_names.add(name)
402 return bindings
405def _build_outputs(
406 output_kinds: Tuple[OutputType, ...], inputs: Tuple[OpInput, ...], mutates_args
407) -> Tuple[OpOutput, ...]:
408 mutations = _parse_mutates_args(mutates_args, inputs, len(output_kinds))
409 used_names = {item.name for item in inputs}
410 output_names = set()
411 outputs = []
412 for index, kind in enumerate(output_kinds):
413 name = mutations.get(index)
414 if name is None:
415 base_name = f"output{index}"
416 name = base_name
417 suffix = 1
418 while name in used_names or name in output_names:
419 name = f"{base_name}_{suffix}"
420 suffix += 1
421 output_names.add(name)
422 outputs.append(OpOutput(name=name, index=index, kind=kind))
423 return tuple(outputs)
426def _build_descriptor(
427 fn: Callable[..., object], op_type: str, mutates_args
428) -> OpProtoDescriptor:
429 signature = inspect.signature(fn)
430 try:
431 type_hints = get_type_hints(fn)
432 except (NameError, TypeError) as exc:
433 raise TypeError(
434 f"failed to resolve annotations for '{fn.__qualname__}': {exc}"
435 ) from exc
437 inputs = []
438 attrs = []
439 for parameter in signature.parameters.values():
440 if parameter.kind in (
441 inspect.Parameter.VAR_POSITIONAL,
442 inspect.Parameter.VAR_KEYWORD,
443 ):
444 raise TypeError(
445 f"register_op does not support variadic parameter '{parameter.name}'"
446 )
447 annotation = type_hints.get(parameter.name, _EMPTY)
448 if parameter.kind in (
449 inspect.Parameter.POSITIONAL_ONLY,
450 inspect.Parameter.POSITIONAL_OR_KEYWORD,
451 ):
452 if annotation is _EMPTY:
453 raise TypeError(f"input '{parameter.name}' must have a type annotation")
454 if parameter.default is not _EMPTY:
455 raise TypeError(
456 f"input '{parameter.name}' must not have a default value"
457 )
458 inputs.append(
459 OpInput(
460 name=parameter.name,
461 index=len(inputs),
462 kind=_parse_input_annotation(parameter.name, annotation),
463 )
464 )
465 continue
466 if parameter.kind is inspect.Parameter.KEYWORD_ONLY:
467 if annotation is _EMPTY:
468 raise TypeError(f"attr '{parameter.name}' must have a type annotation")
469 # Python 3.7 wraps T in Optional[T] when its default value is None.
470 if parameter.default is None:
471 value_type = _get_optional_value_type(annotation)
472 if value_type is not _EMPTY:
473 annotation = value_type
474 ir_type = _parse_attr_annotation(parameter.name, annotation)
475 is_required = parameter.default is _EMPTY
476 default = None
477 if not is_required:
478 _validate_default(parameter.name, annotation, parameter.default)
479 default = parameter.default
480 attrs.append(
481 OpAttr(
482 name=parameter.name,
483 index=len(attrs),
484 type=ir_type,
485 is_required=is_required,
486 default=default,
487 )
488 )
489 continue
490 raise TypeError(f"unsupported parameter kind for '{parameter.name}'")
492 if "return" not in type_hints:
493 raise TypeError("register_op return type annotation is required")
494 input_tuple = tuple(inputs)
495 outputs = _build_outputs(
496 _parse_output_kinds(type_hints["return"]), input_tuple, mutates_args
497 )
498 module_name = fn.__module__
499 func_name = fn.__qualname__
500 return OpProtoDescriptor(
501 descriptor_key=_build_descriptor_key(module_name, func_name, op_type),
502 op_type=op_type,
503 module_name=module_name,
504 func_name=func_name,
505 inputs=input_tuple,
506 attrs=tuple(attrs),
507 outputs=outputs,
508 infer_func=fn,
509 )
512def register_op(*, op_type: str, mutates_args=()) -> callable:
513 """Collect a Python custom op prototype without registering it in C++."""
515 normalized_op_type = _normalize_op_type(op_type)
517 def decorator(fn: Callable[..., object]) -> Callable[..., object]:
518 try:
519 if not inspect.isfunction(fn):
520 raise TypeError("register_op expects a Python function")
521 descriptor = _OP_PROTO_REGISTRY.register(
522 _build_descriptor(fn, normalized_op_type, mutates_args)
523 )
524 except (TypeError, ValueError) as exc:
525 raise type(exc)(
526 f"register_op op_type '{normalized_op_type}' failed: {exc}"
527 ) from exc
528 setattr(fn, "__ge_op_proto_descriptor__", descriptor)
529 return fn
531 return decorator
534def clear_registered_op_protos() -> None:
535 _OP_PROTO_REGISTRY.clear()
538def get_registered_op_protos() -> List[OpProtoDescriptor]:
539 return _OP_PROTO_REGISTRY.get_all()
542def get_registered_op_proto_dicts() -> List[dict]:
543 return [item.to_bridge_dict() for item in get_registered_op_protos()]
546def get_registered_op_proto_by_descriptor_key(
547 descriptor_key: str,
548) -> Optional[OpProtoDescriptor]:
549 return _OP_PROTO_REGISTRY.get_by_descriptor_key(descriptor_key)
552def get_registered_op_proto_by_op_type(op_type: str) -> Optional[OpProtoDescriptor]:
553 return _OP_PROTO_REGISTRY.get_by_op_type(op_type)