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

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

17""" 

18 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_pad_args(x, pad, mode='constant', value=None): 

25 return (x, pad, mode, value), {} 

26 

27 

28class PadDistributedOp(DistributedOp): 

29 """Distributed implementation for Pad operator.""" 

30 

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

32 """ 

33 Preprocess arguments for Pad operator. 

34 

35 Args: 

36 args (tuple): Input tensor followed by pad arguments. 

37 kwargs (dict): Keyword arguments for pad. 

38 

39 Returns: 

40 tuple: (local_args, local_kwargs, cache_values) 

41 """ 

42 args, _ = _normalize_pad_args(*args, **kwargs) 

43 input_tensor = args[0] 

44 pad, mode, value = args[1], args[2], args[3] 

45 

46 local_args = (input_tensor.to_local(), pad, mode, value) 

47 local_kwargs = {} 

48 cache_values = [input_tensor.layout, pad] 

49 return local_args, local_kwargs, cache_values 

50 

51 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221 

52 """ 

53 Infer output layout for Pad operator. 

54 

55 Rules: 

56 1. Input must not have Partial status. 

57 2. pad must be a tuple or list with even length. 

58 3. The number of padded dimensions must not exceed input ndim. 

59 4. Any dimension with non-zero padding must not be sharded. 

60 5. Output layout is identical to the input layout. 

61 

62 Args: 

63 cache_values (list): [input_layout, pad]. 

64 

65 Returns: 

66 tuple: ((output_layout,), None) 

67 

68 Raises: 

69 ValueError: If input has Partial status, pad is invalid, or padding is 

70 attempted on a sharded dimension. 

71 """ 

72 if len(cache_values) != 2: 

73 raise ValueError( 

74 f"For {self.op_name}, cache_values length should be 2, but got {len(cache_values)}" 

75 ) 

76 

77 input_layout, pad = cache_values[0], cache_values[1] 

78 if input_layout is None: 

79 raise ValueError(f"For {self.op_name}, pad requires a valid input tensor layout.") 

80 

81 self._check_partial_inputs([input_layout]) 

82 

83 tensor_map = input_layout.alias_tensor_map 

84 ndim = len(tensor_map) 

85 

86 if not isinstance(pad, (tuple, list)): 

87 raise ValueError( 

88 f"For {self.op_name}, expected pad tuple or list, but got {type(pad)}" 

89 ) 

90 

91 pad_len = len(pad) 

92 

93 if pad_len % 2 != 0: 

94 raise ValueError(f"For {self.op_name}, Pad tuple length must be even, but got {pad_len}") 

95 

96 # Pytorch pad tuple format: (last_dim_left, last_dim_right, 2nd_last_left, 2nd_last_right, ...) 

97 # We need to check if any dimension being padded is currently sharded. 

98 num_padded_dims = pad_len // 2 

99 if num_padded_dims > ndim: 

100 raise ValueError( 

101 f"For {self.op_name}, Padding {num_padded_dims} dimensions but tensor only has {ndim} dimensions." 

102 ) 

103 

104 for i in range(num_padded_dims): 

105 # Calculate the dimension index in the tensor (from 0 to ndim-1) 

106 # pad index 0,1 -> last dimension (ndim - 1) 

107 # pad index 2,3 -> second to last dimension (ndim - 2) 

108 dim_index = ndim - 1 - i 

109 

110 pad_left = pad[2 * i] 

111 pad_right = pad[2 * i + 1] 

112 

113 # If padding is applied on this dimension 

114 if pad_left != 0 or pad_right != 0: 

115 axis_alias = tensor_map[dim_index] 

116 is_sharded = (any(alias != "None" for alias in axis_alias) 

117 if isinstance(axis_alias, (tuple, list)) 

118 else axis_alias != "None") 

119 

120 if is_sharded: 

121 raise ValueError( 

122 f"For {self.op_name}, Distributed Pad operator does not support padding " 

123 f"on a sharded dimension. " 

124 f"Dimension {dim_index} (alias: {axis_alias}) is sharded. " 

125 f"Please redistribute the tensor to Replicate status on this dimension before padding." 

126 ) 

127 

128 # If no sharded dimension is padded, the output layout is identical to the input layout. 

129 # The local tensor shape changes, but the mapping from device mesh to tensor dimensions remains valid. 

130 return ((input_layout,), None) 

131 

132 # Note: get_expand_impl is not overridden because we default to returning None. 

133 # OpDispatcher will use the original function (e.g., torch.nn.functional.pad) on the local tensor. 

134 # Since we ensured the padded dimensions are Replicated, local padding is mathematically correct.