Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / dump / l0.py: 97%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Whole-MegaKernel boundary snapshots and stage-only Tensor storage.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20import os 

21from pathlib import Path 

22from typing import Any 

23 

24import torch 

25 

26from hyper_multicore.diagnostics.dump.layout import _SAFE_FILE_NAME, _stage_names 

27 

28 

29@dataclass(frozen=True) 

30class _TensorPort: 

31 """One internal operator port referring to a logical graph Tensor.""" 

32 

33 operator: str 

34 stage: str 

35 op_type: str 

36 io: str 

37 index: int 

38 

39 

40@dataclass(frozen=True) 

41class _GraphTensor: 

42 """One logical graph Tensor and its L0 capture relationship.""" 

43 

44 name: str 

45 param_position: int 

46 declared_shape: tuple[int, ...] 

47 role: str 

48 producers: tuple[str, ...] 

49 consumers: tuple[str, ...] 

50 ports: tuple[_TensorPort, ...] 

51 

52 

53@dataclass 

54class _TensorSnapshot: 

55 """One Host-owned Tensor value plus export metadata.""" 

56 

57 graph_tensor: _GraphTensor 

58 phase: str 

59 original_dtype: str 

60 stored_dtype: str 

61 shape: tuple[int, ...] 

62 data: torch.Tensor 

63 

64 

65def _operator_type_name(operator: Any) -> str: 

66 op_type = operator.op_type 

67 value = getattr(op_type, "value", op_type) 

68 return str(value) 

69 

70 

71def _describe_graph_tensors(graph: Any) -> tuple[_GraphTensor, ...]: 

72 """Deduplicate ABI-bound graph tensors and classify their producers and consumers.""" 

73 relationships: dict[tuple[str, int], dict[str, Any]] = {} 

74 operators = graph.topological_sort() 

75 stage_names = _stage_names(operators) 

76 for operator in operators: 

77 op_type = _operator_type_name(operator) 

78 for index, tensor in enumerate(operator.inputs): 

79 key = (tensor.name, int(tensor.param_position)) 

80 relationship = relationships.setdefault( 

81 key, 

82 { 

83 "shape": tuple(int(dim) for dim in tensor.shape), 

84 "producers": [], 

85 "consumers": [], 

86 "ports": [], 

87 }, 

88 ) 

89 relationship["consumers"].append(operator.name) 

90 relationship["ports"].append( 

91 _TensorPort(operator.name, stage_names[operator.name], op_type, "input", index) 

92 ) 

93 for index, tensor in enumerate(operator.outputs): 

94 key = (tensor.name, int(tensor.param_position)) 

95 relationship = relationships.setdefault( 

96 key, 

97 { 

98 "shape": tuple(int(dim) for dim in tensor.shape), 

99 "producers": [], 

100 "consumers": [], 

101 "ports": [], 

102 }, 

103 ) 

104 relationship["producers"].append(operator.name) 

105 relationship["ports"].append( 

106 _TensorPort(operator.name, stage_names[operator.name], op_type, "output", index) 

107 ) 

108 

109 descriptions = [] 

110 for (name, param_position), relationship in relationships.items(): 

111 producers = tuple(relationship["producers"]) 

112 consumers = tuple(relationship["consumers"]) 

113 if not producers: 

114 role = "external_input" 

115 elif consumers: 

116 role = "stage_output" 

117 else: 

118 role = "final_output" 

119 descriptions.append( 

120 _GraphTensor( 

121 name=name, 

122 param_position=param_position, 

123 declared_shape=relationship["shape"], 

124 role=role, 

125 producers=producers, 

126 consumers=consumers, 

127 ports=tuple(relationship["ports"]), 

128 ) 

129 ) 

130 return tuple(sorted(descriptions, key=lambda tensor: (tensor.param_position, tensor.name))) 

131 

132 

133def _dtype_name(dtype: Any) -> str: 

134 return str(dtype).removeprefix("torch.") 

135 

136 

137def _snapshot_tensor(tensor: Any, graph_tensor: _GraphTensor, phase: str) -> _TensorSnapshot: 

138 if not torch.is_tensor(tensor): 

139 raise TypeError( 

140 f"dump Tensor {graph_tensor.name!r} at param_position={graph_tensor.param_position} " 

141 f"must map to torch.Tensor, got {type(tensor).__name__}" 

142 ) 

143 host_tensor = tensor.detach().contiguous().cpu().clone() 

144 dtype_name = _dtype_name(host_tensor.dtype) 

145 return _TensorSnapshot( 

146 graph_tensor=graph_tensor, 

147 phase=phase, 

148 original_dtype=dtype_name, 

149 stored_dtype=dtype_name, 

150 shape=tuple(int(dim) for dim in tensor.shape), 

151 data=host_tensor, 

152 ) 

153 

154 

155def _write_l0_stage_files( 

156 snapshot: _TensorSnapshot, 

157 invocation_dir: Path, 

158 window_dir: Path, 

159 safe_name: str, 

160) -> tuple[Path, list[dict[str, Any]]]: 

161 """Save one L0 Tensor under stages and hard-link its remaining ports.""" 

162 metadata = [] 

163 primary_path = None 

164 for port in snapshot.graph_tensor.ports: 

165 safe_stage = _SAFE_FILE_NAME.sub("_", port.stage) 

166 stage_dir = invocation_dir / "stages" / safe_stage 

167 stage_dir.mkdir(parents=True, exist_ok=True) 

168 stage_file_name = f"{port.io}_{port.index:02d}_{safe_name}.pt" 

169 stage_path = stage_dir / stage_file_name 

170 if primary_path is None: 

171 torch.save(snapshot.data, stage_path) 

172 primary_path = stage_path 

173 else: 

174 try: 

175 os.link(primary_path, stage_path) 

176 except OSError as error: 

177 raise OSError( 

178 f"failed to create stage-file hard link {stage_path} for {primary_path}" 

179 ) from error 

180 metadata.append( 

181 { 

182 "operator": port.operator, 

183 "stage": port.stage, 

184 "opType": port.op_type, 

185 "io": port.io, 

186 "index": port.index, 

187 "file": stage_path.relative_to(window_dir).as_posix(), 

188 } 

189 ) 

190 if primary_path is None: 

191 boundary_dir = invocation_dir / "stages" / "__kernel_boundary__" 

192 boundary_dir.mkdir(parents=True, exist_ok=True) 

193 boundary_io = "input" if snapshot.phase == "before" else "output" 

194 file_name = f"{boundary_io}_{snapshot.graph_tensor.param_position:02d}_{safe_name}.pt" 

195 primary_path = boundary_dir / file_name 

196 torch.save(snapshot.data, primary_path) 

197 return primary_path, metadata