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

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 ScatterUpdate operator. 

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_scatter_update_args(input_tensor, indices, updates): 

26 return (input_tensor, indices, updates), {} 

27 

28 

29class ScatterUpdateDistributedOp(DistributedOp): 

30 """Distributed implementation for ScatterUpdate operator.""" 

31 

32 def __init__(self, op_name: str): 

33 super().__init__(op_name) 

34 self._allow_partial_inputs = True 

35 

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

37 """ 

38 Preprocess arguments for ScatterUpdate operator. 

39 

40 Args: 

41 args (tuple): Input arguments containing input_tensor, indices, and updates. 

42 kwargs (dict): Keyword arguments (none expected). 

43 

44 Returns: 

45 tuple: (local_args, local_kwargs, cache_values) 

46 """ 

47 args, kwargs = _normalize_scatter_update_args(*args, **kwargs) 

48 x, indices, updates = args 

49 

50 local_args = ( 

51 x.to_local() if hasattr(x, '_layout') else x, 

52 indices.to_local() if hasattr(indices, '_layout') else indices, 

53 updates.to_local() if hasattr(updates, '_layout') else updates, 

54 ) 

55 

56 cache_values = [ 

57 x.layout if hasattr(x, '_layout') else None, 

58 indices.layout if hasattr(indices, '_layout') else None, 

59 updates.layout if hasattr(updates, '_layout') else None, 

60 ] 

61 local_kwargs = kwargs 

62 return local_args, local_kwargs, cache_values 

63 

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

65 """ 

66 Infer output layout for ScatterUpdate operator. 

67 

68 Rules: 

69 1. All three inputs must be DTensors (layout not None). 

70 2. Input first dimension (indexed/scatter dim) cannot be sharded. 

71 3. Indices must be all-Replicate on every dimension. 

72 4. updates[0:indices_ndim] prefix dimensions cannot be sharded. 

73 5. updates rank must equal indices_ndim + input_ndim - 1. 

74 6. updates[indices_ndim:] sharding must match input[1:] sharding. 

75 7. Output layout inherits input layout (including Partial status). 

76 

77 Args: 

78 cache_values (list): [input_layout, indices_layout, updates_layout] 

79 

80 Returns: 

81 tuple: ((output_layout,), None) 

82 

83 Raises: 

84 ValueError: If any validation rule above is violated. 

85 """ 

86 input_layout, indices_layout, updates_layout = cache_values 

87 

88 if input_layout is None: 

89 raise ValueError( 

90 f"For {self.op_name}, input layout should not be None" 

91 ) 

92 

93 if indices_layout is None: 

94 raise ValueError( 

95 f"For {self.op_name}, indices must be a DTensor when input is a DTensor" 

96 ) 

97 

98 if updates_layout is None: 

99 raise ValueError( 

100 f"For {self.op_name}, updates must be a DTensor when input is a DTensor" 

101 ) 

102 

103 # Partial inputs are intentionally allowed. The output inherits Partial status 

104 # from the input layout (see lines below), making this a Partial-preserving op. 

105 if not self._allow_partial_inputs: 

106 self._check_partial_inputs([input_layout]) 

107 

108 self._validate_strategy(input_layout, indices_layout, updates_layout) 

109 

110 output_layout = Layout( 

111 mesh_shape=input_layout.mesh_shape, 

112 alias_name=input_layout.alias_name, 

113 rank_list=input_layout.rank_list, 

114 ) 

115 output_layout = output_layout(*input_layout.alias_tensor_map) 

116 

117 for i, partial_op in enumerate(input_layout.partial): 

118 if partial_op is not None: 

119 dev_axis_name = input_layout.alias_name[i] 

120 output_layout.set_partial_by_dev_axis(dev_axis_name, partial_op) 

121 

122 return ((output_layout,), None) 

123 

124 def _validate_strategy(self, input_layout, indices_layout, updates_layout): 

125 """Validate sharding strategy for ScatterUpdate.""" 

126 input_map = input_layout.alias_tensor_map 

127 indices_map = indices_layout.alias_tensor_map 

128 updates_map = updates_layout.alias_tensor_map 

129 

130 if not input_map: 

131 raise ValueError( 

132 f"For {self.op_name}, input tensor map should not be empty" 

133 ) 

134 

135 if input_map[0] != "None": 

136 raise ValueError( 

137 f"For {self.op_name}, first dimension of input cannot be sharded, " 

138 f"but it is sharded on '{input_map[0]}'" 

139 ) 

140 

141 for i, axis in enumerate(indices_map): 

142 if axis != "None": 

143 raise ValueError( 

144 f"For {self.op_name}, indices cannot be sharded, " 

145 f"but dimension {i} is sharded on '{axis}'" 

146 ) 

147 

148 indices_ndim = len(indices_map) 

149 for i in range(indices_ndim): 

150 if i >= len(updates_map): 

151 raise ValueError( 

152 f"For {self.op_name}, updates rank is smaller than indices rank" 

153 ) 

154 if updates_map[i] != "None": 

155 raise ValueError( 

156 f"For {self.op_name}, first {indices_ndim} dimensions of updates " 

157 f"cannot be sharded, but dimension {i} is sharded on '{updates_map[i]}'" 

158 ) 

159 

160 expected_updates_ndim = indices_ndim + len(input_map) - 1 

161 if len(updates_map) != expected_updates_ndim: 

162 raise ValueError( 

163 f"For {self.op_name}, updates rank mismatch. " 

164 f"Expected {expected_updates_ndim}, got {len(updates_map)}" 

165 ) 

166 

167 for i in range(1, len(input_map)): 

168 updates_idx = indices_ndim + i - 1 

169 if input_map[i] != updates_map[updates_idx]: 

170 raise ValueError( 

171 f"For {self.op_name}, updates sharding must match input[1:]. " 

172 f"Mismatch at input dim {i}: '{input_map[i]}' != '{updates_map[updates_idx]}'" 

173 ) 

174