Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / hardware.py: 97%

1# Copyright 2025-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"""hardware abstraction""" 

16 

17from __future__ import annotations 

18 

19import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim 

20from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger 

21 

22 

23class Type: 

24 """Machine type""" 

25 

26 name: str 

27 levels: int # levels in hierarchy 

28 level_bound_number: list[int] # devices per level 

29 level_bandwidth: list[int] # bandwidth (GB/s) per level 

30 

31 def __init__(self, name, bounds, bandwidths): 

32 self.name = name 

33 self.level_bound_number = bounds 

34 self.level_bandwidth = bandwidths 

35 if len(bounds) != len(bandwidths): 

36 raise ValueError("bounds and bandwidths must have the same length") 

37 self.levels = len(bounds) 

38 

39 def __str__(self): 

40 return self.name 

41 

42 def __repr__(self): 

43 return str(self) 

44 

45 def devices_below_level(self, level): 

46 """Number of devices below the given hierarchy level""" 

47 devices = 1 

48 for lvl in range(min(level, self.levels)): 

49 devices *= self.level_bound_number[lvl] 

50 return devices 

51 

52 def intra_node_num(self): 

53 """Number of devices in a node""" 

54 return self.devices_below_level(1) 

55 

56 def levels_used(self, device_number): 

57 """Number of hierarchy level used""" 

58 devices = 1 

59 for lvl in range(self.levels): 

60 if self.level_bound_number[lvl]: 

61 devices *= self.level_bound_number[lvl] 

62 if device_number <= devices: 

63 return lvl 

64 else: 

65 return lvl 

66 return self.levels 

67 

68 def level_assign(self, dp=1, tp=1, cp=1, pp=1): 

69 """device assignment of the different parallel dimensions""" 

70 device_number = dp * tp * cp * pp 

71 logger.debug("DP = %d, TP = %d, CP = %d, PP = %d", dp, tp, cp, pp) 

72 assignment = {} 

73 assignment[Dim.TP] = [] 

74 assignment[Dim.CP] = [] 

75 assignment[Dim.DP] = [] 

76 assignment[Dim.PP] = [] 

77 for level in range(self.levels): 

78 bound = self.level_bound_number[level] 

79 if bound: 

80 level_device_number = min(device_number, bound) 

81 device_number = device_number // bound 

82 else: 

83 level_device_number = device_number 

84 remaining_devices = max(level_device_number, 1) 

85 

86 tp_level = min(tp, remaining_devices) 

87 assignment[Dim.TP].append(tp_level) 

88 tp = tp // tp_level 

89 remaining_devices = remaining_devices // tp_level 

90 

91 cp_level = min(cp, remaining_devices) 

92 assignment[Dim.CP].append(cp_level) 

93 cp = cp // cp_level 

94 remaining_devices = remaining_devices // cp_level 

95 

96 dp_level = min(dp, remaining_devices) 

97 assignment[Dim.DP].append(dp_level) 

98 dp = dp // dp_level 

99 remaining_devices = remaining_devices // dp_level 

100 

101 pp_level = min(pp, remaining_devices) 

102 assignment[Dim.PP].append(pp_level) 

103 pp = pp // pp_level 

104 remaining_devices = remaining_devices // pp_level 

105 

106 return assignment 

107 

108 

109# Device_A2 = Machine(devices_per_node=8, inter_node_bw=10, intra_node_bw=50) 

110Device_A2 = Type(name="A2", bounds=[8, None], bandwidths=[50, 10]) 

111Device_A3 = Type( 

112 name="A3", bounds=[16, 24, None], bandwidths=[200, 25, 10] 

113) 

114device_map = { 

115 "A2": Device_A2, 

116 "A3": Device_A3, 

117 "V100": Type(name="V100", bounds=[8, None], bandwidths=[50, 10]), 

118} 

119 

120 

121class Machine: 

122 """Hardware description""" 

123 

124 number: int 

125 device: Type 

126 

127 def __init__(self, number, device): 

128 self.number = number 

129 if isinstance(device, int): 

130 if device == 2: 

131 self.device = Device_A2 

132 elif device == 3: 

133 self.device = Device_A3 

134 else: 

135 raise ValueError(f"Ascend A{device} unknown") 

136 elif isinstance(device, str): 

137 if device not in device_map: 

138 raise ValueError( 

139 f"Device {device} is not supported. " 

140 f"Supported devices: {list(device_map.keys())}" 

141 ) 

142 self.device = device_map[device] 

143 else: 

144 self.device = device 

145 

146 def update_num_if_none(self, num): 

147 """Assign number of device if not already precised""" 

148 if self.number is None: 

149 self.number = num 

150 

151 def pipeline_bound(self): 

152 """Return pipeline bound from hardware topology because as pipeline may currently not cross hierarchy levels""" 

153 max_bound = 1 

154 devices = self.number 

155 while devices > 1: 

156 max_bound = max( 

157 max_bound, 

158 devices 

159 // self.device.devices_below_level( 

160 self.device.levels_used(devices) 

161 ), 

162 ) 

163 devices = devices // 2 

164 # devices = self.devices_below_level(self.levels_used(device_number)) 

165 # return device_number // devices 

166 return max_bound 

167 

168 

169def prime_factors(n): 

170 """Decompose n into a product of prime factors""" 

171 divisor = 2 

172 factors = [] 

173 while n > 1: 

174 while n % divisor != 0: 

175 divisor += 1 

176 factors.append(divisor) 

177 n = n // divisor 

178 return factors 

179 

180 

181def all_factors_combinations(factors): 

182 """Computes all divisors from a prime factor list""" 

183 def rec_factors(n, factors): 

184 combinations = {n} 

185 for u in set(factors): 

186 remaining = factors.copy() 

187 remaining.remove(u) 

188 combinations = combinations.union(rec_factors(n * u, remaining)) 

189 return combinations 

190 return rec_factors(1, factors) 

191 

192 

193def all_divisors(n, reverse=False, min_bound=1, max_bound=float("inf")): 

194 """Computes all divisors of an integer n""" 

195 divisors = sorted( 

196 all_factors_combinations(prime_factors(n)), reverse=reverse 

197 ) 

198 div_in_bound = [] 

199 for d in divisors: 

200 if min_bound <= d <= max_bound: 

201 div_in_bound.append(d) 

202 

203 return div_in_bound 

204 

205 

206def from_prime_factors(factors): 

207 """Compute a number from its prime factor decomposition""" 

208 number = 1 

209 for f in factors: 

210 number *= f 

211 return number 

212 

213 

214def split_node(n, device): 

215 """Split decompositions into intra & inter devices""" 

216 devices_per_node = device.intra_node_num() 

217 nodes = prime_factors(max(1, n // devices_per_node)) 

218 intra = prime_factors(min(n, devices_per_node)) 

219 return [intra, nodes] 

220 

221 

222def unique_factors(factors): 

223 """Remove duplicates. Factors are sorted""" 

224 offset = 0 

225 for i, f in enumerate(factors[:-1]): 

226 j = i - offset 

227 if factors[j + 1] == f: 

228 factors.pop(j) 

229 offset += 1 

230 return factors 

231 

232 

233def highest_power_of_2_divisor(divisor_of): 

234 """Compute the highest number that is both a divisor of 'divisor_of' and a power of 2""" 

235 divisor = 1 

236 factors = prime_factors(divisor_of) 

237 for f in factors: 

238 if f == 2: 

239 divisor *= f 

240 return divisor