Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / boundary_forms.py: 97%
115 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 observed boundary into the forward form the emitter renders.
17Single source of truth shared by regular and recorded-forward emitters. The per-tensor
18transition judgment itself lives in
19``distributed._builder.tp_collective_lowering.classify_tp_transition`` — this
20module only folds a whole observed boundary entry (its ``in_src``/``in_dst``/
21``out_src``/``out_dst`` named placements) into one of four forms:
23- ``region`` — the boundary has a ``local_compute_fn`` injection (EP
24 MoE local region): the forward delegates to ``self.__hp_compute__`` and
25 re-wraps locals per ``out_src`` before the boundary exit.
26- ``identity`` — every declared transition is identity on the plan's
27 active axes AND no inner-wrapper injection is declared: the class forward
28 is not rewritten because the observed contract is a no-op.
29- ``generic`` (forced) — a boundary with an ``inner_wrapper`` injection is
30 generic even when every transition is identity, so the generated CP
31 dispatch remains visible around the original forward.
32- ``tp_collective`` — every declared transition is identity or TP-lowerable
33 (``all_gather`` / ``all_reduce`` / ``reduce_scatter`` on the tp axis): the
34 emitter may sink bare operator calls (``self._hp_tp.all_gather(...)``),
35 subject to source-side structural gates in ``emit/parallel``.
36- ``generic`` — anything else (non-tp axis differences, unsupported
37 placement pairs, expert-mesh routing): the forward is rewritten to call
38 the instance-bound compiled plan (``self._hp_boundary``) with a
39 declarative comment annotation of what each side does.
41The classification is a pure function of the observed entry, active plan
42axes, and optional injection declaration. Runtime-only mesh and backend
43objects never enter this source-emission decision.
44"""
46from __future__ import annotations
48from dataclasses import dataclass, field
49from typing import Any, Optional
51from hyper_parallel.codegen.plan.placements import parse_named_placement
52from hyper_parallel.distributed.recipe_spec import resolve_placements
53from hyper_parallel.distributed._builder.tp_collective_lowering import (
54 classify_tp_transition,
55)
57#: Region form: local-compute injection owns the forward (EP MoE).
58FORM_REGION = "region"
59#: Identity form: prune the rewrite entirely; runtime wrapper covers it.
60FORM_IDENTITY = "identity"
61#: TP-collective form: bare operator calls (source gates permitting).
62FORM_TP_COLLECTIVE = "tp_collective"
63#: Generic form: annotated redistribute through the instance-bound plan.
64FORM_GENERIC = "generic"
66#: Placement fields of a frozen boundary entry, per side.
67_SIDE_FIELDS = (("in_src", "in_dst"), ("out_src", "out_dst"))
70@dataclass(frozen=True)
71class TransitionOp:
72 """One TP-lowerable placement transition, as the emitter renders it."""
74 #: ``all_gather`` / ``all_reduce`` / ``reduce_scatter``.
75 kind: str
76 #: Tensor dim the collective gathers / scatters on (``None`` for all_reduce).
77 tensor_dim: Optional[int] = None
80@dataclass(frozen=True)
81class BoundaryForm:
82 """The classified form of one frozen boundary entry.
84 ``in_ops`` / ``out_ops`` carry the non-identity TP-lowerable transitions
85 (name -> op); the ``*_notes``
86 dicts carry one declarative annotation line per non-identity transition,
87 keyed by name, for the generic form's comment block.
88 """
90 form: str
91 in_ops: dict[str, TransitionOp] = field(default_factory=dict)
92 out_ops: dict[str, TransitionOp] = field(default_factory=dict)
93 in_notes: dict[str, str] = field(default_factory=dict)
94 out_notes: dict[str, str] = field(default_factory=dict)
95 #: Every declared name on each side (identity or not).
96 in_declared: tuple[str, ...] = ()
97 #: The entry's declared output order (``out_names`` or ``out_src`` keys).
98 out_names: tuple[str, ...] = ()
101def classify_boundary_form(
102 entry: dict[str, Any],
103 mesh_dim_names: Optional[Any],
104 injection: Optional[dict[str, Any]] = None,
105) -> BoundaryForm:
106 """Classify one frozen ``param_plan`` boundary entry into its emit form.
108 Args:
109 entry: The frozen boundary entry (placements as canonical strings,
110 e.g. ``{"tp": "S(1)"}``).
111 mesh_dim_names: The plan's active axis names (frozen offline-mesh
112 order). ``None`` / empty means the plan shards on no axis.
113 injection: The boundary's frozen injection rule, if any — a
114 ``local_compute_fn`` makes the form ``region``; an ``inner_wrapper``
115 forces ``generic`` (the CP wrapper mutates the forward at runtime,
116 so the class keeps the rewritten redistribute form).
118 Returns:
119 The :class:`BoundaryForm`; ``form`` is one of ``region`` /
120 ``identity`` / ``tp_collective`` / ``generic``.
121 """
122 # A local-compute region runs its compute fn on local tensors between the
123 # entry/exit redistributes; the sides are still classified below so the
124 # emitter's region template can annotate each transition. An
125 # unlowerable transition does NOT force the generic form there — the
126 # compiled plan's redistribute owns it either way.
127 is_region = injection is not None and injection.get("local_compute_fn") is not None
128 # An inner-wrapper boundary keeps the generic rewrite even when every
129 # transition is identity: ``hp_apply_inner_wrapper`` mutates the target
130 # forward at runtime, and the P1 contract (wrapper weaves into the
131 # rewritten redistribute forward) must not change under pruning.
132 inner_wrap = injection is not None and injection.get("inner_wrapper") is not None
134 axes = tuple(mesh_dim_names or ())
135 if not axes:
136 # No active sharding axes: the runtime routes entries whose ``ep``
137 # placement actually changes to the expert mesh (a runtime object —
138 # not statically decidable) and treats everything else as an exact
139 # no-op passthrough.
140 if is_region:
141 return BoundaryForm(form=FORM_REGION)
142 if inner_wrap or _entry_has_ep_placement(entry):
143 return BoundaryForm(form=FORM_GENERIC)
144 return BoundaryForm(form=FORM_IDENTITY)
146 return _classify_active_axes(entry, axes, is_region=is_region, inner_wrap=inner_wrap)
149def _classify_active_axes(
150 entry: dict[str, Any],
151 axes: tuple[str, ...],
152 *,
153 is_region: bool,
154 inner_wrap: bool,
155) -> BoundaryForm:
156 """Classify an entry whose plan has at least one active axis."""
157 in_ops: dict[str, TransitionOp] = {}
158 in_notes: dict[str, str] = {}
159 in_declared: list[str] = []
160 in_fallback = False
161 for name, result, note in _classify_side(entry, "in_src", "in_dst", axes):
162 in_declared.append(name)
163 if note is not None:
164 in_notes[name] = note
165 if result is None:
166 # An unlowerable transition anywhere on the side forces the
167 # generic form (the compiled plan's DTensor fallback owns it).
168 in_fallback = True
169 elif result["kind"] != "identity":
170 in_ops[name] = TransitionOp(
171 kind=result["kind"], tensor_dim=result.get("tensor_dim")
172 )
174 out_ops: dict[str, TransitionOp] = {}
175 out_notes: dict[str, str] = {}
176 out_fallback = False
177 for name, result, note in _classify_side(entry, "out_src", "out_dst", axes):
178 if note is not None:
179 out_notes[name] = note
180 if result is None:
181 out_fallback = True
182 elif result["kind"] != "identity":
183 out_ops[name] = TransitionOp(
184 kind=result["kind"], tensor_dim=result.get("tensor_dim")
185 )
187 entry_out_names = declared_out_names(entry)
188 if is_region:
189 return BoundaryForm(
190 form=FORM_REGION,
191 in_ops=in_ops,
192 out_ops=out_ops,
193 in_notes=in_notes,
194 out_notes=out_notes,
195 in_declared=tuple(in_declared),
196 out_names=entry_out_names,
197 )
198 if in_fallback or out_fallback or inner_wrap:
199 return BoundaryForm(
200 form=FORM_GENERIC,
201 in_ops=in_ops,
202 out_ops=out_ops,
203 in_notes=in_notes,
204 out_notes=out_notes,
205 in_declared=tuple(in_declared),
206 out_names=entry_out_names,
207 )
208 if not in_ops and not out_ops:
209 return BoundaryForm(
210 form=FORM_IDENTITY,
211 in_ops=in_ops,
212 out_ops=out_ops,
213 in_notes=in_notes,
214 out_notes=out_notes,
215 in_declared=tuple(in_declared),
216 out_names=entry_out_names,
217 )
218 return BoundaryForm(
219 form=FORM_TP_COLLECTIVE,
220 in_ops=in_ops,
221 out_ops=out_ops,
222 in_notes=in_notes,
223 out_notes=out_notes,
224 in_declared=tuple(in_declared),
225 out_names=entry_out_names,
226 )
229def _classify_side(
230 entry: dict[str, Any],
231 src_field: str,
232 dst_field: str,
233 axes: tuple[str, ...],
234):
235 """Yield ``(name, classified, note)`` for every declared name on one side.
237 ``classified`` is ``classify_tp_transition``'s result (``None`` when the
238 transition is not TP-lowerable); ``note`` is the declarative annotation
239 line for non-identity transitions (``None`` for identity / no-op names).
240 """
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]