Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / base.py: 79%

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

16Graph Pass Base - Base Class for Graph Transformation Passes 

17 

18All parallel partitioning passes and optimization passes inherit from this 

19class. The config a pass receives is typed as ``PassConfig`` so passes 

20can rely on ``enable_overlap`` / ``fsdp_enabled`` / ``fsdp_degree`` 

21/ ``tp_size`` being present without each pass re-declaring the contract. 

22""" 

23 

24__all__ = ["GraphPass"] 

25 

26from abc import ABC, abstractmethod 

27from typing import TYPE_CHECKING, Any 

28 

29from ..pass_config import PassConfig 

30 

31if TYPE_CHECKING: 

32 from torch import fx 

33 

34 

35class GraphPass(ABC): 

36 """Graph Optimization Pass Base Class. 

37 

38 All parallel partitioning passes and optimization passes inherit from 

39 this class. 

40 """ 

41 

42 @property 

43 @abstractmethod 

44 def name(self) -> str: 

45 """Pass name.""" 

46 pass # pylint: disable=W0107 

47 

48 @abstractmethod 

49 def run( 

50 self, 

51 graph_module: "fx.GraphModule", 

52 pass_config: PassConfig, 

53 **kwargs: Any, 

54 ) -> "fx.GraphModule": 

55 """Execute graph transformation. 

56 

57 Args: 

58 graph_module: Input graph. 

59 pass_config: Parallel configuration. 

60 

61 Returns: 

62 Transformed graph. 

63 """ 

64 pass # pylint: disable=W0107 

65 

66 def __repr__(self) -> str: 

67 """Return a developer-facing representation of the pass.""" 

68 return f"{self.__class__.__name__}(name={self.name})"