Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / execution / tensor.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"""Shadow tensor. 

16 

17This module provides: :class:`ShadowTensor` — a :class:`torch.Tensor` subclass that 

18resolves a view of its :class:`PhysicalBuffer`'s storage on every dispatch. 

19""" 

20 

21from __future__ import annotations 

22 

23from typing import Any 

24 

25import torch 

26 

27from hyper_parallel.auto_parallel.hyper_offload.runtime.residency import PhysicalBuffer 

28 

29 

30class ShadowTensor(torch.Tensor): 

31 """Tensor subclass that resolves a device view on demand. 

32 

33 The shadow uses :meth:`PhysicalBuffer.device_storage` to obtain the 

34 raw device storage on every :meth:`resolve` call and creates a 

35 fresh tensor view from it. **No cached reference is kept**, so the 

36 shadow never prevents the underlying storage from being freed. 

37 

38 Because :class:`ShadowTensor` is created via 

39 :meth:`torch.Tensor._make_wrapper_subclass`, it carries all tensor 

40 metadata ("dtype", "size", "stride", "storage_offset", 

41 "device") directly. 

42 """ 

43 

44 @staticmethod 

45 def __new__( 

46 cls, 

47 elem: torch.Tensor, 

48 buffer: PhysicalBuffer, 

49 storage_id: int, 

50 ) -> ShadowTensor: 

51 """Create a new instance using *elem*'s metadata.""" 

52 return torch.Tensor._make_wrapper_subclass( 

53 cls, 

54 elem.size(), 

55 strides=elem.stride(), 

56 storage_offset=elem.storage_offset(), 

57 dtype=elem.dtype, 

58 layout=elem.layout, 

59 device=elem.device, 

60 requires_grad=elem.requires_grad, 

61 ) 

62 

63 def __init__( # pylint: disable=unused-argument 

64 self, 

65 elem: torch.Tensor, 

66 buffer: PhysicalBuffer, 

67 storage_id: int, 

68 ) -> None: 

69 self._buffer = buffer 

70 self._storage_id = storage_id 

71 

72 # ------------------------------------------------------------------ 

73 # Public properties 

74 # ------------------------------------------------------------------ 

75 

76 @property 

77 def storage_id(self) -> int: 

78 """Storage ID of the underlying physical block.""" 

79 return self._storage_id 

80 

81 # ------------------------------------------------------------------ 

82 # Resolution (on every call) 

83 # ------------------------------------------------------------------ 

84 

85 def resolve(self) -> torch.Tensor: 

86 """Return a device-resident view of the physical storage. 

87 

88 1. Calls :meth:`PhysicalBuffer.device_storage` to obtain the 

89 raw device storage (demand-paging from host if needed). 

90 2. Builds a fresh tensor view from the shadow's cached metadata. 

91 3. Returns the view — no long-lived reference is kept. 

92 """ 

93 storage = self._buffer.device_storage() 

94 result = torch.empty(0, dtype=self.dtype, device=self.device) 

95 result.set_(storage, self.storage_offset(), self.size(), self.stride()) 

96 return result 

97 

98 # ------------------------------------------------------------------ 

99 # PyTorch dispatch 

100 # ------------------------------------------------------------------ 

101 

102 @classmethod 

103 def __torch_dispatch__(cls, func, types, args=(), kwargs=None): # pylint: disable=unused-argument 

104 """Dispatch a torch operation.""" 

105 kwargs = kwargs or {} 

106 

107 def unwrap(value: Any) -> Any: 

108 if isinstance(value, ShadowTensor): 

109 return value.resolve() 

110 if isinstance(value, tuple): 

111 return tuple(unwrap(v) for v in value) 

112 if isinstance(value, list): 

113 return [unwrap(v) for v in value] 

114 if isinstance(value, dict): 

115 return {k: unwrap(v) for k, v in value.items()} 

116 return value 

117 

118 with torch._C._DisableTorchDispatch(): # pylint: disable=protected-access 

119 return func(*unwrap(args), **unwrap(kwargs))