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

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

16Distributed implementation for Scatter operator. 

17""" 

18 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_scatter_args(input_tensor, dim, index, src): 

25 return (input_tensor, dim, index, src), {} 

26 

27 

28class ScatterDistributedOp(DistributedOp): 

29 """Distributed implementation for torch.scatter.""" 

30 

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

32 """ 

33 Preprocess arguments for Scatter operator. 

34 

35 Args: 

36 args (tuple): Input arguments (input, dim, index, src). 

37 kwargs (dict): Keyword arguments (unused for scatter). 

38 

39 Returns: 

40 tuple: (local_args, local_kwargs, cache_values) where local_args contains 

41 local tensors and cache_values contains layouts plus dim. 

42 """ 

43 args, kwargs = _normalize_scatter_args(*args, **kwargs) 

44 input_tensor, dim, index_tensor, src = args 

45 

46 input_local = input_tensor.to_local() 

47 index_local = index_tensor.to_local() if hasattr(index_tensor, '_layout') else index_tensor 

48 src_local = src.to_local() if hasattr(src, '_layout') else src 

49 local_args = (input_local, dim, index_local, src_local) 

50 local_kwargs = {} 

51 

52 cache_values = [ 

53 input_tensor.layout, 

54 dim, 

55 index_tensor.layout if hasattr(index_tensor, '_layout') else None, 

56 src.layout if hasattr(src, '_layout') else None, 

57 ] 

58 return local_args, local_kwargs, cache_values 

59 

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

61 """ 

62 Infer output layout for Scatter operator. 

63 

64 Rules: 

65 1. Input must not have Partial status. 

66 2. Input must be a DTensor with a valid layout. 

67 3. dim must be an integer within the valid range [-ndim, ndim-1]. 

68 4. The scatter dimension must be replicated (not sharded). 

69 5. Index layout must match input layout (if index is a DTensor). 

70 6. Src layout must match input layout (if src is a DTensor). 

71 7. Output layout is identical to input layout. 

72 

73 Args: 

74 cache_values (list): [input_layout, dim, index_layout, src_layout] 

75 where index_layout and src_layout may be None. 

76 

77 Returns: 

78 tuple: ((output_layout,), None) 

79 

80 Raises: 

81 ValueError: If any rule above is violated. 

82 """ 

83 if not self._allow_partial_inputs: 

84 self._check_partial_inputs([cache_values[0]]) 

85 

86 input_layout = cache_values[0] 

87 if input_layout is None: 

88 raise ValueError( 

89 f"For {self.op_name}, input should be a DTensor with a valid layout, " 

90 f"but got None." 

91 ) 

92 

93 dim = cache_values[1] 

94 if not isinstance(dim, int): 

95 raise ValueError( 

96 f"For {self.op_name}, dim should be an integer, " 

97 f"but got {type(dim)}." 

98 ) 

99 

100 alias_map = input_layout.alias_tensor_map 

101 ndim = len(alias_map) 

102 

103 if dim < 0: 

104 dim += ndim 

105 

106 if dim < 0 or dim >= ndim: 

107 raise ValueError( 

108 f"For {self.op_name}, dim should be in range [{-ndim}, {ndim - 1}], " 

109 f"but got {cache_values[1]}." 

110 ) 

111 

112 # Scatter dimension must be replicated 

113 dim_alias = alias_map[dim] 

114 if dim_alias != "None": 

115 raise ValueError( 

116 f"For {self.op_name}, scatter dim should be replicated, " 

117 f"but dim {cache_values[1]} is mapped to {dim_alias}." 

118 ) 

119 

120 # Index layout must match input layout 

121 index_layout = cache_values[2] 

122 if index_layout is not None: 

123 index_alias = index_layout.alias_tensor_map 

124 if index_alias != alias_map: 

125 raise ValueError( 

126 f"For {self.op_name}, index layout should match input layout, " 

127 f"but got {index_alias} vs {alias_map}." 

128 ) 

129 

130 # Src layout must match input layout 

131 src_layout = cache_values[3] 

132 if src_layout is not None: 

133 src_alias = src_layout.alias_tensor_map 

134 if src_alias != alias_map: 

135 raise ValueError( 

136 f"For {self.op_name}, src layout should match input layout, " 

137 f"but got {src_alias} vs {alias_map}." 

138 ) 

139 

140 return ((input_layout,), None) 

141