Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / _mesh_layout.py: 99%
210 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""Internal layout helpers for DeviceMesh bookkeeping."""
17import math
18from dataclasses import dataclass
19from typing import Any, Iterator, Optional, Union
21import numpy as np
24IntTuple = Union[int, tuple["IntTuple", ...]]
27def _is_int(value: Any) -> bool:
28 return isinstance(value, int) and not isinstance(value, bool)
31def _as_tuple(value: IntTuple) -> tuple[IntTuple, ...]:
32 return value if isinstance(value, tuple) else (value,)
35def _flatten_inttuple(value: IntTuple) -> tuple[int, ...]:
36 if _is_int(value):
37 return (value,)
38 flattened: list[int] = []
39 for item in value:
40 flattened.extend(_flatten_inttuple(item))
41 return tuple(flattened)
44def _match_structure(shape: IntTuple, stride: IntTuple) -> bool:
45 if _is_int(shape):
46 return _is_int(stride)
47 if not isinstance(shape, tuple) or not isinstance(stride, tuple) or len(shape) != len(stride):
48 return False
49 return all(_match_structure(sub_shape, sub_stride) for sub_shape, sub_stride in zip(shape, stride))
52def _numel(value: IntTuple) -> int:
53 return int(np.prod(np.array(_flatten_inttuple(value), dtype=np.int64)))
56def _contiguous_strides(mesh_shape: tuple[int, ...]) -> tuple[int, ...]:
57 if len(mesh_shape) == 0:
58 return ()
59 strides = [1] * len(mesh_shape)
60 for idx in range(len(mesh_shape) - 2, -1, -1):
61 strides[idx] = strides[idx + 1] * mesh_shape[idx + 1]
62 return tuple(strides)
65def _scale_inttuple(value: IntTuple, factor: int) -> IntTuple:
66 if _is_int(value):
67 return int(value) * factor
68 return tuple(_scale_inttuple(item, factor) for item in value)
71def _enumerate_offsets(shape: IntTuple, stride: IntTuple) -> list[int]:
72 if _is_int(shape):
73 return [i * int(stride) for i in range(shape)]
75 offsets = [0]
76 for sub_shape, sub_stride in zip(_as_tuple(shape), _as_tuple(stride)):
77 dim_offsets = _enumerate_offsets(sub_shape, sub_stride)
78 offsets = [base + dim_offset for base in offsets for dim_offset in dim_offsets]
79 return offsets
82def _canonicalize_axis(shape, stride) -> tuple[tuple[int, ...], tuple[int, ...]]:
83 """Normalize one logical axis into a flattened shape/stride pair."""
84 flat_shape = _flatten_inttuple(shape)
85 flat_stride = _flatten_inttuple(stride if stride is not None else _contiguous_strides(flat_shape))
86 if len(flat_shape) != len(flat_stride):
87 raise ValueError(
88 f"shape and stride must have the same length, got {len(flat_shape)} and {len(flat_stride)}"
89 )
91 normalized_shape: list[int] = []
92 normalized_stride: list[int] = []
93 for size, step in zip(flat_shape, flat_stride):
94 if size < 0:
95 raise ValueError(f"shape entries must be non-negative, got {flat_shape}")
96 if size == 1:
97 continue
98 normalized_shape.append(int(size))
99 normalized_stride.append(int(step))
101 coalesced_shape: list[int] = []
102 coalesced_stride: list[int] = []
103 for size, step in zip(normalized_shape, normalized_stride):
104 if coalesced_shape and coalesced_stride[-1] == step * size:
105 coalesced_shape[-1] *= size
106 coalesced_stride[-1] = step
107 else:
108 coalesced_shape.append(size)
109 coalesced_stride.append(step)
110 return tuple(coalesced_shape), tuple(coalesced_stride)
113def _nested_from_flat(value: tuple[int, ...]) -> IntTuple:
114 if len(value) == 1:
115 return value[0]
116 return tuple(value)
119@dataclass(frozen=True)
120class _FlatLayout:
121 """Canonicalized layout for one logical DeviceMesh axis."""
123 shape: tuple[int, ...]
124 stride: tuple[int, ...]
126 def __init__(self, shape: IntTuple, stride: Optional[IntTuple] = None) -> None:
127 """Canonicalize *shape* and *stride* into a frozen flat layout."""
128 flat_shape, flat_stride = _canonicalize_axis(shape, stride)
129 object.__setattr__(self, "shape", flat_shape)
130 object.__setattr__(self, "stride", flat_stride)
132 def numel(self) -> int:
133 """Return the total number of elements in this layout."""
134 return math.prod(self.shape) if len(self.shape) > 0 else 1
136 def cosize(self) -> int:
137 """Return the size of the smallest contiguous block containing all ranks."""
138 ranks = self.all_ranks_from_zero()
139 return max(ranks) + 1 if ranks else 1
141 def check_sorted(self) -> bool:
142 """Return True if strides are in descending order."""
143 return tuple(sorted(self.stride, reverse=True)) == self.stride
145 def check_orthogonal(self) -> bool:
146 """Return True if each axis is independent (strides are orthogonal)."""
147 if len(self.shape) < 2:
148 return True
149 stride, shape = zip(*sorted(zip(self.stride, self.shape), reverse=True))
150 return all(
151 stride[idx] % (stride[idx + 1] * shape[idx + 1]) == 0
152 for idx in range(len(stride) - 1)
153 )
155 def all_ranks_from_zero(self) -> list[int]:
156 """List every rank offset assuming the base is zero."""
157 if len(self.shape) == 0:
158 return [0]
159 return [
160 int(sum(coord[dim] * self.stride[dim] for dim in range(len(self.shape))))
161 for coord in np.ndindex(self.shape)
162 ]
165class _MeshLayout:
166 """Minimal layout helper for DeviceMesh slicing, flattening, and concatenation."""
168 def __init__(
169 self,
170 shape_or_axes: Union[IntTuple, list[_FlatLayout], tuple[_FlatLayout, ...]],
171 stride: Optional[IntTuple] = None,
172 ) -> None:
173 """Build a layout from a nested shape/stride pair or a list of flat axes."""
174 if stride is None and isinstance(shape_or_axes, (list, tuple)) and all(
175 isinstance(axis, _FlatLayout) for axis in shape_or_axes
176 ):
177 axes = list(shape_or_axes)
178 shape = tuple(_nested_from_flat(axis.shape) for axis in axes)
179 stride = tuple(_nested_from_flat(axis.stride) for axis in axes)
180 else:
181 shape = shape_or_axes
182 if not _match_structure(shape, stride):
183 raise ValueError(f"shape {shape} and stride {stride} do not match")
184 self.shape: IntTuple = shape
185 self.stride: IntTuple = stride
187 @classmethod
188 def from_sizes_strides(
189 cls,
190 sizes: tuple[int, ...],
191 strides: Optional[tuple[int, ...]] = None,
192 ) -> "_MeshLayout":
193 """Create a layout from flat sizes and optional strides."""
194 if strides is None:
195 strides = _contiguous_strides(sizes)
196 return cls(sizes, strides)
198 @property
199 def sizes(self) -> IntTuple:
200 """Return the (possibly nested) shape tuple."""
201 return self.shape
203 @property
204 def strides(self) -> IntTuple:
205 """Return the (possibly nested) stride tuple."""
206 return self.stride
208 @property
209 def axes(self) -> tuple[_FlatLayout, ...]:
210 """Return each top-level dimension as a separate flat layout."""
211 return tuple(self[idx].collapse() for idx in range(len(self)))
213 def __len__(self) -> int:
214 """Return the number of top-level dimensions."""
215 return len(self.shape) if isinstance(self.shape, tuple) else 1
217 def __iter__(self) -> Iterator["_MeshLayout"]:
218 """Yield each top-level dimension as its own layout."""
219 for idx in range(len(self)):
220 yield self[idx]
222 def __getitem__(self, idx: int) -> "_MeshLayout":
223 """Select a single top-level dimension by index."""
224 if isinstance(self.shape, tuple):
225 if idx < -len(self.shape) or idx >= len(self.shape):
226 raise IndexError(
227 f"Dim {idx} is out of range for layout with {len(self.shape)} dimensions."
228 )
229 return _MeshLayout(self.shape[idx], self.stride[idx])
230 if idx not in (0, -1):
231 raise IndexError("Dim is out of range for 1D layout.")
232 return _MeshLayout(self.shape, self.stride)
234 def __eq__(self, other: object) -> bool:
235 """Return True if *other* has the same shape and stride."""
236 if not isinstance(other, _MeshLayout):
237 return False
238 return self.shape == other.shape and self.stride == other.stride
240 def __repr__(self) -> str:
241 """Return a printable representation of this layout."""
242 return f"_MeshLayout(shape={self.shape}, stride={self.stride})"
244 def numel(self) -> int:
245 """Return the total number of elements in this layout."""
246 return _numel(self.shape)
248 @property
249 def top_level_sizes(self) -> tuple[int, ...]:
250 """Return the element count of each top-level dimension."""
251 return tuple(self[idx].numel() for idx in range(len(self)))
253 def all_ranks_from_zero(self) -> list[int]:
254 """List every rank offset assuming the base is zero."""
255 return _enumerate_offsets(self.shape, self.stride)
257 def check_non_overlap(self) -> bool:
258 """Return True if no rank appears more than once in this layout."""
259 ranks = self.all_ranks_from_zero()
260 return len(ranks) == len(set(ranks))
262 def coalesce(self) -> "_MeshLayout":
263 """Merge adjacent contiguous axes while preserving the represented layout."""
264 if _is_int(self.shape):
265 return self
267 coalesced_shapes: list[IntTuple] = []
268 coalesced_strides: list[IntTuple] = []
269 for shape, stride in zip(self.shape, self.stride):
270 child = _MeshLayout(shape, stride).coalesce()
271 coalesced_shapes.append(child.shape)
272 coalesced_strides.append(child.stride)
274 merged_shapes: list[IntTuple] = []
275 merged_strides: list[IntTuple] = []
276 for shape, stride in zip(coalesced_shapes, coalesced_strides):
277 can_merge_scalar_axis = (
278 bool(merged_shapes)
279 and _is_int(merged_shapes[-1])
280 and _is_int(merged_strides[-1])
281 and _is_int(shape)
282 and _is_int(stride)
283 )
284 if can_merge_scalar_axis and merged_strides[-1] == stride * shape:
285 merged_shapes[-1] *= shape
286 merged_strides[-1] = stride
287 else:
288 merged_shapes.append(shape)
289 merged_strides.append(stride)
291 if len(merged_shapes) == 1:
292 return _MeshLayout(merged_shapes[0], merged_strides[0])
293 return _MeshLayout(tuple(merged_shapes), tuple(merged_strides))
295 def composition(self, layout: "_MeshLayout") -> "_MeshLayout":
296 """Compose *layout* on top of this axis (unflatten with scaled strides)."""
297 if not _is_int(self.stride):
298 raise NotImplementedError(
299 "Currently, _unflatten only supports unflattening a mesh dim with scalar stride."
300 )
301 return _MeshLayout(layout.shape, _scale_inttuple(layout.stride, int(self.stride)))
303 def nest(self) -> "_MeshLayout":
304 """Wrap all dimensions into a single outer dimension."""
305 if len(self) == 1:
306 return self
307 return _MeshLayout((self.shape,), (self.stride,))
309 def splice(self, start: int, end: int, layout: "_MeshLayout") -> "_MeshLayout":
310 """Replace dimensions [start, end) with the axes of *layout*."""
311 sizes = list(_as_tuple(self.shape))
312 strides = list(_as_tuple(self.stride))
313 sizes[start:end] = list(_as_tuple(layout.shape))
314 strides[start:end] = list(_as_tuple(layout.stride))
315 if len(sizes) == 1:
316 return _MeshLayout(sizes[0], strides[0])
317 return _MeshLayout(tuple(sizes), tuple(strides))
319 def collapse(self) -> _FlatLayout:
320 """Flatten this layout into a single canonicalized flat layout."""
321 return _FlatLayout(self.shape, self.stride)
323 def remap_to_numpy(self, rank_map: Any) -> np.ndarray:
324 """Materialize this layout as a dense numpy mesh over the provided rank map."""
325 rank_map_np = np.asarray(rank_map).reshape(-1)
326 base_offsets = self.all_ranks_from_zero()
327 if len(base_offsets) == 0:
328 raise ValueError("Cannot remap an empty layout.")
330 groups: list[list[int]] = []
331 used: set[int] = set()
332 world_size = rank_map_np.shape[0]
334 for anchor in range(world_size):
335 if anchor in used:
336 continue
337 group = [anchor + offset for offset in base_offsets]
338 if any(index >= world_size for index in group):
339 continue
340 if any(index in used for index in group):
341 continue
342 groups.append(group)
343 used.update(group)
345 if len(used) != world_size:
346 raise ValueError(
347 f"Layout {self} does not form a full partition over rank_map with world size {world_size}."
348 )
350 remapped = rank_map_np[np.array(groups, dtype=np.int64)]
351 remapped = remapped.reshape((len(groups),) + self.top_level_sizes)
352 return remapped