diff --git a/examples/workers/l3/README.md b/examples/workers/l3/README.md index e84ebb1493..2a7d0f40d8 100644 --- a/examples/workers/l3/README.md +++ b/examples/workers/l3/README.md @@ -64,6 +64,10 @@ Two things to know before reading the example: | [`multi_chip_dispatch/`](multi_chip_dispatch/) | Two chips + one SubWorker. An orchestration fn dispatches a `ChipCallable` to each chip, then submits a Python callable to collect/verify results. | | [`child_memory/`](child_memory/) | `orch.malloc` + `ContinuousTensor(child_memory=True)` to load a weight once and reuse it across multiple kernel invocations on the same chip. | | [`allreduce_distributed/`](allreduce_distributed/) | One communication domain allocated inside the orchestration via `orch.allocate_domain`, with PTO-ISA remote reads over the domain window. | +| [`allgather_distributed/`](allgather_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages its slice, synchronizes across ranks, then gathers every rank's window data into a full output. | +| [`reduce_scatter_distributed/`](reduce_scatter_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages all input chunks, synchronizes, then reduces the per-rank chunk across peers into a rank-local output. | +| [`broadcast_distributed/`](broadcast_distributed/) | One communication domain via `orch.allocate_domain`; root stages into the window, synchronizes, then every rank reads the root's scratch slot into its output. | +| [`all_to_all_distributed/`](all_to_all_distributed/) | One communication domain via `orch.allocate_domain`; scratch indexed by destination rank, barrier, then each rank gathers the slice peers sent to it. | | [`ffn_tp_parallel/`](ffn_tp_parallel/) | Local compute followed by one-domain cross-rank reduction through a domain scratch window. | | [`ep_dispatch_combine/`](ep_dispatch_combine/) | MoE-style dispatch/combine over a one-domain communication window. | | [`domain_rank_map/`](domain_rank_map/) | Small two-domain example showing domain-local ranks, missing-domain `KeyError`, separate window slices, and real per-domain allreduce. | diff --git a/examples/workers/l3/all_to_all_distributed/__init__.py b/examples/workers/l3/all_to_all_distributed/__init__.py new file mode 100644 index 0000000000..25708baecc --- /dev/null +++ b/examples/workers/l3/all_to_all_distributed/__init__.py @@ -0,0 +1,9 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Package marker so ``test_*.py`` can do ``from .main import run``.""" diff --git a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp new file mode 100644 index 0000000000..cdde406bd0 --- /dev/null +++ b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp @@ -0,0 +1,139 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * AllToAll kernel — symmetric, 3-phase, HCCL-window scratch pattern. + * + * Phase 1 (stage-in): for dest in 0..nranks-1: input[dest*C..) → scratch[dest*C..) + * Phase 2 (barrier): signal matrix + TWAIT cross-rank sync + * Phase 3 (exchange): for src in 0..nranks-1: TLOAD(peer src scratch[my_rank*C]) → output[src*C..) + * + * Scratch is indexed by destination rank: scratch[dest*C] holds the chunk sent to dest. + * + * args layout: + * tensor(0) = input nranks*COUNT_PER_RANK floats (INPUT) + * tensor(1) = output nranks*COUNT_PER_RANK floats (OUTPUT_EXISTING) + * tensor(2) = scratch HCCL window slot (INOUT) + * scalar(0) = nranks + * scalar(1) = CommContext device pointer + */ + +#include +#include +#include "pto/comm/comm_types.hpp" +#include "pto/comm/pto_comm_inst.hpp" +#include "platform_comm/comm_context.h" +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *localPtr, int pe) { + uint64_t localBase = ctx->windowsIn[ctx->rankId]; + uint64_t offset = (uint64_t)localPtr - localBase; + return (__gm__ T *)(ctx->windowsIn[pe] + offset); +} + +static constexpr size_t COUNT_PER_RANK = 64; +static constexpr int kMaxSupportedRanks = 16; + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *output_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + int nranks = static_cast(args[3]); + __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; + __gm__ float *scratch = + reinterpret_cast<__gm__ float *>(scratch_tensor->buffer.addr) + scratch_tensor->start_offset; + + using ShapeDyn = pto::Shape; + using StrideDyn = pto::Stride; + using Global = pto::GlobalTensor; + using TileData = pto::Tile; + + int my_rank = static_cast(commCtx->rankId); + + if (nranks <= 0 || nranks > kMaxSupportedRanks) { + pipe_barrier(PIPE_ALL); + return; + } + + // signal_base follows the nranks * COUNT_PER_RANK float staging region. + __gm__ int32_t *signal_base = reinterpret_cast<__gm__ int32_t *>(scratch + nranks * COUNT_PER_RANK); + + ShapeDyn shape(1, 1, 1, 1, COUNT_PER_RANK); + StrideDyn stride(COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, 1); + + TileData stageTile(1, COUNT_PER_RANK); + TileData recvTile(1, COUNT_PER_RANK); + TASSIGN(stageTile, 0x0); + TASSIGN(recvTile, 0x10000); + + // ------------------------------------------------------------------ + // Phase 1: stage-in — copy each destination chunk into scratch so + // every peer can read the slice destined for them in Phase 3. + // ------------------------------------------------------------------ + for (int dest = 0; dest < nranks; ++dest) { + Global inputChunkG(input + dest * COUNT_PER_RANK, shape, stride); + Global scratchChunkG(scratch + dest * COUNT_PER_RANK, shape, stride); + TLOAD(stageTile, inputChunkG); + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + TSTORE(scratchChunkG, stageTile); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + } + pipe_barrier(PIPE_ALL); + + // ------------------------------------------------------------------ + // Phase 2: device barrier — notify every peer that stage-in is done, + // then wait until every peer has notified us. + // ------------------------------------------------------------------ + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + __gm__ int32_t *remote_signal = CommRemotePtr(commCtx, signal_base + my_rank, peer); + pto::comm::Signal sig(remote_signal); + pto::comm::TNOTIFY(sig, (int32_t)1, pto::comm::NotifyOp::AtomicAdd); + } + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal sig(signal_base + peer); + pto::comm::TWAIT(sig, (int32_t)1, pto::comm::WaitCmp::GE); + } + pipe_barrier(PIPE_ALL); + + // ------------------------------------------------------------------ + // Phase 3: exchange — read chunk my_rank from every rank's scratch and + // write it into the corresponding slice of the output tensor. + // CommRemotePtr with pe==my_rank returns localPtr unchanged, so the + // self-read goes through the same code path as the remote reads. + // ------------------------------------------------------------------ + for (int src = 0; src < nranks; ++src) { + __gm__ float *remote_chunk = CommRemotePtr(commCtx, scratch + my_rank * COUNT_PER_RANK, src); + Global remoteG(remote_chunk, shape, stride); + Global outputSlotG(output + src * COUNT_PER_RANK, shape, stride); + TLOAD(recvTile, remoteG); + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + TSTORE(outputSlotG, recvTile); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + } + pipe_barrier(PIPE_ALL); +} diff --git a/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp b/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp new file mode 100644 index 0000000000..55973d0fd9 --- /dev/null +++ b/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp @@ -0,0 +1,49 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * AllToAll orchestration shim. + * + * tensor(0) input INPUT (nranks*COUNT_PER_RANK floats) + * tensor(1) output OUTPUT_EXISTING (nranks*COUNT_PER_RANK floats) + * tensor(2) scratch INOUT (HCCL window slot) + * scalar(0) nranks + * scalar(1) CommContext device pointer + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +all_to_all_orchestration_config(const ChipStorageTaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 5, // 3 tensors + 2 scalars + }; +} + +__attribute__((visibility("default"))) void all_to_all_orchestration(const ChipStorageTaskArgs &orch_args) { + Tensor input = from_tensor_arg(orch_args.tensor(0)); + Tensor output = from_tensor_arg(orch_args.tensor(1)); + Tensor scratch = from_tensor_arg(orch_args.tensor(2)); + + Arg params; + params.add_input(input); + params.add_output(output); + params.add_inout(scratch); + params.add_scalar(orch_args.scalar(0)); // nranks + params.add_scalar(orch_args.scalar(1)); // CommContext + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/examples/workers/l3/all_to_all_distributed/main.py b/examples/workers/l3/all_to_all_distributed/main.py new file mode 100644 index 0000000000..ec04cc112f --- /dev/null +++ b/examples/workers/l3/all_to_all_distributed/main.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""End-to-end distributed all-to-all — symmetric 3-phase pattern. + +Each rank owns nranks send chunks; chunk d is payload for rank d. After the +exchange rank r holds chunk s from every source s in output[s*COUNT_PER_RANK]: + + Phase 1 stage-in for d in 0..N-1: input[d*C..) → scratch[d*C..) + Phase 2 device barrier signal matrix cross-rank sync via TNOTIFY/TWAIT + Phase 3 exchange for s in 0..N-1: TLOAD(peer s scratch[my_rank*C]) → output[s*C..) + +Run: + python examples/workers/l3/all_to_all_distributed/main.py -p a2a3sim -d 0-1 + +""" + +from __future__ import annotations + +import argparse +import os +import sys + +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +import torch # noqa: E402 +from simpler.task_interface import ( # noqa: E402 + ArgDirection, + CallConfig, + ChipCallable, + CommBufferSpec, + ContinuousTensor, + CoreCallable, + DataType, + TaskArgs, + TensorArgType, +) +from simpler.worker import Worker # noqa: E402 + +from simpler_setup.elf_parser import extract_text_section # noqa: E402 +from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 +from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 +from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 + +HERE = os.path.dirname(os.path.abspath(__file__)) + +# Must match COUNT_PER_RANK in kernels/aiv/all_to_all_kernel.cpp. +COUNT_PER_RANK = 64 +DTYPE_NBYTES = 4 # float32 +# Signal tail: one int32 slot per rank, bounded by kMaxSupportedRanks. +SIGNAL_TAIL_NBYTES = 16 * 4 # 64 B +# NOTE: the full scratch size depends on nranks (staging area = nranks * COUNT_PER_RANK floats). +# It is computed inside run() once nranks is known. + + +def parse_device_range(spec: str) -> list[int]: + if "-" in spec: + lo, hi = (int(x) for x in spec.split("-")) + ids = list(range(lo, hi + 1)) + else: + ids = [int(spec)] + if not (2 <= len(ids) <= 16): + raise ValueError(f"all_to_all_distributed needs between 2 and 16 devices, got {len(ids)} ({ids})") + return ids + + +def build_chip_callable(platform: str, pto_isa_commit: str | None) -> ChipCallable: + """Compile the AIV all-to-all kernel + its C++ orchestration shim.""" + kc = KernelCompiler(platform=platform) + runtime = "tensormap_and_ringbuffer" + pto_isa_root = ensure_pto_isa_root(commit=pto_isa_commit, clone_protocol="https") + include_dirs = kc.get_orchestration_include_dirs(runtime) + + # src/common — for platform_comm/comm_context.h + kernel_include_dirs = list(include_dirs) + [ + str(kc.project_root / "src" / "common"), + ] + kernel_bytes = kc.compile_incore( + source_path=os.path.join(HERE, "kernels/aiv/all_to_all_kernel.cpp"), + core_type="aiv", + pto_isa_root=pto_isa_root, + extra_include_dirs=kernel_include_dirs, + ) + if not platform.endswith("sim"): + kernel_bytes = extract_text_section(kernel_bytes) + + orch_bytes = kc.compile_orchestration( + runtime_name=runtime, + source_path=os.path.join(HERE, "kernels/orchestration/all_to_all_orch.cpp"), + ) + core_callable = CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT, ArgDirection.INOUT], + binary=kernel_bytes, + ) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT, ArgDirection.INOUT], + func_name="all_to_all_orchestration", + config_name="all_to_all_orchestration_config", + binary=orch_bytes, + children=[(0, core_callable)], + ) + + +def expected_output(nranks: int, rank: int) -> list[float]: + """output[src*C + j] = input_src chunk destined for rank: src*1000 + rank*100 + j.""" + return [float(src * 1000 + rank * 100 + j) for src in range(nranks) for j in range(COUNT_PER_RANK)] + + +def run( + device_ids: list[int], + platform: str = "a2a3", + pto_isa_commit: str | None = None, + build: bool = False, +) -> int: + """Core logic — callable from both CLI and pytest.""" + nranks = len(device_ids) + # Scratch = nranks * COUNT_PER_RANK floats (staging) + signal tail. + scratch_nbytes = nranks * COUNT_PER_RANK * DTYPE_NBYTES + SIGNAL_TAIL_NBYTES + window_size = max(scratch_nbytes, 4 * 1024) + + print(f"[all_to_all] platform={platform} devices={device_ids} nranks={nranks}") + + host_inputs = [ + torch.tensor( + [rank * 1000 + dest * 100 + j for dest in range(nranks) for j in range(COUNT_PER_RANK)], + dtype=torch.float32, + ).share_memory_() + for rank in range(nranks) + ] + host_outputs = [torch.zeros(nranks * COUNT_PER_RANK, dtype=torch.float32).share_memory_() for _ in range(nranks)] + + print("[all_to_all] compiling kernels...") + chip_callable = build_chip_callable(platform, pto_isa_commit) + + worker = Worker( + level=3, + platform=platform, + runtime="tensormap_and_ringbuffer", + device_ids=device_ids, + num_sub_workers=0, + build=build, + ) + chip_cid = worker.register(chip_callable) + + try: + print("[all_to_all] init worker (forks chip children; base comm is lazy)...") + worker.init() + + def orch_fn(orch, _args, cfg): + with orch.allocate_domain( + name="default", + workers=list(range(nranks)), + window_size=window_size, + buffers=[ + CommBufferSpec( + name="scratch", + dtype="float32", + count=nranks * COUNT_PER_RANK, + nbytes=scratch_nbytes, + ) + ], + ) as handle: + for i in range(nranks): + domain = handle[i] + print( + f"[all_to_all] chip {i}: rank={domain.domain_rank}/{domain.domain_size} " + f"window=[0x{domain.local_window_base:x} +{domain.actual_window_size}B] " + f"scratch=0x{domain.buffer_ptrs['scratch']:x}" + ) + chip_args = TaskArgs() + chip_args.add_tensor(make_tensor_arg(host_inputs[i]), TensorArgType.INPUT) + chip_args.add_tensor(make_tensor_arg(host_outputs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_tensor( + ContinuousTensor.make( + data=domain.buffer_ptrs["scratch"], + shapes=(nranks * COUNT_PER_RANK,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INOUT, + ) + chip_args.add_scalar(domain.domain_size) + chip_args.add_scalar(domain.device_ctx) + orch.submit_next_level(chip_cid, chip_args, cfg, worker=i) + + print(f"[all_to_all] running {nranks}-chip all-to-all DAG...") + worker.run(orch_fn, args=None, config=CallConfig()) + + ok = True + for i in range(nranks): + expected = torch.tensor(expected_output(nranks, i), dtype=torch.float32) + max_diff = float(torch.max(torch.abs(host_outputs[i] - expected))) + print(f"[all_to_all] chip {i}: max |out - expected| = {max_diff:.3e}") + if max_diff > 1e-3: + ok = False + for j in range(min(4, nranks * COUNT_PER_RANK)): + print(f" output[{j}]={float(host_outputs[i][j])!r} expected={float(expected[j])!r}") + + if not ok: + print("[all_to_all] golden check FAILED") + return 1 + print("[all_to_all] all ranks matched golden ✅") + return 0 + finally: + worker.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("-p", "--platform", default="a2a3", help="Platform backend, e.g. a2a3 or a2a3sim.") + parser.add_argument( + "-d", "--device", default="0-1", help="Device range, e.g. '0-1' or '0-3'. 2 to 16 chips required." + ) + parser.add_argument( + "--build", action="store_true", help="Rebuild runtime from source instead of using cached libs." + ) + parser.add_argument("--pto-isa-commit", default=None, help="Optional PTO ISA commit/tag to fetch before compiling.") + cli = parser.parse_args() + + return run( + parse_device_range(cli.device), platform=cli.platform, pto_isa_commit=cli.pto_isa_commit, build=cli.build + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/workers/l3/all_to_all_distributed/test_all_to_all.py b/examples/workers/l3/all_to_all_distributed/test_all_to_all.py new file mode 100644 index 0000000000..f45e83ac1a --- /dev/null +++ b/examples/workers/l3/all_to_all_distributed/test_all_to_all.py @@ -0,0 +1,28 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""ST for examples/workers/l3/all_to_all_distributed.""" + +import pytest + +from .main import run + + +@pytest.mark.platforms(["a2a3sim", "a2a3", "a5sim"]) +@pytest.mark.runtime("tensormap_and_ringbuffer") +@pytest.mark.parametrize( + "n_devices", + [ + pytest.param(2, marks=pytest.mark.device_count(2)), + pytest.param(4, marks=pytest.mark.device_count(4)), + ], +) +def test_all_to_all_distributed(st_platform, st_device_ids, n_devices): + assert len(st_device_ids) == n_devices + rc = run([int(d) for d in st_device_ids], platform=st_platform) + assert rc == 0 diff --git a/examples/workers/l3/broadcast_distributed/__init__.py b/examples/workers/l3/broadcast_distributed/__init__.py new file mode 100644 index 0000000000..25708baecc --- /dev/null +++ b/examples/workers/l3/broadcast_distributed/__init__.py @@ -0,0 +1,9 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Package marker so ``test_*.py`` can do ``from .main import run``.""" diff --git a/examples/workers/l3/broadcast_distributed/kernels/aiv/broadcast_kernel.cpp b/examples/workers/l3/broadcast_distributed/kernels/aiv/broadcast_kernel.cpp new file mode 100644 index 0000000000..54a7d2ae20 --- /dev/null +++ b/examples/workers/l3/broadcast_distributed/kernels/aiv/broadcast_kernel.cpp @@ -0,0 +1,138 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Broadcast kernel — symmetric, 3-phase, HCCL-window scratch pattern. + * + * Phase 1 (stage-in): root only: input → scratch + * Phase 2 (barrier): signal matrix + TWAIT cross-rank sync + * Phase 3 (broadcast): TLOAD(root scratch) → TSTORE(output) + * + * args layout: + * tensor(0) = input COUNT_PER_RANK floats (INPUT) + * tensor(1) = output COUNT_PER_RANK floats (OUTPUT_EXISTING) + * tensor(2) = scratch HCCL window slot (INOUT) + * scalar(0) = nranks + * scalar(1) = root + * scalar(2) = CommContext device pointer + */ + +#include +#include +#include "pto/comm/comm_types.hpp" +#include "pto/comm/pto_comm_inst.hpp" +#include "platform_comm/comm_context.h" +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *localPtr, int pe) { + uint64_t localBase = ctx->windowsIn[ctx->rankId]; + uint64_t offset = (uint64_t)localPtr - localBase; + return (__gm__ T *)(ctx->windowsIn[pe] + offset); +} + +static constexpr size_t COUNT_PER_RANK = 64; +static constexpr int kMaxSupportedRanks = 16; + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *output_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + int nranks = static_cast(args[3]); + int root = static_cast(args[4]); + __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[5]); + + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; + __gm__ float *scratch = + reinterpret_cast<__gm__ float *>(scratch_tensor->buffer.addr) + scratch_tensor->start_offset; + // Signal area: nranks int32 slots at the tail of the scratch buffer. + // Peer r writes into my_rank's signal[r] when its stage-in is done. + __gm__ int32_t *signal_base = reinterpret_cast<__gm__ int32_t *>(scratch + COUNT_PER_RANK); + + using ShapeDyn = pto::Shape; + using StrideDyn = pto::Stride; + using Global = pto::GlobalTensor; + using TileData = pto::Tile; + + int my_rank = static_cast(commCtx->rankId); + + if (nranks <= 0 || nranks > kMaxSupportedRanks || root < 0 || root >= nranks) { + pipe_barrier(PIPE_ALL); + return; + } + + ShapeDyn shape(1, 1, 1, 1, COUNT_PER_RANK); + StrideDyn stride(COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, 1); + + TileData stageTile(1, COUNT_PER_RANK); + TileData recvTile(1, COUNT_PER_RANK); + TASSIGN(stageTile, 0x0); + TASSIGN(recvTile, 0x10000); + + Global inputG(input, shape, stride); + Global scratchG(scratch, shape, stride); + Global outputG(output, shape, stride); + + // ------------------------------------------------------------------ + // Phase 1: stage-in — root copies its input into the HCCL window so + // every peer can TLOAD it in Phase 3. + // ------------------------------------------------------------------ + if (my_rank == root) { + TLOAD(stageTile, inputG); + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + TSTORE(scratchG, stageTile); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + } + pipe_barrier(PIPE_ALL); + + // ------------------------------------------------------------------ + // Phase 2: device barrier — notify every peer that stage-in is done, + // then wait until every peer has notified us. + // ------------------------------------------------------------------ + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + __gm__ int32_t *remote_signal = CommRemotePtr(commCtx, signal_base + my_rank, peer); + pto::comm::Signal sig(remote_signal); + pto::comm::TNOTIFY(sig, (int32_t)1, pto::comm::NotifyOp::AtomicAdd); + } + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal sig(signal_base + peer); + pto::comm::TWAIT(sig, (int32_t)1, pto::comm::WaitCmp::GE); + } + pipe_barrier(PIPE_ALL); + + // ------------------------------------------------------------------ + // Phase 3: broadcast — read the root scratch slot and write it into + // the local output. CommRemotePtr with pe==root returns localPtr when + // my_rank==root, so the root follows the same code path as receivers. + // ------------------------------------------------------------------ + __gm__ float *root_scratch = CommRemotePtr(commCtx, scratch, root); + Global rootG(root_scratch, shape, stride); + TLOAD(recvTile, rootG); + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + TSTORE(outputG, recvTile); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + + pipe_barrier(PIPE_ALL); +} diff --git a/examples/workers/l3/broadcast_distributed/kernels/orchestration/broadcast_orch.cpp b/examples/workers/l3/broadcast_distributed/kernels/orchestration/broadcast_orch.cpp new file mode 100644 index 0000000000..337980eb33 --- /dev/null +++ b/examples/workers/l3/broadcast_distributed/kernels/orchestration/broadcast_orch.cpp @@ -0,0 +1,51 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Broadcast orchestration shim. + * + * tensor(0) input INPUT (COUNT_PER_RANK floats) + * tensor(1) output OUTPUT_EXISTING (COUNT_PER_RANK floats) + * tensor(2) scratch INOUT (HCCL window slot) + * scalar(0) nranks + * scalar(1) root + * scalar(2) CommContext device pointer + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +broadcast_orchestration_config(const ChipStorageTaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 6, // 3 tensors + 3 scalars + }; +} + +__attribute__((visibility("default"))) void broadcast_orchestration(const ChipStorageTaskArgs &orch_args) { + Tensor input = from_tensor_arg(orch_args.tensor(0)); + Tensor output = from_tensor_arg(orch_args.tensor(1)); + Tensor scratch = from_tensor_arg(orch_args.tensor(2)); + + Arg params; + params.add_input(input); + params.add_output(output); + params.add_inout(scratch); + params.add_scalar(orch_args.scalar(0)); // nranks + params.add_scalar(orch_args.scalar(1)); // root + params.add_scalar(orch_args.scalar(2)); // CommContext + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/examples/workers/l3/broadcast_distributed/main.py b/examples/workers/l3/broadcast_distributed/main.py new file mode 100644 index 0000000000..d10229e2d2 --- /dev/null +++ b/examples/workers/l3/broadcast_distributed/main.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""End-to-end distributed broadcast — symmetric 3-phase pattern. + +Root rank stages its input into the HCCL window; after barrier every rank +reads the root's scratch slot into its local output: + + Phase 1 stage-in root: input → scratch + Phase 2 device barrier signal matrix cross-rank sync via TNOTIFY/TWAIT + Phase 3 broadcast TLOAD(root scratch) → TSTORE(output) + +Run: + python examples/workers/l3/broadcast_distributed/main.py -p a2a3sim -d 0-1 + +""" + +from __future__ import annotations + +import argparse +import os +import sys + +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +import torch # noqa: E402 +from simpler.task_interface import ( # noqa: E402 + ArgDirection, + CallConfig, + ChipCallable, + CommBufferSpec, + ContinuousTensor, + CoreCallable, + DataType, + TaskArgs, + TensorArgType, +) +from simpler.worker import Worker # noqa: E402 + +from simpler_setup.elf_parser import extract_text_section # noqa: E402 +from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 +from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 +from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 + +HERE = os.path.dirname(os.path.abspath(__file__)) + +# Must match COUNT_PER_RANK in kernels/aiv/broadcast_kernel.cpp. +COUNT_PER_RANK = 64 +DTYPE_NBYTES = 4 # float32 +BUFFER_NBYTES = COUNT_PER_RANK * DTYPE_NBYTES # 256 B per rank's scratch slot +# Signal tail: one int32 slot per rank, bounded by kMaxSupportedRanks. +SIGNAL_TAIL_NBYTES = 16 * 4 # 64 B +SCRATCH_NBYTES = BUFFER_NBYTES + SIGNAL_TAIL_NBYTES # 320 B +ROOT_RANK = 0 + + +def parse_device_range(spec: str) -> list[int]: + if "-" in spec: + lo, hi = (int(x) for x in spec.split("-")) + ids = list(range(lo, hi + 1)) + else: + ids = [int(spec)] + if not (2 <= len(ids) <= 16): + raise ValueError(f"broadcast_distributed needs between 2 and 16 devices, got {len(ids)} ({ids})") + return ids + + +def build_chip_callable(platform: str, pto_isa_commit: str | None) -> ChipCallable: + """Compile the AIV broadcast kernel + its C++ orchestration shim.""" + kc = KernelCompiler(platform=platform) + runtime = "tensormap_and_ringbuffer" + pto_isa_root = ensure_pto_isa_root(commit=pto_isa_commit, clone_protocol="https") + include_dirs = kc.get_orchestration_include_dirs(runtime) + + # src/common — for platform_comm/comm_context.h + kernel_include_dirs = list(include_dirs) + [ + str(kc.project_root / "src" / "common"), + ] + kernel_bytes = kc.compile_incore( + source_path=os.path.join(HERE, "kernels/aiv/broadcast_kernel.cpp"), + core_type="aiv", + pto_isa_root=pto_isa_root, + extra_include_dirs=kernel_include_dirs, + ) + if not platform.endswith("sim"): + kernel_bytes = extract_text_section(kernel_bytes) + + orch_bytes = kc.compile_orchestration( + runtime_name=runtime, + source_path=os.path.join(HERE, "kernels/orchestration/broadcast_orch.cpp"), + ) + core_callable = CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT, ArgDirection.INOUT], + binary=kernel_bytes, + ) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT, ArgDirection.INOUT], + func_name="broadcast_orchestration", + config_name="broadcast_orchestration_config", + binary=orch_bytes, + children=[(0, core_callable)], + ) + + +def expected_output(root: int) -> list[float]: + """Every rank receives the root payload: output[i] = root*100 + i.""" + return [float(root * 100 + i) for i in range(COUNT_PER_RANK)] + + +def run( + device_ids: list[int], + platform: str = "a2a3", + pto_isa_commit: str | None = None, + build: bool = False, + root: int = ROOT_RANK, +) -> int: + """Core logic — callable from both CLI and pytest.""" + nranks = len(device_ids) + if root < 0 or root >= nranks: + raise ValueError(f"root must be in [0, {nranks}), got {root}") + window_size = max(SCRATCH_NBYTES, 4 * 1024) + + print(f"[broadcast] platform={platform} devices={device_ids} nranks={nranks} root={root}") + + host_inputs = [ + torch.tensor( + [i + rank * 100 for i in range(COUNT_PER_RANK)] if rank == root else [0.0] * COUNT_PER_RANK, + dtype=torch.float32, + ).share_memory_() + for rank in range(nranks) + ] + host_outputs = [torch.zeros(COUNT_PER_RANK, dtype=torch.float32).share_memory_() for _ in range(nranks)] + + print("[broadcast] compiling kernels...") + chip_callable = build_chip_callable(platform, pto_isa_commit) + + worker = Worker( + level=3, + platform=platform, + runtime="tensormap_and_ringbuffer", + device_ids=device_ids, + num_sub_workers=0, + build=build, + ) + chip_cid = worker.register(chip_callable) + + try: + print("[broadcast] init worker (forks chip children; base comm is lazy)...") + worker.init() + + def orch_fn(orch, _args, cfg): + with orch.allocate_domain( + name="default", + workers=list(range(nranks)), + window_size=window_size, + buffers=[CommBufferSpec(name="scratch", dtype="float32", count=COUNT_PER_RANK, nbytes=SCRATCH_NBYTES)], + ) as handle: + for i in range(nranks): + domain = handle[i] + print( + f"[broadcast] chip {i}: rank={domain.domain_rank}/{domain.domain_size} " + f"window=[0x{domain.local_window_base:x} +{domain.actual_window_size}B] " + f"scratch=0x{domain.buffer_ptrs['scratch']:x}" + ) + chip_args = TaskArgs() + chip_args.add_tensor(make_tensor_arg(host_inputs[i]), TensorArgType.INPUT) + chip_args.add_tensor(make_tensor_arg(host_outputs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_tensor( + ContinuousTensor.make( + data=domain.buffer_ptrs["scratch"], + shapes=(COUNT_PER_RANK,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INOUT, + ) + chip_args.add_scalar(domain.domain_size) + chip_args.add_scalar(root) + chip_args.add_scalar(domain.device_ctx) + orch.submit_next_level(chip_cid, chip_args, cfg, worker=i) + + print(f"[broadcast] running {nranks}-chip broadcast DAG...") + worker.run(orch_fn, args=None, config=CallConfig()) + + expected = torch.tensor(expected_output(root), dtype=torch.float32) + ok = True + for i in range(nranks): + max_diff = float(torch.max(torch.abs(host_outputs[i] - expected))) + print(f"[broadcast] chip {i}: max |out - expected| = {max_diff:.3e}") + if max_diff > 1e-3: + ok = False + for j in range(min(4, COUNT_PER_RANK)): + print(f" output[{j}]={float(host_outputs[i][j])!r} expected={float(expected[j])!r}") + + if not ok: + print("[broadcast] golden check FAILED") + return 1 + print("[broadcast] all ranks matched golden ✅") + return 0 + finally: + worker.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("-p", "--platform", default="a2a3", help="Platform backend, e.g. a2a3 or a2a3sim.") + parser.add_argument( + "-d", "--device", default="0-1", help="Device range, e.g. '0-1' or '0-3'. 2 to 16 chips required." + ) + parser.add_argument( + "--build", action="store_true", help="Rebuild runtime from source instead of using cached libs." + ) + parser.add_argument("--pto-isa-commit", default=None, help="Optional PTO ISA commit/tag to fetch before compiling.") + cli = parser.parse_args() + + return run( + parse_device_range(cli.device), platform=cli.platform, pto_isa_commit=cli.pto_isa_commit, build=cli.build + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/workers/l3/broadcast_distributed/test_broadcast.py b/examples/workers/l3/broadcast_distributed/test_broadcast.py new file mode 100644 index 0000000000..12c1d07aa2 --- /dev/null +++ b/examples/workers/l3/broadcast_distributed/test_broadcast.py @@ -0,0 +1,28 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""ST for examples/workers/l3/broadcast_distributed.""" + +import pytest + +from .main import run + + +@pytest.mark.platforms(["a2a3sim", "a2a3", "a5sim"]) +@pytest.mark.runtime("tensormap_and_ringbuffer") +@pytest.mark.parametrize( + "n_devices", + [ + pytest.param(2, marks=pytest.mark.device_count(2)), + pytest.param(4, marks=pytest.mark.device_count(4)), + ], +) +def test_broadcast_distributed(st_platform, st_device_ids, n_devices): + assert len(st_device_ids) == n_devices + rc = run([int(d) for d in st_device_ids], platform=st_platform) + assert rc == 0