Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/utils/ge_utils.py: 98%

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"""GE public utility APIs.""" 

14 

15import ctypes 

16from typing import List, Tuple 

17 

18from ge._capi.pyge_utils_wrapper import ge_utils_lib 

19from ge.graph import Graph, Node 

20 

21 

22class GeUtils: 

23 """Public GE utility APIs.""" 

24 

25 @staticmethod 

26 def infer_shape(graph: "Graph", input_shapes: List[List[int]]) -> None: 

27 """Infer shapes for a graph with the given input shapes. 

28 

29 This API only runs shape inference and does not apply other graph 

30 optimizations, such as constant folding or dead edge elimination. 

31 

32 Args: 

33 graph: Graph object to run shape inference on. 

34 input_shapes: Input shape list. Each element describes one graph input shape. 

35 

36 Raises: 

37 TypeError: If input_shapes is not a list of integer shape lists. 

38 RuntimeError: If shape inference fails. 

39 """ 

40 flat_dims, shape_ranks = GeUtils._normalize_input_shapes(input_shapes) 

41 dims_num = len(flat_dims) 

42 shape_num = len(shape_ranks) 

43 dims_arr = (ctypes.c_int64 * dims_num)(*flat_dims) if dims_num > 0 else None 

44 shape_ranks_arr = (ctypes.c_size_t * shape_num)(*shape_ranks) if shape_num > 0 else None 

45 ret = ge_utils_lib.GeApiWrapper_GeUtils_InferShape( 

46 graph._handle, dims_arr, dims_num, shape_ranks_arr, shape_num 

47 ) 

48 if ret != 0: 

49 raise RuntimeError("Failed to infer shape") 

50 

51 @staticmethod 

52 def check_node_support_on_aicore(node: "Node") -> Tuple[bool, str]: 

53 """Check whether a node is supported on AICore. 

54 

55 Args: 

56 node: Node object to check. 

57 

58 Returns: 

59 Tuple of (is_supported, unsupported_reason). 

60 

61 Raises: 

62 RuntimeError: If AICore support checking fails. 

63 """ 

64 is_supported = ctypes.c_bool(False) 

65 unsupported_reason = ctypes.POINTER(ctypes.c_char)() 

66 ret = ge_utils_lib.GeApiWrapper_GeUtils_CheckNodeSupportOnAicore( 

67 node._handle, ctypes.byref(is_supported), ctypes.byref(unsupported_reason) 

68 ) 

69 try: 

70 if ret != 0: 

71 raise RuntimeError("Failed to check node support on AICore") 

72 reason = ctypes.string_at(unsupported_reason).decode("utf-8") if unsupported_reason else "" 

73 return bool(is_supported.value), reason 

74 finally: 

75 if unsupported_reason: 

76 ge_utils_lib.GeApiWrapper_GeUtils_FreeString(unsupported_reason) 

77 

78 @staticmethod 

79 def _normalize_input_shapes( 

80 input_shapes: List[List[int]], 

81 ) -> Tuple[List[int], List[int]]: 

82 if not isinstance(input_shapes, list): 

83 raise TypeError("input_shapes must be a list of shape lists") 

84 

85 flat_dims: List[int] = [] 

86 shape_ranks: List[int] = [] 

87 for shape in input_shapes: 

88 if not isinstance(shape, list): 

89 raise TypeError("input_shapes must be a list of shape lists") 

90 dims: List[int] = [] 

91 for dim in shape: 

92 if not isinstance(dim, int): 

93 raise TypeError("each shape dim must be an integer") 

94 dims.append(dim) 

95 shape_ranks.append(len(dims)) 

96 flat_dims.extend(dims) 

97 return flat_dims, shape_ranks