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 dataclasses import replace 

18from typing import Callable, Optional, Tuple 

19 

20from .registry import ( 

21 PARSE_NODE, 

22 PARSE_OPERATOR, 

23 OnnxPluginDescriptor, 

24 register_onnx_plugin, 

25 replace_registered_onnx_plugin, 

26) 

27 

28 

29def _normalize_name( 

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

31) -> str: 

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

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

34 if reject_origin_separator and "::" in value: 

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

36 return value 

37 

38 

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

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

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

42 if not opsets: 

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

44 normalized = set() 

45 for opset in opsets: 

46 if type(opset) is not int: 

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

48 if opset <= 0: 

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

50 normalized.add(opset) 

51 return tuple(sorted(normalized)) 

52 

53 

54class OnnxPlugin: 

55 """ONNX source-to-target descriptor for parser callbacks.""" 

56 

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

58 

59 def __init__( 

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

61 ) -> None: 

62 self._source = source 

63 self._domain = domain 

64 self._opsets = opsets 

65 self._target = target 

66 self._descriptor: Optional[OnnxPluginDescriptor] = None 

67 

68 def _bind_callback( 

69 self, fn: Callable[..., None], callback_kind: str 

70 ) -> Callable[..., None]: 

71 if not inspect.isfunction(fn): 

72 raise TypeError(f"OnnxPlugin {callback_kind} expects a Python function") 

73 

74 if self._descriptor is not None: 

75 descriptor = self._descriptor 

76 if callback_kind in descriptor.callback_kinds: 

77 raise ValueError(f"OnnxPlugin {callback_kind} is already bound") 

78 descriptor = replace( 

79 descriptor, 

80 parser_node=fn 

81 if callback_kind == PARSE_NODE 

82 else descriptor.parser_node, 

83 parser_operator=( 

84 fn 

85 if callback_kind == PARSE_OPERATOR 

86 else descriptor.parser_operator 

87 ), 

88 ) 

89 replace_registered_onnx_plugin(descriptor) 

90 if descriptor.parser_node is not None: 

91 setattr( 

92 descriptor.parser_node, "__ge_onnx_plugin_descriptor__", descriptor 

93 ) 

94 if descriptor.parser_operator is not None: 

95 setattr( 

96 descriptor.parser_operator, 

97 "__ge_onnx_plugin_descriptor__", 

98 descriptor, 

99 ) 

100 self._descriptor = descriptor 

101 return fn 

102 

103 module_name = fn.__module__ 

104 callback_name = fn.__qualname__ 

105 origin_types = tuple( 

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

107 ) 

108 descriptor = register_onnx_plugin( 

109 OnnxPluginDescriptor( 

110 descriptor_key=( 

111 f"{module_name}:{callback_name}:{callback_kind}:" 

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

113 ), 

114 source=self._source, 

115 domain=self._domain, 

116 opsets=self._opsets, 

117 target=self._target, 

118 origin_types=origin_types, 

119 module_name=module_name, 

120 parser_node=fn if callback_kind == PARSE_NODE else None, 

121 parser_operator=fn if callback_kind == PARSE_OPERATOR else None, 

122 ) 

123 ) 

124 self._descriptor = descriptor 

125 setattr(fn, "__ge_onnx_plugin_descriptor__", descriptor) 

126 return fn 

127 

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

129 return self._bind_callback(fn, PARSE_NODE) 

130 

131 def parse_operator(self, fn: Callable[..., None]) -> Callable[..., None]: 

132 return self._bind_callback(fn, PARSE_OPERATOR) 

133 

134 

135def onnx_plugin( 

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

137) -> OnnxPlugin: 

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

139 

140 return OnnxPlugin( 

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

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

143 opsets=_normalize_opsets(opsets), 

144 target=_normalize_name(target, "target"), 

145 )