diff --git a/docs/en/dev/distributed_ops.md b/docs/en/dev/distributed_ops.md index 32ed5d972a..7133dabbb7 100644 --- a/docs/en/dev/distributed_ops.md +++ b/docs/en/dev/distributed_ops.md @@ -38,6 +38,78 @@ There are **fifteen ops** and **four ABI enums**: The seven side-effect-only ops produce [`UnknownType`](ir/02-types.md): they exist for their cross-rank effect, not for an SSA value a consumer reads. +## Ergonomic collective API — auto-managed signals + +The `pld.tensor.*` collectives require an explicit window-bound INT32 signal +buffer per call. The ergonomic short forms (`pld.all_reduce`, `pld.all_gather`, +`pld.reduce_scatter`, `pld.broadcast`, `pld.all_to_all`, `pld.all_to_all_v`, +`pld.barrier`) allocate a **fresh, correctly shaped** signal automatically and +delegate to the corresponding `pld.tensor.*` HOST builtin: + +```python +data = pld.all_reduce(data, op=pld.ReduceOp.Sum) # mesh: no signal needed (host synthesis) +data = pld.all_reduce(data, mode="ring", nranks=2) # ring: [2*(NR-1)+1, NR] signal auto-allocated +data = pld.all_gather(local, target) +data = pld.reduce_scatter(target, op=pld.ReduceOp.Sum) +data = pld.broadcast(target, root=0) +data = pld.all_to_all(input, target) +data = pld.all_to_all_v(input, target, send_counts, recv_counts, nranks=NR) +sig = pld.barrier(sig) # requires an explicit, covered signal +``` + +Semantics and constraints: + +- **HOST-orchestration only.** The wrappers build on the host-only + `alloc_window_buffer` / `window` / `world_size` primitives; operands are + window-bound `DistributedTensor`s, exactly as for `pld.tensor.*`. +- **Signal shape** is chosen per op and matches the HOST builtin's requirement: + `broadcast` / `reduce_scatter` use a rank-1 `[world_size]` signal; + `all_gather` / `all_to_all` use a rank-2 `[world_size, 1]` signal; + `all_to_all_v` uses a rank-2 `[nranks, 1]` signal (static `nranks` required); + `allreduce mode="mesh"` needs no signal (the compiler synthesizes + it), while `mode="ring"` needs a static `[2*(NR-1)+1, NR]` signal, so + `nranks` is required. +- **Fresh per call.** Signals self-clear under the credit-barrier protocol + (pypto #2175, merged), so they are safe to reuse across back-to-back calls; + the wrappers still allocate a fresh buffer per call, which remains correct + and is the simplest safe default. +- **`pld.barrier(sig)` needs an explicit signal** with comm-domain coverage — a + barrier has no data buffer from which coverage could be inherited, so an + auto-allocated signal would be rejected by `MaterializeCommDomainScopes`. + Pass an INT32 window that is also consumed by a device-tagged dispatch (see + `tests/st/distributed/test_l3_host_tensor_barrier.py`); a zero-arg auto + barrier lands with pypto #2243 (plan 65). +- **Loops (HOST rail).** The wrappers are loop-safe: the HOST builtin kernels self-clear their barrier cells after every call (#2279 / plan 83), and for the mesh `all_reduce` the signal is the compiler-synthesized shared signal (#2504 / plan 88) — an implicit-signal allreduce inside `for` / `while` loops is accepted. Ring and sibling wrappers auto-allocate a reusable signal. +- **`pld.all_to_all_v`** (HOST) requires the `builtin.tensor.all_to_all_v` + rail (pypto #2243, plan 65); until it merges, the HOST path is rejected. + +### Putting it together: publish → collective → consume + +The wrappers auto-manage the **signal** only; the **data** window, the +per-rank publish dispatches, and the read-back are still explicit. A complete +HOST-orchestrator allreduce looks like: + +```python +data_buf = pld.alloc_window_buffer(64 * pl.FP32.get_byte()) +for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, 64], dtype=pl.FP32) + self.publish_orch(inputs[r], data, device=r) # user InCore publish step +data = pld.window(data_buf, [1, 64], dtype=pl.FP32) +data = pld.all_reduce(data, op=pld.ReduceOp.Sum) # mesh: signal auto-synthesized +for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[r], device=r) # user InCore consume step +``` + +Only the `pld.all_reduce` line is collective-specific; the publish / consume +steps are plain cross-scope dispatches you write once per kernel (full program: +`tests/st/distributed/test_l3_ergonomic_api.py`). + +**Mesh vs ring:** the default `mode="mesh"` is a direct all-to-all exchange — +simplest, best for small payloads. `mode="ring"` streams data in `2*(NR-1)` +pipelined steps with a smaller signal footprint and usually wins for large +payloads / high NR; it requires `nranks` (a static world size) and currently +supports `ReduceOp.Sum` + FP32 only. + ## Namespacing: why `tile.*` vs `tensor.*` vs `system.*` The namespace encodes the IR level the op lives at, not an arbitrary grouping: diff --git a/docs/zh/dev/distributed_ops.md b/docs/zh/dev/distributed_ops.md index 0bd83394eb..fda74e952b 100644 --- a/docs/zh/dev/distributed_ops.md +++ b/docs/zh/dev/distributed_ops.md @@ -36,6 +36,70 @@ TPUT/TGET 在该侧只需要一段可读/可写的*本地* GM 区域。窗口绑 [`UnknownType`](ir/02-types.md)：它们因跨 rank 副作用而存在，而非为消费者读取的 SSA 值而存在。 +## 便捷集合通信 API —— 自动管理信号 + +`pld.tensor.*` 集合通信每次调用都需要显式提供窗口绑定的 INT32 信号缓冲。 +便捷短形式（`pld.all_reduce`、`pld.all_gather`、`pld.reduce_scatter`、 +`pld.broadcast`、`pld.all_to_all`、`pld.all_to_all_v`、`pld.barrier`）会自动 +分配**全新且形状正确**的信号，并委托给对应的 `pld.tensor.*` HOST 内建算子： + +```python +data = pld.all_reduce(data, op=pld.ReduceOp.Sum) # mesh：无需信号（host 综合） +data = pld.all_reduce(data, mode="ring", nranks=2) # ring：自动分配 [2*(NR-1)+1, NR] 信号 +data = pld.all_gather(local, target) +data = pld.reduce_scatter(target, op=pld.ReduceOp.Sum) +data = pld.broadcast(target, root=0) +data = pld.all_to_all(input, target) +data = pld.all_to_all_v(input, target, send_counts, recv_counts, nranks=NR) +sig = pld.barrier(sig) # 需要显式且已覆盖的信号 +``` + +语义与约束： + +- **仅限 HOST 编排。** 这些包装基于仅 HOST 可用的 `alloc_window_buffer` / + `window` / `world_size` 原语；操作数为窗口绑定的 `DistributedTensor`，与 + `pld.tensor.*` 完全一致。 +- **信号形状**按算子选择，并与 HOST 内建算子的要求一致：`broadcast` / + `reduce_scatter` 使用一维 `[world_size]` 信号；`all_gather` / `all_to_all` + 使用二维 `[world_size, 1]` 信号；`all_to_all_v` 使用二维 `[nranks, 1]` + 信号（必须提供静态 `nranks`）；`allreduce mode="mesh"` 无需信号（编译器自动 + 综合），而 `mode="ring"` 需要静态的 `[2*(NR-1)+1, NR]` 信号，因此必须提供 + `nranks`。 +- **每次调用全新分配。** 信用屏障协议（pypto #2175，已合并）使信号自清零，可安全 + 地在连续调用间复用；包装仍每次分配新缓冲，这依然正确且是最简单的安全默认。 +- **`pld.barrier(sig)` 需要显式、已覆盖的信号**——屏障没有可继承通信域覆盖的 + 数据缓冲，自动分配的信号会被 `MaterializeCommDomainScopes` 拒绝。请传入一个 + 同时被设备标记派发消费的 INT32 窗口（参见 + `tests/st/distributed/test_l3_host_tensor_barrier.py`）；零参自动 barrier + 随 pypto #2243（计划 65）落地。 +- **循环（HOST 轨道）。** 包装是循环安全的：HOST 内建 kernel 会在每次调用后自清零屏障 cell（#2279 / plan 83），而 mesh `all_reduce` 的信号是编译器合成的共享信号（#2504 / plan 88）——`for` / `while` 循环内的隐式信号 allreduce 已被接受。ring 及兄弟包装会分配可复用的信号。 +- **`pld.all_to_all_v`**（HOST）需要 `builtin.tensor.all_to_all_v` 轨道 + （pypto #2243，计划 65）；在合并之前，HOST 路径会被拒绝。 + +### 完整流程：发布 → 集合通信 → 消费 + +包装仅自动管理**信号**；**数据**窗口、各 rank 的发布派发与读回仍需显式编写。 +一个完整的 HOST 编排 allreduce 如下： + +```python +data_buf = pld.alloc_window_buffer(64 * pl.FP32.get_byte()) +for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, 64], dtype=pl.FP32) + self.publish_orch(inputs[r], data, device=r) # 用户 InCore 发布步骤 +data = pld.window(data_buf, [1, 64], dtype=pl.FP32) +data = pld.all_reduce(data, op=pld.ReduceOp.Sum) # mesh：信号自动综合 +for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[r], device=r) # 用户 InCore 消费步骤 +``` + +只有 `pld.all_reduce` 一行与集合通信相关；发布 / 消费步骤是每个内核编写一次的 +普通跨作用域派发（完整程序见 `tests/st/distributed/test_l3_ergonomic_api.py`）。 + +**mesh 与 ring：** 默认的 `mode="mesh"` 是直接的全对全交换 —— 最简单，适合小 +数据量。`mode="ring"` 以 `2*(NR-1)` 步流水传输数据，信号占用更小，通常在大 +数据量 / 高 NR 时更优；它需要 `nranks`（静态世界大小），目前仅支持 +`ReduceOp.Sum` + FP32。 + ## 命名空间：为何区分 `tile.*` / `tensor.*` / `system.*` 命名空间编码的是算子所在的 IR 层级，而非随意分组： @@ -128,7 +192,7 @@ deducer 会校验打包的 `int` 落在枚举范围内,使 codegen 无需二次 ## 屏障-信号协议 每个 `pld.tensor.*` 集合通信算子（`allreduce`、`barrier`、`broadcast`、 -`reduce_scatter`、`allgather`、`all_to_all`）都使用同一个**自清理信用屏障** +`reduce_scatter`、`allgather`、`all_to_all`、`all_to_all_v`）都使用同一个**自清理信用屏障** （self-clearing credit barrier）进行同步，该屏障由 `pld.system.notify` / `pld.system.wait` 构建： diff --git a/python/pypto/language/distributed/__init__.py b/python/pypto/language/distributed/__init__.py index e74f72bc46..b245252055 100644 --- a/python/pypto/language/distributed/__init__.py +++ b/python/pypto/language/distributed/__init__.py @@ -38,10 +38,17 @@ from pypto.pypto_core.ir import AtomicType, NotifyOp, ReduceOp, WaitCmp from .op import ( + all_gather, + all_reduce, + all_to_all, + all_to_all_v, alloc_window_buffer, + barrier, + broadcast, get_comm_ctx, nranks, rank, + reduce_scatter, remote_load, remote_store, system, @@ -59,10 +66,17 @@ "NotifyOp", "ReduceOp", "WaitCmp", + "all_gather", + "all_reduce", + "all_to_all", + "all_to_all_v", "alloc_window_buffer", + "barrier", + "broadcast", "get_comm_ctx", "nranks", "rank", + "reduce_scatter", "remote_load", "remote_store", "system", diff --git a/python/pypto/language/distributed/op/__init__.py b/python/pypto/language/distributed/op/__init__.py index 79bfff96bd..80592f66ad 100644 --- a/python/pypto/language/distributed/op/__init__.py +++ b/python/pypto/language/distributed/op/__init__.py @@ -29,10 +29,17 @@ from . import tensor_ops as tensor from . import tile_ops as tile from .unified_ops import ( + all_gather, + all_reduce, + all_to_all, + all_to_all_v, alloc_window_buffer, + barrier, + broadcast, get_comm_ctx, nranks, rank, + reduce_scatter, remote_load, remote_store, window, @@ -40,10 +47,17 @@ ) __all__ = [ + "all_gather", + "all_reduce", + "all_to_all", + "all_to_all_v", "alloc_window_buffer", + "barrier", + "broadcast", "get_comm_ctx", "nranks", "rank", + "reduce_scatter", "remote_load", "remote_store", "system", diff --git a/python/pypto/language/distributed/op/collective_api.py b/python/pypto/language/distributed/op/collective_api.py new file mode 100644 index 0000000000..ff76b2a40c --- /dev/null +++ b/python/pypto/language/distributed/op/collective_api.py @@ -0,0 +1,283 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Ergonomic ``pld.*`` collective wrappers — auto-managed barrier signals. + +Each wrapper allocates a **fresh**, correctly shaped INT32 signal window (via +``pld.tensor.alloc_window_buffer(..., name=...)`` + ``pld.tensor.window``) and +delegates to the corresponding ``pld.tensor.*`` HOST builtin, removing the +per-call signal-buffer boilerplate. Two exceptions: ``all_reduce(mode="mesh")`` +uses a compiler-synthesized signal (no explicit window), and ``barrier`` takes +a caller-provided, comm-domain-covered signal. + +HOST-orchestration only: the underlying collective ops require window-bound +:class:`pld.DistributedTensor` operands and host-only allocation primitives +(``alloc_window_buffer`` / ``window`` / ``world_size``). + +Signals self-clear under the credit-barrier protocol (pypto #2175, merged), so +they are safe to reuse across calls. The wrappers still allocate a **fresh** +buffer per call, which remains correct and is the simplest safe default. +""" + +import itertools +from collections.abc import Sequence +from typing import Literal, TypeGuard, overload + +from pypto.language.typing import IntLike, Tensor +from pypto.pypto_core import DataType +from pypto.pypto_core import ir as _ir +from pypto.pypto_core.ir import ReduceOp + +from ..typing.distributed_tensor import DistributedTensor +from . import tensor_ops as _tensor +from ._utils import _unwrap +from .system_ops import world_size + +__all__ = [ + "all_gather", + "all_reduce", + "all_to_all", + "all_to_all_v", + "barrier", + "broadcast", + "reduce_scatter", +] + +_SIGNAL_COUNTER = itertools.count() + + +def _is_static_positive_int(value: object) -> TypeGuard[int]: + """True when ``value`` is a plain positive int (bool excluded).""" + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def _fresh_signal(op_name: str, shape: Sequence[IntLike]) -> DistributedTensor: + """Allocate a fresh, correctly shaped INT32 signal window for a collective call. + + Args: + op_name: Short op name used to derive a unique buffer identifier + (``__auto__``). The parser injects ``alloc_window_buffer`` + names only from program-body assignments; an explicit generated name + is required here because the call lives inside helper Python. + shape: Per-rank signal shape (e.g. ``[world_size(), 1]`` or a static + ring shape). + + Returns: + A window-bound :class:`pld.DistributedTensor` INT32 signal. + """ + name = f"__auto_{op_name}_{next(_SIGNAL_COUNTER)}" + buf = _tensor.alloc_window_buffer(shape, dtype=DataType.INT32, name=name) + return _tensor.window(buf, shape, dtype=DataType.INT32) + + +@overload +def all_reduce( + target: DistributedTensor, + *, + op: ReduceOp = ReduceOp.Sum, + mode: Literal["mesh"] = "mesh", + nranks: None = None, +) -> DistributedTensor: ... + + +@overload +def all_reduce( + target: DistributedTensor, + *, + op: ReduceOp = ReduceOp.Sum, + mode: Literal["ring"], + nranks: int, +) -> DistributedTensor: ... + + +def all_reduce( + target: DistributedTensor, + *, + op: ReduceOp = ReduceOp.Sum, + mode: str = "mesh", + nranks: int | None = None, +) -> DistributedTensor: + """In-place cross-rank allreduce, auto-managing the barrier signal. + + HOST-orchestration only; ``target`` is a window-bound + :class:`pld.DistributedTensor` holding each rank's partial result. + + * ``mode="mesh"`` (default): no signal is allocated — the compiler + synthesizes a private INT32 signal (``[world_size, 1]``). + * ``mode="ring"``: a fresh signal of shape ``[2*(nranks-1)+1, nranks]`` is + allocated; ``nranks`` (the static world size) is required because the + ring signal shape is a compile-time constant. + + Args: + target: Window-bound :class:`pld.DistributedTensor`, reduced in place. + op: :class:`pld.ReduceOp` (``Sum`` / ``Max`` / ``Min`` / ``Prod``). + mode: ``"mesh"`` (default) or ``"ring"``. ``"ring"`` currently supports + ``ReduceOp.Sum`` + FP32 only. + nranks: Static world size, required for ``mode="ring"`` and rejected + otherwise (the ring signal shape is a compile-time constant). + + Returns: + The rebound ``target`` (window-as-result). + """ + if mode not in ("mesh", "ring"): + raise ValueError(f'pld.all_reduce mode must be "mesh" or "ring", got {mode!r}') + if mode == "ring": + if not _is_static_positive_int(nranks): + raise ValueError( + "pld.all_reduce(mode='ring') requires a positive static int `nranks` " + "(the ring signal shape [2*(NR-1)+1, NR] is a compile-time constant)" + ) + if op != ReduceOp.Sum: + raise ValueError( + "pld.all_reduce(mode='ring') supports only ReduceOp.Sum " + "(the HOST ring schedule implements Sum only)" + ) + # Guarded read: a non-tensor arg (e.g. an int) falls through _unwrap, so + # skip the dtype check and let the delegation raise its TypeError. + target_type = getattr(_unwrap(target), "type", None) + if isinstance(target_type, _ir.ShapedType) and target_type.dtype != DataType.FP32: + raise ValueError( + "pld.all_reduce(mode='ring') supports only FP32 targets (the HOST ring schedule is FP32-only)" + ) + signal = _fresh_signal("allreduce_ring", [2 * (nranks - 1) + 1, nranks]) + return _tensor.allreduce(target, signal, op=op, mode="ring") + if nranks is not None: + raise ValueError( + "pld.all_reduce nranks is only used with mode='ring'; " + "omit it for mesh allreduce (or pass mode='ring')" + ) + return _tensor.allreduce(target, op=op) + + +def all_gather(local_data: DistributedTensor, target: DistributedTensor) -> DistributedTensor: + """All-gather ``local_data`` into ``target`` (push-based), auto-managing the signal. + + Args: + local_data: Window-bound :class:`pld.DistributedTensor` ``[1, SIZE]`` + holding this rank's chunk. Must differ from ``target``. + target: Window-bound :class:`pld.DistributedTensor` ``[NR, SIZE]`` + result window. + + Returns: + The ``target`` :class:`pld.DistributedTensor` (window-as-result). + """ + signal = _fresh_signal("allgather", [world_size(), 1]) + return _tensor.allgather(local_data, target, signal) + + +def reduce_scatter( + target: DistributedTensor, + *, + op: ReduceOp = ReduceOp.Sum, +) -> DistributedTensor: + """Reduce-scatter ``target`` in place, auto-managing the signal. + + Args: + target: Window-bound :class:`pld.DistributedTensor` ``[NR, SIZE]`` + (each rank stages all NR chunks, one per row). + op: :class:`pld.ReduceOp`. HOST ``reduce_scatter`` supports ``Sum`` + only; the wrapper rejects other ops up front. + + Returns: + The rebound ``target`` :class:`pld.DistributedTensor`. + """ + # The HOST builtin.tensor.reduce_scatter requires a rank-1 [world_size] signal. + if op != ReduceOp.Sum: + raise ValueError( + "pld.reduce_scatter supports only ReduceOp.Sum (the HOST builtin implements Sum only)" + ) + signal = _fresh_signal("reduce_scatter", [world_size()]) + return _tensor.reduce_scatter(target, signal, op=op) + + +def broadcast(target: DistributedTensor, *, root: int) -> DistributedTensor: + """Broadcast ``root``'s data to all ranks, auto-managing the signal. + + Args: + target: Window-bound :class:`pld.DistributedTensor`; root must stage + its data before the call. + root: Root rank index (int). + + Returns: + The rebound ``target`` :class:`pld.DistributedTensor`. + """ + # The HOST builtin.tensor.broadcast requires a rank-1 [world_size] signal. + signal = _fresh_signal("broadcast", [world_size()]) + return _tensor.broadcast(target, signal, root=root) + + +def barrier(signal: DistributedTensor) -> DistributedTensor: + """Cross-rank barrier using an explicit, comm-domain-covered signal. + + Unlike the other collectives, ``pld.barrier`` has **no data buffer** from + which ``MaterializeCommDomainScopes`` can inherit comm-domain coverage for + an auto-allocated signal, and a barrier-only signal must be consumed by a + device-tagged chip dispatch on ``main``. The caller therefore passes a + signal that already carries coverage (e.g. an INT32 window also passed to a + publish/dispatch call). A zero-arg auto-signal ``pld.barrier()`` becomes + possible once the coverage fallback lands (pypto #2243, plan 65). + + Args: + signal: Window-bound INT32 :class:`pld.DistributedTensor` with + comm-domain coverage. + + Returns: + The rebound signal :class:`pld.DistributedTensor`. + """ + return _tensor.barrier(signal) + + +def all_to_all(input: Tensor | DistributedTensor, target: DistributedTensor) -> DistributedTensor: + """Symmetric all-to-all (push-based), auto-managing the signal. + + Args: + input: ``[NR, SIZE]`` :class:`pl.Tensor` or :class:`pld.DistributedTensor` + with per-destination chunks, distinct from ``target``. + target: Window-bound :class:`pld.DistributedTensor` ``[NR, SIZE]`` + result window. + + Returns: + The ``target`` :class:`pld.DistributedTensor` (window-as-result). + """ + signal = _fresh_signal("all_to_all", [world_size(), 1]) + return _tensor.all_to_all(input, target, signal) + + +def all_to_all_v( + input: Tensor | DistributedTensor, + target: DistributedTensor, + send_counts: Tensor | DistributedTensor, + recv_counts: DistributedTensor, + *, + nranks: int, +) -> DistributedTensor: + """Variable-size all-to-all (push-based), auto-managing the signal. + + Args: + input: Flat ``[NR*MAX_RECV, SIZE]`` send buffer. + target: Flat window-bound :class:`pld.DistributedTensor` + ``[NR*MAX_RECV, SIZE]`` staging/result window. + send_counts: Local INT32 ``[NR]`` per-peer send row counts. + recv_counts: Published INT32 :class:`pld.DistributedTensor` ``[NR, 1]``. + nranks: Static world size (int), required — the signal's first + dimension and ``MAX_RECV = target.shape[0] // NR`` must be + compile-time constants, so a dynamic ``world_size()`` is rejected + by the ``pld.tensor.all_to_all_v`` type deducer. + + Returns: + The ``target`` :class:`pld.DistributedTensor` (window-as-result). + """ + if not _is_static_positive_int(nranks): + raise ValueError( + "pld.all_to_all_v requires a positive static int `nranks` " + "(the signal shape [nranks, 1] and MAX_RECV = target.shape[0] // NR " + "need a compile-time NR)" + ) + signal = _fresh_signal("all_to_all_v", [nranks, 1]) + return _tensor.all_to_all_v(input, target, signal, send_counts, recv_counts) diff --git a/python/pypto/language/distributed/op/unified_ops.py b/python/pypto/language/distributed/op/unified_ops.py index df978a3d5a..6044388f2c 100644 --- a/python/pypto/language/distributed/op/unified_ops.py +++ b/python/pypto/language/distributed/op/unified_ops.py @@ -31,15 +31,31 @@ from ..typing.distributed_tensor import DistributedTensor from . import tensor_ops as _tensor from . import tile_ops as _tile +from .collective_api import ( + all_gather, + all_reduce, + all_to_all, + all_to_all_v, + barrier, + broadcast, + reduce_scatter, +) from .system_ops import get_comm_ctx, nranks, rank, world_size from .tensor_ops import alloc_window_buffer, window from .tile_ops import remote_load __all__ = [ "alloc_window_buffer", + "all_gather", + "all_reduce", + "all_to_all", + "all_to_all_v", + "barrier", + "broadcast", "get_comm_ctx", "nranks", "rank", + "reduce_scatter", "remote_load", "remote_store", "window", diff --git a/tests/st/distributed/test_l3_ergonomic_api.py b/tests/st/distributed/test_l3_ergonomic_api.py new file mode 100644 index 0000000000..89585212c4 --- /dev/null +++ b/tests/st/distributed/test_l3_ergonomic_api.py @@ -0,0 +1,393 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""L3 distributed ST: ergonomic ``pld.*`` collective wrappers (plan 58). + +Each test drives the HOST auto-signal wrapper (``pld.all_reduce`` / +``pld.broadcast`` / ``pld.all_gather`` / ``pld.barrier``) through the real +compiler and run path, with per-rank inputs so the goldens validate the +cross-rank exchange (not just local pass-through). The underlying collective +semantics are additionally covered by the ``test_l3_host_tensor_*.py`` siblings; +here the point is the auto-managed signal + parser resolution end-to-end. + +``all_to_all_v`` is intentionally not run end-to-end here: the HOST +``builtin.tensor.all_to_all_v`` rail is plan 65 / #2243 (not on ``main`` yet). +Its wrapper delegation is unit-tested in ``tests/ut/language/test_collective_api.py``. +""" + +import pypto.language as pl +import pypto.language.distributed as pld +import pytest +import torch +from pypto import ir +from pypto.ir.distributed_compiled_program import DistributedConfig + +SIZE = 64 +NR = 2 + + +def _make_rank_inputs(n_ranks: int, size: int = SIZE) -> torch.Tensor: + """Per-rank distinct inputs: rank r is [r*100, r*100+size).""" + rows = [ + torch.arange(r * 100.0, r * 100.0 + size, dtype=torch.float32).reshape(1, size) + for r in range(n_ranks) + ] + return torch.stack(rows) + + +def _compile(program, test_config, device_ids): + return ir.compile( + program, + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:NR], + num_sub_workers=0, + ), + ) + + +def _assert_close(outputs, expected, label): + assert torch.allclose(outputs, expected), ( + f"{label} mismatch: max diff = {(outputs - expected).abs().max().item()}" + ) + + +@pl.program +class ErgonomicAllReduce: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(data, [0, 0], [1, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.all_reduce(data, op=pld.ReduceOp.Sum) # mesh: signal auto-synthesized + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[r], device=r) + return outputs + + +@pl.program +class ErgonomicAllReduceRing: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(data, [0, 0], [1, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + for r in pl.range(NR): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.all_reduce(data, op=pld.ReduceOp.Sum, mode="ring", nranks=NR) + for r in pl.range(NR): + self.consume_orch(data, outputs[r], device=r) + return outputs + + +@pl.program +class ErgonomicBroadcast: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + if my_rank == 0: # root stages only + local = pl.load(inp, [0, 0], [1, SIZE]) + return pl.store(local, [0, 0], data) + return data + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data, my_rank) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(data, [0, 0], [1, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], data, r, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.broadcast(data, root=0) + for r in pl.range(NR): + self.consume_orch(data, outputs[r], device=r) + return outputs + + +@pl.program +class ErgonomicAllGather: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + local: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], local) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + local: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, local) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + target: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: + # Copy the whole gathered target so every rank's chunk is validated. + return pl.store(pl.load(target, [0, 0], [NR, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + target: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: + return self.consume_step(target, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[NR, NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, NR, SIZE], pl.FP32]: + local_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + target_buf = pld.alloc_window_buffer(NR * SIZE * pl.FP32.get_byte()) + for r in pl.range(NR): + local = pld.window(local_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], local, device=r) + local = pld.window(local_buf, [1, SIZE], dtype=pl.FP32) + target = pld.window(target_buf, [NR, SIZE], dtype=pl.FP32) + target = pld.all_gather(local, target) + for r in pl.range(NR): + self.consume_orch(target, outputs[r], device=r) + return outputs + + +@pl.program +class ErgonomicBarrier: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + sig: pld.DistributedTensor[[NR], pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + # ``sig`` is declared so the dispatch tags the signal window with + # comm-domain coverage (a barrier-only signal cannot be auto-covered + # on ``main`` — see ``pld.barrier``'s docstring). + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + peer: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + recv = pld.tile.remote_load(data, peer=peer, offsets=[0, 0], shape=[1, SIZE]) + return pl.store(recv, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + peer: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out, peer) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], data, signal, device=r) + + signal = pld.barrier(signal) # covered signal; all ranks must reach it + + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + peer = (r + 1) % pld.world_size() + self.consume_orch(data, outputs[r], peer, device=r) + return outputs + + +class TestErgonomicCollectiveApi: + def test_all_reduce_mesh(self, test_config, device_ids): + if len(device_ids) < NR: + pytest.skip(f"ergonomic allreduce P={NR} needs {NR} devices, got {device_ids}") + compiled = _compile(ErgonomicAllReduce, test_config, device_ids) + inputs = _make_rank_inputs(NR) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + expected = torch.stack([inputs.sum(dim=0)] * NR) + _assert_close(outputs, expected, "mesh AR") + + def test_all_reduce_ring(self, test_config, device_ids): + if len(device_ids) < NR: + pytest.skip(f"ergonomic ring allreduce P={NR} needs {NR} devices, got {device_ids}") + compiled = _compile(ErgonomicAllReduceRing, test_config, device_ids) + inputs = _make_rank_inputs(NR) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + expected = torch.stack([inputs.sum(dim=0)] * NR) + _assert_close(outputs, expected, "ring AR") + + def test_broadcast(self, test_config, device_ids): + if len(device_ids) < NR: + pytest.skip(f"ergonomic broadcast P={NR} needs {NR} devices, got {device_ids}") + compiled = _compile(ErgonomicBroadcast, test_config, device_ids) + inputs = _make_rank_inputs(NR) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + expected = torch.stack([inputs[0]] * NR) + _assert_close(outputs, expected, "broadcast") + + def test_all_gather(self, test_config, device_ids): + if len(device_ids) < NR: + pytest.skip(f"ergonomic all_gather P={NR} needs {NR} devices, got {device_ids}") + compiled = _compile(ErgonomicAllGather, test_config, device_ids) + inputs = _make_rank_inputs(NR) + outputs = torch.zeros((NR, NR, SIZE), dtype=inputs.dtype, device=inputs.device) + compiled(inputs, outputs) + # Every rank reads the full gathered target: output[r] holds all rank chunks. + expected = torch.stack([inputs[:, 0, :]] * NR) + _assert_close(outputs, expected, "all_gather") + + def test_barrier(self, test_config, device_ids): + if len(device_ids) < NR: + pytest.skip(f"ergonomic barrier P={NR} needs {NR} devices, got {device_ids}") + compiled = _compile(ErgonomicBarrier, test_config, device_ids) + inputs = _make_rank_inputs(NR) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + # Each rank remote-loads its peer's published data after the barrier — + # validates the barrier actually synchronized the ranks. + expected = torch.stack([inputs[(r + 1) % NR] for r in range(NR)]) + assert torch.equal(outputs, expected) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/ut/language/test_collective_api.py b/tests/ut/language/test_collective_api.py new file mode 100644 index 0000000000..d6c8cd7c77 --- /dev/null +++ b/tests/ut/language/test_collective_api.py @@ -0,0 +1,327 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Unit tests for the ergonomic ``pld.*`` collective wrappers (``collective_api.py``). + +Covers auto-signal allocation (shape per op, fresh-per-call unique names), +mesh-vs-ring signal handling, kwarg passthrough, and parser resolution of the +``pld.`` short forms inside a host-orchestration program body. +""" + +from typing import Any, cast + +import pypto.language as pl +import pypto.language.distributed as pld +import pytest +from pypto import ir +from pypto.pypto_core import ir as _ir +from pypto.pypto_core.ir import Call, ConstInt, ShapedType + +SIZE = 8 + + +def _window(shape, dtype, name) -> pld.DistributedTensor: + """Build a window-bound DistributedTensor in pure Python (explicit alloc name).""" + buf = pld.tensor.alloc_window_buffer(shape, dtype=dtype, name=name) + return pld.tensor.window(buf, shape, dtype=dtype) + + +def _as_call(result) -> Call: + """Return the underlying IR Call of a wrapper result (narrowed from Expr).""" + return cast(Call, result._expr) + + +def _shape_of(arg) -> list[int]: + """Return the per-rank shape of a DistributedTensor-typed argument as ints.""" + dist_type = cast(ShapedType, cast(Call, arg).type) + return [int(cast(ConstInt, d).value) for d in dist_type.shape] + + +def _signal_alloc_name(call: Call, signal_arg_index: int) -> str: + """Extract the alloc_window_buffer name backing a signal argument.""" + window_call = cast(Call, call.args[signal_arg_index]) + alloc_call = cast(Call, window_call.args[0]) # window(buf, ...) -> buf is the alloc call + return cast(str, alloc_call.kwargs["name"]) + + +class TestAllReduceSignal: + def test_mesh_emits_no_explicit_signal(self): + target = _window([1, SIZE], pl.FP32, "t") + call = _as_call(pld.all_reduce(target, op=pld.ReduceOp.Sum)) + assert call.op.name == _ir.get_op("pld.tensor.allreduce").name + # Host synthesis path: no explicit signal argument. + assert len(call.args) == 1 + assert call.kwargs["op"] == int(pld.ReduceOp.Sum) + + def test_ring_allocates_static_signal_shape(self): + target = _window([1, SIZE], pl.FP32, "t") + call = _as_call(pld.all_reduce(target, mode="ring", nranks=2)) + assert call.op.name == _ir.get_op("pld.tensor.allreduce").name + # Ring signal is [2*(NR-1)+1, NR] = [3, 2] for NR=2. + assert len(call.args) == 2 + assert _shape_of(call.args[1]) == [3, 2] + assert cast(ShapedType, call.args[1].type).dtype == pl.INT32 + assert call.kwargs["mode"] == "ring" + + def test_ring_requires_static_nranks(self): + target = _window([1, SIZE], pl.FP32, "t") + # Widened to Any: the Literal overloads make mode="ring" without nranks a + # static type error — this exercises the runtime guard for DSL code that + # is not type-checked. + with pytest.raises(ValueError, match="nranks"): + cast(Any, pld.all_reduce)(target, mode="ring") + + def test_mesh_rejects_nranks(self): + target = _window([1, SIZE], pl.FP32, "t") + with pytest.raises(ValueError, match="nranks"): + cast(Any, pld.all_reduce)(target, nranks=2) + + def test_invalid_mode_rejected(self): + target = _window([1, SIZE], pl.FP32, "t") + with pytest.raises(ValueError, match="mesh"): + cast(Any, pld.all_reduce)(target, mode="tree") + + def test_ring_rejects_non_sum(self): + target = _window([1, SIZE], pl.FP32, "t") + with pytest.raises(ValueError, match="Sum"): + pld.all_reduce(target, mode="ring", nranks=2, op=pld.ReduceOp.Max) + + def test_ring_rejects_non_fp32(self): + target = _window([1, SIZE], pl.FP16, "t") + with pytest.raises(ValueError, match="FP32"): + pld.all_reduce(target, mode="ring", nranks=2) + + def test_ring_rejects_non_positive_nranks(self): + target = _window([1, SIZE], pl.FP32, "t") + with pytest.raises(ValueError, match="positive"): + cast(Any, pld.all_reduce)(target, mode="ring", nranks=0) + + +class TestPassthrough: + def test_broadcast_root_passthrough(self): + target = _window([1, SIZE], pl.FP32, "t") + call = _as_call(pld.broadcast(target, root=0)) + assert call.op.name == _ir.get_op("pld.tensor.broadcast").name + assert call.kwargs["root"] == 0 + # target + auto signal; broadcast builtin requires a rank-1 signal. + assert len(call.args) == 2 + assert len(cast(ShapedType, call.args[1].type).shape) == 1 + + def test_reduce_scatter_op_passthrough(self): + target = _window([2, SIZE], pl.FP32, "t") + call = _as_call(pld.reduce_scatter(target, op=pld.ReduceOp.Sum)) + assert call.op.name == _ir.get_op("pld.tensor.reduce_scatter").name + assert call.kwargs["op"] == int(pld.ReduceOp.Sum) + # target + auto signal; reduce_scatter builtin requires a rank-1 signal. + assert len(call.args) == 2 + assert len(cast(ShapedType, call.args[1].type).shape) == 1 + + def test_reduce_scatter_rejects_non_sum(self): + target = _window([2, SIZE], pl.FP32, "t") + with pytest.raises(ValueError, match="Sum"): + pld.reduce_scatter(target, op=pld.ReduceOp.Max) + + def test_all_to_all_and_all_gather_delegate(self): + inp = _window([2, SIZE], pl.FP32, "inp") + target = _window([2, SIZE], pl.FP32, "tgt") + a2a = _as_call(pld.all_to_all(inp, target)) + assert a2a.op.name == _ir.get_op("pld.tensor.all_to_all").name + assert len(a2a.args) == 3 # input + target + signal + + local = _window([1, SIZE], pl.FP32, "loc") + gather = _as_call(pld.all_gather(local, target)) + assert gather.op.name == _ir.get_op("pld.tensor.allgather").name + assert len(gather.args) == 3 + + def test_all_to_all_v_delegates_with_static_nranks(self): + inp = _window([2 * SIZE, SIZE], pl.FP32, "inp") + target = _window([2 * SIZE, SIZE], pl.FP32, "tgt") + send = _window([2], pl.INT32, "send") + recv = _window([2, 1], pl.INT32, "recv") + call = _as_call(pld.all_to_all_v(inp, target, send, recv, nranks=2)) + assert call.op.name == _ir.get_op("pld.tensor.all_to_all_v").name + # arg order: input, target, signal, send_counts, recv_counts. + assert len(call.args) == 5 + # signal is [nranks, 1] = [2, 1] INT32. + assert _shape_of(call.args[2]) == [2, 1] + assert cast(ShapedType, call.args[2].type).dtype == pl.INT32 + + def test_all_to_all_v_rejects_non_positive_nranks(self): + inp = _window([2 * SIZE, SIZE], pl.FP32, "inp") + target = _window([2 * SIZE, SIZE], pl.FP32, "tgt") + send = _window([2], pl.INT32, "send") + recv = _window([2, 1], pl.INT32, "recv") + with pytest.raises(ValueError, match="positive"): + pld.all_to_all_v(inp, target, send, recv, nranks=0) + + def test_barrier_requires_covered_signal(self): + signal = _window([2], pl.INT32, "sig") + call = _as_call(pld.barrier(signal)) + assert call.op.name == _ir.get_op("pld.tensor.barrier").name + assert len(call.args) == 1 # the user-provided signal only + + +class TestFreshSignal: + def test_signal_allocated_fresh_per_call(self): + names = [] + for i in range(2): + local = _window([1, SIZE], pl.FP32, f"loc{i}") + target = _window([2, SIZE], pl.FP32, f"tgt{i}") + call = _as_call(pld.all_gather(local, target)) + names.append(_signal_alloc_name(call, signal_arg_index=2)) + assert len(set(names)) == 2, "signals must be fresh (unique) per call" + assert all(n.startswith("__auto_") for n in names) + + def test_signal_shape_is_rank2_world_size(self): + local = _window([1, SIZE], pl.FP32, "loc") + target = _window([2, SIZE], pl.FP32, "tgt") + call = _as_call(pld.all_gather(local, target)) + # signal = [world_size(), 1] — dynamic NR, static second extent 1. + sig_type = cast(ShapedType, call.args[2].type) + assert int(cast(ConstInt, sig_type.shape[1]).value) == 1 + + +@pl.program +class _HostAllReduce: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(data, [0, 0], [1, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[2, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[2, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[2, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.all_reduce(data, op=pld.ReduceOp.Sum) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[r], device=r) + return outputs + + +@pl.program +class _HostAllGather: + """Program exercising a signal-bearing wrapper (pld.all_gather).""" + + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + local: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(inp, [0, 0], [1, SIZE]), [0, 0], local) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + local: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, local) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + target: pld.DistributedTensor[[2, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return pl.store(pl.load(target, [0, 0], [1, SIZE]), [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + target: pld.DistributedTensor[[2, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(target, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[2, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[2, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[2, 1, SIZE], pl.FP32]: + local_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + target_buf = pld.alloc_window_buffer(2 * SIZE * pl.FP32.get_byte()) + for r in pl.range(2): + local = pld.window(local_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[r], local, device=r) + local = pld.window(local_buf, [1, SIZE], dtype=pl.FP32) + target = pld.window(target_buf, [2, SIZE], dtype=pl.FP32) + target = pld.all_gather(local, target) # auto signal: __auto_allgather_ + for r in pl.range(2): + self.consume_orch(target, outputs[r], device=r) + return outputs + + +class TestParserResolution: + def test_pld_all_reduce_resolves_in_host_body(self): + """The parser resolves the pld.all_reduce short form and round-trips.""" + # as_python() parses the program body; an unresolved pld.all_reduce + # would raise "Unknown distributed operation" here. The printer + # expands the wrapper to the canonical pld.tensor.allreduce IR op. + printed = _HostAllReduce.as_python() + assert "pld.tensor.allreduce(" in printed + reparsed = pl.parse_program(printed) + assert isinstance(reparsed, ir.Program) + ir.assert_structural_equal(_HostAllReduce, reparsed) + + @pytest.mark.xfail( + strict=True, + reason=( + "as_python() prints the auto signal as an inline " + "window(alloc_window_buffer(...)) with the name dropped; parse_program " + "requires alloc_window_buffer as the RHS of a simple assignment, so " + "signal-bearing wrappers cannot round-trip yet (python_printer hoist fix)." + ), + ) + def test_signal_bearing_wrapper_round_trips(self): + """A wrapper with an auto-allocated signal must survive print→reparse.""" + printed = _HostAllGather.as_python() + assert "pld.tensor.allgather(" in printed + reparsed = pl.parse_program(printed) + assert isinstance(reparsed, ir.Program) + ir.assert_structural_equal(_HostAllGather, reparsed) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])