Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_infer_meta.py: 96%

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"""Python infer_meta callback helpers for GE custom operators.""" 

14 

15from typing import Optional 

16 

17from ._ir_types import InputType, OutputType 

18from ._signature import _get_runtime_attr_spec 

19from .proto import get_registered_op_proto_by_op_type 

20 

21 

22def _build_infer_inputs(ctx, ir_inputs: list) -> list: 

23 args = [] 

24 for ir_index, item in enumerate(ir_inputs): 

25 kind = item["kind"] 

26 if kind == InputType.REQUIRED: 

27 args.append(ctx.get_required_input_tensor(ir_index)) 

28 elif kind == InputType.OPTIONAL: 

29 args.append(ctx.get_optional_input_tensor(ir_index)) 

30 elif kind == InputType.DYNAMIC: 

31 instance_num = ctx.get_dynamic_input_num(ir_index) 

32 descs = [] 

33 for relative_index in range(instance_num): 

34 descs.append(ctx.get_dynamic_input_tensor(ir_index, relative_index)) 

35 args.append(descs) 

36 return args 

37 

38 

39def _read_infer_attr(attrs, index: int, ir_type: str): 

40 getter_name, _ = _get_runtime_attr_spec(ir_type, index) 

41 return getattr(attrs, getter_name)(index) 

42 

43 

44def _build_infer_attrs(ctx, ir_attrs: list) -> dict: 

45 if not ir_attrs: 

46 return {} 

47 attrs = ctx.get_attrs() 

48 return { 

49 item["name"]: _read_infer_attr(attrs, index, item["type"]) 

50 for index, item in enumerate(ir_attrs) 

51 } 

52 

53 

54def _validate_tensor_desc(desc, output_index: int) -> None: 

55 from ge.runtime import TensorDesc 

56 

57 if not isinstance(desc, TensorDesc): 

58 raise TypeError( 

59 f"infer_meta output[{output_index}] must be TensorDesc, " 

60 f"got {type(desc).__name__}" 

61 ) 

62 

63 

64def _flatten_infer_outputs(ir_outputs: list, result) -> tuple: 

65 if not ir_outputs: 

66 return [], [] 

67 

68 if len(ir_outputs) == 1: 

69 kind = ir_outputs[0]["kind"] 

70 if kind == OutputType.REQUIRED: 

71 _validate_tensor_desc(result, 0) 

72 return [result], [1] 

73 if kind == OutputType.DYNAMIC: 

74 if not isinstance(result, (list, tuple)): 

75 raise TypeError( 

76 f"infer_meta output[0] is dynamic, must return list, " 

77 f"got {type(result).__name__}" 

78 ) 

79 for i, desc in enumerate(result): 

80 _validate_tensor_desc(desc, i) 

81 flattened = list(result) 

82 return flattened, [len(flattened)] 

83 

84 if not isinstance(result, (list, tuple)): 

85 raise TypeError( 

86 "infer_meta must return list/tuple for multiple outputs, " 

87 f"got {type(result).__name__}" 

88 ) 

89 if len(result) != len(ir_outputs): 

90 raise TypeError( 

91 f"infer_meta return count {len(result)} != output count {len(ir_outputs)}" 

92 ) 

93 flattened = [] 

94 slot_sizes = [] 

95 for ir_index, (item, desc) in enumerate(zip(ir_outputs, result)): 

96 kind = item["kind"] 

97 if kind == OutputType.REQUIRED: 

98 _validate_tensor_desc(desc, ir_index) 

99 flattened.append(desc) 

100 slot_sizes.append(1) 

101 elif kind == OutputType.DYNAMIC: 

102 if not isinstance(desc, (list, tuple)): 

103 raise TypeError( 

104 f"infer_meta output[{ir_index}] is dynamic, must return list" 

105 ) 

106 for i, d in enumerate(desc): 

107 _validate_tensor_desc(d, ir_index) 

108 flattened.extend(desc) 

109 slot_sizes.append(len(desc)) 

110 return flattened, slot_sizes 

111 

112 

113def call_infer_meta(op_type: str, ir_meta: Optional[dict], ctx) -> list: 

114 try: 

115 proto = get_registered_op_proto_by_op_type(op_type) 

116 infer_func = proto.infer_func 

117 args = _build_infer_inputs(ctx, ir_meta["inputs"]) 

118 kwargs = _build_infer_attrs(ctx, ir_meta["attrs"]) 

119 result = infer_func(*args, **kwargs) 

120 flattened, slot_sizes = _flatten_infer_outputs(ir_meta["outputs"], result) 

121 for ir_index, item in enumerate(ir_meta["outputs"]): 

122 kind = item["kind"] 

123 if kind == OutputType.DYNAMIC: 

124 instance_num = ctx.get_dynamic_output_num(ir_index) 

125 actual_num = slot_sizes[ir_index] 

126 if instance_num != actual_num: 

127 raise TypeError( 

128 f"infer_meta dynamic output[{ir_index}] instance count mismatch: " 

129 f"expected {instance_num}, got {actual_num}" 

130 ) 

131 return [ 

132 ( 

133 list(desc.shape.origin_shape.dims), 

134 list(desc.shape.storage_shape.dims), 

135 int(desc.data_type), 

136 ) 

137 for desc in flattened 

138 ] 

139 finally: 

140 ctx._invalidate()