Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_split.py: 50%

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""" 

16Distributed implementation for Split operator. 

17""" 

18 

19import copy 

20import math 

21from typing import Tuple 

22 

23from .parallel_ops import DistributedOp 

24 

25 

26def _normalize_split_with_size_args(x, split_sections, dim): 

27 return (x, split_sections, dim), {} 

28 

29 

30def _normalize_split_args(x, split_size_or_sections, dim=0): 

31 return (x, split_size_or_sections, dim), {} 

32 

33 

34def _normalize_split_tensor_args(x, split_size, dim): 

35 return (x, split_size, dim), {} 

36 

37 

38def _normalize_tensor_split_args(x, indices_or_sections, dim=0): 

39 return (x, indices_or_sections, dim), {} 

40 

41 

42class SplitWithSizeDistributedOp(DistributedOp): 

43 """Distributed implementation for SplitWithSize operator.""" 

44 

45 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

46 """ 

47 Preprocess arguments for SplitWithSize operator. 

48 

49 Args: 

50 args (tuple): Input arguments (input, split_sections, dim). 

51 kwargs (dict): Keyword arguments (empty for this operator). 

52 

53 Returns: 

54 tuple: (local_args, local_kwargs, cache_values) 

55 """ 

56 args, kwargs = _normalize_split_with_size_args(*args, **kwargs) 

57 input_tensor, split_sections, dim = args 

58 output_num = len(split_sections) 

59 local_args = (input_tensor.to_local(), split_sections, dim) 

60 local_kwargs = {} 

61 cache_values = [input_tensor.layout, dim, output_num] 

62 return local_args, local_kwargs, cache_values 

63 

64 # pylint: disable=W0237 

65 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

66 """ 

67 Infer output layouts for SplitWithSize operator. 

68 

69 Rules: 

70 1. Input must not have Partial status. 

71 2. The split dimension must not be sharded. 

72 

73 Args: 

74 cache_values (list): [input_layout, dim, output_num] 

75 

76 Returns: 

77 tuple: ((output_layouts,), None) 

78 

79 Raises: 

80 ValueError: If any rule above is violated. 

81 """ 

82 layout = cache_values[0] 

83 dim = cache_values[1] 

84 output_num = cache_values[2] 

85 

86 if not self._allow_partial_inputs: 

87 self._check_partial_inputs([layout]) 

88 

89 in_tensor_map = layout.alias_tensor_map 

90 ndim = len(in_tensor_map) 

91 

92 if dim < 0: 

93 dim = ndim + dim 

94 if not 0 <= dim < ndim: 

95 raise ValueError( 

96 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

97 f"but got {dim}." 

98 ) 

99 

100 if in_tensor_map[dim] != "None": 

101 raise ValueError( 

102 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

103 f"but got layout: {layout}." 

104 ) 

105 

106 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

107 

108 

109class SplitWithSizeViewDistributedOp(DistributedOp): 

110 """Distributed implementation for SplitWithSizeView operator.""" 

111 

112 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

113 """ 

114 Preprocess arguments for SplitWithSizeView operator. 

115 

116 Args: 

117 args (tuple): Input arguments (input, split_sections, dim). 

118 kwargs (dict): Keyword arguments (empty for this operator). 

119 

120 Returns: 

121 tuple: (local_args, local_kwargs, cache_values) 

122 """ 

123 args, kwargs = _normalize_split_with_size_args(*args, **kwargs) 

124 input_tensor, split_sections, dim = args 

125 output_num = len(split_sections) 

126 local_args = (input_tensor.to_local(), split_sections, dim) 

127 local_kwargs = {} 

128 cache_values = [input_tensor.layout, dim, output_num] 

129 return local_args, local_kwargs, cache_values 

130 

131 # pylint: disable=W0237 

132 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

133 """ 

134 Infer output layouts for SplitWithSizeView operator. 

135 

136 Rules: 

137 1. Input must not have Partial status. 

138 2. The split dimension must not be sharded. 

139 

140 Args: 

141 cache_values (list): [input_layout, dim, output_num] 

142 

143 Returns: 

144 tuple: ((output_layouts,), None) 

145 

146 Raises: 

147 ValueError: If any rule above is violated. 

148 """ 

149 layout = cache_values[0] 

150 dim = cache_values[1] 

151 output_num = cache_values[2] 

152 

153 if not self._allow_partial_inputs: 

154 self._check_partial_inputs([layout]) 

155 

156 in_tensor_map = layout.alias_tensor_map 

157 ndim = len(in_tensor_map) 

158 

159 if dim < 0: 

160 dim = ndim + dim 

161 if not 0 <= dim < ndim: 

162 raise ValueError( 

163 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

164 f"but got {dim}." 

165 ) 

166 

167 if in_tensor_map[dim] != "None": 

