Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_ops.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""" 

16Distributed operator implementation. 

17""" 

18 

19from typing import Optional 

20 

21from .parallel_ops_register import register_distributed_op 

22 

23 

24class DistributedOp: 

25 """ 

26 Base class for distributed operator implementations. 

27 

28 This class provides default implementations for distributed operators. 

29 Subclasses should override methods as needed for specific operators. 

30 

31 Args: 

32 op_name (str): Name of the operator to register. 

33 """ 

34 def __init__(self, op_name): 

35 self.op_name = op_name 

36 register_distributed_op(op_name, self) 

37 self._allow_partial_inputs = False 

38 

39 def _check_partial_inputs(self, layouts): 

40 """ 

41 Check if any input layout has partial status and raise an error if not allowed. 

42 

43 This method can be called by subclasses to enforce that partial inputs 

44 are not supported for a particular operator. Subclasses that support 

45 partial inputs should not call this method. 

46 

47 Args: 

48 layouts (tuple): Layouts of input tensor. 

49 

50 Raises: 

51 ValueError: If any input layout has partial status. 

52 """ 

53 for i, layout in enumerate(layouts): 

54 if layout is not None and layout.is_partial(): 

55 raise ValueError( 

56 f"For {self.op_name}, input {i} with {layout} has Partial status which is not allowed. " 

57 f"Should be without Partial status for this operation." 

58 ) 

59 

60 # pylint: disable=W0613 

61 def preprocess(self, args: tuple, kwargs: dict) -> Optional[tuple]: 

62 """ 

63 Unified preprocessing: parameter parsing + to_local + cache_values construction. 

64 

65 Subclasses override this to participate in the new dispatch flow. 

66 

67 Args: 

68 args (tuple): Positional arguments passed to the operator call. 

69 kwargs (dict): Keyword arguments passed to the operator call. 

70 

71 Returns: 

72 None: Fall back to legacy dispatch (default). 

73 tuple: (local_args, local_kwargs, cache_values) 

74 - local_args: Local tensor positional arguments (DTensors already to_local'd). 

75 - local_kwargs: Local tensor keyword arguments (DTensors already to_local'd). 

76 - cache_values: Values affecting layout inference (fixed order). 

77 Contains Layout objects (with compact_str) and raw values (int, bool, tuple, etc.). 

78 """ 

79 return None 

80 

81 def infer_layout(self, cache_values: list) -> Optional[tuple]: 

82 """ 

83 Infer output layouts based on cache_values built by preprocess. 

84 

85 Default implementation extracts the first Layout from cache_values and 

86 returns it as the output layout (element-wise default). Subclasses should 

87 override this method to provide custom layout inference logic. 

88 

89 Args: 

90 cache_values (list): Values built by preprocess that affect layout inference. 

91 Contains Layout objects and non-layout values (shapes, scalars, etc.). 

92 

93 Returns: 

94 tuple: ((output_layouts,), None) or None if no layouts found. 

95 """ 

96 if not self._allow_partial_inputs: 

97 self._check_partial_inputs(cache_values) 

98 

99 if cache_values: 

100 return (cache_values[0],) 

101 return None 

102 

103 # pylint: disable=W0613 

104 def get_expand_impl( 

105 self, 

106 func: Optional[callable], 

107 infer_result: tuple, 

108 cache_values: list, 

109 ) -> Optional[callable]: 

110 """ 

111 Get expand implementation for the operator. 

112 

113 Args: 

114 func (Optional[callable]): The underlying operator function. 

115 infer_result (tuple): Result returned by infer_layout (output_layouts, extra_info). 

116 cache_values (list): Values built by preprocess, forwarded from the dispatch layer. 

117 

118 Returns: 

119 Optional[callable]: A closure that wraps the operator call with extra logic, 

120 or None if no expansion is needed. 

121 """ 

122 return None 

123 

124 @staticmethod 

125 def wrap_output(py_output, output_layouts): 

126 """Wrap local outputs into DTensors according to inferred layouts. 

127 

128 Subclasses may override this when a specific operator needs custom 

129 packing semantics for certain output slots. 

130 """ 

131 # pylint: disable=C0415 

132 from hyper_parallel.core.dtensor.dtensor import DTensor 

133 

134 if isinstance(py_output, (tuple, list)): 

135 if len(py_output) != len(output_layouts): 

136 raise RuntimeError( 

137 f"Output tuple size ({len(py_output)}) " 

138 f"does not match layout tuple size ({len(output_layouts)})") 

139 return tuple( 

140 DTensor.from_local(item, layout.mesh, layout.alias_placements) 

141 for item, layout in zip(py_output, output_layouts) 

142 ) 

143 

144 if isinstance(output_layouts, (tuple, list)): 

145 if len(output_layouts) != 1: 

146 raise RuntimeError( 

147 f"Scalar output expects a single layout, but got {len(output_layouts)} layouts" 

148 ) 

149 output_layout = output_layouts[0] 

150 else: 

151 output_layout = output_layouts 

152 

153 return DTensor.from_local( 

154 py_output, output_layout.mesh, output_layout.alias_placements 

155 )