diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_consumer.cpp b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_consumer.cpp new file mode 100644 index 0000000000..c0d698bcfa --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_consumer.cpp @@ -0,0 +1,64 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +#include + +#include "tensor.h" + +using namespace pto; + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *src_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + + __gm__ float *src = reinterpret_cast<__gm__ float *>(src_tensor->buffer.addr) + src_tensor->start_offset; + __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset; + + constexpr int kTotalRows = 128; + constexpr int kRows = 64; + constexpr int kCols = 128; + constexpr int kIters = kTotalRows / kRows; + using DynShapeDim5 = Shape<1, 1, 1, kRows, kCols>; + using DynStrideDim5 = pto::Stride<1, 1, 1, kCols, 1>; + using GlobalData = GlobalTensor; + using TileData = Tile; + + TileData src_tile(kRows, kCols); + TileData result_tile(kRows, kCols); + TASSIGN(src_tile, 0x0); + TASSIGN(result_tile, 0x10000); + + constexpr int kChunkElems = kRows * kCols; + for (int iter = 0; iter < kIters; ++iter) { + GlobalData src_global(src + iter * kChunkElems); + GlobalData result_global(result + iter * kChunkElems); + TLOAD(src_tile, src_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TADDS(result_tile, src_tile, 1.0f); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + TSTORE(result_global, result_tile); + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); + } +} diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp new file mode 100644 index 0000000000..c5001ab01a --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp @@ -0,0 +1,81 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +#ifdef MEMORY_BASE +#undef MEMORY_BASE +#endif +#ifndef REGISTER_BASE +#define REGISTER_BASE +#endif + +#include + +#include "backend/urma/urma_completion_kernel.h" +#include "platform_comm/comm_context.h" +#include "tensor.h" + +using namespace pto; + +namespace { + +constexpr int kElems = 128 * 128; + +template +static inline __aicore__ __gm__ T *tensor_data(__gm__ Tensor *tensor) { + return reinterpret_cast<__gm__ T *>(tensor->buffer.addr) + tensor->start_offset; +} + +} // namespace + +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 *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ CommContext *comm_ctx = reinterpret_cast<__gm__ CommContext *>(args[2]); + + // workSpace == 0 means the URMA overlay is not built in + // (SIMPLER_ENABLE_PTO_URMA_WORKSPACE=OFF, see docs/a5-sdma-overlay.md + // #1315): self-skip rather than dereferencing a null workspace. + if (comm_ctx == nullptr || comm_ctx->rankNum != 2 || comm_ctx->rankId >= comm_ctx->rankNum || + comm_ctx->workSpace == 0 || comm_ctx->windowsIn[comm_ctx->rankId] == 0) { + pipe_barrier(PIPE_ALL); + return; + } + + __gm__ float *local_input = tensor_data(input_tensor); + __gm__ float *local_out = tensor_data(out_tensor); + uint32_t peer_rank = 1u - comm_ctx->rankId; + uint64_t input_offset = reinterpret_cast(local_input) - comm_ctx->windowsIn[comm_ctx->rankId]; + __gm__ float *remote_input = pto2::urma_backend::peer_mr_ptr( + reinterpret_cast<__gm__ uint8_t *>(comm_ctx->workSpace), peer_rank, input_offset + ); + + using FlatShape = Shape<1, 1, 1, 1, kElems>; + using FlatStride = pto::Stride; + using GlobalData = GlobalTensor; + + GlobalData remote_global(remote_input); + GlobalData local_global(local_out); + + AsyncCtx async_ctx = get_async_ctx(args); + (void)send_request_entry( + async_ctx, + UrmaTget(local_global, remote_global, reinterpret_cast<__gm__ uint8_t *>(comm_ctx->workSpace), peer_rank) + ); +} diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/orchestration/urma_deferred_completion_orch.cpp b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/orchestration/urma_deferred_completion_orch.cpp new file mode 100644 index 0000000000..ea92905c60 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/orchestration/urma_deferred_completion_orch.cpp @@ -0,0 +1,52 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "platform_comm/comm_context.h" +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +urma_deferred_completion_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{.expected_arg_count = 4}; +} + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + return urma_deferred_completion_orchestration_config(orch_args); +} + +__attribute__((visibility("default"))) void urma_deferred_completion_orchestration(const L2TaskArgs &orch_args) { + if (orch_args.tensor_count() != 3 || orch_args.scalar_count() != 1) { + LOG_ERROR("urma_deferred_completion_demo: expected 3 tensors and 1 scalar"); + return; + } + + const Tensor &input = orch_args.tensor(0).ref(); + const Tensor &out = orch_args.tensor(1).ref(); + const Tensor &result = orch_args.tensor(2).ref(); + auto *comm_ctx = reinterpret_cast(static_cast(orch_args.scalar(0))); + + L0TaskArgs producer_args; + producer_args.add_input(input); + producer_args.add_output(out); + producer_args.add_scalar(reinterpret_cast(comm_ctx)); + rt_submit_aiv_task(0, producer_args); + + L0TaskArgs consumer_args; + consumer_args.add_input(out); + consumer_args.add_output(result); + rt_submit_aiv_task(1, consumer_args); +} + +} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py new file mode 100644 index 0000000000..ab3e6df166 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py @@ -0,0 +1,233 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""URMA deferred completion smoke test for onboard a5. + +Each rank stages its input inside the HCCL/URMA communication window. The +producer TGET_ASYNCs the peer rank's input into local ``out`` and registers the +URMA AsyncEvent through the deferred completion path. The consumer depends on +that producer output and writes ``result = out + 1``. Correct ``out`` and +``result`` validate both URMA completion polling and deferred-release dependency +handling. +""" + +from __future__ import annotations + +import argparse +import os + +import pytest +import torch +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CommBufferSpec, + CoreCallable, + DataType, + TaskArgs, + Tensor, + TensorArgType, +) +from simpler.worker import Worker + +from simpler_setup.elf_parser import extract_text_section +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.pto_isa import ensure_pto_isa_root +from simpler_setup.torch_interop import make_tensor_arg + +HERE = os.path.dirname(os.path.abspath(__file__)) +N = 128 * 128 +DTYPE_NBYTES = 4 +URMA_DATA_OFFSET_NBYTES = 64 * 4 +_URMA_WORKSPACE_ENV = "SIMPLER_ENABLE_PTO_URMA_WORKSPACE" +_WORKSPACE_TRUTHY = {"1", "ON", "TRUE", "YES"} + + +def _urma_workspace_enabled() -> bool: + return os.environ.get(_URMA_WORKSPACE_ENV, "").upper() in _WORKSPACE_TRUTHY + + +def _require_urma_workspace_enabled() -> None: + if _urma_workspace_enabled(): + return + raise RuntimeError( + "urma_deferred_completion_demo requires host runtime built with " + f"{_URMA_WORKSPACE_ENV}=ON; set it before rebuilding simpler." + ) + + +def parse_device_range(spec: str) -> list[int]: + if "," in spec: + return [int(x) for x in spec.split(",") if x] + if "-" in spec: + lo, hi = (int(x) for x in spec.split("-")) + return list(range(lo, hi + 1)) + return [int(spec)] + + +def build_chip_callable(platform: str) -> ChipCallable: + kc = KernelCompiler(platform=platform) + runtime = "tensormap_and_ringbuffer" + pto_isa_root = ensure_pto_isa_root() + include_dirs = kc.get_orchestration_include_dirs(runtime) + extra_includes = list(include_dirs) + [str(kc.project_root / "src" / "common")] + + children = [] + for func_id, rel, signature in [ + ( + 0, + "kernels/aiv/kernel_urma_tget_async.cpp", + [ArgDirection.IN, ArgDirection.OUT, ArgDirection.IN], + ), + ( + 1, + "kernels/aiv/kernel_consumer.cpp", + [ArgDirection.IN, ArgDirection.OUT], + ), + ]: + kernel = kc.compile_incore( + source_path=os.path.join(HERE, rel), + core_type="aiv", + pto_isa_root=pto_isa_root, + extra_include_dirs=extra_includes, + ) + if not platform.endswith("sim"): + kernel = extract_text_section(kernel) + children.append((func_id, CoreCallable.build(signature=signature, binary=kernel))) + + orch = kc.compile_orchestration( + runtime_name=runtime, + source_path=os.path.join(HERE, "kernels/orchestration/urma_deferred_completion_orch.cpp"), + extra_include_dirs=[str(kc.project_root / "src" / "common")], + ) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT, ArgDirection.OUT, ArgDirection.IN], + func_name="urma_deferred_completion_orchestration", + config_name="urma_deferred_completion_orchestration_config", + binary=orch, + children=children, + ) + + +def run(platform: str = "a5", device_ids: list[int] | None = None) -> int: + _require_urma_workspace_enabled() + if device_ids is None: + device_ids = [0, 1] + nranks = len(device_ids) + if nranks != 2: + raise ValueError(f"urma_deferred_completion_demo needs exactly 2 devices, got {device_ids}") + if platform != "a5": + raise ValueError("urma_deferred_completion_demo requires onboard a5 hardware") + + input_nbytes = N * DTYPE_NBYTES + window_size = max(URMA_DATA_OFFSET_NBYTES + input_nbytes, 4 * 1024 * 1024) + + # `inputs` must live in shared memory: `orch.copy_to` stages each rank's + # data into its HCCL window from the forked chip child, which reads `src` + # out of its own address space. + inputs = [ + torch.tensor([float(rank * 1000 + (i % 251)) / 10.0 for i in range(N)], dtype=torch.float32).share_memory_() + for rank in range(nranks) + ] + out = [torch.zeros(N, dtype=torch.float32).share_memory_() for _ in range(nranks)] + result = [torch.zeros(N, dtype=torch.float32).share_memory_() for _ in range(nranks)] + + chip_callable = build_chip_callable(platform) + worker = Worker( + level=3, + platform=platform, + runtime="tensormap_and_ringbuffer", + device_ids=device_ids, + num_sub_workers=0, + ) + chip_cid = worker.register(chip_callable) + try: + worker.init() + + def orch_fn(orch, _args, cfg): + with orch.allocate_domain( + name="urma_deferred_completion", + workers=list(range(nranks)), + window_size=window_size, + buffers=[ + CommBufferSpec( + name="urma_reserved", + dtype="int32", + count=URMA_DATA_OFFSET_NBYTES // 4, + nbytes=URMA_DATA_OFFSET_NBYTES, + ), + CommBufferSpec(name="input_window", dtype="float32", count=N, nbytes=input_nbytes), + ], + ) as handle: + # Stage every rank's input window before submitting any kernel: + # each producer TGET_ASYNCs the *peer* rank's window, so all + # windows must hold real data before execution begins. + for rank in range(nranks): + orch.copy_to( + rank, + dst=handle[rank].buffer_ptrs["input_window"], + src=inputs[rank].data_ptr(), + size=input_nbytes, + ) + for rank in range(nranks): + domain = handle[rank] + args = TaskArgs() + args.add_tensor( + Tensor.make( + data=domain.buffer_ptrs["input_window"], + shapes=(N,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_scalar(domain.device_ctx) + orch.submit_next_level(chip_cid, args, cfg, worker=rank) + + worker.run(orch_fn, args=None, config=CallConfig()) + + ok = True + for rank in range(nranks): + peer = 1 - rank + expected_out = inputs[peer] + expected_result = expected_out + 1.0 + max_out = float(torch.max(torch.abs(out[rank] - expected_out))) + max_result = float(torch.max(torch.abs(result[rank] - expected_result))) + print(f"[urma_deferred_completion_demo] rank {rank}: max_out={max_out:.3e} max_result={max_result:.3e}") + ok = ok and max_out <= 1e-3 and max_result <= 1e-3 + return 0 if ok else 1 + finally: + worker.close() + + +@pytest.mark.platforms(["a5"]) +@pytest.mark.runtime("tensormap_and_ringbuffer") +@pytest.mark.device_count(2) +@pytest.mark.skipif( + not _urma_workspace_enabled(), + reason="URMA workspace overlay not enabled (set SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON to run). " + "See docs/a5-sdma-overlay.md (#1315).", +) +def test_urma_deferred_completion_demo(st_device_ids, st_platform) -> None: + assert run(st_platform, [int(d) for d in st_device_ids]) == 0 + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("-p", "--platform", default="a5") + parser.add_argument("-d", "--device", default="0-1") + args = parser.parse_args() + return run(args.platform, parse_device_range(args.device)) + + +if __name__ == "__main__": + raise SystemExit(main())