168 raise ValueError( 

169 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

170 f"but got layout: {layout}." 

171 ) 

172 

173 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

174 

175 

176class SplitDistributedOp(DistributedOp): 

177 """Distributed implementation for Split operator (MindSpore Split and torch.split).""" 

178 

179 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

180 """ 

181 Preprocess arguments for Split operator. 

182 

183 Args: 

184 args (tuple): Input arguments (input, split_size_or_sections[, dim]). 

185 kwargs (dict): Keyword arguments (may contain dim). 

186 

187 Returns: 

188 tuple: (local_args, local_kwargs, cache_values) 

189 """ 

190 args, kwargs = _normalize_split_args(*args, **kwargs) 

191 input_tensor, split_size_or_sections, dim = args 

192 

193 if isinstance(split_size_or_sections, int): 

194 output_num = math.ceil(input_tensor.shape[dim] / split_size_or_sections) 

195 else: 

196 output_num = len(split_size_or_sections) 

197 

198 local_args = (input_tensor.to_local(), split_size_or_sections, dim) 

199 local_kwargs = {} 

200 cache_values = [input_tensor.layout, dim, output_num] 

201 return local_args, local_kwargs, cache_values 

202 

203 # pylint: disable=W0237 

204 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

205 """ 

206 Infer output layouts for Split operator. 

207 

208 Rules: 

209 1. Input must not have Partial status. 

210 2. The split dimension must not be sharded. 

211 3. dim must be within the valid range of input dimensions. 

212 

213 Args: 

214 cache_values (list): [input_layout, dim, output_num] 

215 

216 Returns: 

217 tuple: ((output_layouts,), None) 

218 

219 Raises: 

220 ValueError: If any rule above is violated. 

221 """ 

222 layout = cache_values[0] 

223 dim = cache_values[1] 

224 output_num = cache_values[2] 

225 

226 if not self._allow_partial_inputs: 

227 self._check_partial_inputs([layout]) 

228 

229 in_tensor_map = layout.alias_tensor_map 

230 ndim = len(in_tensor_map) 

231 

232 if dim < 0: 

233 dim = ndim + dim 

234 if not 0 <= dim < ndim: 

235 raise ValueError( 

236 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

237 f"but got {dim}." 

238 ) 

239 

240 if in_tensor_map[dim] != "None": 

241 raise ValueError( 

242 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

243 f"but got layout: {layout}." 

244 ) 

245 

246 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

247 

248 

249class SplitTensorDistributedOp(DistributedOp): 

250 """Distributed implementation for SplitTensor operator.""" 

251 

252 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

253 """ 

254 Preprocess arguments for SplitTensor operator. 

255 

256 Args: 

257 args (tuple): Input arguments (input, split_size, dim). 

258 kwargs (dict): Keyword arguments (empty for this operator). 

259 

260 Returns: 

261 tuple: (local_args, local_kwargs, cache_values) 

262 """ 

263 args, kwargs = _normalize_split_tensor_args(*args, **kwargs) 

264 input_tensor, split_size, dim = args 

265 output_num = math.ceil(input_tensor.shape[dim] / split_size) 

266 local_args = (input_tensor.to_local(), split_size, dim) 

267 local_kwargs = {} 

268 cache_values = [input_tensor.layout, dim, output_num] 

269 return local_args, local_kwargs, cache_values 

270 

271 # pylint: disable=W0237 

272 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

273 """ 

274 Infer output layouts for SplitTensor operator. 

275 

276 Rules: 

277 1. Input must not have Partial status. 

278 2. The split dimension must not be sharded. 

279 

280 Args: 

281 cache_values (list): [input_layout, dim, output_num] 

282 

283 Returns: 

284 tuple: ((output_layouts,), None) 

285 

286 Raises: 

287 ValueError: If any rule above is violated. 

288 """ 

289 layout = cache_values[0] 

290 dim = cache_values[1] 

291 output_num = cache_values[2] 

292 

293 if not self._allow_partial_inputs: 

294 self._check_partial_inputs([layout]) 

295 

296 in_tensor_map = layout.alias_tensor_map 

297 ndim = len(in_tensor_map) 

298 

299 if dim < 0: 

300 dim = ndim + dim 

301 if not 0 <= dim < ndim: 

302 raise ValueError( 

303 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

304 f"but got {dim}." 

305 ) 

306 

307 if in_tensor_map[dim] != "None": 

308 raise ValueError( 

309 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

310 f"but got layout: {layout}." 

311 ) 

312 

313 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

314 

315 

316class SplitTensorViewDistributedOp(DistributedOp): 

317 """Distributed implementation for SplitTensorView operator.""" 

318 

319 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

