Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / boundary_forms.py: 27%
113 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Classify each frozen boundary into the forward form the emitter renders.
17Single source of truth shared by generation (``emit/parallel``), install-time
18validation (``runtime``), and preflight (``check/preflight``). The per-tensor
19transition judgment itself lives in
20``distributed._builder.tp_collective_lowering.classify_tp_transition`` — this
21module only folds a whole frozen boundary entry (its ``in_src``/``in_dst``/
22``out_src``/``out_dst`` named placements) into one of four forms:
24- ``region`` — the boundary has a ``local_compute_fn`` injection (EP
25 MoE local region): the forward delegates to ``self.__hp_compute__`` and
26 re-wraps locals per ``out_src`` before the boundary exit.
27- ``identity`` — every declared transition is identity on the plan's
28 active axes AND no inner-wrapper injection is declared: the class forward
29 is NOT rewritten (pruned); the runtime install path covers the exact no-op
30 / to-local semantics with its generic wrapper.
31- ``generic`` (forced) — a boundary with an ``inner_wrapper`` injection is
32 generic even when every transition is identity: the CP wrapper is applied
33 to the module at runtime by ``hp_apply_inner_wrapper``, so the class
34 must keep the rewritten redistribute forward that P1 always emitted.
35- ``tp_collective`` — every declared transition is identity or TP-lowerable
36 (``all_gather`` / ``all_reduce`` / ``reduce_scatter`` on the tp axis): the
37 emitter may sink bare operator calls (``self._hp_tp.all_gather(...)``),
38 subject to source-side structural gates (see ``emit/parallel``); the
39 runtime re-validates against the live mesh at install time and falls back
40 to the generic engine when the compiled plan disagrees.
41- ``generic`` — anything else (non-tp axis differences, unsupported
42 placement pairs, expert-mesh routing): the forward is rewritten to call
43 the instance-bound compiled plan (``self._hp_boundary``) with a
44 declarative comment annotation of what each side does.
46The classification is a pure function of the frozen entry, the plan's active
47axis names, and the (optional) injection rule — the same inputs the drift
48check re-emits from, so generation, install, and preflight always agree.
50The runtime-conditional facts (live mesh axes, rank order, backend, expert
51mesh) are deliberately NOT classified here: the emit-time decision is
52optimistic and the runtime's install-time validation is the safety net.
53"""
55from __future__ import annotations
57from dataclasses import dataclass, field
58from typing import Any, Optional
60#: Region form: local-compute injection owns the forward (EP MoE).
61FORM_REGION = "region"
62#: Identity form: prune the rewrite entirely; runtime wrapper covers it.
63FORM_IDENTITY = "identity"
64#: TP-collective form: bare operator calls (source gates permitting).
65FORM_TP_COLLECTIVE = "tp_collective"
66#: Generic form: annotated redistribute through the instance-bound plan.
67FORM_GENERIC = "generic"
69#: Placement fields of a frozen boundary entry, per side.
70_SIDE_FIELDS = (("in_src", "in_dst"), ("out_src", "out_dst"))
73@dataclass(frozen=True)
74class TransitionOp:
75 """One TP-lowerable placement transition, as the emitter renders it."""
77 #: ``all_gather`` / ``all_reduce`` / ``reduce_scatter``.
78 kind: str
79 #: Tensor dim the collective gathers / scatters on (``None`` for all_reduce).
80 tensor_dim: Optional[int] = None
83@dataclass(frozen=True)
84class BoundaryForm:
85 """The classified form of one frozen boundary entry.
87 ``in_ops`` / ``out_ops`` carry the non-identity TP-lowerable transitions
88 (name -> op); the ``*_notes``
89 dicts carry one declarative annotation line per non-identity transition,
90 keyed by name, for the generic form's comment block.
91 """
93 form: str
94 in_ops: dict[str, TransitionOp] = field(default_factory=dict)
95 out_ops: dict[str, TransitionOp] = field(default_factory=dict)
96 in_notes: dict[str, str] = field(default_factory=dict)
97 out_notes: dict[str, str] = field(default_factory=dict)
98 #: Every declared name on each side (identity or not).
99 in_declared: tuple[str, ...] = ()
100 #: The entry's declared output order (``out_names`` or ``out_src`` keys).
101 out_names: tuple[str, ...] = ()
104def classify_boundary_form(
105 entry: dict[str, Any],
106 mesh_dim_names: Optional[Any],
107 injection: Optional[dict[str, Any]] = None,
108) -> BoundaryForm:
109 """Classify one frozen ``param_plan`` boundary entry into its emit form.
111 Args:
112 entry: The frozen boundary entry (placements as canonical strings,
113 e.g. ``{"tp": "S(1)"}``).
114 mesh_dim_names: The plan's active axis names (frozen offline-mesh
115 order). ``None`` / empty means the plan shards on no axis.
116 injection: The boundary's frozen injection rule, if any — a
117 ``local_compute_fn`` makes the form ``region``; an ``inner_wrapper``
118 forces ``generic`` (the CP wrapper mutates the forward at runtime,
119 so the class keeps the rewritten redistribute form).
121 Returns:
122 The :class:`BoundaryForm`; ``form`` is one of ``region`` /
123 ``identity`` / ``tp_collective`` / ``generic``.
124 """
125 # A local-compute region runs its compute fn on local tensors between the
126 # entry/exit redistributes; the sides are still classified below so the
127 # emitter's region template can annotate each transition. An
128 # unlowerable transition does NOT force the generic form there — the
129 # compiled plan's redistribute owns it either way.
130 is_region = injection is not None and injection.get("local_compute_fn") is not None
131 # An inner-wrapper boundary keeps the generic rewrite even when every
132 # transition is identity: ``hp_apply_inner_wrapper`` mutates the target
133 # forward at runtime, and the P1 contract (wrapper weaves into the
134 # rewritten redistribute forward) must not change under pruning.
135 inner_wrap = injection is not None and injection.get("inner_wrapper") is not None
137 axes = tuple(mesh_dim_names or ())
138 if not axes:
139 # No active sharding axes: the runtime routes entries whose ``ep``
140 # placement actually changes to the expert mesh (a runtime object —
141 # not statically decidable) and treats everything else as an exact
142 # no-op passthrough.
143 if is_region:
144 return BoundaryForm(form=FORM_REGION)
145 if inner_wrap or _entry_has_ep_placement(entry):
146 return BoundaryForm(form=FORM_GENERIC)
147 return BoundaryForm(form=FORM_IDENTITY)
149 in_ops: dict[str, TransitionOp] = {}
150 in_notes: dict[str, str] = {}
151 in_declared: list[str] = []
152 in_fallback = False
153 for name, result, note in _classify_side(entry, "in_src", "in_dst", axes):
154 in_declared.append(name)
155 if note is not None:
156 in_notes[name] = note
157 if result is None:
158 # An unlowerable transition anywhere on the side forces the
159 # generic form (the compiled plan's DTensor fallback owns it).
160 in_fallback = True
161 elif result["kind"] != "identity":
162 in_ops[name] = TransitionOp(
163 kind=result["kind"], tensor_dim=result.get("tensor_dim")
164 )
166 out_ops: dict[str, TransitionOp] = {}
167 out_notes: dict[str, str] = {}
168 out_fallback = False
169 for name, result, note in _classify_side(entry, "out_src", "out_dst", axes):
170 if note is not None:
171 out_notes[name] = note
172 if result is None:
173 out_fallback = True
174 elif result["kind"] != "identity":
175 out_ops[name] = TransitionOp(
176 kind=result["kind"], tensor_dim=result.get("tensor_dim")
177 )
179 entry_out_names = declared_out_names(entry)
180 if is_region:
181 return BoundaryForm(
182 form=FORM_REGION,
183 in_ops=in_ops,
184 out_ops=out_ops,
185 in_notes=in_notes,
186 out_notes=out_notes,
187 in_declared=tuple(in_declared),
188 out_names=entry_out_names,
189 )
190 if in_fallback or out_fallback or inner_wrap:
191 return BoundaryForm(
192 form=FORM_GENERIC,
193 in_ops=in_ops,
194 out_ops=out_ops,
195 in_notes=in_notes,
196 out_notes=out_notes,
197 in_declared=tuple(in_declared),
198 out_names=entry_out_names,
199 )
200 if not in_ops and not out_ops:
201 return BoundaryForm(
202 form=FORM_IDENTITY,
203 in_ops=in_ops,
204 out_ops=out_ops,
205 in_notes=in_notes,
206 out_notes=out_notes,
207 in_declared=tuple(in_declared),
208 out_names=entry_out_names,
209 )
210 return BoundaryForm(
211 form=FORM_TP_COLLECTIVE,
212 in_ops=in_ops,
213 out_ops=out_ops,
214 in_notes=in_notes,
215 out_notes=out_notes,
216 in_declared=tuple(in_declared),
217 out_names=entry_out_names,
218 )
221def _classify_side(
222 entry: dict[str, Any],
223 src_field: str,
224 dst_field: str,
225 axes: tuple[str, ...],
226):
227 """Yield ``(name, classified, note)`` for every declared name on one side.
229 ``classified`` is ``classify_tp_transition``'s result (``None`` when the
230 transition is not TP-lowerable); ``note`` is the declarative annotation
231 line for non-identity transitions (``None`` for identity / no-op names).
232 """
233 # Lazy imports: plan modules must not pull the distributed package (and
234 # its torch-dependent __init__) at codegen import time.
235 from hyper_parallel.codegen.plan.freeze import parse_named_placement
236 from hyper_parallel.distributed.recipe_spec import resolve_placements
237 from hyper_parallel.distributed._builder.tp_collective_lowering import (
238 classify_tp_transition,
239 )
241 src_named = parse_named_placement(entry.get(src_field) or {})
242 dst_named = parse_named_placement(entry.get(dst_field) or {})
243 src_raw = entry.get(src_field) or {}
244 dst_raw = entry.get(dst_field) or {}
245 for name in sorted(set(src_named) | set(dst_named)):
246 src_p = tuple(resolve_placements(src_named.get(name) or {}, axes))
247 dst_p = tuple(resolve_placements(dst_named.get(name) or {}, axes))
248 result = classify_tp_transition(src_p, dst_p, axes)
249 note = _transition_note(
250 name, src_raw.get(name) or {}, dst_raw.get(name) or {}, axes, result
251 )
252 yield name, result, note
255def _transition_note(
256 name: str,
257 src_raw: dict[str, Any],
258 dst_raw: dict[str, Any],
259 axes: tuple[str, ...],
260 result: Optional[dict],
261) -> Optional[str]:
262 """Render one transition's declarative annotation, or ``None`` for identity.
264 Only axes the plan actually shards on participate (an ``ep`` key is
265 dropped by the dense-mesh compile, so claiming it here would lie); the
266 raw frozen strings are shown verbatim — the annotation is the plan's own
267 declaration, not a re-interpretation.
268 """
269 diffs = [
270 f"{axis} {src_raw.get(axis, 'R')} -> {dst_raw.get(axis, 'R')}"
271 for axis in axes
272 if src_raw.get(axis, "R") != dst_raw.get(axis, "R")
273 ]
274 if not diffs:
275 return None
276 head = f"{name}: " + ", ".join(diffs)
277 if result is None:
278 return f"{head} ==> dtensor redistribute"
279 return f"{head} ==> {_op_text(result)}"
282def _op_text(result: dict) -> str:
283 """Compact text of one classified op, as it appears in annotations/calls."""
284 kind = result["kind"]
285 if kind == "all_reduce":
286 return "all_reduce(sum)"
287 return f"{kind}(dim={result['tensor_dim']})"
290def _entry_has_ep_placement(entry: dict[str, Any]) -> bool:
291 """Whether any declared name's ``ep`` placement actually changes.
293 Every frozen entry declares full per-axis dicts, so an identity ``ep`` key
294 (``R -> R``) is the plan's way of saying "not EP-dependent" — not a
295 dependency on the runtime expert mesh. Under a degenerate topology
296 (``tp=cp=ep=1``) the plan carries no active axes and every ``ep`` key is
297 ``R -> R``; reading key *existence* as EP dependence would force the
298 generic form on boundaries whose semantics are an exact no-op, while the
299 install path is skipped entirely for such meshes — the rewritten
300 forward's ``_hp_boundary`` reference would never be bound.
301 """
302 for src_field, dst_field in _SIDE_FIELDS:
303 src_named = entry.get(src_field) or {}
304 dst_named = entry.get(dst_field) or {}
305 for name in sorted(set(src_named) | set(dst_named)):
306 if _ep_axis(src_named.get(name)) != _ep_axis(dst_named.get(name)):
307 return True
308 return False
311def _ep_axis(placement: Any) -> str:
312 """One named placement's canonical ``ep`` value (absent / non-dict = ``R``)."""
313 if not isinstance(placement, dict):
314 return "R"
315 value = placement.get("ep")
316 return value if isinstance(value, str) and value else "R"
319def declared_out_names(entry: dict[str, Any]) -> tuple[str, ...]:
320 """The entry's declared output order (``out_names`` or ``out_src`` keys).
322 Output names are not retained on the frozen entry, so a local-region
323 forward wraps its local results per ``out_src`` before the boundary exit
324 and reads the declared order off the same place the rewrap does.
325 """
326 declared = entry.get("out_src") or {}
327 out_names = entry.get("out_names") or declared.keys()
328 return tuple(out_names)
331def is_boundary_entry(entry: Any) -> bool:
332 """Whether a frozen param-plan entry is a boundary.
334 Single definition for both the emitter (which sinks the entry's forward
335 into its boundary form) and the installer (which wraps the module): an
336 entry carrying any layout contract is a boundary, an explicit
337 ``is_boundary`` wins over that inference. Two copies of this predicate
338 had drifted apart -- they disagreed on an entry carrying only
339 ``out_names`` -- so both consumers read it from here.
340 """
341 if not isinstance(entry, dict):
342 return False
343 if entry.get("is_boundary") is not None:
344 return bool(entry["is_boundary"])
345 return any(entry.get(field) for field in _ENTRY_META_FIELDS)
348#: Frozen-entry keys that make an entry a boundary: the in/out placement
349#: contracts plus the declared output order a local-region forward wraps by.
350_ENTRY_META_FIELDS = ("in_src", "in_dst", "out_src", "out_dst", "out_names")
353__all__ = [
354 "BoundaryForm",
355 "FORM_GENERIC",
356 "FORM_IDENTITY",
357 "FORM_REGION",
358 "FORM_TP_COLLECTIVE",
359 "TransitionOp",
360 "classify_boundary_form",
361 "declared_out_names",
362 "is_boundary_entry",
363]