From d5774be16a16ddee2114a5aafaf9e2fd5888b9ba Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 28 Jul 2026 20:09:08 -0700 Subject: [PATCH] test(a5/tmr): port a2a3 TMR scenes and examples to a5 Port the remaining a2a3-only tensormap_and_ringbuffer coverage to a5 for differential validation across architectures. Add the scene and L3 tests for alternating matmul/add, batch and multi-round paged attention, 4D unroll, dynamic registration, fanin lookup, empty-ring heap rebase, SPMD batch dispatch, SPMD paged attention, and L3 dependency, group, and host-buffer behavior. Add the scalar-data, paged-attention ringbuffer, and merge-pipeline-barrier examples, plus the a5 InCore orchestration guide. Adapt a5 APIs and worker topology where they differ from a2a3, and guard alternating orchestration against invalid or non-divisible task geometry. Keep spmd_paged_attention_highperf out of scope until its CCE implementation has a dedicated dav-c310 port. --- .../docs/INCORE_ORCHESTRATION_GUIDE.md | 83 ++ .../kernels/aiv/merge_pipeline.cpp | 213 +++++ .../kernels/orchestration/merge_orch.cpp | 50 ++ .../test_merge_pipeline_barrier.py | 99 +++ .../test_paged_attention_ringbuffer.py | 114 +++ .../kernels/aiv/kernel_add.cpp | 90 ++ .../kernels/aiv/kernel_noop.cpp | 33 + .../orchestration/scalar_data_orch.cpp | 265 ++++++ .../scalar_data_test/test_scalar_data.py | 86 ++ .../orchestration/alternating_orch.cpp | 13 + .../kernels/aic/kernel_matmul.cpp | 133 +++ .../kernels/aiv/kernel_add.cpp | 93 ++ .../orchestration/alternating_orch.cpp | 139 +++ .../test_alternating_matmul_add.py | 132 +++ .../kernels/aic/aic_pv_matmul.cpp | 137 +++ .../kernels/aic/aic_qk_matmul.cpp | 144 ++++ .../kernels/aiv/aiv_online_update.cpp | 221 +++++ .../kernels/aiv/aiv_softmax_prepare.cpp | 191 ++++ .../orchestration/paged_attention_orch.cpp | 215 +++++ .../test_batch_paged_attention.py | 213 +++++ .../dynamic_register/test_dynamic_register.py | 448 ++++++++++ .../aic/kernel_write_const_visible.cpp | 55 ++ .../orchestration/fanin_lookup_perf_orch.cpp | 93 ++ .../test_fanin_lookup_perf.py | 86 ++ .../kernels/aic/kernel_copy_first.cpp | 48 ++ .../kernels/aic/kernel_write_const.cpp | 45 + .../heap_empty_ring_rebase_orch.cpp | 66 ++ .../test_heap_empty_ring_rebase.py | 85 ++ .../test_multi_round_paged_attention.py | 155 ++++ .../kernels/aic/aic_pv_matmul.cpp | 171 ++++ .../kernels/aic/aic_qk_matmul.cpp | 135 +++ .../kernels/aiv/aiv_online_update.cpp | 251 ++++++ .../kernels/aiv/aiv_softmax_prepare.cpp | 265 ++++++ .../orchestration/paged_attention_orch.cpp | 186 ++++ .../test_paged_attention_unroll_4dims.py | 144 ++++ .../kernels/aic/kernel_write.cpp | 61 ++ .../kernels/aiv/kernel_write.cpp | 63 ++ .../spmd_batch_dispatch_oob_orch.cpp | 64 ++ .../test_spmd_batch_dispatch_oob.py | 88 ++ .../kernels/mix/paged_attention_parallel.cpp | 816 ++++++++++++++++++ .../spmd_paged_attention_orch.cpp | 159 ++++ .../test_spmd_paged_attention.py | 130 +++ .../test_l3_dependency.py | 106 +++ .../tensormap_and_ringbuffer/test_l3_group.py | 120 +++ .../test_l3_host_buffer_registration.py | 138 +++ 45 files changed, 6642 insertions(+) create mode 100644 examples/a5/tensormap_and_ringbuffer/docs/INCORE_ORCHESTRATION_GUIDE.md create mode 100644 examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/aiv/merge_pipeline.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/orchestration/merge_orch.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/test_merge_pipeline_barrier.py create mode 100644 examples/a5/tensormap_and_ringbuffer/paged_attention_ringbuffer/test_paged_attention_ringbuffer.py create mode 100644 examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_add.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_noop.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/orchestration/scalar_data_orch.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/scalar_data_test/test_scalar_data.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_copy_first.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_write_const.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/orchestration/heap_empty_ring_rebase_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/test_heap_empty_ring_rebase.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py diff --git a/examples/a5/tensormap_and_ringbuffer/docs/INCORE_ORCHESTRATION_GUIDE.md b/examples/a5/tensormap_and_ringbuffer/docs/INCORE_ORCHESTRATION_GUIDE.md new file mode 100644 index 0000000000..6acc69e87d --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/docs/INCORE_ORCHESTRATION_GUIDE.md @@ -0,0 +1,83 @@ +# InCore Orchestration Guide: tensormap_and_ringbuffer + +## Goal + +In tensormap_and_ringbuffer, the orchestration function runs on AICPU and builds the graph directly on device. Dependencies are discovered automatically by TensorMap based on tensor overlap, and task memory is allocated from ring buffers. + +## Where To Put Orchestration Code + +- Each example keeps orchestration sources under `examples/a5/tensormap_and_ringbuffer//kernels/orchestration/`. +- `examples/a5/tensormap_and_ringbuffer//kernels/kernel_config.py` selects the orchestration source and the runtime `tensormap_and_ringbuffer`. + +## Required Exports + +Your orchestration shared object must export: + +```cpp +extern "C" PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args); +extern "C" void aicpu_orchestration_entry(const L2TaskArgs &orch_args); +``` + +Both symbols are loaded by AICPU via `dlopen` in `src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp`. + +## Argument Layout + +Arguments are constructed by `SceneTestCase` (via `TaskArgsBuilder`, see `simpler_setup/scene_test.py`) and passed through host init into `Runtime::orch_args` as device pointers or scalars. For the default `TENSOR_ORDER` flow, the layout is: + +```text +[ptr_0, ptr_1, ..., ptr_n, nbytes_0, nbytes_1, ..., nbytes_n, element_count] +``` + +Validate `arg_count` in `aicpu_orchestration_config` and interpret pointers as device addresses. + +## Building The Graph + +1. Wrap orchestration in scopes with `PTO2_SCOPE()` to control tensor lifetimes. +2. Use `make_tensor_external` for existing device buffers and `TensorCreateInfo` + `add_output(...)` for runtime-created intermediates. +3. Use `add_inout(...)` for existing tensors that a kernel writes. +4. Build `L0TaskArgs` with `add_input`, `add_output`, `add_inout` for tensors and `add_scalar` for scalars. + > **Constraint**: All tensor parameters (`add_input` / `add_output` / `add_inout`) **must** be added before any scalar parameters (`add_scalar` / `add_scalars`). Violating this order will trigger an assertion failure. This is because the runtime dispatches tensor arguments first in kernel args, followed by scalars, and the layout must match. + + ```cpp + // Correct + L0TaskArgs p; + p.add_input(a); + p.add_inout(b); + p.add_scalar(val); // scalars after all tensors + + // Wrong — triggers assertion + L0TaskArgs p; + p.add_scalar(val); // scalar added too early + p.add_input(a); // assertion: "scalar must add after all tensor added" + ``` + +5. Submit tasks with one of: + - `rt_submit_aic_task(kernel_id, args)` — AIC (CUBE) task + - `rt_submit_aiv_task(kernel_id, args)` — AIV (VECTOR) task + - `rt_submit_task(mixed_kernels, args)` — mixed task with a `MixedKernels` struct + +Dependencies are inferred by TensorMap from input/inout/output tensors, so you do not add explicit edges. + +## Submit API And Kernel IDs + +- Submit helpers are defined in `pto_orchestration_api.h`. +- `rt_submit_aic_task` and `rt_submit_aiv_task` are convenience wrappers around `rt_submit_task` with a `MixedKernels` struct. +- For mixed AIC+AIV tasks, construct a `MixedKernels` struct directly: + + ```cpp + MixedKernels mk; + mk.aic_kernel_id = FUNC_QK; + mk.aiv0_kernel_id = FUNC_SF; + rt_submit_task(mk, args); + ``` + +- Kernel `func_id` values are defined in `kernels/kernel_config.py` under `KERNELS`. + +## Completion Semantics + +Do not call `rt_orchestration_done` yourself in device mode. The executor wraps the entry call in an outer scope and signals completion after `aicpu_orchestration_entry` returns. + +## Examples + +- `examples/a5/tensormap_and_ringbuffer/vector_example/kernels/orchestration/example_orchestration.cpp` (AIV-only tasks) +- `examples/a5/tensormap_and_ringbuffer/bgemm/kernels/orchestration/bgemm_orch.cpp` (mixed AIC + AIV tasks) diff --git a/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/aiv/merge_pipeline.cpp b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/aiv/merge_pipeline.cpp new file mode 100644 index 0000000000..30fc6875c9 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/aiv/merge_pipeline.cpp @@ -0,0 +1,213 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * The whole 3-stage pipeline (x+1 -> *2 -> +1) as ONE block_num=8 SPMD task, + * with a uniform cross-core barrier between segments instead of inter-task + * deps: stage A does 8 tiles on block 0 only, stage B 2 tiles each on blocks + * 0..3, stage C 1 tile each on all 8. Every block hits every barrier (idle + * blocks still arrive), so each barrier waits for block_num. Bulk data uses + * tile DMA; the barrier uses ld_dev/st_dev (see MERGE_BARRIER_COUNTER below). + * + * args: [0]=x, [1]=sync (block_num int32 slots, zero-init), [2]=s1, [3]=s2, + * [4]=out, [5]=timing (per-block get_sys_cnt gaps, 32-int32 stride) + */ + +#include +#include + +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +#if defined(__CCE_AICORE__) && !defined(__CPU_SIM) && !defined(__COSTMODEL) +#define MERGE_USE_DEV_INTRIN 1 +#else +#define MERGE_USE_DEV_INTRIN 0 +#endif + +static __aicore__ inline void st_i32_dev(__gm__ int32_t *p, int32_t v) { +#if MERGE_USE_DEV_INTRIN + st_dev(static_cast(v), reinterpret_cast<__gm__ uint32_t *>(p), 0); +#else + // CPU_SIM fallback is thread-based; volatile stops the spin loop hoisting. + *reinterpret_cast(p) = v; +#endif +} + +static __aicore__ inline int32_t ld_i32_dev(__gm__ int32_t *p) { +#if MERGE_USE_DEV_INTRIN + return static_cast(ld_dev(reinterpret_cast<__gm__ uint32_t *>(p), 0)); +#else + return *reinterpret_cast(p); +#endif +} + +static __aicore__ inline void ddr_fence() { +#if MERGE_USE_DEV_INTRIN + dsb(DSB_DDR); +#endif +} + +// Per-slot cross-core barrier over `n` blocks; slot advances monotonically per +// epoch. Single-writer per slot (no atomic), ld_dev poll (no dcci). Reader polls +// all n slots -> O(n) non-cacheable reads per barrier. +static __aicore__ inline void barrier_slots(__gm__ int32_t *sync, int32_t n, int32_t i, int32_t epoch) { + ddr_fence(); + st_i32_dev(&sync[i], epoch); + for (int32_t j = 0; j < n; ++j) { + while (ld_i32_dev(&sync[j]) < epoch) {} + } + ddr_fence(); +} + +// Scalar hardware atomic-add to a GM int32 (no UB/DMA). The dcci writes the line +// out atomically; the reader polls with ld_dev, so no dcci on the hot path. +static __aicore__ inline void atomic_add_i32(__gm__ int32_t *p, int32_t v) { +#if MERGE_USE_DEV_INTRIN + set_st_atomic_cfg(ATOMIC_S32, ATOMIC_SUM); + dcci(reinterpret_cast<__gm__ void *>(p), SINGLE_CACHE_LINE); + st_atomic(v, p); + dcci(reinterpret_cast<__gm__ void *>(p), SINGLE_CACHE_LINE); + dsb(DSB_DDR); + set_st_atomic_cfg(ATOMIC_NONE, ATOMIC_SUM); +#else + volatile int32_t *vp = reinterpret_cast(p); + *vp = *vp + v; +#endif +} + +// Single shared counter barrier: atomic-add arrival, poll ONE counter with +// ld_dev until it reaches epoch*n -> O(1) read per barrier. +static __aicore__ inline void barrier_counter(__gm__ int32_t *counter, int32_t n, int32_t epoch) { + ddr_fence(); + atomic_add_i32(counter, 1); + while (ld_i32_dev(counter) < epoch * n) {} + ddr_fence(); +} + +// Select the barrier at compile time: 0 = per-slot st_dev, 1 = atomic counter. +#ifndef MERGE_BARRIER_COUNTER +#define MERGE_BARRIER_COUNTER 1 +#endif +static __aicore__ inline void barrier(__gm__ int32_t *sync, int32_t n, int32_t i, int32_t epoch) { +#if MERGE_BARRIER_COUNTER + (void)i; + barrier_counter(sync, n, epoch); +#else + barrier_slots(sync, n, i, epoch); +#endif +} + +constexpr int32_t kR = 1; +constexpr int32_t kC = 128; +using GData = GlobalTensor, Stride<1, 1, 1, kC, 1>>; +using VTile = Tile; + +static __aicore__ inline void wait_store_done() { + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); +} + +// tile compute: dst[t] = op0 ? (src[t] + b) : (src[t] + src[t]) +static __aicore__ inline void +stage_tile(__gm__ float *src, __gm__ float *dst, int32_t t, int32_t op, float b, VTile &t0, VTile &t1) { + GData ig(src + t * kC); + GData og(dst + t * kC); + TLOAD(t0, ig); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + if (op == 0) { + TADDS(t1, t0, b); + } else { + TADD(t1, t0, t0); + } + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(og, t1); + // Drain this store before the caller reuses t0/t1 for the next tile. + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID6); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID6); +} + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *x_t = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *sync_t = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *s1_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *s2_t = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *out_t = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *tm_t = reinterpret_cast<__gm__ Tensor *>(args[5]); + + __gm__ float *x = reinterpret_cast<__gm__ float *>(x_t->buffer.addr) + x_t->start_offset; + __gm__ int32_t *sync = reinterpret_cast<__gm__ int32_t *>(sync_t->buffer.addr) + sync_t->start_offset; + __gm__ float *s1 = reinterpret_cast<__gm__ float *>(s1_t->buffer.addr) + s1_t->start_offset; + __gm__ float *s2 = reinterpret_cast<__gm__ float *>(s2_t->buffer.addr) + s2_t->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_t->buffer.addr) + out_t->start_offset; + __gm__ int32_t *tm = reinterpret_cast<__gm__ int32_t *>(tm_t->buffer.addr) + tm_t->start_offset; + + int32_t N = get_block_num(args); // 8 + int32_t i = get_block_idx(args); + + VTile t0(kR, kC), t1(kR, kC); + TASSIGN(t0, 0x0); + TASSIGN(t1, 0x10000); + + // get_sys_cnt() = 50 MHz counter (20 ns/tick). Deltas below are in ticks. + int64_t c0 = get_sys_cnt(); + + // Segment A (stage 1): block 0 does all 8 tiles: s1 = x + 1 + if (i == 0) { + for (int32_t t = 0; t < 8; ++t) { + stage_tile(x, s1, t, 0, 1.0f, t0, t1); + } + wait_store_done(); + } + int64_t c1 = get_sys_cnt(); // arrived at barrier 1 + barrier(sync, N, i, 1); + int64_t c2 = get_sys_cnt(); // released from barrier 1 + + // Segment B (stage 2): blocks 0..3 do 2 tiles each: s2 = s1 * 2 + if (i < 4) { + for (int32_t t = i * 2; t < i * 2 + 2; ++t) { + stage_tile(s1, s2, t, 1, 0.0f, t0, t1); + } + wait_store_done(); + } + int64_t c3 = get_sys_cnt(); // arrived at barrier 2 + barrier(sync, N, i, 2); + int64_t c4 = get_sys_cnt(); // released from barrier 2 + + // Segment C (stage 3): all 8 blocks do 1 tile each: out = s2 + 1 + stage_tile(s2, out, i, 0, 1.0f, t0, t1); + int64_t c5 = get_sys_cnt(); // end + + // Per-block timing (ticks): [segA, barrier1_gap, segB, barrier2_gap, segC, total]. + // Stride each block by 32 int32 (one 128B cacheline): scalar st_dev only + // commits the first 8 int32 of a cacheline, so keep all 6 writes in that head. + st_i32_dev(&tm[i * 32 + 0], static_cast(c1 - c0)); + st_i32_dev(&tm[i * 32 + 1], static_cast(c2 - c1)); + st_i32_dev(&tm[i * 32 + 2], static_cast(c3 - c2)); + st_i32_dev(&tm[i * 32 + 3], static_cast(c4 - c3)); + st_i32_dev(&tm[i * 32 + 4], static_cast(c5 - c4)); + st_i32_dev(&tm[i * 32 + 5], static_cast(c5 - c0)); + ddr_fence(); // flush timing writes to HBM before the task completes / copy-back + + pipe_sync(); +} diff --git a/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/orchestration/merge_orch.cpp b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/orchestration/merge_orch.cpp new file mode 100644 index 0000000000..c1a210163f --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/kernels/orchestration/merge_orch.cpp @@ -0,0 +1,50 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Merge mode: one block_num=8 SPMD task (merge_pipeline, func 0) runs all three + * stages with an intra-task barrier between segments. s1/s2 are runtime scratch; + * sync is a caller-provided zero-initialized barrier buffer. + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{.expected_arg_count = 4}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_x = orch_args.tensor(0).ref(); + const Tensor &ext_sync = orch_args.tensor(1).ref(); + const Tensor &ext_out = orch_args.tensor(2).ref(); + const Tensor &ext_timing = orch_args.tensor(3).ref(); + + uint32_t SIZE = ext_x.shapes[0]; + uint32_t sh[1] = {SIZE}; + TensorCreateInfo ci(sh, 1, DataType::FLOAT32); + + L0TaskArgs p; + p.add_input(ext_x); // args[0] + p.add_input(ext_sync); // args[1] + p.add_output(ci); // args[2] s1 scratch + p.add_output(ci); // args[3] s2 scratch + p.add_output(ext_out); // args[4] + p.add_output(ext_timing); // args[5] per-block timing (ticks) + p.launch_spec.set_block_num(8); + rt_submit_aiv_task(0, p); +} + +} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/test_merge_pipeline_barrier.py b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/test_merge_pipeline_barrier.py new file mode 100644 index 0000000000..92d9cba0bc --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/test_merge_pipeline_barrier.py @@ -0,0 +1,99 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Merged 3-stage pipeline as ONE SPMD task with an intra-task cross-core barrier. + +Three stages (x+1 -> *2 -> +1 = 2x+3) run as a single block_num=8 AIV task; the +inter-stage ordering is an intra-task cross-core barrier instead of separate +scheduled tasks. Stage A does 8 tiles on block 0 only, stage B 2 tiles each on +blocks 0..3, stage C 1 tile each on all 8; every block hits every barrier (idle +blocks still arrive), so each barrier waits for block_num. + +The barrier has two compile-time variants (MERGE_BARRIER_COUNTER in +merge_pipeline.cpp): a per-slot st_dev poll (O(n) non-cacheable reads) and a +single-counter st_atomic+dcci arrival + one-ld_dev poll (O(1) read). The kernel +records each block's per-barrier gap with get_sys_cnt; the `compare_outputs` +override prints them (ticks at 50 MHz -> us) and leaves the diagnostic 'timing' +output out of the golden comparison. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +TILES = 8 +M = 128 +SIZE = TILES * M # 1024 +NBLK = 8 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestMergePipelineBarrier(SceneTestCase): + """out = 2x + 3 across 3 stages, ordered by an intra-task cross-core barrier.""" + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/merge_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/merge_pipeline.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT, D.OUT, D.OUT, D.OUT], + }, + ], + } + + CASES = [ + { + # Onboard only: the busy-wait barrier needs all block_num logical + # blocks co-resident, so config block_dim (24) > block_num (8, set + # per-task via launch_spec.set_block_num in merge_orch.cpp). + "name": "merge", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("x", torch.arange(SIZE, dtype=torch.float32) / 7.0), + Tensor("sync", torch.zeros(NBLK, dtype=torch.int32)), # barrier slots / counter (zero-init) + Tensor("out", torch.zeros(SIZE, dtype=torch.float32)), + Tensor("timing", torch.zeros(NBLK * 32, dtype=torch.int32)), # per-block gaps (32-int32 stride) + ) + + def compute_golden(self, args, params): + args.out[:] = 2.0 * args.x + 3.0 + + def compare_outputs(self, test_args, golden_args, output_names, params): + # 'timing' is a diagnostic output with no golden: print the per-block + # gaps, then compare everything else. Each block's slots are strided by + # 32 int32 (one cacheline) because scalar st_dev only commits the first + # 8 int32 of a 128B line. + tm = test_args.timing.reshape(NBLK, 32)[:, :6].to(torch.float64) / 50.0 + labels = ["segA", "bar1", "segB", "bar2", "segC", "total"] + print("[GAP] per-block us (50 MHz counter):") + print(" blk " + " ".join(f"{lbl:>7}" for lbl in labels)) + for b in range(NBLK): + print(f" {b:>3} " + " ".join(f"{tm[b, k]:7.2f}" for k in range(6))) + print( + f"[GAP] barrier1: max={tm[:, 1].max():.2f} min={tm[:, 1].min():.2f} us | " + f"barrier2: max={tm[:, 3].max():.2f} min={tm[:, 3].min():.2f} us" + ) + super().compare_outputs(test_args, golden_args, [n for n in output_names if n != "timing"], params) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/paged_attention_ringbuffer/test_paged_attention_ringbuffer.py b/examples/a5/tensormap_and_ringbuffer/paged_attention_ringbuffer/test_paged_attention_ringbuffer.py new file mode 100644 index 0000000000..5ee9e02406 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/paged_attention_ringbuffer/test_paged_attention_ringbuffer.py @@ -0,0 +1,114 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention with small ring buffer sizes — stress test for ring rotation/reclamation. + +Drives per-case ring sizing through ``config.runtime_env`` (ring_task_window / +ring_heap / ring_dep_pool) rather than the process-global PTO2_RING_* env, plus +INOUT tensors, bfloat16, and AIC+AIV mixed execution. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden # noqa: PLC0415 +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs # noqa: PLC0415 + +PA_KERNELS = "../../../../tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels" + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPagedAttentionRingbuffer(SceneTestCase): + """Paged attention with small ring buffer sizes for stress testing.""" + + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": f"{PA_KERNELS}/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{PA_KERNELS}/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{PA_KERNELS}/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{PA_KERNELS}/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 3, + "source": f"{PA_KERNELS}/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "ringbuffer_stress", + "platforms": ["a5"], + # ring_heap is bytes per ring. Non power-of-2 sizes are accepted, + # but 4 MiB keeps the small-ring stress intent compact. + "config": { + "aicpu_thread_num": 4, + "runtime_env": { + "ring_task_window": 64, + "ring_heap": 4 * 1024 * 1024, + "ring_dep_pool": 256, + }, + }, + "params": { + "batch": 32, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 4096, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + inputs = _pa_generate_inputs(params) + specs = [] + for name, val in inputs: + if isinstance(val, torch.Tensor): + specs.append(Tensor(name, val)) + else: + specs.append(Scalar(name, val)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_add.cpp b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_add.cpp new file mode 100644 index 0000000000..8a119554db --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_add.cpp @@ -0,0 +1,90 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Element-wise Tensor Addition Kernel + * + * Implements: out[i] = src0[i] + src1[i] + * + * This kernel performs element-wise addition of two tensors. It's compiled + * separately as a standalone kernel and linked with the dispatcher using + * function pointers, demonstrating the separation pattern used in production + * systems where kernel binaries are loaded dynamically. + */ + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +/** + * Element-wise addition kernel implementation + * + * Unified signature: all arguments passed via int64_t array + * @param args Argument array: + * args[0] = src0 pointer (first input tensor) + * args[1] = src1 pointer (second input tensor) + * args[2] = out pointer (output tensor) + * args[3] = size (number of elements) + */ +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack arguments (Tensor* pointers from runtime) + __gm__ Tensor *src0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *src1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *src0 = reinterpret_cast<__gm__ float *>(src0_tensor->buffer.addr) + src0_tensor->start_offset; + __gm__ float *src1 = reinterpret_cast<__gm__ float *>(src1_tensor->buffer.addr) + src1_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + // Configuration: float, 128, 128, 128, 128 + constexpr int kTRows_ = 128; + constexpr int kTCols_ = 128; + constexpr int vRows = 128; + constexpr int vCols = 128; + + using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>; + using DynStridDim5 = Stride<1, 1, 1, kTCols_, 1>; + using GlobalData = GlobalTensor; + using TileData = Tile; + + TileData src0Tile(vRows, vCols); + TileData src1Tile(vRows, vCols); + TileData dstTile(vRows, vCols); + TASSIGN(src0Tile, 0x0); + TASSIGN(src1Tile, 0x10000); + TASSIGN(dstTile, 0x20000); + + GlobalData src0Global(src0); + GlobalData src1Global(src1); + GlobalData dstGlobal(out); + + TLOAD(src0Tile, src0Global); + TLOAD(src1Tile, src1Global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(dstTile, src0Tile, src1Tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dstGlobal, dstTile); + + pipe_sync(); +} diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_noop.cpp b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_noop.cpp new file mode 100644 index 0000000000..8187197c4c --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/aiv/kernel_noop.cpp @@ -0,0 +1,33 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * No-op Kernel + * + * Empty kernel used to trigger runtime allocation for tensors passed + * as OUTPUT/INOUT via add_inout(). The runtime allocates HeapRing memory + * and writes initial values before dispatching this task; the kernel + * itself does not read or modify any data. + */ + +#include +#include + +using namespace pto; + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; } diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/orchestration/scalar_data_orch.cpp b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/orchestration/scalar_data_orch.cpp new file mode 100644 index 0000000000..a0a8ed7d88 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/kernels/orchestration/scalar_data_orch.cpp @@ -0,0 +1,265 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Scalar Data Dependency Test Orchestration + * + * End-to-end test for get_tensor_data, set_tensor_data, and add_inout + * with runtime-created outputs and initial value support. + * + * Flow: + * 1. c = a + b (kernel_add, runtime-created tensor) + * 2. get_tensor_data(c, {0}) → check[0] = 2.0 + * 3. get_tensor_data(c, {100}) → check[1] = 102.0 + * 4. scalar_tensor = add_output(TensorCreateInfo, 77.0f), submit noop + * 5. get_tensor_data(scalar_tensor, {0}) → check[2] = 77.0 + * 6. add_inout(scalar_tensor) (INOUT path), submit noop + * 7. get_tensor_data(scalar_tensor, {0}) → check[3] = 77.0 + * 8. check[4] = 2.0 + 77.0 = 79.0 (orchestration arithmetic) + * 9. set_tensor_data(scalar_tensor, {0}, 42.0), get_tensor_data → check[5] = 42.0 + * 10. Orch set_tensor_data(d, {0}, 10.0) → kernel_add(d, a) → check[6] = 12.0 + * 11. WAW+WAR: kernel_add reads c → set_tensor_data(c, 88.0) auto-waits → check[7] = 88.0 + * 12. External WAR with INOUT: noop(ext_b as INOUT) → set_tensor_data(ext_b) → check[8] = 55.0 + * 13. result = a + b (kernel_add, external output via INOUT) + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_ADD 0 +#define FUNC_NOOP 1 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 4, // a, b, result, check + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // External tensors from golden.py + const Tensor &ext_a = orch_args.tensor(0).ref(); + const Tensor &ext_b = orch_args.tensor(1).ref(); + const Tensor &ext_result = orch_args.tensor(2).ref(); + const Tensor &ext_check = orch_args.tensor(3).ref(); + + uint32_t SIZE = orch_args.tensor(0).ref().shapes[0]; + LOG_INFO_V0("scalar_data_test: SIZE=%u, check_size=%u", SIZE, orch_args.tensor(3).ref().shapes[0]); + + uint32_t inter_shapes[1] = {SIZE}; + TensorCreateInfo inter_ci(inter_shapes, 1, DataType::FLOAT32); + + // ========================================================= + // Step 1: c = a + b (runtime-created tensor, kernel_add) + // ========================================================= + L0TaskArgs params_c; + params_c.add_input(ext_a); + params_c.add_input(ext_b); + params_c.add_output(inter_ci); + TaskOutputTensors c_outs = rt_submit_aiv_task(FUNC_ADD, params_c); + const Tensor &c = c_outs.get_ref(0); + + // ========================================================= + // Step 2: get_tensor_data(c, {0}) → check[0] + // Tests TensorMap lookup + spin-wait for kernel completion + // ========================================================= + uint32_t idx[1] = {0}; + float c0_val = get_tensor_data(c, 1, idx); + LOG_INFO_V0("get_tensor_data(c, {0}) = %f (expected 2.0)", static_cast(c0_val)); + + uint32_t check_idx[1] = {0}; + set_tensor_data(ext_check, 1, check_idx, c0_val); + + // ========================================================= + // Step 3: get_tensor_data(c, {100}) → check[1] + // Tests flat offset calculation for non-zero index + // ========================================================= + idx[0] = 100; + float c100_val = get_tensor_data(c, 1, idx); + LOG_INFO_V0("get_tensor_data(c, {100}) = %f (expected 102.0)", static_cast(c100_val)); + + check_idx[0] = 1; + set_tensor_data(ext_check, 1, check_idx, c100_val); + + // ========================================================= + // Step 4: Runtime-created scalar output with initial value + // Runtime allocates HeapRing buffer, writes 77.0 to element [0] + // ========================================================= + uint32_t scalar_shapes[1] = {1}; + TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32); + scalar_ci.set_initial_value(77.0f); + TaskOutputTensors scalar_alloc_outs = alloc_tensors(scalar_ci); + const Tensor &scalar_tensor = scalar_alloc_outs.get_ref(0); + + // ========================================================= + // Step 5: get_tensor_data(scalar_tensor, {0}) → check[2] + // Verifies initial value was written correctly + // ========================================================= + idx[0] = 0; + float s0_val = get_tensor_data(scalar_tensor, 1, idx); + LOG_INFO_V0("get_tensor_data(scalar_tensor, {0}) after init = %f (expected 77.0)", static_cast(s0_val)); + + check_idx[0] = 2; + set_tensor_data(ext_check, 1, check_idx, s0_val); + + // ========================================================= + // Step 6: add_inout(scalar_tensor) second use → INOUT path + // Buffer already exists, so the noop just registers dependency + // ========================================================= + { + L0TaskArgs args; + args.add_inout(scalar_tensor); + rt_submit_aiv_task(FUNC_NOOP, args); + } + + // ========================================================= + // Step 7: get_tensor_data(scalar_tensor, {0}) → check[3] + // Value should be preserved (noop kernel didn't modify it) + // ========================================================= + float s1_val = get_tensor_data(scalar_tensor, 1, idx); + LOG_INFO_V0("get_tensor_data(scalar_tensor, {0}) after 2nd noop = %f (expected 77.0)", static_cast(s1_val)); + + check_idx[0] = 3; + set_tensor_data(ext_check, 1, check_idx, s1_val); + + // ========================================================= + // Step 8: set_tensor_data with orchestration-computed value → check[4] + // Tests set_tensor_data write + orchestration arithmetic + // ========================================================= + float combined = c0_val + s0_val; // 2.0 + 77.0 = 79.0 + LOG_INFO_V0( + "Orchestration arithmetic: %f + %f = %f", static_cast(c0_val), static_cast(s0_val), + static_cast(combined) + ); // NOLINT(whitespace/line_length) + + check_idx[0] = 4; + set_tensor_data(ext_check, 1, check_idx, combined); + + // ========================================================= + // Step 9: Orch set→get round-trip on internal tensor + // Validates that set_tensor_data writes are visible to get_tensor_data + // on the same tensor. Uses scalar_tensor (currently 77.0), overwrites to 42.0. + // ========================================================= + set_tensor_data(scalar_tensor, 1, idx, 42.0f); + float rw_val = get_tensor_data(scalar_tensor, 1, idx); + LOG_INFO_V0("set_tensor_data→get_tensor_data round-trip = %f (expected 42.0)", static_cast(rw_val)); + + check_idx[0] = 5; + set_tensor_data(ext_check, 1, check_idx, rw_val); + + // ========================================================= + // Step 10: Orch→AICore RAW (set_tensor_data → kernel reads) + // Orchestration writes d[0]=10.0 via set_tensor_data, then + // kernel_add reads d as input: e[0] = d[0] + a[0] = 12.0 + // ========================================================= + TaskOutputTensors d_alloc_outs = alloc_tensors(inter_ci); + const Tensor &d = d_alloc_outs.get_ref(0); + + idx[0] = 0; + set_tensor_data(d, 1, idx, 10.0f); + + L0TaskArgs params_e; + params_e.add_input(d); + params_e.add_input(ext_a); + params_e.add_output(inter_ci); + TaskOutputTensors e_outs = rt_submit_aiv_task(FUNC_ADD, params_e); + const Tensor &e = e_outs.get_ref(0); + + float e0_val = get_tensor_data(e, 1, idx); + LOG_INFO_V0("Orch→AICore RAW: e[0] = %f (expected 12.0)", static_cast(e0_val)); + + check_idx[0] = 6; + set_tensor_data(ext_check, 1, check_idx, e0_val); + + // ========================================================= + // Step 11: WAW + WAR on internal tensor + // c was written by Step 1 (kernel_add, TensorMap has producer entry). + // Submit a new kernel that reads c as INPUT (creates consumer dep). + // Then set_tensor_data(c) — no manual get_tensor_data sync. + // set_tensor_data internally waits for: + // - WAW: producer (Step 1) COMPLETED + // - WAR: consumer (this kernel) done (fanout_refcount check) + // + // NOTE on external tensors: ext_a was read by Step 1 as INPUT, + // but TensorMap has no producer entry for ext_a (only consumers). + // set_tensor_data(ext_a) would NOT detect the reader — data race. + // To ensure WAR safety on external tensors, use add_inout() + // instead of add_input() so TensorMap tracks the access chain. + // ========================================================= + { + L0TaskArgs args; + args.add_input(c); + args.add_input(ext_b); + args.add_output(inter_ci); + (void)rt_submit_aiv_task(FUNC_ADD, args); // NOLINT(readability/casting) + } + + // set_tensor_data auto-waits for producer + consumer before writing + idx[0] = 0; + set_tensor_data(c, 1, idx, 88.0f); + float waw_val = get_tensor_data(c, 1, idx); + LOG_INFO_V0("WAW+WAR: set_tensor_data(c, 88.0) after consumer = %f (expected 88.0)", static_cast(waw_val)); + + check_idx[0] = 7; + set_tensor_data(ext_check, 1, check_idx, waw_val); + + // ========================================================= + // Step 12: External tensor WAR — must use add_output or add_inout, not add_input + // + // For external tensors, using add_input() does NOT create a + // TensorMap entry. set_tensor_data would then write immediately + // without waiting for the reader kernel — a WAR data race. + // + // Using add_output() (or add_inout()) creates a TensorMap entry, + // enabling set_tensor_data to detect the producer via TensorMap lookup + // and wait for fanout_refcount (all consumers done). + // + // Here we submit noop with ext_b as write-only output (noop doesn't + // read data), then set_tensor_data overwrites ext_b[0] = 55.0. + // set_tensor_data auto-waits for the noop to complete. + // ========================================================= + { + L0TaskArgs args; + args.add_output(ext_b); // write-only: creates TensorMap entry (not add_input!) + rt_submit_aiv_task(FUNC_NOOP, args); + } + + idx[0] = 0; + set_tensor_data(ext_b, 1, idx, 55.0f); + float ext_war_val = get_tensor_data(ext_b, 1, idx); + LOG_INFO_V0( + "External WAR (INOUT): set_tensor_data(ext_b, 55.0) = %f (expected 55.0)", static_cast(ext_war_val) + ); + + check_idx[0] = 8; + set_tensor_data(ext_check, 1, check_idx, ext_war_val); + + // Restore ext_b[0] for final result comparison + set_tensor_data(ext_b, 1, idx, 0.0f); + + // ========================================================= + // Step 13: result = a + b (external output via add_output, kernel_add) + // ========================================================= + { + L0TaskArgs args; + args.add_input(ext_a); + args.add_input(ext_b); + args.add_output(ext_result); + rt_submit_aiv_task(FUNC_ADD, args); + } + + LOG_INFO_V0("scalar_data_test: orchestration complete"); +} + +} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data_test/test_scalar_data.py b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/test_scalar_data.py new file mode 100644 index 0000000000..1e5722a2c7 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/test_scalar_data.py @@ -0,0 +1,86 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Scalar data dependency test: GetTensorData, SetTensorData, add_inout. + +Tests orchestration-level data manipulation: scalar initialization, +Get/Set round-trips, WAW+WAR dependency auto-wait, and external tensor WAR. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestScalarData(SceneTestCase): + """Scalar data dependency: Get/SetTensorData, add_inout with initial value.""" + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/scalar_data_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/kernel_noop.cpp", + "core_type": "aiv", + "signature": [], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + SIZE = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((SIZE,), 2.0, dtype=torch.float32)), + Tensor("b", torch.arange(SIZE, dtype=torch.float32)), + Tensor("result", torch.zeros(SIZE, dtype=torch.float32)), + # Exactly 9 slots: the orchestration writes check[0..8] via + # set_tensor_data. Output-tensor slots are not seeded from the host, + # so any extra slot reads undefined device memory. + Tensor("check", torch.zeros(9, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + # result = a + b (computed by kernel_add) + args.result[:] = args.a + args.b + + # check values written by orchestration via SetTensorData + args.check[0] = 2.0 # GetTensorData(c, {0}): c = a + b, c[0] = 2.0+0.0 + args.check[1] = 102.0 # GetTensorData(c, {100}): c[100] = 2.0+100.0 + args.check[2] = 77.0 # runtime-created scalar output initialized to 77.0 + args.check[3] = 77.0 # second noop via add_inout preserves the value + args.check[4] = 79.0 # orchestration arithmetic: 2.0 + 77.0 + args.check[5] = 42.0 # Orch set->get round-trip: SetTensorData then GetTensorData + args.check[6] = 12.0 # Orch->AICore RAW: SetTensorData(d,10.0) + kernel_add(d,a) -> 10.0+2.0 + args.check[7] = 88.0 # WAW+WAR: kernel reads c, SetTensorData(c,88.0) auto-waits + args.check[8] = 55.0 # External WAR: noop(ext_b INOUT) -> SetTensorData(ext_b,55.0) auto-waits + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp index d08f7645bc..49ccd773cd 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp +++ b/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp @@ -69,6 +69,19 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta int total_matmul_tasks = batch * M; int total_add_tasks = batch * N; + if (matmul_batch <= 0 || add_batch <= 0) { + LOG_ERROR( + "[alternating_orch] batch sizes must be positive (matmul_batch=%d, add_batch=%d)", matmul_batch, add_batch + ); + return; + } + if (total_matmul_tasks % matmul_batch != 0 || total_add_tasks % add_batch != 0) { + LOG_ERROR( + "[alternating_orch] task counts must divide evenly (matmul %d/%d, add %d/%d)", total_matmul_tasks, + matmul_batch, total_add_tasks, add_batch + ); + return; + } int num_matmul_groups = total_matmul_tasks / matmul_batch; int num_add_groups = total_add_tasks / add_batch; diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp new file mode 100644 index 0000000000..607e5d657f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp @@ -0,0 +1,133 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Matrix Multiplication Kernel (Cube Core) + * + * Computes: C = A @ B (TILE x TILE x TILE matmul) + * Uses TMATMUL instruction + * + * Args (Tensor*): + * args[0] = A (INPUT) - TILE x TILE + * args[1] = B (INPUT) - TILE x TILE + * args[2] = C (OUTPUT) - TILE x TILE + */ + +#include +#include +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +AICORE constexpr inline T CeilAlign(T num_1, T num_2) { + if (num_2 == 0) { + return 0; + } + return (num_1 + num_2 - 1) / num_2 * num_2; +} + +static __aicore__ inline int get_num_tiles(__gm__ Tensor *tensor, uint64_t tile_elems) { + uint64_t total_elems = tensor->shapes[0]; + return static_cast(total_elems / tile_elems); +} + +template +static __aicore__ void matmul_impl(__gm__ float *input_a, __gm__ float *input_b, __gm__ float *output) { + constexpr int blockAlign = C0_SIZE_BYTE / sizeof(float); + constexpr int M = CeilAlign(TILE, 16); + constexpr int K = CeilAlign(TILE, blockAlign); + constexpr int N = CeilAlign(TILE, blockAlign); + + using GlobalDataA = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + using GlobalDataB = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + using GlobalDataC = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + + GlobalDataA src0Global(input_a); + GlobalDataB src1Global(input_b); + GlobalDataC dstGlobal(output); + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + TLOAD(aMatTile, src0Global); + TLOAD(bMatTile, src1Global); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(dstGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_a = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *input_b = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *output = reinterpret_cast<__gm__ Tensor *>(args[2]); + + constexpr uint64_t TILE_ELEMS = 128 * 128; + int num_tiles = get_num_tiles(input_a, TILE_ELEMS); + + __gm__ float *base_a = reinterpret_cast<__gm__ float *>(input_a->buffer.addr) + input_a->start_offset; + __gm__ float *base_b = reinterpret_cast<__gm__ float *>(input_b->buffer.addr) + input_b->start_offset; + __gm__ float *base_c = reinterpret_cast<__gm__ float *>(output->buffer.addr) + output->start_offset; + + for (int tile_idx = 0; tile_idx < num_tiles; tile_idx++) { + __gm__ float *a_ptr = base_a + (tile_idx * TILE_ELEMS); + __gm__ float *b_ptr = base_b + (tile_idx * TILE_ELEMS); + __gm__ float *c_ptr = base_c + (tile_idx * TILE_ELEMS); + + matmul_impl<128>(a_ptr, b_ptr, c_ptr); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp new file mode 100644 index 0000000000..ba0cdfa3b4 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp @@ -0,0 +1,93 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Element-wise Tensor Addition Kernel + * + * Implements: out[i] = src0[i] + src1[i] + * Tile size: ROWS x COLS + * + * Args (Tensor*): + * args[0] = src0 (INPUT) - ROWS x COLS + * args[1] = src1 (INPUT) - ROWS x COLS + * args[2] = out (OUTPUT) - ROWS x COLS + */ + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static __aicore__ inline int get_num_tiles(__gm__ Tensor *tensor, uint64_t tile_elems) { + uint64_t total_elems = tensor->shapes[0]; + return static_cast(total_elems / tile_elems); +} + +template +static __aicore__ void add_impl(__gm__ float *src0, __gm__ float *src1, __gm__ float *out) { + using DynShapeDim5 = Shape<1, 1, 1, ROWS, COLS>; + using DynStridDim5 = pto::Stride<1, 1, 1, COLS, 1>; + using GlobalData = GlobalTensor; + using TileData = Tile; + + TileData src0Tile(ROWS, COLS); + TileData src1Tile(ROWS, COLS); + TileData dstTile(ROWS, COLS); + TASSIGN(src0Tile, 0x0); + TASSIGN(src1Tile, 0x10000); + TASSIGN(dstTile, 0x20000); + + GlobalData src0Global(src0); + GlobalData src1Global(src1); + GlobalData dstGlobal(out); + + TLOAD(src0Tile, src0Global); + TLOAD(src1Tile, src1Global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(dstTile, src0Tile, src1Tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dstGlobal, dstTile); + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *src0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *src1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + constexpr uint64_t TILE_ELEMS = 128 * 128; + int num_tiles = get_num_tiles(src0_tensor, TILE_ELEMS); + + __gm__ float *base_src0 = reinterpret_cast<__gm__ float *>(src0_tensor->buffer.addr) + src0_tensor->start_offset; + __gm__ float *base_src1 = reinterpret_cast<__gm__ float *>(src1_tensor->buffer.addr) + src1_tensor->start_offset; + __gm__ float *base_out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + for (int tile_idx = 0; tile_idx < num_tiles; tile_idx++) { + __gm__ float *src0_ptr = base_src0 + (tile_idx * TILE_ELEMS); + __gm__ float *src1_ptr = base_src1 + (tile_idx * TILE_ELEMS); + __gm__ float *out_ptr = base_out + (tile_idx * TILE_ELEMS); + + add_impl<128, 128>(src0_ptr, src1_ptr, out_ptr); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp new file mode 100644 index 0000000000..49ccd773cd --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Alternating Matmul-Add Orchestration Function (tensormap_and_ringbuffer Runtime) + * + * Submits independent matmul and add tasks per batch. + * + * Configuration read from scalar args: + * - batch: Number of batches + * - M: Number of matmul tasks per batch + * - N: Number of add tasks per batch + * - matmul_batch: Number of matmul tiles per task group + * - add_batch: Number of add tiles per task group + * + * Task pattern: interleaved [matmul_0, add_0, matmul_1, add_1, ...] + * All tasks are completely independent (no dependencies). + * + * Arg layout: [A, B, C, X, Y, Z, batch, M_val, N_val, matmul_batch, add_batch] + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_MATMUL 0 +#define FUNC_ADD 1 + +static constexpr uint64_t MATMUL_ELEMS = 128 * 128; +static constexpr uint64_t ADD_ELEMS = 128 * 128; + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 11, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Tensor args + const Tensor &ext_A = orch_args.tensor(0).ref(); + const Tensor &ext_B = orch_args.tensor(1).ref(); + const Tensor &ext_C = orch_args.tensor(2).ref(); + const Tensor &ext_X = orch_args.tensor(3).ref(); + const Tensor &ext_Y = orch_args.tensor(4).ref(); + const Tensor &ext_Z = orch_args.tensor(5).ref(); + + // Scalar config args + int batch = static_cast(orch_args.scalar(0)); + int M = static_cast(orch_args.scalar(1)); + int N = static_cast(orch_args.scalar(2)); + int matmul_batch = static_cast(orch_args.scalar(3)); + int add_batch = static_cast(orch_args.scalar(4)); + + LOG_INFO_V0( + "[alternating_orch] Batch: %d, M: %d, N: %d, matmul_batch: %d, add_batch: %d", batch, M, N, matmul_batch, + add_batch + ); + + int total_matmul_tasks = batch * M; + int total_add_tasks = batch * N; + if (matmul_batch <= 0 || add_batch <= 0) { + LOG_ERROR( + "[alternating_orch] batch sizes must be positive (matmul_batch=%d, add_batch=%d)", matmul_batch, add_batch + ); + return; + } + if (total_matmul_tasks % matmul_batch != 0 || total_add_tasks % add_batch != 0) { + LOG_ERROR( + "[alternating_orch] task counts must divide evenly (matmul %d/%d, add %d/%d)", total_matmul_tasks, + matmul_batch, total_add_tasks, add_batch + ); + return; + } + int num_matmul_groups = total_matmul_tasks / matmul_batch; + int num_add_groups = total_add_tasks / add_batch; + + int total_matmul = 0; + int total_add = 0; + + int max_groups = num_matmul_groups > num_add_groups ? num_matmul_groups : num_add_groups; + + // Interleaved submit: matmul and add groups alternate + for (int group_idx = 0; group_idx < max_groups; group_idx++) { + if (group_idx < num_matmul_groups) { + int start_task_idx = group_idx * matmul_batch; + uint64_t offset = static_cast(start_task_idx) * MATMUL_ELEMS; + uint64_t group_size = static_cast(matmul_batch) * MATMUL_ELEMS; + + uint32_t matmul_group_shapes[1] = {static_cast(group_size)}; + uint32_t view_offsets[1] = {static_cast(offset)}; + + Tensor A_view = ext_A.view(matmul_group_shapes, view_offsets); + Tensor B_view = ext_B.view(matmul_group_shapes, view_offsets); + Tensor C_view = ext_C.view(matmul_group_shapes, view_offsets); + + L0TaskArgs params_matmul; + params_matmul.add_input(A_view); + params_matmul.add_input(B_view); + params_matmul.add_output(C_view); + rt_submit_aic_task(FUNC_MATMUL, params_matmul); + total_matmul++; + } + + if (group_idx < num_add_groups) { + int start_task_idx = group_idx * add_batch; + uint64_t offset = static_cast(start_task_idx) * ADD_ELEMS; + uint64_t group_size = static_cast(add_batch) * ADD_ELEMS; + + uint32_t add_group_shapes[1] = {static_cast(group_size)}; + uint32_t view_offsets[1] = {static_cast(offset)}; + + Tensor X_view = ext_X.view(add_group_shapes, view_offsets); + Tensor Y_view = ext_Y.view(add_group_shapes, view_offsets); + Tensor Z_view = ext_Z.view(add_group_shapes, view_offsets); + + L0TaskArgs params_add; + params_add.add_input(X_view); + params_add.add_input(Y_view); + params_add.add_output(Z_view); + rt_submit_aiv_task(FUNC_ADD, params_add); + total_add++; + } + } + + LOG_INFO_V9("[alternating_orch] Submitted %d matmul groups and %d add groups", total_matmul, total_add); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py new file mode 100644 index 0000000000..9fc261b4d8 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py @@ -0,0 +1,132 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Alternating matmul + add: interleaved AIC (matmul 128x128) and AIV (add 128x128) tasks. + +Tests AIC+AIV mixed execution with scalar parameters and batched task submission. +C[b,m] = A[b,m] @ B[b,m], Z[b,n] = X[b,n] + Y[b,n]. +""" + +import ctypes + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAlternatingMatmulAdd(SceneTestCase): + """Alternating matmul + add with scalar parameters.""" + + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/alternating_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/kernel_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 1, "M": 1, "N": 1, "matmul_batch": 1, "add_batch": 1}, + }, + { + "name": "Case1", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 500, "M": 4, "N": 4, "matmul_batch": 4, "add_batch": 4}, + "manual": True, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 512, "M": 2, "N": 5, "matmul_batch": 4, "add_batch": 5}, + "manual": True, + }, + ] + + def generate_args(self, params): + batch = params["batch"] + M = params["M"] + N = params["N"] + matmul_batch = params.get("matmul_batch", 1) + add_batch = params.get("add_batch", 1) + matmul_size = 128 + add_rows = 128 + add_cols = 128 + + torch.manual_seed(42) + A = torch.randn(batch, M, matmul_size, matmul_size, dtype=torch.float32) * 0.01 + B = torch.randn(batch, M, matmul_size, matmul_size, dtype=torch.float32) * 0.01 + C = torch.zeros(batch, M, matmul_size, matmul_size, dtype=torch.float32) + X = torch.randn(batch, N, add_rows, add_cols, dtype=torch.float32) * 0.01 + Y = torch.randn(batch, N, add_rows, add_cols, dtype=torch.float32) * 0.01 + Z = torch.zeros(batch, N, add_rows, add_cols, dtype=torch.float32) + + return TaskArgsBuilder( + Tensor("A", A.flatten()), + Tensor("B", B.flatten()), + Tensor("C", C.flatten()), + Tensor("X", X.flatten()), + Tensor("Y", Y.flatten()), + Tensor("Z", Z.flatten()), + Scalar("batch", ctypes.c_int64(batch)), + Scalar("M_val", ctypes.c_int64(M)), + Scalar("N_val", ctypes.c_int64(N)), + Scalar("matmul_batch", ctypes.c_int64(matmul_batch)), + Scalar("add_batch", ctypes.c_int64(add_batch)), + ) + + def compute_golden(self, args, params): + batch = params["batch"] + M = params["M"] + N = params["N"] + matmul_size = 128 + add_rows = 128 + add_cols = 128 + + A = args.A.reshape(batch, M, matmul_size, matmul_size) + B = args.B.reshape(batch, M, matmul_size, matmul_size) + C = args.C.reshape(batch, M, matmul_size, matmul_size) + X = args.X.reshape(batch, N, add_rows, add_cols) + Y = args.Y.reshape(batch, N, add_rows, add_cols) + Z = args.Z.reshape(batch, N, add_rows, add_cols) + + for b in range(batch): + for m in range(M): + C[b, m] = torch.matmul(A[b, m], B[b, m]) + for b in range(batch): + for n in range(N): + Z[b, n] = X[b, n] + Y[b, n] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..5dcd577ea6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,137 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched PV Matmul Kernel: for each batch b, pij(M, K) @ vj(K, N) -> oi_new(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// Template: M=q_tile, K=block_size, N=head_dim + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void pv_matmul_batch_impl( + __gm__ Tensor *pij_batch, __gm__ Tensor *value_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *oi_new_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t block_num, uint64_t batch_start +) { + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ bfloat16_t *val_base = reinterpret_cast<__gm__ bfloat16_t *>(value_cache->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + using GlobalOut = GlobalTensor, pto::Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *pij_addr = pij_base + b * M * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *vj_addr = val_base + static_cast(phys_block) * K * N; + __gm__ float *oi_addr = oi_base + b * M * N; + + GlobalA pijGlobal(pij_addr); + GlobalB vjGlobal(vj_addr); + GlobalOut oiGlobal(oi_addr); + + TLOAD(aMatTile, pijGlobal); + TLOAD(bMatTile, vjGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(oiGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *value_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t block_num = static_cast(args[6]); + uint64_t batch_start = static_cast(args[7]); + + uint64_t q_tile_size = static_cast(pij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + pv_matmul_batch_impl<16, 16, 16>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else if (q_tile_size == 16) { + pv_matmul_batch_impl<16, 128, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else { + pv_matmul_batch_impl<64, 64, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..64f4147cc3 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,144 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched QK Matmul Kernel: for each batch b, qi(M, K) @ kj.T(K, N) -> sij(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// Template: M=q_tile, K=head_dim, N=block_size + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void qk_matmul_batch_impl( + __gm__ Tensor *query, __gm__ Tensor *key_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *sij_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t q_offset, uint64_t block_num, uint64_t num_heads, + uint64_t batch_start +) { + __gm__ bfloat16_t *query_base = reinterpret_cast<__gm__ bfloat16_t *>(query->buffer.addr); + __gm__ bfloat16_t *key_base = reinterpret_cast<__gm__ bfloat16_t *>(key_cache->buffer.addr); + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride, Layout::DN>; + using GlobalOut = GlobalTensor, pto::Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *qi_addr = query_base + ((batch_start + b) * num_heads + q_offset) * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *kj_addr = key_base + static_cast(phys_block) * N * K; + __gm__ float *sij_addr = sij_base + b * M * N; + + GlobalA qiGlobal(qi_addr); + GlobalB kjGlobal(kj_addr); + GlobalOut sijGlobal(sij_addr); + + TLOAD(aMatTile, qiGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TLOAD(bMatTile, kjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *query = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t q_offset = static_cast(args[6]); + uint64_t block_num = static_cast(args[7]); + uint64_t num_heads = static_cast(args[8]); + uint64_t batch_start = static_cast(args[9]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(sij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + qk_matmul_batch_impl<16, 16, 16>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else if (q_tile_size == 16) { + qk_matmul_batch_impl<16, 128, 128>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else { + qk_matmul_batch_impl<64, 128, 64>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..f55220b1cc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,221 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Batched Online Softmax Update + Normalize Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b, updates accumulators mi/li/oi with new block's mij/lij/oi_new. +// On is_last, normalizes and writes to the output tensor at the correct batch offset. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, head_dim=128 +// Case2: (64, 128) -- q_tile=64, head_dim=128 +// +// Scalar layout strategy: +// M scalar floats stored contiguously in GM can be loaded as either: +// - ND (kScalarRows, kScalarCols) RowMajor for element-wise ops +// - DN (kAlignedRows, 1) ColMajor for row-broadcast ops +// Conversion between layouts uses TRESHAPE (UB-internal, zero GM access). + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_batch_impl( + __gm__ Tensor *mij_batch, __gm__ Tensor *lij_batch, __gm__ Tensor *oi_new_batch, __gm__ Tensor *mi_batch, + __gm__ Tensor *li_batch, __gm__ Tensor *oi_batch, __gm__ Tensor *out, uint64_t is_first, uint64_t is_last, + uint64_t batch_count, uint64_t q_offset, uint64_t num_heads, uint64_t batch_start +) { + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ float *oi_new_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ float *mi_base = reinterpret_cast<__gm__ float *>(mi_batch->buffer.addr); + __gm__ float *li_base = reinterpret_cast<__gm__ float *>(li_batch->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_batch->buffer.addr); + __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out->buffer.addr); + + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + using TileDataMxN = Tile; + using TileScalarND = + Tile; + using TileScalarDN = Tile; + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarNDBytes = kScalarRows * kScalarCols * sizeof(float); + + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + TileScalarND mijND, lijND, miND, liND; + TileScalarND miNewND, alphaND, betaND, tmpND; + + TileScalarDN alphaDN, betaDN, liDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijND, 2 * kDataBytes); + TASSIGN(lijND, 2 * kDataBytes + kScalarNDBytes); + TASSIGN(miND, 2 * kDataBytes + 2 * kScalarNDBytes); + TASSIGN(liND, 2 * kDataBytes + 3 * kScalarNDBytes); + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarNDBytes); + TASSIGN(alphaND, 2 * kDataBytes + 5 * kScalarNDBytes); + TASSIGN(betaND, 2 * kDataBytes + 6 * kScalarNDBytes); + TASSIGN(tmpND, 2 * kDataBytes + 7 * kScalarNDBytes); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ float *mij_ptr = mij_base + b * M; + __gm__ float *lij_ptr = lij_base + b * M; + __gm__ float *oi_new_ptr = oi_new_base + b * M * N; + __gm__ float *mi_ptr = mi_base + b * M; + __gm__ float *li_ptr = li_base + b * M; + __gm__ float *oi_ptr = oi_base + b * M * N; + __gm__ float *dst_ptr = out_base + ((batch_start + b) * num_heads + q_offset) * N; + + GlobalDataMxN oiNewGlobal(oi_new_ptr); + GlobalDataMxN oiGlobal(oi_ptr); + GlobalDataMxN dstGlobal(dst_ptr); + + GlobalScalarND mijGlobalND(mij_ptr); + GlobalScalarND lijGlobalND(lij_ptr); + GlobalScalarND miGlobalND(mi_ptr); + GlobalScalarND liGlobalND(li_ptr); + + if (is_first) { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); + TSTORE(liGlobalND, lijND); + TSTORE(oiGlobal, oiNewTile); + + if (is_last) { + TRESHAPE(liDN, lijND); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TROWEXPANDDIV(oiNewTile, oiNewTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + TLOAD(miND, miGlobalND); + TLOAD(liND, liGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TMAX(miNewND, miND, mijND); + TSUB(alphaND, miND, miNewND); + TSUB(betaND, mijND, miNewND); + TEXP(alphaND, alphaND); + TEXP(betaND, betaND); + TMUL(liND, alphaND, liND); + TMUL(tmpND, betaND, lijND); + TADD(liND, liND, tmpND); + + TRESHAPE(alphaDN, alphaND); + TRESHAPE(betaDN, betaND); + if (is_last) { + TRESHAPE(liDN, liND); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); + TSTORE(liGlobalND, liND); + + TROWEXPANDMUL(oiTile, oiTile, alphaDN); + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); + TADD(oiTile, oiTile, oiNewTile); + + if (is_last) { + TROWEXPANDDIV(oiTile, oiTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiTile); + } else { + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(oiGlobal, oiTile); + } + } + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi_batch = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *out = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t batch_count = static_cast(args[9]); + uint64_t q_offset = static_cast(args[10]); + uint64_t num_heads = static_cast(args[11]); + uint64_t batch_start = static_cast(args[12]); + + uint64_t q_tile_size = static_cast(mij_batch->shapes[0] / batch_count); + uint64_t head_dim = static_cast(oi_new_batch->shapes[1]); + + if (q_tile_size == 16 && head_dim <= 16) { + online_update_batch_impl<16, 16>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else if (q_tile_size == 16) { + online_update_batch_impl<16, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else { + online_update_batch_impl<64, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..0dd0b24256 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,191 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched Softmax Preparation Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b at block_idx bn: +// valid_len = min(N, context_lens[b] - bn * N) +// sij_masked = pad(sij[b], valid_len, -inf) +// sij_scale = sij_masked * scale +// mij[b] = row_max(sij_scale) +// pij[b] = exp(sij_scale - mij[b]) (truncated to bf16 then back) +// lij[b] = row_sum(pij[b]) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, block_size=128 +// Case2: (64, 64) -- q_tile=64, block_size=64 + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void softmax_prepare_batch_impl( + __gm__ Tensor *sij_batch, __gm__ Tensor *context_lens_t, __gm__ Tensor *pij_batch, __gm__ Tensor *mij_batch, + __gm__ Tensor *lij_batch, float scale_value, uint64_t batch_count, uint64_t block_idx, uint64_t batch_start +) { + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ int32_t *ctx_lens = reinterpret_cast<__gm__ int32_t *>(context_lens_t->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + + using TileSijDyn = Tile; + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + TASSIGN(sijTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + for (uint64_t b = 0; b < batch_count; b++) { + int32_t cur_seq = ctx_lens[batch_start + b]; + uint64_t start = block_idx * N; + uint64_t valid_len = 0; + if (start < static_cast(cur_seq)) { + uint64_t remaining = static_cast(cur_seq) - start; + valid_len = (remaining < N) ? remaining : N; + } + + __gm__ float *sij_addr = sij_base + b * M * N; + __gm__ bfloat16_t *pij_addr = pij_base + b * M * N; + __gm__ float *mij_addr = mij_base + b * M; + __gm__ float *lij_addr = lij_base + b * M; + + GlobalDataMxN sijGlobal(sij_addr); + GlobalDataMxN_bf16 pijGlobal(pij_addr); + GlobalScalarDN mijGlobal(mij_addr); + GlobalScalarDN lijGlobal(lij_addr); + + if (valid_len == 0) { + // Block entirely beyond sequence: write mij=-1e30, lij=0, pij=0 + // Use -1e30 instead of -inf to avoid NaN in online_update (exp(-inf - (-inf)) = NaN) + constexpr float NEG_LARGE = -1e30f; + for (int i = 0; i < kAlignedRows; i++) { + maxTile.SetValue(i, NEG_LARGE); + sumTile.SetValue(i, 0.0f); + } + for (int i = 0; i < M * N; i++) { + pijBf16Tile.SetValue(i, static_cast(0.0f)); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(mijGlobal, maxTile); + TSTORE(lijGlobal, sumTile); + TSTORE(pijGlobal, pijBf16Tile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + continue; + } + + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TileSijDyn sijDynTile(static_cast(valid_len)); + TASSIGN(sijDynTile, 0x0); + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + + TMULS(sijTile, sijTile, scale_value); + TROWMAX(maxTile, sijTile, tmpTile); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first, then compute lij from truncated values (matches golden) + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + TROWSUM(sumTile, pijTile, tmpTile); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + TSTORE(mijGlobal, maxTile); + TSTORE(lijGlobal, sumTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *context_lens_t = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[5]); + float scale_value = scale_conv.f; + uint64_t batch_count = static_cast(args[6]); + uint64_t block_idx = static_cast(args[7]); + uint64_t batch_start = static_cast(args[8]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + softmax_prepare_batch_impl<16, 16>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else if (q_tile_size == 16) { + softmax_prepare_batch_impl<16, 128>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else { + softmax_prepare_batch_impl<64, 64>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..1717ebc48a --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,215 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Batch Paged Attention Orchestration Function - Production Scale + * + * Chunked batched architecture: the full batch is split into chunks of + * IN_CORE_BATCH size. Each chunk's QK/SF/PV/UP tasks are independent + * and can be scheduled to different cores in parallel. + * + * Task count = num_chunks * (1 + max_bn * 4), where + * num_chunks = ceil(batch / IN_CORE_BATCH) + * + * For batch <= IN_CORE_BATCH, behavior is identical to the non-chunked version. + * + * Memory Layout: + * Query: (batch * num_heads, head_dim) bf16 + * Key: (total_blocks, block_size, head_dim) bf16 (stored as K^T for QK) + * Value: (total_blocks, block_size, head_dim) bf16 + * + * Per-chunk intermediate tensors (contiguous across chunk_bc dimension): + * sij: (chunk_bc * q_tile, block_size) fp32 + * pij: (chunk_bc * q_tile, block_size) bf16 + * mij/lij: (chunk_bc * q_tile) fp32 + * oi_new: (chunk_bc * q_tile, head_dim) fp32 + * oi: (chunk_bc * q_tile, head_dim) fp32 accumulator + * mi/li: (chunk_bc * q_tile) fp32 accumulator + * + * Kernels receive global tensors + scalar metadata (including batch_start) + * and compute per-batch addresses internally. + */ + +#include +#include + +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + uint64_t elem_size = get_element_size(data_type); + + LOG_INFO_V0("batch_paged_attention: batch=%" PRIu64 ", num_heads=%" PRIu64, batch, num_heads); + + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + uint64_t max_bn = 0; + for (uint64_t b = 0; b < batch; b++) { + uint32_t cl_idx[1] = {static_cast(b)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_b = (cur_seq + block_size - 1) / block_size; + if (bn_b > max_bn) max_bn = bn_b; + } + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + uint64_t kv_total_rows = total_blocks_count * block_size; + uint32_t key_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t value_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32, true); + + constexpr uint64_t IN_CORE_BATCH = 16; + uint64_t num_chunks = (batch + IN_CORE_BATCH - 1) / IN_CORE_BATCH; + + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + uint64_t q_offset = q_idx * q_tile; + + for (uint64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + uint64_t chunk_bc = batch - chunk_idx * IN_CORE_BATCH; + if (chunk_bc > IN_CORE_BATCH) chunk_bc = IN_CORE_BATCH; + uint64_t batch_start = chunk_idx * IN_CORE_BATCH; + + PTO2_SCOPE() { + uint32_t oi_acc_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(head_dim)}; + uint32_t scalar_acc_shapes[1] = {static_cast(chunk_bc * q_tile)}; + TensorCreateInfo oi_batch_ci(oi_acc_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_acc_ci(scalar_acc_shapes, 1, DataType::FLOAT32); + TaskOutputTensors alloc_outs = alloc_tensors(oi_batch_ci, scalar_acc_ci, scalar_acc_ci); + const Tensor &oi_batch = alloc_outs.get_ref(0); + const Tensor &li_batch = alloc_outs.get_ref(1); + const Tensor &mi_batch = alloc_outs.get_ref(2); + + // Inner-loop create infos: shapes are loop-invariant, hoist out of bn loop + uint32_t sij_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(block_size)}; + uint32_t vec_shapes[1] = {static_cast(chunk_bc * q_tile)}; + uint32_t oi_new_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(head_dim)}; + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_ci(sij_shapes, 2, data_type); + TensorCreateInfo vec_ci(vec_shapes, 1, DataType::FLOAT32); + TensorCreateInfo oi_new_ci(oi_new_shapes, 2, DataType::FLOAT32); + + for (uint64_t bn = 0; bn < max_bn; bn++) { + PTO2_SCOPE() { + L0TaskArgs params_qk; + params_qk.add_input(query); + params_qk.add_input(key_cache); + params_qk.add_input(block_table); + params_qk.add_output(sij_ci); + params_qk.add_scalar(chunk_bc); + params_qk.add_scalar(bn); + params_qk.add_scalar(q_offset); + params_qk.add_scalar(block_num); + params_qk.add_scalar(num_heads); + params_qk.add_scalar(batch_start); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij_b = qk_outs.get_ref(0); + + L0TaskArgs params_sf; + params_sf.add_input(sij_b); + params_sf.add_input(context_lens); + params_sf.add_output(pij_ci); + params_sf.add_output(vec_ci); + params_sf.add_output(vec_ci); + params_sf.add_scalar(scale_value); + params_sf.add_scalar(chunk_bc); + params_sf.add_scalar(bn); + params_sf.add_scalar(batch_start); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_b = sf_outs.get_ref(0); + const Tensor &mij_b = sf_outs.get_ref(1); + const Tensor &lij_b = sf_outs.get_ref(2); + + L0TaskArgs params_pv; + params_pv.add_input(pij_b); + params_pv.add_input(value_cache); + params_pv.add_input(block_table); + params_pv.add_output(oi_new_ci); + params_pv.add_scalar(chunk_bc); + params_pv.add_scalar(bn); + params_pv.add_scalar(block_num); + params_pv.add_scalar(batch_start); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_new_b = pv_outs.get_ref(0); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == max_bn - 1) ? 1 : 0; + L0TaskArgs params_up; + params_up.add_input(mij_b); + params_up.add_input(lij_b); + params_up.add_input(oi_new_b); + params_up.add_inout(mi_batch); + params_up.add_inout(li_batch); + params_up.add_inout(oi_batch); + params_up.add_inout(out); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + params_up.add_scalar(chunk_bc); + params_up.add_scalar(q_offset); + params_up.add_scalar(num_heads); + params_up.add_scalar(batch_start); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + } + } + } + } + } + + LOG_INFO_V0( + "batch_paged_attention: %" PRIu64 " tasks (batch=%" PRIu64 ", max_bn=%" PRIu64 ", chunks=%" PRIu64 + ", IN_CORE_BATCH=%" PRIu64 ")", + static_cast(num_chunks * (1 + max_bn * 4)), batch, max_bn, num_chunks, IN_CORE_BATCH + ); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py new file mode 100644 index 0000000000..e7392531cb --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py @@ -0,0 +1,213 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Batch paged attention: batched online softmax with AIC/AIV subgraph splitting (bfloat16).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestBatchPagedAttention(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case3", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 256, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 31, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall3", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py b/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py new file mode 100644 index 0000000000..e5ba8404e7 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -0,0 +1,448 @@ +#!/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 ST for post-start Worker.register(ChipCallable) at L3. + +Exercises the _CTRL_REGISTER IPC path end-to-end: parent stages a +ChipCallable in shared memory after child startup, broadcasts CTRL_REGISTER to +every chip child, the child mmaps + prepares, and the resulting +CallableHandle is indistinguishable from a pre-start preparation when used +in run(). + +The UT suite (tests/ut/py/test_worker/test_host_worker.py) already covers +the facade-level paths (lock guard, capacity overflow, lambda rejection, run +race detection, shm name generator). This file's job is to prove the +bytes actually traverse shm to the chip child and prepare succeeds — +which only a real (sim or device) chip child can confirm. +""" + +import os + +import pytest +import torch +from _task_interface import MAX_REGISTERED_CALLABLE_IDS # pyright: ignore[reportMissingImports] +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import CallConfig, ChipCallable +from simpler.worker import Worker + +from simpler_setup import TaskArgsBuilder, Tensor +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.scene_test import _build_l3_task_args + +_RUNTIME = "tensormap_and_ringbuffer" +_HERE = os.path.dirname(os.path.abspath(__file__)) +_KERNELS = os.path.join( + _HERE, + "..", + "..", + "..", + "..", + "..", + "examples", + "a5", + "tensormap_and_ringbuffer", + "vector_example", + "kernels", +) +_ORCH_SRC = os.path.join(_KERNELS, "orchestration", "example_orchestration.cpp") +_AIV_ADD = os.path.join(_KERNELS, "aiv", "kernel_add.cpp") +_AIV_ADD_SCALAR = os.path.join(_KERNELS, "aiv", "kernel_add_scalar.cpp") +_AIV_MUL = os.path.join(_KERNELS, "aiv", "kernel_mul.cpp") + +_ORCH_SIG = [D.IN, D.IN, D.OUT] + + +def _build_vector_callable(platform: str, *, extra_unused_child: bool = False) -> ChipCallable: + """Compile the vector_example orchestration + 3 AIV kernels. + + Mirrors how SceneTestCase._compile_chip_callable_from_spec assembles + a ChipCallable, but inline so the test can call register_callable() on it both + before and after init(). + """ + from simpler.task_interface import CoreCallable # noqa: PLC0415 + + from simpler_setup.elf_parser import extract_text_section # noqa: PLC0415 + from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: PLC0415 + + kc = KernelCompiler(platform=platform) + pto_isa_root = ensure_pto_isa_root() + inc_dirs = kc.get_orchestration_include_dirs(_RUNTIME) + + orch_bytes = kc.compile_orchestration(runtime_name=_RUNTIME, source_path=_ORCH_SRC) + + def _aiv(path: str) -> bytes: + raw = kc.compile_incore(path, core_type="aiv", pto_isa_root=pto_isa_root, extra_include_dirs=inc_dirs) + return raw if platform.endswith("sim") else extract_text_section(raw) + + add = CoreCallable.build(signature=[D.IN, D.IN, D.OUT], binary=_aiv(_AIV_ADD)) + add_scalar = CoreCallable.build(signature=[D.IN, D.OUT], binary=_aiv(_AIV_ADD_SCALAR)) + mul = CoreCallable.build(signature=[D.IN, D.IN, D.OUT], binary=_aiv(_AIV_MUL)) + + children = [(0, add), (1, add_scalar), (2, mul)] + if extra_unused_child: + children.append((99, add)) + + return ChipCallable.build( + signature=_ORCH_SIG, + func_name="aicpu_orchestration_entry", + binary=orch_bytes, + children=children, + ) + + +def _unique_py_callable(index: int): + def fn(args, _index=index): + return _index + + return fn + + +def _make_args(a: float, b: float) -> TaskArgsBuilder: + size = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((size,), a, dtype=torch.float32).share_memory_()), + Tensor("b", torch.full((size,), b, dtype=torch.float32).share_memory_()), + Tensor("f", torch.zeros(size, dtype=torch.float32).share_memory_()), + ) + + +def _golden(a: float, b: float) -> float: + # Matches the orchestration: f = (a+b+1) * (a+b+2) + (a+b) + s = a + b + return (s + 1) * (s + 2) + s + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_new_identity_after_start_then_run(st_platform, st_device_ids): + """Happy path: prepare one identity pre-start and another post-start. + + Proves the post-start control path delivers a usable handle for a + previously unseen hashid. Both identities execute equivalent kernels and + must produce numerically identical outputs. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + # Pre-allocate both runs' tensors BEFORE Worker.init() so the + # share_memory_() mappings are inherited by the forked chip child. + # share_memory_ regions created after fork in the parent are not visible + # to the chip child, so dispatch on those would segfault. + a, b = 2.0, 3.0 + expected = _golden(a, b) + args_pre = _make_args(a, b) + args_post = _make_args(a, b) + chip_args_pre, output_names_pre = _build_l3_task_args(args_pre, _ORCH_SIG) + chip_args_post, output_names_post = _build_l3_task_args(args_post, _ORCH_SIG) + assert output_names_pre == ["f"] and output_names_post == ["f"] + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + # 1. Run pre_handle once against the snapshot-prepared chip callable. + # init() already forked the chip children and prepared the startup + # snapshot, so they are in _run_chip_main_loop — the only state in + # which the CTRL_REGISTER broadcast in step 2 can be ACKed. + def orch_pre(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + got_pre = args_pre.f + assert torch.allclose(got_pre, torch.full_like(got_pre, expected), rtol=1e-5, atol=1e-5), ( + f"pre_handle={pre_handle.hashid}: expected {expected}, got {got_pre[:4].tolist()}..." + ) + + # 2. Now do the post-start dynamic prepare. The parent stages bytes + # in shm and broadcasts CTRL_REGISTER; the child mmaps and calls + # prepare_callable_from_blob. post_handle is unknown to the + # CoW-inherited registry on the child side — only the IPC path + # can deliver it. + post_handle = worker.register(post_callable) + assert post_handle.hashid != pre_handle.hashid + + # 3. Run with post_handle. If CTRL_REGISTER delivered correctly, the + # child has the identity prepared; otherwise dispatch will fail. + def orch_post(o, _args, _cfg): + o.submit_next_level(post_handle, chip_args_post, config, worker=0) + + worker.run(orch_post) + got_post = args_post.f + assert torch.allclose(got_post, torch.full_like(got_post, expected), rtol=1e-5, atol=1e-5), ( + f"post_handle={post_handle.hashid}: expected {expected}, got {got_post[:4].tolist()}..." + ) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(2) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_new_identity_after_start_parallel_broadcast(st_platform, st_device_ids): + """Two chip children, post-start prepare broadcasts to both in parallel. + + Asserts that the prepared handle runs successfully on each chip — proving + the C++ broadcast (one std::thread per WorkerThread) delivers the bytes + to every chip's mailbox and each prepare_callable_from_blob runs without + racing against the others. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + device_ids = [int(d) for d in st_device_ids[:2]] + worker = Worker( + level=3, + device_ids=device_ids, + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + a, b = 2.0, 3.0 + expected = _golden(a, b) + # Use one first-run bundle to start the control-capable child loop, then one + # post-start bundle per chip so the dynamically prepared handle is executed + # by both broadcast targets. + args_pre = _make_args(a, b) + args_post = [_make_args(a, b) for _ in device_ids] + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) + chip_args_post = [_build_l3_task_args(args, _ORCH_SIG)[0] for args in args_post] + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_pre(o, _a, _c): + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + assert torch.allclose(args_pre.f, torch.full_like(args_pre.f, expected), rtol=1e-5, atol=1e-5) + + # Now broadcast CTRL_REGISTER to BOTH chip mailboxes in parallel. + post_handle = worker.register(post_callable) + + def orch_post(o, _a, _c): + o.submit_next_level_group(post_handle, chip_args_post, config, workers=[0, 1]) + + worker.run(orch_post) + for args in args_post: + assert torch.allclose(args.f, torch.full_like(args.f, expected), rtol=1e-5, atol=1e-5) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_capacity_overflow_post_start(st_platform, st_device_ids): + """Saturate callable capacity pre-start, then verify post-start prepare hits + the same ``MAX_REGISTERED_CALLABLE_IDS`` ceiling for a new hashid. + + Confirms the public capacity guard is shared between pre-start preparation + and the post-start control path (and that the error message is + protocol-aware so the operator sees the same diagnostic in both paths). + """ + chip_callable = _build_vector_callable(st_platform) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + # A sub worker gives the LOCAL_PYTHON fillers below a valid resolver, so + # they exercise the shared capacity ceiling as eligible registrations + # (not inert ones) alongside the chip callable. + num_sub_workers=1, + platform=st_platform, + runtime=_RUNTIME, + ) + # Fill the registry pre-start with distinct sub fn identities (cheap, no + # device cost). + for i in range(MAX_REGISTERED_CALLABLE_IDS - 1): + worker.register(_unique_py_callable(i)) + chip_handle = worker.register(chip_callable) # final capacity entry + + a, b = 2.0, 3.0 + args_pre = _make_args(a, b) + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_pre(o, _a, _c): + o.submit_next_level(chip_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + + # The very next dynamic prepare of a new identity hits the capacity + # ceiling. Re-preparing ``chip_callable`` itself would only create + # another handle to the existing identity. + with pytest.raises(RuntimeError, match="MAX_REGISTERED_CALLABLE_IDS"): + worker.register(_build_vector_callable(st_platform, extra_unused_child=True)) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_duplicate_prepare_same_hashid_survives_one_unregister(st_platform, st_device_ids): + """prepare same hashid twice, unregister one handle, run the other. + + This is the hashid-specific post-start path: the second + ``register_callable(same_chip_callable)`` must return a distinct handle for + the same hashid. Unregistering one handle must only drop that public + handle; the remaining handle must still dispatch successfully. + """ + chip_callable = _build_vector_callable(st_platform) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + a, b = 2.0, 3.0 + expected = _golden(a, b) + # Two runs total — preallocate both args bundles BEFORE init() so + # the share_memory_ mappings are inherited by the forked chip child. + args_one = _make_args(a, b) + args_two = _make_args(a, b) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + # 1. Trigger fork via pre_handle to put the chip child into the main loop. + def orch_one(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) + + worker.run(orch_one) + assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) + + # 2. Prepare the same callable after start. This returns another + # public handle for the same hashid, not a new identity. + duplicate_handle = worker.register(chip_callable) + assert duplicate_handle.hashid == pre_handle.hashid + assert duplicate_handle.digest == pre_handle.digest + assert duplicate_handle._handle_id != pre_handle._handle_id + + # 3. Drop the first handle. The child must keep the prepared identity + # alive for duplicate_handle. + worker.unregister(pre_handle) + + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(pre_handle, chip_args_one, config, worker=0)) + + def orch_two(o, _args, _cfg): + o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0) + + worker.run(orch_two) + assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) + + # 4. Dropping the final handle invalidates it through the public API. + worker.unregister(duplicate_handle) + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0)) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_unregister_last_handle_allows_reprepare_same_hashid(st_platform, st_device_ids): + """prepare → run → unregister final handle → prepare same identity again. + + Proves the IPC unregister path works end-to-end: after CTRL_UNREGISTER + propagates to the chip child, the old handle is invalid and a subsequent + post-start prepare of that identity materializes a usable handle again. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + a, b = 2.0, 3.0 + expected = _golden(a, b) + args_one = _make_args(a, b) + args_two = _make_args(a, b) + args_three = _make_args(a, b) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) + chip_args_three, _ = _build_l3_task_args(args_three, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_one(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) + + worker.run(orch_one) + assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) + + dyn_handle = worker.register(post_callable) + + def orch_two(o, _args, _cfg): + o.submit_next_level(dyn_handle, chip_args_two, config, worker=0) + + worker.run(orch_two) + assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) + + worker.unregister(dyn_handle) + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(dyn_handle, chip_args_two, config, worker=0)) + + # Re-prepare the same hashid after its final handle was dropped. + again_handle = worker.register(post_callable) + assert again_handle.hashid == dyn_handle.hashid + assert again_handle.digest == dyn_handle.digest + assert again_handle._handle_id != dyn_handle._handle_id + + def orch_three(o, _args, _cfg): + o.submit_next_level(again_handle, chip_args_three, config, worker=0) + + worker.run(orch_three) + assert torch.allclose(args_three.f, torch.full_like(args_three.f, expected), rtol=1e-5, atol=1e-5) + finally: + worker.close() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp new file mode 100644 index 0000000000..543cb83618 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp @@ -0,0 +1,55 @@ +/* + * 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 + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + // Keep the swimlane bars visible. Lookup-only timing uses dummy tasks, so + // this spin does not affect the fanin lookup cost measurement. + volatile uint32_t spin = 0; + for (uint32_t i = 0; i < 4096; i++) { + spin += i; + } + + out[0] = 42.0f; + if (spin == 0xffffffffu) { + out[0] = 43.0f; + } + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp new file mode 100644 index 0000000000..566e213cde --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp @@ -0,0 +1,93 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Explicit 64x64 fanin DAG validation scene. + * + * Args layout: + * tensor[0]: producer outputs, split into disjoint producer slices + * tensor[1]: consumer outputs, split into disjoint consumer slices + * scalar[0]: producer_count + * scalar[1]: consumer_count + * scalar[2]: use_real_kernels + * + * The scene submits producer_count independent producers, then consumer_count + * independent consumers where every consumer explicitly depends on every + * producer. Real-kernel mode writes disjoint tensor slices so tensormap + * auto-deps do not add producer chains or consumer chains. + * + * When use_real_kernels is false, the same dependency shape is submitted with + * dummy tasks to isolate orchestrator fanin lookup cost. + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +static constexpr int32_t MAX_PRODUCERS = 64; +static constexpr int32_t MAX_CONSUMERS = 64; +static constexpr uint32_t SLOT_ELEMS = 16; + +#define FUNC_WRITE_CONST 0 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 5, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &producer_outputs = orch_args.tensor(0).ref(); + const Tensor &consumer_outputs = orch_args.tensor(1).ref(); + int32_t producer_count = static_cast(orch_args.scalar(0)); + int32_t consumer_count = static_cast(orch_args.scalar(1)); + bool use_real_kernels = orch_args.scalar(2) != 0; + if (producer_count < 1 || producer_count > MAX_PRODUCERS || consumer_count < 1 || consumer_count > MAX_CONSUMERS) { + rt_report_fatal( + PTO2_ERROR_INVALID_ARGS, + "producer_count=%d consumer_count=%d exceed supported range producers=[1, %d] consumers=[1, %d]", + producer_count, consumer_count, MAX_PRODUCERS, MAX_CONSUMERS + ); + return; + } + + PTO2TaskId producer_ids[MAX_PRODUCERS]; + uint32_t slot_shape[1] = {SLOT_ELEMS}; + for (int32_t i = 0; i < producer_count; i++) { + L0TaskArgs args; + if (use_real_kernels) { + uint32_t offset[1] = {static_cast(i) * SLOT_ELEMS}; + Tensor producer_out = producer_outputs.view(slot_shape, offset); + args.add_inout(producer_out); + producer_ids[i] = rt_submit_aic_task(FUNC_WRITE_CONST, args).task_id(); + } else { + producer_ids[i] = rt_submit_dummy_task(args).task_id(); + } + } + + for (int32_t c = 0; c < consumer_count; c++) { + L0TaskArgs args; + args.set_dependencies(producer_ids, static_cast(producer_count)); + if (use_real_kernels) { + uint32_t offset[1] = {static_cast(c) * SLOT_ELEMS}; + Tensor consumer_out = consumer_outputs.view(slot_shape, offset); + args.add_inout(consumer_out); + rt_submit_aic_task(FUNC_WRITE_CONST, args); + } else { + rt_submit_dummy_task(args); + } + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py new file mode 100644 index 0000000000..155894272a --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py @@ -0,0 +1,86 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""fanin_lookup_perf: validate a 64x64 explicit fanin DAG. + +The orchestration submits 64 independent producers and 64 independent +consumers. Each consumer explicitly depends on all 64 producers. The real +kernel case uses disjoint tensor slices so tensormap auto-deps do not add +producer chains or consumer chains. +""" + +import ctypes + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestFaninLookupPerf(SceneTestCase): + """Validate a wide 64-producer/64-consumer explicit fanin DAG.""" + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/fanin_lookup_perf_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "WRITE_CONST", + "source": "kernels/aic/kernel_write_const_visible.cpp", + "core_type": "aic", + # Single-AIC task with one INOUT tensor at payload slot 0. + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "LookupOnlyProducers64Consumers64", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"producer_count": 64, "consumer_count": 64, "use_real_kernels": 0}, + }, + { + "name": "SwimlaneProducers64Consumers64", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"producer_count": 64, "consumer_count": 64, "use_real_kernels": 1}, + }, + ] + + def generate_args(self, params): + slot_elems = 16 + producer_count = int(params["producer_count"]) + consumer_count = int(params["consumer_count"]) + return TaskArgsBuilder( + Tensor("producer_out", torch.full((producer_count * slot_elems,), -1.0, dtype=torch.float32)), + Tensor("consumer_out", torch.full((consumer_count * slot_elems,), -1.0, dtype=torch.float32)), + Scalar("producer_count", ctypes.c_int64(producer_count)), + Scalar("consumer_count", ctypes.c_int64(consumer_count)), + Scalar("use_real_kernels", ctypes.c_int64(int(params["use_real_kernels"]))), + ) + + def compute_golden(self, args, params): + if not params["use_real_kernels"]: + return + slot_elems = 16 + for i in range(int(params["producer_count"])): + args.producer_out[i * slot_elems] = 42.0 + for c in range(int(params["consumer_count"])): + args.consumer_out[c * slot_elems] = 42.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_copy_first.cpp b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_copy_first.cpp new file mode 100644 index 0000000000..1d4af93d9f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_copy_first.cpp @@ -0,0 +1,48 @@ +/* + * 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 + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *in_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + + __gm__ float *in = reinterpret_cast<__gm__ float *>(in_tensor->buffer.addr) + in_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = in[0]; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_write_const.cpp b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_write_const.cpp new file mode 100644 index 0000000000..a3d593bc2f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/aic/kernel_write_const.cpp @@ -0,0 +1,45 @@ +/* + * 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 + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = 42.0f; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/orchestration/heap_empty_ring_rebase_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/orchestration/heap_empty_ring_rebase_orch.cpp new file mode 100644 index 0000000000..e81af56d59 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/kernels/orchestration/heap_empty_ring_rebase_orch.cpp @@ -0,0 +1,66 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Allocates one 1 MiB scratch in each of two consecutive scopes on ring 1. + * Ring 1 has 1.5 MiB of heap, so the second allocation requires the drained + * ring parked at offset 1 MiB to rebase to zero. + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_FILL_CONST 0 +#define FUNC_COPY_FIRST 1 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 2, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_Y1 = orch_args.tensor(0).ref(); + const Tensor &ext_Y2 = orch_args.tensor(1).ref(); + + uint32_t scratch_shapes[2] = {1024, 256}; + TensorCreateInfo scratch_ci(scratch_shapes, 2, DataType::FLOAT32); + + PTO2_SCOPE() { + L0TaskArgs fill1; + fill1.add_output(scratch_ci); + TaskOutputTensors t1_outs = rt_submit_aic_task(FUNC_FILL_CONST, fill1); + const Tensor &t1 = t1_outs.get_ref(0); + + L0TaskArgs copy1; + copy1.add_input(t1); + copy1.add_inout(ext_Y1); + rt_submit_aic_task(FUNC_COPY_FIRST, copy1); + } + + PTO2_SCOPE() { + L0TaskArgs fill2; + fill2.add_output(scratch_ci); + TaskOutputTensors t2_outs = rt_submit_aic_task(FUNC_FILL_CONST, fill2); + const Tensor &t2 = t2_outs.get_ref(0); + + L0TaskArgs copy2; + copy2.add_input(t2); + copy2.add_inout(ext_Y2); + rt_submit_aic_task(FUNC_COPY_FIRST, copy2); + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/test_heap_empty_ring_rebase.py b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/test_heap_empty_ring_rebase.py new file mode 100644 index 0000000000..022d3609e6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/test_heap_empty_ring_rebase.py @@ -0,0 +1,85 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Verify that a drained heap ring can reuse its full capacity. + +Each scope allocates a 1 MiB scratch on ring 1. The first scope drains the ring +at offset 1 MiB; the second allocation therefore requires an empty-ring rebase +when ring 1 is limited to 1.5 MiB. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +SENTINEL = 42.0 +INIT_VAL = -1.0 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestHeapEmptyRingRebase(SceneTestCase): + """A drained ring parked away from zero accepts a whole-span allocation.""" + + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/heap_empty_ring_rebase_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "FILL_CONST", + "source": "kernels/aic/kernel_write_const.cpp", + "core_type": "aic", + "signature": [D.OUT], + }, + { + "func_id": 1, + "name": "COPY_FIRST", + "source": "kernels/aic/kernel_copy_first.cpp", + "core_type": "aic", + "signature": [D.IN, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "EmptyRingRebase", + "platforms": ["a5sim", "a5"], + "config": { + "aicpu_thread_num": 2, + "runtime_env": { + "ring_heap": [268435456, 1572864, 268435456, 268435456], + }, + }, + "params": {}, + }, + ] + + def generate_args(self, params): + y1 = torch.full((16,), INIT_VAL, dtype=torch.float32) + y2 = torch.full((16,), INIT_VAL, dtype=torch.float32) + return TaskArgsBuilder( + Tensor("y1", y1), + Tensor("y2", y2), + ) + + def compute_golden(self, args, params): + args.y1[0] = SENTINEL + args.y2[0] = SENTINEL + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py b/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py new file mode 100644 index 0000000000..28a54467fc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py @@ -0,0 +1,155 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Multi-round paged attention: benchmark multi-round execution (default 10 rounds). + +Run with --rounds 10 --skip-golden for benchmarking. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + +_PA_KERNELS = "../../../../../examples/a5/tensormap_and_ringbuffer/paged_attention/kernels" + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestMultiRoundPagedAttention(SceneTestCase): + RTOL = 1e-2 + ATOL = 1e-2 + + CALLABLE = { + "orchestration": { + "source": f"{_PA_KERNELS}/orchestration/paged_attention_orch.cpp", + "function_name": "build_paged_attention_graph", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": f"{_PA_KERNELS}/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": f"{_PA_KERNELS}/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": f"{_PA_KERNELS}/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": f"{_PA_KERNELS}/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..ed14f900b6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,171 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// SplitK PV Matmul Kernel: Accumulated P @ V across n_blocks +// +// Processes n_blocks blocks using SplitK accumulation pattern: +// Block 0: TMATMUL(C, A, B) — initialize accumulator +// Block i: TMATMUL_ACC(C, C, A, B) — accumulate into same C +// +// Per-block pij addresses: contiguous slices of pij_buf (n_blocks * M * K) +// Per-block vj addresses: value_cache base + block_indices lookup +// Single output: oi_new (M, N) fp32 = sum of P_i @ V_i across all blocks +// +// Optimizations: +// - Double-buffered L1 tiles (ping/pong for A and B via MTE2) +// - Double-buffered L0 tiles (ping/pong for L0A and L0B via MTE1) +// - TLOAD(next) overlaps with TMATMUL(current) via MTE2/M-pipe parallelism +// - Canonical 3-stage pipeline: TLOAD(MTE2) → TMOV(MTE1) → TMATMUL(M) +// - Reverse-dependency events ensure buffer safety across iterations +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// pij is bfloat16 (from softmax_prepare TCVT). +// vj is stored as (K, N) = (block_size, head_dim) in row-major (ND) layout. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void pv_matmul_n_impl( + __gm__ Tensor *pij_buf, __gm__ Tensor *value_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *oi_new, + uint64_t n_blocks, uint64_t bt_offset +) { + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_buf->buffer.addr) + pij_buf->start_offset; + __gm__ bfloat16_t *val_base = reinterpret_cast<__gm__ bfloat16_t *>(value_cache->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr) + oi_new->start_offset; + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride<1, M * K, M * K, K, 1>>; + using GlobalB = GlobalTensor, pto::Stride>; + using GlobalOut = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + // L1 memory layout: double-buffered A and B tiles (tightly packed) + constexpr int kATileBytes = M * K * static_cast(sizeof(bfloat16_t)); + constexpr int kBTileBytes = K * N * static_cast(sizeof(bfloat16_t)); + + TileMatA aMatTile[2]; + TileMatB bMatTile[2]; + TASSIGN(aMatTile[0], 0x0); + TASSIGN(aMatTile[1], kATileBytes); + TASSIGN(bMatTile[0], 2 * kATileBytes); + TASSIGN(bMatTile[1], 2 * kATileBytes + kBTileBytes); + + // L0 memory layout: double-buffered L0A and L0B, single accumulator L0C + LeftTile aTile[2]; + RightTile bTile[2]; + AccTile cTile; + TASSIGN(aTile[0], 0x0); + TASSIGN(aTile[1], kATileBytes); + TASSIGN(bTile[0], 0x0); + TASSIGN(bTile[1], kBTileBytes); + TASSIGN(cTile, 0x0); + + GlobalOut oiGlobal(oi_base); + + // Seed reverse-dependency flags: all ping/pong buffers initially free + // PIPE_MTE1 → PIPE_MTE2: L1 buffer [0/1] safe for TLOAD to overwrite + // PIPE_M → PIPE_MTE1: L0 buffer [0/1] safe for TMOV to overwrite + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + + for (uint64_t i = 0; i < n_blocks; i++) { + int cur = static_cast(i % 2); + GlobalA pijGlobal(pij_base + i * M * K); + GlobalB vjGlobal(val_base + bt[bt_offset + i] * K * N); + + // Stage 1: TLOAD (MTE2: GM → L1[cur]) + // Wait for MTE1 to release L1[cur] (reverse dep from previous iteration) + wait_flag(PIPE_MTE1, PIPE_MTE2, static_cast<::event_t>(cur)); + TLOAD(aMatTile[cur], pijGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // forward: A in L1 ready + TLOAD(bMatTile[cur], vjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // forward: B in L1 ready + + // Stage 2: TMOV (MTE1: L1[cur] → L0[cur]) + // Wait for M-pipe to release L0[cur] (reverse dep from previous iteration) + wait_flag(PIPE_M, PIPE_MTE1, static_cast<::event_t>(cur)); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // forward: wait A loaded + TMOV(aTile[cur], aMatTile[cur]); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // forward: wait B loaded + TMOV(bTile[cur], bMatTile[cur]); + set_flag(PIPE_MTE1, PIPE_MTE2, static_cast<::event_t>(cur)); // reverse: release L1[cur] + + // Stage 3: TMATMUL (M-pipe: L0A[cur] × L0B[cur] → L0C) + set_flag(PIPE_MTE1, PIPE_M, static_cast<::event_t>(cur)); // forward: L0[cur] ready + wait_flag(PIPE_MTE1, PIPE_M, static_cast<::event_t>(cur)); + if (i == 0) { + TMATMUL(cTile, aTile[cur], bTile[cur]); + } else { + TMATMUL_ACC(cTile, cTile, aTile[cur], bTile[cur]); + } + set_flag(PIPE_M, PIPE_MTE1, static_cast<::event_t>(cur)); // reverse: release L0[cur] + } + + // Drain outstanding reverse-dependency flags + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE(oiGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij_buf = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *value_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t n_blocks = static_cast(args[4]); + uint64_t bt_offset = static_cast(args[5]); + + // pij_buf is 4D (1, 1, q_tile, n_blocks*block_size) to match qk's 4D output. + uint64_t q_tile_size = static_cast(pij_buf->shapes[2]); + + if (q_tile_size == 16) { + pv_matmul_n_impl<16, 128, 128>(pij_buf, value_cache, block_table_t, oi_new, n_blocks, bt_offset); + } else { + pv_matmul_n_impl<64, 64, 128>(pij_buf, value_cache, block_table_t, oi_new, n_blocks, bt_offset); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..02ccd5987f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,135 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Multi-block QK Matmul Kernel: qi(M, K) @ kj.T(K, N) -> sij(M, N) for each block +// +// Processes n_blocks blocks in a single kernel invocation. +// Per-block kj addresses computed from key_cache base + block_indices lookup. +// qi is shared across all blocks (same query head against different key blocks). +// +// Output layout: n_blocks contiguous (M, N) tiles stacked vertically. +// Block i occupies sij[i*M : (i+1)*M, 0:N]. +// +// Optimizations: +// - qi TLOAD hoisted before the loop (constant across all iterations) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// Template: M=q_tile, K=head_dim, N=block_size + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void qk_matmul_n_impl( + __gm__ Tensor *qi, __gm__ Tensor *key_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *sij_buf, + uint64_t n_blocks, uint64_t bt_offset +) { + // Decode 4D query view: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + __gm__ bfloat16_t *qi_base = reinterpret_cast<__gm__ bfloat16_t *>(qi->buffer.addr) + qi->start_offset; + __gm__ bfloat16_t *key_base = reinterpret_cast<__gm__ bfloat16_t *>(key_cache->buffer.addr); + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_buf->buffer.addr) + sij_buf->start_offset; + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride<1, M * K, M * K, K, 1>>; + using GlobalB = GlobalTensor, pto::Stride, Layout::DN>; + using GlobalOut = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Hoist qi TLOAD before the loop (qi is constant across all blocks) + GlobalA qiGlobal(qi_base); + TLOAD(aMatTile, qiGlobal); + + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalB kjGlobal(key_base + bt[bt_offset + i] * N * K); + GlobalOut sijGlobal(sij_base + i * M * N); + + // Load only B each iteration (qi already in L1 from hoist) + TLOAD(bMatTile, kjGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + // TMOV qi from L1→L0A (re-copy since TMATMUL consumed L0A) and kj from L1→L0B + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + if (i + 1 < n_blocks) { + pipe_barrier(PIPE_ALL); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *qi = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *sij_buf = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t n_blocks = static_cast(args[4]); + uint64_t bt_offset = static_cast(args[5]); + + // qi is a 4D view (batch, q_len, num_heads_tile, head_dim); decode fixes batch=q_len=1. + uint64_t q_tile_size = static_cast(qi->shapes[2]); + + if (q_tile_size == 16) { + qk_matmul_n_impl<16, 128, 128>(qi, key_cache, block_table_t, sij_buf, n_blocks, bt_offset); + } else { + qk_matmul_n_impl<64, 128, 64>(qi, key_cache, block_table_t, sij_buf, n_blocks, bt_offset); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..b3f51fc11c --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,251 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + // 4D data views (1, 1, q_tile, head_dim) — oi, oi_new, dst. + using GlobalData4D = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalData4D oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalData4D oiGlobal(oi_ptr + oi->start_offset); + GlobalData4D dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to same UB as DN tiles for ND-format store + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs as DN (ColMajor) + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + TMUL(tmpRow, alphaRow, liRow); // alpha * li + TMUL(liNewRow, betaRow, lijRow); // beta * lij + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + // mij is 3D (1, 1, q_tile) to match softmax's 3D scalar output. + uint64_t q_tile_size = static_cast(mij->shapes[2]); + + if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8c05ea8a00 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,265 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Two-Pass Softmax Kernel (AIV) for n_blocks tiles +// +// Input: sij_buf (n_blocks * M, N) fp32 — QK results stacked vertically +// Output: pij_buf (n_blocks * M, N) bf16 — attention weights per block +// mij (M,) fp32 — global row max across all blocks +// lij (M,) fp32 — total row sum across all blocks +// +// Pass 1: Iterate over n_blocks tiles, apply scale, mask last block, +// find global m = max over all blocks of rowmax(S_i * scale) +// Uses TRESHAPE for DN↔Row conversion to keep globalMax in UB +// (eliminates 63 × 4 GM round-trip operations). +// Pass 2: Iterate again, compute P_i = exp(S_i * scale - m) -> bf16, +// accumulate l = sum over all blocks of rowsum(P_i) +// Uses double-buffered sij tiles to overlap TLOAD with computation. +// +// Two-pass ensures all P_i tiles share the same scale (global max), +// enabling direct TMATMUL_ACC accumulation in the PV kernel. +// +// Supports two tile configurations via runtime dispatch: +// Case1: M=16, N=128 (q_tile=16, block_size=128) +// Case2: M=64, N=64 (q_tile=64, block_size=64) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_n_impl( + __gm__ Tensor *sij_buf, __gm__ Tensor *pij_buf, __gm__ Tensor *mij, __gm__ Tensor *lij, float scale_value, + uint64_t n_blocks, uint64_t valid_len_last +) { + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_buf->buffer.addr) + sij_buf->start_offset; + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_buf->buffer.addr) + pij_buf->start_offset; + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr) + mij->start_offset; + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr) + lij->start_offset; + + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + + // --- GlobalTensor types --- + // 4D data views (1, 1, q_tile, n_blocks*block_size) for sij/pij. + using GlobalDataMxN = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + using GlobalDataMxN_bf16 = + GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + // DN/ND scalar globals stay 2D: scalar vectors only need per-element layout. + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + // --- Tile types --- + using TileSijDyn = Tile; + using TileSijPad = Tile; + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + using TileScalarND = + Tile; + // RowMajor (1, M) tile for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // --- UB memory layout (double-buffered sij) --- + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Double-buffered sij tiles + TileVecMxN sijTile_A; + TileSijPad sijPadTile_A; + TileVecMxN sijTile_B; + TileSijPad sijPadTile_B; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileVecMxN sumAccTile; + TileScalarDN localMaxDN; + TileScalarDN globalMaxDN; + TileScalarDN sumDN; + TileVecMxN_bf16 pijBf16Tile; + + // TRESHAPE aliases (same UB address as their DN counterparts) + TileScalarRow localMaxRow; + TileScalarRow globalMaxRow; + + // ND alias for storing globalMax to GM + TileScalarND globalMaxND; + + TASSIGN(sijTile_A, 0x0); + TASSIGN(sijPadTile_A, 0x0); + TASSIGN(sijTile_B, kDataBytes); + TASSIGN(sijPadTile_B, kDataBytes); + TASSIGN(pijTile, 2 * kDataBytes); + TASSIGN(tmpTile, 3 * kDataBytes); + TASSIGN(sumAccTile, 4 * kDataBytes); + int scalarBase = 5 * kDataBytes; + TASSIGN(localMaxDN, scalarBase); + TASSIGN(localMaxRow, scalarBase); // alias: same UB as localMaxDN + TASSIGN(globalMaxDN, scalarBase + kScalarDNBytes); + TASSIGN(globalMaxRow, scalarBase + kScalarDNBytes); // alias: same UB as globalMaxDN + TASSIGN(globalMaxND, scalarBase + kScalarDNBytes); // alias: same UB as globalMaxDN + TASSIGN(sumDN, scalarBase + 2 * kScalarDNBytes); + TASSIGN(pijBf16Tile, scalarBase + 3 * kScalarDNBytes); + + // GM aliases (mij/lij output buffers) + GlobalScalarND mijGlobalND(mij_addr); + GlobalScalarDN lijGlobalDN(lij_addr); + + // ======== Pass 1: Find global row max via TRESHAPE (no GM round-trip) ======== + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalDataMxN sijGlobal(sij_base + i * M * N); + TLOAD(sijTile_A, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if (i == n_blocks - 1 && valid_len_last < static_cast(N)) { + TileSijDyn sijDynTile(static_cast(valid_len_last)); + TASSIGN(sijDynTile, 0x0); + TFILLPAD_INPLACE(sijPadTile_A, sijDynTile); + } + + TMULS(sijTile_A, sijTile_A, scale_value); + TROWMAX(localMaxDN, sijTile_A, tmpTile); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise TMAX + TRESHAPE(localMaxRow, localMaxDN); + if (i == 0) { + TMAX(globalMaxRow, localMaxRow, localMaxRow); + } else { + TMAX(globalMaxRow, globalMaxRow, localMaxRow); + } + } + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for Pass 2's TROWEXPANDSUB + TRESHAPE(globalMaxDN, globalMaxRow); + + // Store final global max to mij for online_update to consume + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(mijGlobalND, globalMaxND); + + // ======== Pass 2: Compute softmax with double-buffered sij ======== + // globalMaxDN is already in UB from TRESHAPE — no reload needed. + // Sync MTE3→MTE2 to ensure the mij TSTORE completed before first sij TLOAD. + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + + // Pre-load first sij tile into buffer A + GlobalDataMxN sijGlobal_0(sij_base); + TLOAD(sijTile_A, sijGlobal_0); + + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalDataMxN_bf16 pijGlobal(pij_base + i * M * N); + + // Wait for current tile's TLOAD to complete + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TFILLPAD on current buffer if last block with partial valid length + if (i == n_blocks - 1 && valid_len_last < static_cast(N)) { + TileSijDyn curSijDyn(static_cast(valid_len_last)); + if (i % 2 == 0) { + TASSIGN(curSijDyn, 0x0); + TFILLPAD_INPLACE(sijPadTile_A, curSijDyn); + } else { + TASSIGN(curSijDyn, static_cast(kDataBytes)); + TFILLPAD_INPLACE(sijPadTile_B, curSijDyn); + } + } + + // Compute on current buffer (select A or B based on iteration parity) + if (i % 2 == 0) { + TMULS(sijTile_A, sijTile_A, scale_value); + TROWEXPANDSUB(pijTile, sijTile_A, globalMaxDN); + } else { + TMULS(sijTile_B, sijTile_B, scale_value); + TROWEXPANDSUB(pijTile, sijTile_B, globalMaxDN); + } + TEXP(pijTile, pijTile); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + + if (i == 0) { + TMULS(sumAccTile, pijTile, 1.0f); + } else { + TADD(sumAccTile, sumAccTile, pijTile); + } + + // Store pij (must complete before next iteration's TCVT overwrites pijBf16Tile) + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Prefetch next sij into alternate buffer (after TSTORE to avoid UB race) + if (i + 1 < n_blocks) { + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + GlobalDataMxN sijGlobal_next(sij_base + (i + 1) * M * N); + if (i % 2 == 0) { + TLOAD(sijTile_B, sijGlobal_next); + } else { + TLOAD(sijTile_A, sijGlobal_next); + } + } + } + + // Compute final row sum from accumulated pij values + TROWSUM(sumDN, sumAccTile, tmpTile); + + // Store lij (total sum). mij already stored after Pass 1. + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(lijGlobalDN, sumDN); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij_buf = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij_buf = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t n_blocks = static_cast(args[5]); + uint64_t valid_len_last = static_cast(args[6]); + + // sij_buf is 4D (1, 1, q_tile, n_blocks*block_size) to match qk's 4D output semantic. + uint64_t q_tile_size = static_cast(sij_buf->shapes[2]); + + if (q_tile_size == 16) { + softmax_prepare_n_impl<16, 128>(sij_buf, pij_buf, mij, lij, scale_value, n_blocks, valid_len_last); + } else { + softmax_prepare_n_impl<64, 64>(sij_buf, pij_buf, mij, lij, scale_value, n_blocks, valid_len_last); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..fb5b2015ed --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,186 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Paged Attention Orchestration - 4D input shapes, N_UNROLL=64, 4 Tasks Per Group + * + * Batches up to N_UNROLL blocks per group. Each group submits exactly 4 tasks: + * 1. QK matmul: qi @ K^T for n_blocks → sij_buf (1, 1, q_tile, n_blocks * block_size) + * 2. Softmax: two-pass over sij_buf → pij_buf, mi, li + * 3. PV matmul: SplitK accumulated P @ V → oi_new (1, 1, q_tile, head_dim) + * 4. Update: online softmax accumulation with group-level mi, li, oi_new + * + * Memory Layout (4D throughout): + * Query: (batch, seq_len=1, num_heads, head_dim) bf16 + * Key: (total_blocks, block_size, kv_head_num, head_dim) bf16 + * Value: (total_blocks, block_size, kv_head_num, head_dim) bf16 + * Out: (batch, seq_len=1, num_heads, head_dim) fp32 + */ + +#include +#include + +#include "pto_orchestration_api.h" + +#define N_UNROLL 64 + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 + +extern "C" { +/** + * Orchestration config — the executor reads these values to set up + * shared memory and runtime before calling aicpu_orchestration_entry. + */ +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Read dimensions from tensor metadata + // query: shape=[batch, seq_len, num_heads, head_dim] + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[2]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[3]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + // key_cache: shape=[total_blocks, block_size, kv_head_num, head_dim] + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + + // block_table: shape=[batch, max_num_blocks_per_req] + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + // scale from scalar arg + uint64_t scale_value = orch_args.scalar(0); + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + + // External 4D tensors inherit shape/dtype from TaskArg (golden provides 4D). + const Tensor &query = orch_args.tensor(0).ref(); + const Tensor &key_cache = orch_args.tensor(1).ref(); + const Tensor &value_cache = orch_args.tensor(2).ref(); + const Tensor &block_table = orch_args.tensor(3).ref(); + const Tensor &out = orch_args.tensor(5).ref(); + + int *host_context_lens = orch_args.tensor(4).ref().data_as(); + + // Loop-invariant shape descriptors: 4D data tiles (1, 1, q_tile, head_dim), + // 3D scalar vectors (1, 1, q_tile). + uint32_t tile4d_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)head_dim}; + uint32_t scalar_shapes[3] = {1, 1, (uint32_t)q_tile}; + TensorCreateInfo tile4d_ci(tile4d_shapes, 4, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 3, DataType::FLOAT32); + + // Prefetch first block host_context_lens data into cache + __builtin_prefetch(&host_context_lens[0], 0, 3); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint64_t cur_seq = host_context_lens[b_idx]; + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + + // Prefetch next block host_context_lens data while processing current batch + if (b_idx + 1 < batch) { + __builtin_prefetch(&host_context_lens[b_idx + 1], 0, 3); + } + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE() { + // 4D views into query/out, matching (1, 1, q_tile, head_dim). + uint32_t view_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)head_dim}; + uint32_t view_offsets[4] = {(uint32_t)b_idx, 0, (uint32_t)(q_idx * q_tile), 0}; + Tensor qi = query.view(view_shapes, view_offsets); + Tensor out_view = out.view(view_shapes, view_offsets, true); + + // Per-group accumulators: oi (4D data), mi_update/li_update (3D scalars). + TaskOutputTensors alloc_outs = alloc_tensors(tile4d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + + // Reusable Arg objects — reset() before each use avoids + // repeated stack-frame construction in the inner loop. + L0TaskArgs params_qk, params_sf, params_pv, params_up; + + for (uint64_t bn = 0; bn < bn_this_batch; bn += N_UNROLL) { + uint64_t n_blocks = std::min((uint64_t)N_UNROLL, bn_this_batch - bn); + + // Valid length for last block in this group + uint64_t last_block_seq_start = (bn + n_blocks - 1) * block_size; + uint64_t valid_len_last = std::min(block_size, cur_seq - last_block_seq_start); + + // === Task 1: Batched QK matmul — produces 4D sij_buf === + uint32_t sij_buf_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)(n_blocks * block_size)}; + TensorCreateInfo sij_buf_ci(sij_buf_shapes, 4, DataType::FLOAT32); + + params_qk.reset(); + params_qk.add_input(qi); + params_qk.add_input(key_cache); + params_qk.add_input(block_table); + params_qk.add_output(sij_buf_ci); + params_qk.add_scalar(n_blocks); + params_qk.add_scalar(b_idx * block_num + bn); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij_buf = qk_outs.get_ref(0); + + // === Task 2: Two-pass softmax — produces 4D pij_buf, 3D mi, li === + uint32_t pij_buf_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)(n_blocks * block_size)}; + TensorCreateInfo pij_buf_ci(pij_buf_shapes, 4, data_type); + + params_sf.reset(); + params_sf.add_input(sij_buf); + params_sf.add_output(pij_buf_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + params_sf.add_scalar(scale_value); + params_sf.add_scalar(n_blocks); + params_sf.add_scalar(valid_len_last); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_buf = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + + // === Task 3: SplitK PV matmul — produces 4D oi_new === + params_pv.reset(); + params_pv.add_input(pij_buf); + params_pv.add_input(value_cache); + params_pv.add_input(block_table); + params_pv.add_output(tile4d_ci); + params_pv.add_scalar(n_blocks); + params_pv.add_scalar(b_idx * block_num + bn); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_new = pv_outs.get_ref(0); + + // === Task 4: Online update (per-group) === + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn + n_blocks >= bn_this_batch) ? 1 : 0; + + params_up.reset(); + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_new); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + } + } + } + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py new file mode 100644 index 0000000000..3eb8d49691 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py @@ -0,0 +1,144 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention unroll with 4D input shapes (batch, seq_len, num_heads, head_dim). + +Query and output tensors use 4D format instead of the standard 3D. +6 kernels: QK/PV matmul (AIC), softmax_prepare/online_update (AIV). +Orchestration with N_UNROLL=64, 4 tasks per group, online softmax accumulation. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPagedAttentionUnroll4dims(SceneTestCase): + """Paged attention unroll with 4D query/out shapes.""" + + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + "manual": True, + }, + { + "name": "Case3", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 256, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + "manual": True, + }, + ] + + def generate_args(self, params): + inputs = _pa_generate_inputs(params) + batch = params["batch"] + num_heads = params["num_heads"] + head_dim = params["head_dim"] + specs = [] + for name, val in inputs: + if isinstance(val, torch.Tensor): + if name in ("query", "out"): + val = val.reshape(batch, 1, num_heads, head_dim) + specs.append(Tensor(name, val)) + else: + specs.append(Scalar(name, val)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + batch = params["batch"] + num_heads = params["num_heads"] + head_dim = params["head_dim"] + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + # Reshape 4D out to 3D for shared golden, then restore + out_4d = tensors["out"] + tensors["out"] = out_4d.reshape(batch, num_heads, head_dim) + _pa_compute_golden(tensors, params) + tensors["out"] = out_4d + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp new file mode 100644 index 0000000000..e9fb327157 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp @@ -0,0 +1,61 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * AIC kernel: writes float(block_idx) at cache line (base_cl + block_idx*3 + 0). + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 0) * FLOATS_PER_CACHE_LINE; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp new file mode 100644 index 0000000000..1a7fe065bf --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp @@ -0,0 +1,63 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * AIV kernel: writes float(block_idx) at cache line + * (base_cl + block_idx*3 + 1 + sub_block_id). + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t block_idx = get_block_idx(args); + int32_t sub_block_id = get_sub_block_id(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 1 + sub_block_id) * FLOATS_PER_CACHE_LINE; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp new file mode 100644 index 0000000000..366dabbdd1 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Regression test for batch dispatch OOB (issue #565). + * + * Submits two MIX tasks with block_num=48 back-to-back so they are both + * in the ready queue when the scheduler runs pop_ready_tasks_batch. + * + * Args layout: [output] + */ + +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_AIC 0 +#define FUNC_AIV0 1 +#define FUNC_AIV1 2 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +static void submit_spmd_mix(const Tensor &out, int16_t block_num, int64_t base_cl) { + MixedKernels mk; + mk.aic_kernel_id = FUNC_AIC; + mk.aiv0_kernel_id = FUNC_AIV0; + mk.aiv1_kernel_id = FUNC_AIV1; + + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.launch_spec.set_core_num(block_num); + rt_submit_task(mk, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + + // Two back-to-back tasks with block_num=48 (2x cluster count). + // Both land in the ready queue simultaneously, triggering got=2 in + // pop_ready_tasks_batch — the scenario that causes OOB without the fix. + submit_spmd_mix(ext_output, 48, 0); + submit_spmd_mix(ext_output, 48, 144); + + LOG_INFO_V9("[spmd_batch_dispatch_oob] Submitted 2 MIX tasks: block_num=48,48"); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py new file mode 100644 index 0000000000..0e8dadeac9 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py @@ -0,0 +1,88 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Regression test for batch dispatch OOB (issue #565). + +Submits two back-to-back MIX tasks each with block_num=48 (>> 24 clusters). +When both tasks enter the ready queue simultaneously, pop_ready_tasks_batch +returns got=2. Without the fix, the first task's do-while drains all idle +clusters, and the second task's do-while calls pop_first() on an empty mask, +returning -1 as cluster_offset — an out-of-bounds array index. + +Each block writes float(block_idx) at 3 cache lines (AIC, AIV0, AIV1). +Output tensor: 2 * 48 * 3 = 288 cache lines = 4608 float32. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +SLOTS_PER_BLOCK = 3 +TASKS = [(48, 0), (48, 144)] +TOTAL_CL = sum(bn * SLOTS_PER_BLOCK for bn, _ in TASKS) + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdBatchDispatchOob(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/kernel_write.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/kernel_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "source": "kernels/aiv/kernel_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {}, + } + ] + + def generate_args(self, params): + return TaskArgsBuilder(Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32))) + + def compute_golden(self, args, params): + out = args.output + for block_num, base_cl in TASKS: + for block_idx in range(block_num): + for slot in range(SLOTS_PER_BLOCK): + cl = base_cl + block_idx * SLOTS_PER_BLOCK + slot + out[cl * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp new file mode 100644 index 0000000000..c922bd4236 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp @@ -0,0 +1,816 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Paged Attention MIX Kernel — AIC + AIV in single source via TPUSH/TPOP + * + * Hardware block_num is fixed at 24. Each hardware block strides over + * total_logical_blocks = batch * q_loop logical work items: + * for (block_idx = hw_block_idx; block_idx < total_logical_blocks; block_idx += 24) + * Each logical block_idx encodes one (batch_idx, q_tile_idx) position. + * + * q_tile adapts to num_heads at runtime: q_tile = min(num_heads, MAX_Q_TILE). + * When num_heads <= MAX_Q_TILE, q_loop = 1 and each block processes all heads. + * Two q_tile shapes are statically dispatched: 16 (default) and 64. + * + * Compiled twice: once with __DAV_CUBE__ (AIC), once with __DAV_VEC__ (AIV). + * AIC and AIV cooperate via 3 GM-backed FIFO pipes (one set per hardware block, + * reused across stride-loop iterations): + * - sij_pipe (C2V): QK scores (Q_TILE, block_size) fp32, TILE_UP_DOWN + * - pij_pipe (V2C): softmax probs (Q_TILE, block_size) bf16, TILE_UP_DOWN + * - oi_pipe (C2V): PV output (Q_TILE, head_dim) fp32, TILE_UP_DOWN + * + * Per-block pipeline: + * AIC: QK matmul → TPUSH(sij) → TPOP(pij) → PV matmul → TPUSH(oi_new) + * AIV: TPOP(sij) → online softmax → TPUSH(pij) → TPOP(oi_new) → online update + * + * MixedKernels args: + * args[0] = query Tensor* (batch*num_heads, head_dim) bf16 + * args[1] = key_cache Tensor* (kv_total_rows, head_dim) bf16 + * args[2] = value_cache Tensor* (kv_total_rows, head_dim) bf16 + * args[3] = block_table Tensor* (batch, max_blocks_per_req) int32 + * args[4] = context_lens Tensor* (batch,) int32 + * args[5] = out Tensor* (batch*num_heads, head_dim) float32 [output] + * args[6] = sij_fifo Tensor* GM ring buffer for sij pipe + * args[7] = pij_fifo Tensor* GM ring buffer for pij pipe + * args[8] = oi_fifo Tensor* GM ring buffer for oi_new pipe + * args[9] = scale_value scalar (float bits in uint64) + * args[10] = num_heads scalar + * args[11] = head_dim scalar + * args[12] = block_size scalar + * args[13] = max_num_blocks_per_req scalar + * args[14] = q_loop scalar + * args[15] = total_logical_blocks scalar (= batch * q_loop) + * args[16] = q_tile scalar (16 or 64) + */ + +#include +// NOLINTBEGIN(clang-diagnostic-error,bugprone-reserved-identifier,bugprone-easily-swappable-parameters,modernize-use-auto) +#include +#include + +#include "tensor.h" + +using pto::BLayout; +using pto::Direction; +using pto::GlobalTensor; +using pto::Layout; +using pto::PadValue; +using pto::RoundMode; +using pto::Shape; +using pto::SLayout; +using pto::Stride; +using pto::Tile; +using pto::TileAcc; +using pto::TileLeft; +using pto::TileRight; +using pto::TileSplitAxis; +using pto::TileType; +using pto::TPipe; + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#ifdef __DAV_CUBE__ +constexpr bool DAV_CUBE = true; +#else +constexpr bool DAV_CUBE = false; +#endif + +#ifdef __DAV_VEC__ +constexpr bool DAV_VEC = true; +#else +constexpr bool DAV_VEC = false; +#endif + +#include "intrinsic.h" + +static constexpr int MAX_Q_TILE = 64; +static constexpr int HEAD_DIM = 128; +static constexpr int MAX_BLOCK_SIZE = 128; + +// TPUSH/TPOP pipe flag IDs (each consumes 2 consecutive IDs: data + backpressure) +static constexpr uint16_t SIJ_FLAG_ID = 0; +static constexpr uint16_t PIJ_FLAG_ID = 2; +static constexpr uint16_t OI_FLAG_ID = 4; +static constexpr uint8_t FIFO_DEPTH = 2; + +// Per-q_tile compile-time configuration: pipe types, slot sizes, UB/L1 layouts. +// QT must be 16 or 64. SUB_QT = QT / 2 (each of AIV0/AIV1 handles half the rows). +template +struct PAConfig { + static constexpr int Q_TILE = QT; + static constexpr int SUB_QT = QT / 2; + + // GM FIFO slot sizes (full tile per slot, sized for max block_size to allow + // the same FIFO to host both block_size=64 and block_size=128 cases). + static constexpr uint32_t SIJ_SLOT_SIZE = QT * MAX_BLOCK_SIZE * sizeof(float); + static constexpr uint32_t PIJ_SLOT_SIZE = QT * MAX_BLOCK_SIZE * sizeof(bfloat16_t); + static constexpr uint32_t OI_SLOT_SIZE = QT * HEAD_DIM * sizeof(float); + + using SijPipeT = TPipe; + using PijPipeT = TPipe; + using OiPipeT = TPipe; + + // AIV UB consumer buffer layout (sized for SUB_QT rows per AIV lane) + static constexpr uint32_t SIJ_UB_BASE = 0x0; + static constexpr uint32_t SIJ_UB_SIZE = 2 * SUB_QT * MAX_BLOCK_SIZE * sizeof(float); + static constexpr uint32_t OI_UB_BASE = SIJ_UB_BASE + SIJ_UB_SIZE; + static constexpr uint32_t OI_UB_SIZE = 2 * SUB_QT * HEAD_DIM * sizeof(float); + static constexpr uint32_t WORK_UB_BASE = OI_UB_BASE + OI_UB_SIZE; + + // AIC L1 consumer buffer for V2C pij pipe (full QT * MAX_BLOCK_SIZE rows) + static constexpr uint32_t PIJ_L1_BASE = 0x40000; + static constexpr uint32_t PIJ_L1_SIZE = 2 * QT * MAX_BLOCK_SIZE * sizeof(bfloat16_t); +}; + +// ============================================================================ +// AIC (Cube) processing — QK-first offset-loop software pipeline +// +// QK-first order: each steady-state iteration does QK[i] then PV[i-1]. +// This maximizes overlap by hiding AIV's softmax behind AIC's QK matmul: +// while AIC computes QK[i], AIV concurrently processes SF[i-1]. +// By the time AIC finishes QK[i] and needs pij[i-1], SF[i-1] is done. +// FIFO_DEPTH=2 supports the 2-deep sij buffering (sij[i-1] + sij[i]). +// +// Timeline (steady state): +// AIC: QK[i] → TPUSH(sij[i]) → TPOP(pij[i-1]) → PV[i-1] → TPUSH(oi[i-1]) +// AIV: TPOP(sij[i-1]) → SF[i-1] → TPUSH(pij[i-1]) → TPOP(oi[i-2]) → UP[i-2] +// ────────────────────────────────────────────────────────────────────────── +// QK[i] overlaps with SF[i-1] (Cube compute ∥ Vector softmax) +// PV[i-1] overlaps with UP[i-2] (Cube compute ∥ Vector online update) +// ============================================================================ + +// Helper: QK matmul for block i — load key, move to L0, matmul, TPUSH sij +template < + int M, int K, int N, typename SijPipeT, typename GlobalB_QK, typename TileMatA_QK, typename TileMatB_QK, + typename LeftTile_QK, typename RightTile_QK, typename AccTile_QK> +static __aicore__ void aic_qk_step( + __gm__ bfloat16_t *key_base, uint64_t kv_block_id, uint64_t i, TileMatA_QK &aMatTile_QK, TileMatB_QK &bMatTile_QK_A, + TileMatB_QK &bMatTile_QK_B, LeftTile_QK &aTile_QK, RightTile_QK &bTile_QK, AccTile_QK &cTile_QK, SijPipeT &sij_pipe, + bool current_loaded = false, bool has_next = false, uint64_t next_kv_block_id = 0 +) { + if (!current_loaded) { + GlobalB_QK kjGlobal(key_base + kv_block_id * N * K); + if (i % 2 == 0) { + TLOAD(bMatTile_QK_A, kjGlobal); + } else { + TLOAD(bMatTile_QK_B, kjGlobal); + } + } + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile_QK, aMatTile_QK); + if (i % 2 == 0) { + TMOV(bTile_QK, bMatTile_QK_A); + } else { + TMOV(bTile_QK, bMatTile_QK_B); + } + + if (has_next) { + GlobalB_QK kjGlobalNext(key_base + next_kv_block_id * N * K); + if ((i + 1) % 2 == 0) { + TLOAD(bMatTile_QK_A, kjGlobalNext); + } else { + TLOAD(bMatTile_QK_B, kjGlobalNext); + } + } + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile_QK, aTile_QK, bTile_QK); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TPUSH(sij_pipe, cTile_QK); + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); +} + +// Helper: PV matmul for block i — TPOP pij, load value, move to L0, matmul, TPUSH oi +template < + int M, int K, int N, typename PijPipeT, typename OiPipeT, typename GlobalB_PV, typename PijMatTile, + typename TileMatB_PV, typename LeftTile_PV, typename RightTile_PV, typename AccTile_PV> +static __aicore__ void aic_pv_step( + __gm__ bfloat16_t *val_base, uint64_t kv_block_id, uint64_t i, PijMatTile &pijMatTile, TileMatB_PV &bMatTile_PV_A, + TileMatB_PV &bMatTile_PV_B, LeftTile_PV &aTile_PV, RightTile_PV &bTile_PV, AccTile_PV &cTile_PV, PijPipeT &pij_pipe, + OiPipeT &oi_pipe, bool current_loaded = false, bool has_next = false, uint64_t next_kv_block_id = 0 +) { + if (!current_loaded) { + GlobalB_PV vjGlobal(val_base + kv_block_id * N * K); + if (i % 2 == 0) { + TLOAD(bMatTile_PV_A, vjGlobal); + } else { + TLOAD(bMatTile_PV_B, vjGlobal); + } + } + + TPOP(pij_pipe, pijMatTile); + + // PV step uses EVENT_ID1 (QK step uses EVENT_ID0) to avoid flag aliasing + // when pipe_barrier(PIPE_ALL) is removed between steps. + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + + TMOV(aTile_PV, pijMatTile); + if (i % 2 == 0) { + TMOV(bTile_PV, bMatTile_PV_A); + } else { + TMOV(bTile_PV, bMatTile_PV_B); + } + + if (has_next) { + GlobalB_PV vjGlobalNext(val_base + next_kv_block_id * N * K); + if ((i + 1) % 2 == 0) { + TLOAD(bMatTile_PV_A, vjGlobalNext); + } else { + TLOAD(bMatTile_PV_B, vjGlobalNext); + } + } + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + + TMATMUL(cTile_PV, aTile_PV, bTile_PV); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID1); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1); + + TPUSH(oi_pipe, cTile_PV); + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); +} + +template +static __aicore__ void aic_process_blocks( + __gm__ bfloat16_t *qi_base, __gm__ bfloat16_t *key_base, __gm__ bfloat16_t *val_base, __gm__ int32_t *bt, + uint64_t bt_offset, uint64_t n_blocks, typename Cfg::SijPipeT &sij_pipe, typename Cfg::PijPipeT &pij_pipe, + typename Cfg::OiPipeT &oi_pipe +) { + constexpr int M = Cfg::Q_TILE; + using SijPipeT = typename Cfg::SijPipeT; + using PijPipeT = typename Cfg::PijPipeT; + using OiPipeT = typename Cfg::OiPipeT; + + using GlobalA_QK = GlobalTensor, Stride>; + using GlobalB_QK = GlobalTensor, Stride, Layout::DN>; + using TileMatA_QK = Tile; + using TileMatB_QK = Tile; + using LeftTile_QK = TileLeft; + using RightTile_QK = TileRight; + using AccTile_QK = TileAcc; + + using GlobalB_PV = GlobalTensor, Stride>; + using TileMatB_PV = Tile; + using PijMatTile = Tile; + using LeftTile_PV = TileLeft; + using RightTile_PV = TileRight; + using AccTile_PV = TileAcc; + + constexpr int kQKBBytes = K * N * static_cast(sizeof(bfloat16_t)); + constexpr int kPVBBytes = N * K * static_cast(sizeof(bfloat16_t)); + + TileMatA_QK aMatTile_QK; + TileMatB_QK bMatTile_QK_A, bMatTile_QK_B; + TASSIGN(aMatTile_QK, 0x0); + TASSIGN(bMatTile_QK_A, 0x20000); + TASSIGN(bMatTile_QK_B, 0x20000 + kQKBBytes); + + LeftTile_QK aTile_QK; + RightTile_QK bTile_QK; + AccTile_QK cTile_QK; + TASSIGN(aTile_QK, 0x0); + TASSIGN(bTile_QK, 0x0); + TASSIGN(cTile_QK, 0x0); + + PijMatTile pijMatTile; + TileMatB_PV bMatTile_PV_A, bMatTile_PV_B; + TASSIGN(bMatTile_PV_A, Cfg::PIJ_L1_BASE + Cfg::PIJ_L1_SIZE); + TASSIGN(bMatTile_PV_B, Cfg::PIJ_L1_BASE + Cfg::PIJ_L1_SIZE + kPVBBytes); + + LeftTile_PV aTile_PV; + RightTile_PV bTile_PV; + AccTile_PV cTile_PV; + TASSIGN(aTile_PV, 0x0); + TASSIGN(bTile_PV, 0x0); + TASSIGN(cTile_PV, 0x0); + + GlobalA_QK qiGlobal(qi_base); + TLOAD(aMatTile_QK, qiGlobal); + + if (n_blocks == 1) { + // Degenerate case: no pipeline overlap possible + uint64_t block_id = static_cast(bt[bt_offset]); + aic_qk_step( + key_base, block_id, 0, aMatTile_QK, bMatTile_QK_A, bMatTile_QK_B, aTile_QK, bTile_QK, cTile_QK, sij_pipe + ); + aic_pv_step( + val_base, block_id, 0, pijMatTile, bMatTile_PV_A, bMatTile_PV_B, aTile_PV, bTile_PV, cTile_PV, pij_pipe, + oi_pipe + ); + } else { + // Prologue: QK[0] — produces sij[0] for AIV to start SF[0] + uint64_t prev_block_id = static_cast(bt[bt_offset]); + uint64_t next_block_id = static_cast(bt[bt_offset + 1]); + aic_qk_step( + key_base, prev_block_id, 0, aMatTile_QK, bMatTile_QK_A, bMatTile_QK_B, aTile_QK, bTile_QK, cTile_QK, + sij_pipe, false, true, next_block_id + ); + // Steady state: QK[i] then PV[i-1] (QK-first order). + for (uint64_t i = 1; i < n_blocks; i++) { + uint64_t block_id = static_cast(bt[bt_offset + i]); + uint64_t next_block_id = (i + 1 < n_blocks) ? static_cast(bt[bt_offset + i + 1]) : 0; + aic_qk_step( + key_base, block_id, i, aMatTile_QK, bMatTile_QK_A, bMatTile_QK_B, aTile_QK, bTile_QK, cTile_QK, + sij_pipe, true, i + 1 < n_blocks, next_block_id + ); + aic_pv_step( + val_base, prev_block_id, i - 1, pijMatTile, bMatTile_PV_A, bMatTile_PV_B, aTile_PV, bTile_PV, cTile_PV, + pij_pipe, oi_pipe, i > 1, i < n_blocks, block_id + ); + prev_block_id = block_id; + } + + // Epilogue: PV[n-1] — consume last pij + aic_pv_step( + val_base, prev_block_id, n_blocks - 1, pijMatTile, bMatTile_PV_A, bMatTile_PV_B, aTile_PV, bTile_PV, + cTile_PV, pij_pipe, oi_pipe, n_blocks > 1 + ); + } +} + +// ============================================================================ +// AIV (Vector) processing — SF-first offset-loop software pipeline +// +// SF-first order: each steady-state iteration does SF[i] then UP[i-1]. +// This ensures pij[i] is produced as early as possible so AIC's TPOP(pij) +// never stalls behind a pending UP computation. Combined with AIC's +// QK-first order, SF[i] overlaps with AIC's PV[i-1] Cube matmul. +// ============================================================================ + +// Helper: softmax step for block i — TPOP sij, compute softmax, TPUSH pij +// +// globalMaxRow is used as a running accumulator: on entry it holds the max +// from the previous iteration (or is undefined when i==0). SF updates it +// in-place to max(globalMaxRow, localMaxRow_i * scale). The caller must +// save globalMaxRow before calling SF if the old value is still needed. +template < + typename Cfg, int TM, int TN, typename SijVecTile, typename TileSijPad, typename TileVecMxN, + typename PijVecBf16Tile, typename TileScalarDN, typename TileScalarRow> +static __aicore__ void aiv_sf_step( + uint64_t i, bool is_last_partial, uint64_t valid_len_last, float scale_value, SijVecTile &sijTile, + TileSijPad &sijPadTile, TileVecMxN &pijTile, TileVecMxN &tmpTile, PijVecBf16Tile &pijBf16Tile, + TileScalarDN &localMaxDN, TileScalarDN &globalMaxDN, TileScalarDN &llDN, TileScalarRow &localMaxRow, + TileScalarRow &globalMaxRow, typename Cfg::SijPipeT &sij_pipe, typename Cfg::PijPipeT &pij_pipe +) { + using TileSijDyn = Tile; + using SijPipeT = typename Cfg::SijPipeT; + using PijPipeT = typename Cfg::PijPipeT; + + TPOP(sij_pipe, sijTile); + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if (is_last_partial) { + int sij_addr = Cfg::SIJ_UB_BASE + static_cast((i % 2) * TM * TN * static_cast(sizeof(float))); + TASSIGN(sijPadTile, sij_addr); + TileSijDyn sijDynTile(static_cast(valid_len_last)); + TASSIGN(sijDynTile, sij_addr); + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + } + + TROWMAX(localMaxDN, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TRESHAPE(localMaxRow, localMaxDN); + + if (i == 0) { + TMULS(globalMaxRow, localMaxRow, scale_value); + } else { + TMULS(localMaxRow, localMaxRow, scale_value); + pipe_barrier(PIPE_V); + TMAX(globalMaxRow, globalMaxRow, localMaxRow); + } + TRESHAPE(globalMaxDN, globalMaxRow); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, globalMaxDN); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + pipe_barrier(PIPE_V); + + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + + TROWSUM(llDN, pijTile, tmpTile); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TPUSH(pij_pipe, pijBf16Tile); +} + +// Helper: online update step for block i — TPOP oi, merge with accumulators +// +// curMaxRow = M[i] = running max over blocks 0..i (mij in FlashAttention notation) +// prevMaxRow = M[i-1] = running max over blocks 0..i-1 (dm / old max) +// llDN_i = row-sum of pij for block i +// +// alpha = exp(prevMaxRow - curMaxRow), used to rescale accumulated go and gl. +template < + typename Cfg, int TM, int TN, typename OiVecTile, typename TileDataMxHD, typename TileScalarDN, + typename TileScalarND, typename TileScalarRow> +static __aicore__ void aiv_up_step( + uint64_t i, OiVecTile &oiNewTile, TileDataMxHD &goTile, TileScalarDN &alphaDN_dn, TileScalarDN &llDN_i, + TileScalarND &glND, TileScalarND &alphaND, TileScalarND &llND, TileScalarND &dmND, TileScalarND &mijND, + TileScalarRow &curMaxRow, TileScalarRow &prevMaxRow, typename Cfg::OiPipeT &oi_pipe +) { + using OiPipeT = typename Cfg::OiPipeT; + TPOP(oi_pipe, oiNewTile); + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + + if (i == 0) { + TMULS(goTile, oiNewTile, 1.0f); + TRESHAPE(llND, llDN_i); + pipe_barrier(PIPE_V); + TMULS(glND, llND, 1.0f); + } else { + TRESHAPE(llND, llDN_i); + TRESHAPE(mijND, curMaxRow); + TRESHAPE(dmND, prevMaxRow); + + TSUB(alphaND, dmND, mijND); + pipe_barrier(PIPE_V); + TEXP(alphaND, alphaND); + pipe_barrier(PIPE_V); + + TRESHAPE(alphaDN_dn, alphaND); + TROWEXPANDMUL(goTile, goTile, alphaDN_dn); + pipe_barrier(PIPE_V); + TADD(goTile, goTile, oiNewTile); + + TMUL(glND, glND, alphaND); + pipe_barrier(PIPE_V); + TADD(glND, glND, llND); + } + + pipe_barrier(PIPE_V); +} + +template +static __aicore__ void aiv_process_blocks( + float scale_value, uint64_t n_blocks, uint64_t valid_len_last, __gm__ float *dst_ptr, + typename Cfg::SijPipeT &sij_pipe, typename Cfg::PijPipeT &pij_pipe, typename Cfg::OiPipeT &oi_pipe +) { + constexpr int TM = Cfg::SUB_QT; + constexpr int HD = HEAD_DIM; + constexpr int kAlignedRows = ((TM * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = TM / kScalarCols; + + using SijVecTile = Tile; + using PijVecBf16Tile = Tile; + using OiVecTile = Tile; + + using TileVecMxN = Tile; + using TileSijPad = + Tile; + using TileScalarDN = Tile; + using TileScalarND = + Tile; + using TileScalarRow = Tile; + using TileDataMxHD = Tile; + using GlobalDataMxHD = GlobalTensor, Stride<1, 1, 1, HD, 1>>; + + constexpr int kSijBytes = TM * TN * sizeof(float); + constexpr int kPijBf16Bytes = TM * TN * sizeof(bfloat16_t); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + constexpr int kScalarNDBytes = kScalarRows * kScalarCols * sizeof(float); + + SijVecTile sijTile; + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + PijVecBf16Tile pijBf16Tile; + TileScalarDN localMaxDN, globalMaxDN; + TileScalarDN alphaDN_dn, llDN, glDN; + TileScalarDN savedLlDN; + TileScalarND gmND, glND, alphaND, llND, dmND, miNewND, mijND; + TileScalarRow localMaxRow, globalMaxRow; + TileScalarRow savedMaxRow, prevMaxRow; + OiVecTile oiNewTile; + TileDataMxHD goTile; + + int ub = Cfg::WORK_UB_BASE; + TASSIGN(pijTile, ub); + ub += kSijBytes; + TASSIGN(pijBf16Tile, ub); + ub += kPijBf16Bytes; + TASSIGN(tmpTile, ub); + ub += kSijBytes; + + int sb = ub; + TASSIGN(localMaxDN, sb); + TASSIGN(localMaxRow, sb); + sb += kScalarDNBytes; + TASSIGN(globalMaxDN, sb); + TASSIGN(globalMaxRow, sb); + sb += kScalarDNBytes; + TASSIGN(gmND, sb); + TASSIGN(savedMaxRow, sb); + sb += kScalarDNBytes; + TASSIGN(glND, sb); + TASSIGN(glDN, sb); + sb += kScalarDNBytes; + TASSIGN(alphaND, sb); + TASSIGN(alphaDN_dn, sb); + sb += kScalarDNBytes; + TASSIGN(llND, sb); + TASSIGN(llDN, sb); + sb += kScalarDNBytes; + TASSIGN(dmND, sb); + sb += kScalarNDBytes; + TASSIGN(miNewND, sb); + sb += kScalarNDBytes; + TASSIGN(mijND, sb); + sb += kScalarNDBytes; + TASSIGN(prevMaxRow, sb); + sb += kScalarDNBytes; + TASSIGN(savedLlDN, sb); + sb += kScalarDNBytes; + + TASSIGN(goTile, sb); + + GlobalDataMxHD dstGlobal(dst_ptr); + + bool last_partial = (valid_len_last < static_cast(TN)); + + if (n_blocks == 1) { + aiv_sf_step( + 0, last_partial, valid_len_last, scale_value, sijTile, sijPadTile, pijTile, tmpTile, pijBf16Tile, + localMaxDN, globalMaxDN, llDN, localMaxRow, globalMaxRow, sij_pipe, pij_pipe + ); + aiv_up_step( + 0, oiNewTile, goTile, alphaDN_dn, llDN, glND, alphaND, llND, dmND, mijND, globalMaxRow, globalMaxRow, + oi_pipe + ); + } else { + // Prologue: SF[0] — not the last block + aiv_sf_step( + 0, false, valid_len_last, scale_value, sijTile, sijPadTile, pijTile, tmpTile, pijBf16Tile, localMaxDN, + globalMaxDN, llDN, localMaxRow, globalMaxRow, sij_pipe, pij_pipe + ); + + // Steady state: SF[i] then UP[i-1] (SF-first order). + for (uint64_t i = 1; i < n_blocks; i++) { + // Shift max history: prevMaxRow ← savedMaxRow (M[i-2]) + // Save current: savedMaxRow ← globalMaxRow (M[i-1]) + TMULS(prevMaxRow, savedMaxRow, 1.0f); + TMULS(savedMaxRow, globalMaxRow, 1.0f); + TMULS(savedLlDN, llDN, 1.0f); + pipe_barrier(PIPE_V); + + bool cur_last_partial = (i == n_blocks - 1) && last_partial; + aiv_sf_step( + i, cur_last_partial, valid_len_last, scale_value, sijTile, sijPadTile, pijTile, tmpTile, pijBf16Tile, + localMaxDN, globalMaxDN, llDN, localMaxRow, globalMaxRow, sij_pipe, pij_pipe + ); + + aiv_up_step( + i - 1, oiNewTile, goTile, alphaDN_dn, savedLlDN, glND, alphaND, llND, dmND, mijND, savedMaxRow, + prevMaxRow, oi_pipe + ); + } + + // Epilogue: UP[n-1] — uses live globalMaxRow (M[n-1]) and savedMaxRow (M[n-2]) + aiv_up_step( + n_blocks - 1, oiNewTile, goTile, alphaDN_dn, llDN, glND, alphaND, llND, dmND, mijND, globalMaxRow, + savedMaxRow, oi_pipe + ); + } + + // Final normalization: output = goTile / glDN + TRESHAPE(glDN, glND); + pipe_barrier(PIPE_V); + TROWEXPANDDIV(goTile, goTile, glDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dstGlobal, goTile); + + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7); +} + +// ============================================================================ +// Per-config dispatch: builds pipes from per-hw-block FIFO bases, then runs +// the AIC or AIV stride loop over total_logical_blocks. +// ============================================================================ + +template +static __aicore__ void run_aic( + __gm__ int64_t *args, __gm__ int32_t *ctx_ptr, int32_t hw_block_idx, int32_t hw_block_num, + int64_t total_logical_blocks, int64_t num_heads, int64_t head_dim, int64_t block_size, int64_t max_blocks_per_req, + int64_t q_loop, __gm__ void *sij_fifo_base, __gm__ void *pij_fifo_base, __gm__ void *oi_fifo_base +) { + typename Cfg::SijPipeT sij_pipe(sij_fifo_base, Cfg::SIJ_UB_BASE, 0U); + typename Cfg::PijPipeT pij_pipe(pij_fifo_base, 0U, Cfg::PIJ_L1_BASE); + typename Cfg::OiPipeT oi_pipe(oi_fifo_base, Cfg::OI_UB_BASE, 0U); + + __gm__ Tensor *query_t = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache_t = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *value_cache_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[3]); + + __gm__ bfloat16_t *query_base = reinterpret_cast<__gm__ bfloat16_t *>(query_t->buffer.addr) + query_t->start_offset; + __gm__ bfloat16_t *key_base = + reinterpret_cast<__gm__ bfloat16_t *>(key_cache_t->buffer.addr) + key_cache_t->start_offset; + __gm__ bfloat16_t *val_base = + reinterpret_cast<__gm__ bfloat16_t *>(value_cache_t->buffer.addr) + value_cache_t->start_offset; + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr) + block_table_t->start_offset; + + for (int32_t block_idx = hw_block_idx; block_idx < total_logical_blocks; block_idx += hw_block_num) { + int64_t batch_idx = block_idx / q_loop; + int64_t q_tile_idx = block_idx % q_loop; + + int64_t cur_seq = static_cast(ctx_ptr[batch_idx]); + int64_t n_blocks = (cur_seq + block_size - 1) / block_size; + if (n_blocks <= 0) continue; + + int64_t q_offset = (batch_idx * num_heads + q_tile_idx * Cfg::Q_TILE) * head_dim; + __gm__ bfloat16_t *qi_base = query_base + q_offset; + uint64_t bt_offset = static_cast(batch_idx * max_blocks_per_req); + + if (block_size == 128) { + aic_process_blocks( + qi_base, key_base, val_base, bt, bt_offset, static_cast(n_blocks), sij_pipe, pij_pipe, oi_pipe + ); + } else { + aic_process_blocks( + qi_base, key_base, val_base, bt, bt_offset, static_cast(n_blocks), sij_pipe, pij_pipe, oi_pipe + ); + } + } +} + +template +static __aicore__ void run_aiv( + __gm__ int64_t *args, __gm__ int32_t *ctx_ptr, int32_t hw_block_idx, int32_t hw_block_num, + int64_t total_logical_blocks, int64_t num_heads, int64_t head_dim, int64_t block_size, int64_t q_loop, + __gm__ void *sij_fifo_base, __gm__ void *pij_fifo_base, __gm__ void *oi_fifo_base +) { + typename Cfg::SijPipeT sij_pipe(sij_fifo_base, Cfg::SIJ_UB_BASE, 0U); + typename Cfg::PijPipeT pij_pipe(pij_fifo_base, 0U, Cfg::PIJ_L1_BASE); + typename Cfg::OiPipeT oi_pipe(oi_fifo_base, Cfg::OI_UB_BASE, 0U); + + __gm__ Tensor *out_t = reinterpret_cast<__gm__ Tensor *>(args[5]); + float scale_value = from_u64(static_cast(args[9])); + + int32_t sub_block_id = get_sub_block_id(args); + int64_t row_offset = sub_block_id * Cfg::SUB_QT; + + // pto-isa TPUSH/TPOP add `get_subblockid() * sub_rows * cols * elem_bytes` internally, but the CCE + // `get_subblockid()` register is 0 for both lanes under simpler onboard MIX dispatch; add the lane split + // explicitly from GlobalContext.sub_block_id (block_size wide for sij/pij, HEAD_DIM for oi). + sij_pipe.cons.setEntryOffset( + sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(float)) + ); + pij_pipe.prod.setEntryOffset( + sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(bfloat16_t)) + ); + oi_pipe.cons.setEntryOffset(sub_block_id * Cfg::SUB_QT * HEAD_DIM * static_cast(sizeof(float))); + + __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out_t->buffer.addr) + out_t->start_offset; + + for (int32_t block_idx = hw_block_idx; block_idx < total_logical_blocks; block_idx += hw_block_num) { + int64_t batch_idx = block_idx / q_loop; + int64_t q_tile_idx = block_idx % q_loop; + + int64_t cur_seq = static_cast(ctx_ptr[batch_idx]); + int64_t n_blocks = (cur_seq + block_size - 1) / block_size; + + int64_t out_offset = (batch_idx * num_heads + q_tile_idx * Cfg::Q_TILE + row_offset) * head_dim; + __gm__ float *dst = out_base + out_offset; + + if (n_blocks <= 0) { + using ZeroTile = + Tile; + using ZeroGlobal = GlobalTensor, Stride<1, 1, 1, HEAD_DIM, 1>>; + ZeroTile zeroTile; + TASSIGN(zeroTile, Cfg::WORK_UB_BASE); + TEXPANDS(zeroTile, 0.0f); + pipe_barrier(PIPE_V); + ZeroGlobal dstZero(dst); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dstZero, zeroTile); + pipe_barrier(PIPE_MTE3); + continue; + } + + int64_t last_block_seq = (n_blocks - 1) * block_size; + int64_t remaining = cur_seq - last_block_seq; + uint64_t valid_len_last = (remaining >= block_size) ? static_cast(block_size) : + (remaining > 0 ? static_cast(remaining) : 0); + + if (block_size == 128) { + aiv_process_blocks( + scale_value, static_cast(n_blocks), valid_len_last, dst, sij_pipe, pij_pipe, oi_pipe + ); + } else { + aiv_process_blocks( + scale_value, static_cast(n_blocks), valid_len_last, dst, sij_pipe, pij_pipe, oi_pipe + ); + } + } +} + +// ============================================================================ +// Entry point — shared by AIC and AIV via DAV_CUBE / DAV_VEC guards +// ============================================================================ + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *context_lens_t = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *sij_fifo_t = reinterpret_cast<__gm__ Tensor *>(args[6]); + __gm__ Tensor *pij_fifo_t = reinterpret_cast<__gm__ Tensor *>(args[7]); + __gm__ Tensor *oi_fifo_t = reinterpret_cast<__gm__ Tensor *>(args[8]); + + int64_t num_heads = static_cast(args[10]); + int64_t head_dim = static_cast(args[11]); + int64_t block_size = static_cast(args[12]); + int64_t max_blocks_per_req = static_cast(args[13]); + int64_t q_loop = static_cast(args[14]); + int64_t total_logical_blocks = static_cast(args[15]); + int64_t q_tile = static_cast(args[16]); + + int32_t hw_block_idx = get_block_idx(args); + int32_t hw_block_num = get_block_num(args); + + __gm__ int32_t *ctx_ptr = + reinterpret_cast<__gm__ int32_t *>(context_lens_t->buffer.addr) + context_lens_t->start_offset; + + // GM FIFO buffer per hardware block (reused across stride-loop iterations). + // Slot stride is sized for max(Q_TILE) so the same offset works for both q_tile=16 and 64. + constexpr uint32_t SIJ_HW_STRIDE = PAConfig::SIJ_SLOT_SIZE * FIFO_DEPTH; + constexpr uint32_t PIJ_HW_STRIDE = PAConfig::PIJ_SLOT_SIZE * FIFO_DEPTH; + constexpr uint32_t OI_HW_STRIDE = PAConfig::OI_SLOT_SIZE * FIFO_DEPTH; + + __gm__ void *sij_fifo_base = reinterpret_cast<__gm__ void *>( + reinterpret_cast<__gm__ uint8_t *>(sij_fifo_t->buffer.addr) + hw_block_idx * SIJ_HW_STRIDE + ); + __gm__ void *pij_fifo_base = reinterpret_cast<__gm__ void *>( + reinterpret_cast<__gm__ uint8_t *>(pij_fifo_t->buffer.addr) + hw_block_idx * PIJ_HW_STRIDE + ); + __gm__ void *oi_fifo_base = reinterpret_cast<__gm__ void *>( + reinterpret_cast<__gm__ uint8_t *>(oi_fifo_t->buffer.addr) + hw_block_idx * OI_HW_STRIDE + ); + + if constexpr (DAV_CUBE) { + if (q_tile == 16) { + run_aic>( + args, ctx_ptr, hw_block_idx, hw_block_num, total_logical_blocks, num_heads, head_dim, block_size, + max_blocks_per_req, q_loop, sij_fifo_base, pij_fifo_base, oi_fifo_base + ); + } else { + run_aic>( + args, ctx_ptr, hw_block_idx, hw_block_num, total_logical_blocks, num_heads, head_dim, block_size, + max_blocks_per_req, q_loop, sij_fifo_base, pij_fifo_base, oi_fifo_base + ); + } + } + + if constexpr (DAV_VEC) { + if (q_tile == 16) { + run_aiv>( + args, ctx_ptr, hw_block_idx, hw_block_num, total_logical_blocks, num_heads, head_dim, block_size, + q_loop, sij_fifo_base, pij_fifo_base, oi_fifo_base + ); + } else { + run_aiv>( + args, ctx_ptr, hw_block_idx, hw_block_num, total_logical_blocks, num_heads, head_dim, block_size, + q_loop, sij_fifo_base, pij_fifo_base, oi_fifo_base + ); + } + } +} +// NOLINTEND(clang-diagnostic-error,bugprone-reserved-identifier,bugprone-easily-swappable-parameters,modernize-use-auto) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp new file mode 100644 index 0000000000..07e90500b8 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp @@ -0,0 +1,159 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * SPMD Paged Attention Orchestration with TPUSH/TPOP (fixed block_num=24) + * + * Submits a single MixedKernels task with hardware block_num fixed at 24. + * total_logical_blocks = batch * q_loop logical work items are distributed + * across the 24 hardware blocks via a stride loop inside the kernel: + * for (block_idx = hw_block_idx; block_idx < total_logical_blocks; block_idx += 24) + * + * q_tile adapts to num_heads at runtime: q_tile = min(num_heads, MAX_Q_TILE). + * When num_heads <= MAX_Q_TILE (=64), q_loop = 1 and each block processes all heads. + * + * Each iteration of the stride loop processes one (batch_idx, q_tile_idx) logical + * position, running the full AIC/AIV cooperative pipeline via TPUSH/TPOP pipes: + * AIC: QK matmul -> TPUSH(sij) -> TPOP(pij) -> PV matmul -> TPUSH(oi_new) + * AIV: TPOP(sij) -> online softmax -> TPUSH(pij) -> TPOP(oi_new) -> online update + * + * GM FIFO buffers for TPUSH/TPOP are sized for the 24 hardware blocks (not for + * total_logical_blocks). Each hardware block owns its own FIFO slots and reuses + * them across stride-loop iterations. + */ + +#include +#include + +#include + +#include "pto_orchestration_api.h" + +#define FUNC_PA_AIC 0 +#define FUNC_PA_AIV 1 + +static constexpr uint64_t MAX_Q_TILE = 64; +static constexpr uint64_t HEAD_DIM = 128; +static constexpr uint64_t MAX_BLOCK_SIZE = 128; +static constexpr int16_t SPMD_BLOCK_NUM = 24; + +// GM FIFO slot sizes (must match kernel's PAConfig constants). +// Sized for the maximum (q_tile, block_size) so the same FIFO layout works +// for both q_tile=16 and q_tile=64 dispatch paths inside the kernel. +static constexpr uint32_t SIJ_SLOT_SIZE = MAX_Q_TILE * MAX_BLOCK_SIZE * sizeof(float); +static constexpr uint32_t PIJ_SLOT_SIZE = MAX_Q_TILE * MAX_BLOCK_SIZE * sizeof(uint16_t); +static constexpr uint32_t OI_SLOT_SIZE = MAX_Q_TILE * HEAD_DIM * sizeof(float); +static constexpr uint32_t FIFO_DEPTH = 2; + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t max_num_blocks_per_req = orch_args.tensor(3).ref().shapes[1]; + uint64_t scale_value = orch_args.scalar(0); + + // q_tile adapts to num_heads: use 64 when num_heads >= 64, else 16. + // The kernel statically dispatches on q_tile == 16 vs 64. + uint64_t q_tile = (num_heads >= MAX_Q_TILE) ? MAX_Q_TILE : 16; + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + int64_t total_logical_blocks = static_cast(batch * q_loop); + + LOG_INFO_V0( + "SPMD PA TPUSH/TPOP: batch=%" PRIu64 " heads=%" PRIu64 " hd=%" PRIu64 " bs=%" PRIu64 " q_tile=%" PRIu64 + " q_loop=%" PRIu64 " hw_blocks=%d logical_blocks=%" PRId64, + batch, num_heads, head_dim, block_size, q_tile, q_loop, SPMD_BLOCK_NUM, total_logical_blocks + ); + + // Wrap host tensors + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint64_t total_kv_blocks = orch_args.tensor(1).ref().shapes[0]; + uint64_t kv_total_rows = total_kv_blocks * block_size; + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint32_t kv_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, kv_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, kv_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(max_num_blocks_per_req)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + // GM FIFO buffers for TPUSH/TPOP (one set of slots per hardware block) + uint32_t sij_fifo_total = static_cast(SPMD_BLOCK_NUM) * SIJ_SLOT_SIZE * FIFO_DEPTH; + uint32_t pij_fifo_total = static_cast(SPMD_BLOCK_NUM) * PIJ_SLOT_SIZE * FIFO_DEPTH; + uint32_t oi_fifo_total = static_cast(SPMD_BLOCK_NUM) * OI_SLOT_SIZE * FIFO_DEPTH; + + // Allocate as 1D byte tensors (using INT32 for 4-byte alignment, divide by 4) + uint32_t sij_fifo_shapes[1] = {sij_fifo_total / sizeof(int32_t)}; + uint32_t pij_fifo_shapes[1] = {pij_fifo_total / sizeof(int32_t)}; + uint32_t oi_fifo_shapes[1] = {oi_fifo_total / sizeof(int32_t)}; + + TensorCreateInfo sij_fifo_ci(sij_fifo_shapes, 1, DataType::INT32); + TensorCreateInfo pij_fifo_ci(pij_fifo_shapes, 1, DataType::INT32); + TensorCreateInfo oi_fifo_ci(oi_fifo_shapes, 1, DataType::INT32); + + PTO2_SCOPE() { + L0TaskArgs args; + args.add_input(query); + args.add_input(key_cache); + args.add_input(value_cache); + args.add_input(block_table); + args.add_input(context_lens); + args.add_inout(out); + args.add_output(sij_fifo_ci); + args.add_output(pij_fifo_ci); + args.add_output(oi_fifo_ci); + args.add_scalar(scale_value); + args.add_scalar(static_cast(num_heads)); + args.add_scalar(static_cast(head_dim)); + args.add_scalar(static_cast(block_size)); + args.add_scalar(static_cast(max_num_blocks_per_req)); + args.add_scalar(static_cast(q_loop)); + args.add_scalar(total_logical_blocks); + args.add_scalar(static_cast(q_tile)); + args.launch_spec.set_block_num(SPMD_BLOCK_NUM); + + MixedKernels mk; + mk.aic_kernel_id = FUNC_PA_AIC; + mk.aiv0_kernel_id = FUNC_PA_AIV; + mk.aiv1_kernel_id = FUNC_PA_AIV; + rt_submit_task(mk, args); + } + + LOG_INFO_V0( + "SPMD PA TPUSH/TPOP: submitted 1 MixedKernels task, hw_blocks=%d logical=%" PRId64, + static_cast(SPMD_BLOCK_NUM), total_logical_blocks + ); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py new file mode 100644 index 0000000000..e47e4f5073 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py @@ -0,0 +1,130 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention unroll with TPUSH/TPOP: MIX kernel AIC+AIV cooperative pipeline.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPagedAttentionUnrollTpushPop(SceneTestCase): + # Tolerances relaxed (2e-3 -> 5e-3 in #825, then 5e-3 -> 1e-2 for #848) + # to absorb hardware numerical drift in the AIC/AIV cooperative TPUSH/TPOP + # pipeline; observed max_diff ~5.5e-3. + RTOL = 1e-2 + ATOL = 1e-2 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "PA_AIC", + "source": "kernels/mix/paged_attention_parallel.cpp", + "core_type": "aic", + # Cooperative mix: AIC and AIV share one 9-tensor args[]. Each + # half declares the shared payload (task-level directions); the + # dump records each tensor once per declaring subtask, under its + # own func_id. Consumed only by the dump; dispatch ignores it. + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.INOUT, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 1, + "name": "PA_AIV", + "source": "kernels/mix/paged_attention_parallel.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.INOUT, D.OUT, D.OUT, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + # Intra-core trace target only (--case SmallCase1; manual -> not in + # the default onboard CI sweep). batch=24 == the orchestration's + # hardcoded SPMD_BLOCK_NUM, so every hw block gets one logical block + # (fewer stalls in the AIC<->AIV handshake). Same q_tile=16 path as + # Case1; passes golden at context_len=8192. + "name": "SmallCase1", + "platforms": ["a5"], + "manual": True, + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 24, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py new file mode 100644 index 0000000000..d27e9ba3a5 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py @@ -0,0 +1,106 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""L3 ChipTask → SubTask dependency via TensorMap. + +Worker(level=3) submits a ChipTask then a SubTask that depends on it. +Verifies: TensorMap dependency inference, cross-fork data visibility, +SubWorker reads result produced by ChipWorker. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test +from simpler_setup.scene_test import _build_l3_task_args + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + + +def verify(args): + """SubCallable — dependency target, runs after ChipTask completes.""" + + +def run_dag(orch, callables, task_args, config): + """L3 orchestration: ChipTask → SubTask dependency.""" + # ChipTask: tags inside chip_args drive deps (INPUT → lookup; OUTPUT_EXISTING → insert). + chip_args, _ = _build_l3_task_args(task_args, callables.vector_kernel_sig) + callables.keep(chip_args) # prevent GC before drain + + orch.submit_next_level(callables.vector_kernel, chip_args, config, worker=0) + + # SubTask: tag the chip output as INPUT — Orchestrator wires the dep via TensorMap. + sub_args = TaskArgs() + sub_args.add_tensor(make_tensor_arg(task_args.f), TensorArgType.INPUT) + orch.submit_sub(callables.verify, sub_args) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3Dependency(SceneTestCase): + """L3: ChipTask produces output, SubTask depends on it via TensorMap.""" + + CALLABLE = { + "orchestration": run_dag, + "callables": [ + { + "name": "vector_kernel", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "verify", "callable": verify}, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"device_count": 1, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + SIZE = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + ) + + def compute_golden(self, args, params): + args.f[:] = (args.a + args.b + 1) * (args.a + args.b + 2) + (args.a + args.b) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py new file mode 100644 index 0000000000..f628b0a4c6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py @@ -0,0 +1,120 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""L3 group task — 2 ChipWorkers (process-isolated) on 1 DAG node. + +Each chip runs the same kernel with its own args (different tensors). +A downstream SubTask depends on the group output. +Verifies: fork+shm process isolation, 2-chip concurrent execution, +group completion aggregation, downstream dependency waits for group. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + + +def verify(args): + """SubCallable — runs after group completes.""" + + +def _chip_args(in_a, in_b, out_f) -> TaskArgs: + """Build per-chip TaskArgs with INPUT/INPUT/OUTPUT_EXISTING tags.""" + a = TaskArgs() + a.add_tensor(make_tensor_arg(in_a), TensorArgType.INPUT) + a.add_tensor(make_tensor_arg(in_b), TensorArgType.INPUT) + a.add_tensor(make_tensor_arg(out_f), TensorArgType.OUTPUT_EXISTING) + return a + + +def run_dag(orch, callables, task_args, config): + """L3 orchestration: group of 2 chips → SubTask dependency.""" + args0 = _chip_args(task_args.a0, task_args.b0, task_args.f0) + args1 = _chip_args(task_args.a1, task_args.b1, task_args.f1) + callables.keep(args0, args1) # prevent GC before drain + + orch.submit_next_level_group(callables.vector_kernel, [args0, args1], config, workers=[0, 1]) + + # SubTask depends on both group outputs (f0, f1) — tag both as INPUT. + sub_args = TaskArgs() + sub_args.add_tensor(make_tensor_arg(task_args.f0), TensorArgType.INPUT) + sub_args.add_tensor(make_tensor_arg(task_args.f1), TensorArgType.INPUT) + orch.submit_sub(callables.verify, sub_args) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3Group(SceneTestCase): + """L3: Group of 2 ChipWorkers as 1 DAG node, SubTask depends on group.""" + + CALLABLE = { + "orchestration": run_dag, + "callables": [ + { + "name": "vector_kernel", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "verify", "callable": verify}, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"device_count": 2, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + SIZE = 128 * 128 + return TaskArgsBuilder( + Tensor("a0", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b0", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f0", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + Tensor("a1", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b1", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f1", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + ) + + def compute_golden(self, args, params): + args.f0[:] = (args.a0 + args.b0 + 1) * (args.a0 + args.b0 + 2) + (args.a0 + args.b0) + args.f1[:] = (args.a1 + args.b1 + 1) * (args.a1 + args.b1 + 2) + (args.a1 + args.b1) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py new file mode 100644 index 0000000000..4e88784efc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py @@ -0,0 +1,138 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""L3 post-fork zero-copy host buffers (issue #1027 #1). + +A host tensor created *after* the chip children are forked (lazily on the +first ``Worker.run()``) is not visible to those children: the orch fn runs in +the parent and carries a raw parent VA that is unmapped (or stale) in the child. +``Worker.create_host_buffer`` hands back born-shared memory already attached into +every chip child, so a tensor built over it with ``torch.frombuffer`` round-trips +with **no per-run copy** — the child reads and writes the same physical pages the +parent sees. + +Covers the mechanism end-to-end (allocate a post-fork buffer, fill it in place, +run, read the result back). The host-side staging (born-shared bytes need no +copy; an in-range view is validated to fit) is unit-tested in +``tests/ut/py/test_worker/test_host_buffer_registration.py``. + +a5sim: ``create_host_buffer`` is pure host-side (POSIX shm + a control +broadcast to the forked chip children) with no platform branching, so the sim +backend exercises the full mechanism without needing a device. The +vector_example orchestration kernels exist only for a5. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import CallConfig, TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, make_tensor_arg, scene_test + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + +SIZE = 128 * 128 +DTYPE = torch.float32 + + +def _golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + s = a + b + return (s + 1) * (s + 2) + s + + +def _one_task_orch(chip_handle, a, b, out): + def orch_fn(orch, _args, cfg): + ta = TaskArgs() + ta.add_tensor(make_tensor_arg(a), TensorArgType.INPUT) + ta.add_tensor(make_tensor_arg(b), TensorArgType.INPUT) + ta.add_tensor(make_tensor_arg(out), TensorArgType.OUTPUT_EXISTING) + orch.submit_next_level(chip_handle, ta, cfg, worker=0) + + return orch_fn + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestPostForkHostBufferZeroCopy(SceneTestCase): + """Post-fork zero-copy host buffers on a single L3 worker (issue #1027 #1).""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + ], + } + + CASES = [ + {"name": "post_fork_zero_copy", "platforms": ["a5sim"]}, + ] + + def test_run(self, st_worker): + """Zero-copy: buffers allocated AFTER the fork via ``create_host_buffer``, + filled in place, run, and read back — all without a per-run copy. + + ``Worker.init()`` is eager, so the chip child is already forked when the + buffers are created; a born-shared ``create_host_buffer`` is the mapped + path a post-init host tensor takes to reach that child. + """ + worker = st_worker + chip_handle = type(self)._st_chip_handles["vector"] + + nbytes = SIZE * DTYPE.itemsize # element count × dtype size, not a magic 4 + ba = worker.create_host_buffer(nbytes) + bb = worker.create_host_buffer(nbytes) + bout = worker.create_host_buffer(nbytes) + a = b = out = None + result = False + try: + a = torch.frombuffer(ba.buffer, dtype=DTYPE, count=SIZE) + b = torch.frombuffer(bb.buffer, dtype=DTYPE, count=SIZE) + out = torch.frombuffer(bout.buffer, dtype=DTYPE, count=SIZE) + a.fill_(5.0) # in place → lands directly in the child-visible shm + b.fill_(7.0) + out.zero_() + worker.run(_one_task_orch(chip_handle, a, b, out), args=None, config=CallConfig()) + result = torch.allclose(out, _golden(a, b), rtol=self.RTOL, atol=self.ATOL) + finally: + # Drop the views before freeing so the shm releases promptly, and do it + # in finally so a run()/assert failure above still cleans up all three + # buffers instead of leaving live views warning on the first free. + del a, b, out + worker.free_host_buffer(ba) + worker.free_host_buffer(bb) + worker.free_host_buffer(bout) + assert result + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__)