320 """ 

321 Preprocess arguments for SplitTensorView operator. 

322 

323 Args: 

324 args (tuple): Input arguments (input, split_size, dim). 

325 kwargs (dict): Keyword arguments (empty for this operator). 

326 

327 Returns: 

328 tuple: (local_args, local_kwargs, cache_values) 

329 """ 

330 args, kwargs = _normalize_split_tensor_args(*args, **kwargs) 

331 input_tensor, split_size, dim = args 

332 output_num = math.ceil(input_tensor.shape[dim] / split_size) 

333 local_args = (input_tensor.to_local(), split_size, dim) 

334 local_kwargs = {} 

335 cache_values = [input_tensor.layout, dim, output_num] 

336 return local_args, local_kwargs, cache_values 

337 

338 # pylint: disable=W0237 

339 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

340 """ 

341 Infer output layouts for SplitTensorView operator. 

342 

343 Rules: 

344 1. Input must not have Partial status. 

345 2. The split dimension must not be sharded. 

346 

347 Args: 

348 cache_values (list): [input_layout, dim, output_num] 

349 

350 Returns: 

351 tuple: ((output_layouts,), None) 

352 

353 Raises: 

354 ValueError: If any rule above is violated. 

355 """ 

356 layout = cache_values[0] 

357 dim = cache_values[1] 

358 output_num = cache_values[2] 

359 

360 if not self._allow_partial_inputs: 

361 self._check_partial_inputs([layout]) 

362 

363 in_tensor_map = layout.alias_tensor_map 

364 ndim = len(in_tensor_map) 

365 

366 if dim < 0: 

367 dim = ndim + dim 

368 if not 0 <= dim < ndim: 

369 raise ValueError( 

370 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

371 f"but got {dim}." 

372 ) 

373 

374 if in_tensor_map[dim] != "None": 

375 raise ValueError( 

376 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

377 f"but got layout: {layout}." 

378 ) 

379 

380 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

381 

382 

383class TensorSplitDistributedOp(DistributedOp): 

384 """Distributed implementation for tensor_split operator.""" 

385 

386 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

387 """ 

388 Preprocess arguments for tensor_split operator. 

389 

390 Args: 

391 args (tuple): Input arguments (input, indices_or_sections[, dim]). 

392 kwargs (dict): Keyword arguments (may contain dim). 

393 

394 Returns: 

395 tuple: (local_args, local_kwargs, cache_values) 

396 """ 

397 args, kwargs = _normalize_tensor_split_args(*args, **kwargs) 

398 input_tensor, indices_or_sections, dim = args 

399 

400 if isinstance(indices_or_sections, int): 

401 output_num = indices_or_sections 

402 elif isinstance(indices_or_sections, (list, tuple)): 

403 output_num = len(indices_or_sections) + 1 

404 elif hasattr(indices_or_sections, "shape") and len(indices_or_sections.shape) == 1: 

405 output_num = indices_or_sections.shape[0] + 1 

406 else: 

407 raise TypeError( 

408 f"For {self.op_name}, indices_or_sections must be an integer, " 

409 f"list, tuple, or 1D tensor." 

410 ) 

411 

412 local_indices = indices_or_sections 

413 if hasattr(indices_or_sections, "_layout"): 

414 local_indices = indices_or_sections.to_local() 

415 

416 local_args = (input_tensor.to_local(), local_indices, dim) 

417 local_kwargs = {} 

418 cache_values = [input_tensor.layout, dim, output_num] 

419 return local_args, local_kwargs, cache_values 

420 

421 # pylint: disable=W0237 

422 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

423 """ 

424 Infer output layouts for tensor_split operator. 

425 

426 Rules: 

427 1. Input must not have Partial status. 

428 2. The split dimension must not be sharded. 

429 3. dim must be within the valid range of input dimensions. 

430 

431 Args: 

432 cache_values (list): [input_layout, dim, output_num] 

433 

434 Returns: 

435 tuple: ((output_layouts,), None) 

436 

437 Raises: 

438 ValueError: If any rule above is violated. 

439 """ 

440 layout = cache_values[0] 

441 dim = cache_values[1] 

442 output_num = cache_values[2] 

443 

444 if not self._allow_partial_inputs: 

445 self._check_partial_inputs([layout]) 

446 

447 in_tensor_map = layout.alias_tensor_map 

448 ndim = len(in_tensor_map) 

449 

450 if dim < 0: 

451 dim = ndim + dim 

452 if not 0 <= dim < ndim: 

453 raise ValueError( 

454 f"For {self.op_name}, dimension should be in range [0, {ndim}), " 

455 f"but got {dim}." 

456 ) 

457 

458 if in_tensor_map[dim] != "None": 

459 raise ValueError( 

460 f"For {self.op_name}, can not split tensor at sharded axis[{dim}], " 

461 f"but got layout: {layout}." 

462 ) 

463 

464 return (tuple(copy.deepcopy(layout) for _ in range(output_num)), None) 

465