Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/onnx_plugin/plugin.py: 100%

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# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Public descriptor and callback decorator for Python ONNX Plugins.""" 

14 

15import inspect 

16from collections.abc import Collection 

17from typing import Callable, Optional, Tuple 

18 

19from .registry import OnnxPluginDescriptor, register_onnx_plugin 

20 

21 

22def _normalize_name( 

23 value: str, field_name: str, *, reject_origin_separator=False 

24) -> str: 

25 if not isinstance(value, str) or not value: 

26 raise TypeError(f"onnx_plugin {field_name} must be a non-empty string") 

27 if reject_origin_separator and "::" in value: 

28 raise TypeError(f"onnx_plugin {field_name} must not contain '::'") 

29 return value 

30 

31 

32def _normalize_opsets(opsets: Collection[int]) -> Tuple[int, ...]: 

33 if isinstance(opsets, (str, bytes)) or not isinstance(opsets, Collection): 

34 raise TypeError("onnx_plugin opsets must be a collection of positive integers") 

35 if not opsets: 

36 raise ValueError("onnx_plugin opsets must not be empty") 

37 normalized = set() 

38 for opset in opsets: 

39 if type(opset) is not int: 

40 raise TypeError("onnx_plugin opsets must contain only integers") 

41 if opset <= 0: 

42 raise ValueError("onnx_plugin opsets must contain only positive integers") 

43 normalized.add(opset) 

44 return tuple(sorted(normalized)) 

45 

46 

47class OnnxPlugin: 

48 """ONNX source-to-target descriptor awaiting a parse_node callback.""" 

49 

50 __slots__ = ("_descriptor", "_domain", "_opsets", "_source", "_target") 

51 

52 def __init__( 

53 self, *, source: str, domain: str, opsets: Tuple[int, ...], target: str 

54 ) -> None: 

55 self._source = source 

56 self._domain = domain 

57 self._opsets = opsets 

58 self._target = target 

59 self._descriptor: Optional[OnnxPluginDescriptor] = None 

60 

61 def parse_node(self, fn: Callable[..., None]) -> Callable[..., None]: 

62 if self._descriptor is not None: 

63 raise ValueError("OnnxPlugin parse_node is already bound") 

64 if not inspect.isfunction(fn): 

65 raise TypeError("OnnxPlugin parse_node expects a Python function") 

66 

67 module_name = fn.__module__ 

68 parser_node_name = fn.__qualname__ 

69 origin_types = tuple( 

70 f"{self._domain}::{opset}::{self._source}" for opset in self._opsets 

71 ) 

72 descriptor = register_onnx_plugin( 

73 OnnxPluginDescriptor( 

74 descriptor_key=( 

75 f"{module_name}:{parser_node_name}:{self._domain}:" 

76 f"{self._source}:{','.join(map(str, self._opsets))}" 

77 ), 

78 source=self._source, 

79 domain=self._domain, 

80 opsets=self._opsets, 

81 target=self._target, 

82 origin_types=origin_types, 

83 module_name=module_name, 

84 parser_node=fn, 

85 ) 

86 ) 

87 self._descriptor = descriptor 

88 setattr(fn, "__ge_onnx_plugin_descriptor__", descriptor) 

89 return fn 

90 

91 

92def onnx_plugin( 

93 *, source: str, domain: str, opsets: Collection[int], target: str 

94) -> OnnxPlugin: 

95 """Create an ONNX Plugin descriptor for binding a parse_node callback.""" 

96 

97 return OnnxPlugin( 

98 source=_normalize_name(source, "source", reject_origin_separator=True), 

99 domain=_normalize_name(domain, "domain", reject_origin_separator=True), 

100 opsets=_normalize_opsets(opsets), 

101 target=_normalize_name(target, "target"), 

102 )