From fc9a3c3d5f71ba9c3f0eb5f4c203e61e3cd4e4e1 Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 28 Jul 2026 20:09:08 -0700 Subject: [PATCH] test(a5/tmr): port single-device TMR scenes and examples to a5 Port the remaining single-device a2a3 tensormap_and_ringbuffer coverage to a5 for differential validation across architectures. Add scenes for alternating matmul/add, batch and multi-round paged attention, 4D unroll, dynamic registration, fanin lookup, empty-ring heap rebase, SPMD batch dispatch, and L3 dependency and host-buffer behavior. Add the scalar-data, paged-attention ringbuffer, and merge-pipeline-barrier examples, plus the a5 InCore orchestration guide. Adapt the ports to current a5 launch, logging, stride, and runtime APIs, and keep scene compile-cache entries distinct across modules. Exclude multi-device coverage and spmd_paged_attention from this port. --- conftest.py | 2 +- .../docs/INCORE_ORCHESTRATION_GUIDE.md | 83 ++++ .../kernels/aiv/merge_pipeline.cpp | 213 ++++++++++ .../kernels/orchestration/merge_orch.cpp | 50 +++ .../test_merge_pipeline_barrier.py | 98 +++++ .../test_paged_attention_ringbuffer.py | 113 ++++++ .../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 | 85 ++++ simpler_setup/scene_test.py | 10 +- .../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 | 383 ++++++++++++++++++ .../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 | 84 ++++ .../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 ++++ .../test_l3_dependency.py | 106 +++++ .../test_l3_host_buffer_registration.py | 138 +++++++ tests/ut/py/test_scene_test_cache.py | 32 +- 44 files changed, 5385 insertions(+), 7 deletions(-) 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/test_l3_dependency.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py diff --git a/conftest.py b/conftest.py index b63465540b..f7ddf233d2 100644 --- a/conftest.py +++ b/conftest.py @@ -1500,7 +1500,7 @@ def st_worker(request, st_platform, device_pool, _l2_worker_pool, _l2_poisoned): ) name = entry["name"] - cache_key = l3_compile_cache_key(cls.__qualname__, name, st_platform, runtime) + cache_key = l3_compile_cache_key(cls.__module__, cls.__qualname__, name, st_platform, runtime) chip = _compile_chip_callable_from_spec(entry, st_platform, runtime, cache_key) handle = w.register(chip) chip_handles[name] = handle 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..02f2c59fe1 --- /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, pto::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..85c1a4a7f0 --- /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_core_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..57f01ab602 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/merge_pipeline_barrier/test_merge_pipeline_barrier.py @@ -0,0 +1,98 @@ +#!/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_core_num in merge_orch.cpp). + "name": "merge", + "platforms": ["a5"], + "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..584f50dd91 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/paged_attention_ringbuffer/test_paged_attention_ringbuffer.py @@ -0,0 +1,113 @@ +#!/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": { + "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..9e139f9d6c --- /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 = pto::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..15e020241c --- /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("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("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("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("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("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( + "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("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("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("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( + "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("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..469a7954b3 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data_test/test_scalar_data.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. +# ----------------------------------------------------------------------------------------------------------- +"""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"], + "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/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index 0e6a8593f0..38babfc3cf 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -51,7 +51,7 @@ def _pto_isa_compile_cache_token() -> str: return read_pto_isa_pin() -def l3_compile_cache_key(qualname: str, name: str, platform: str, runtime: str) -> tuple: +def l3_compile_cache_key(module: str, qualname: str, name: str, platform: str, runtime: str) -> tuple: """Session compile-cache key for an L3 orchestration entry. Every L3 compile path (scene-test ``test_run``, standalone worker, the @@ -60,7 +60,7 @@ def l3_compile_cache_key(qualname: str, name: str, platform: str, runtime: str) same kernels once per path. The pin token pins the key to the current pto-isa revision. """ - return (qualname, name, platform, runtime, _pto_isa_compile_cache_token()) + return (module, qualname, name, platform, runtime, _pto_isa_compile_cache_token()) def clear_compile_cache() -> None: @@ -1111,7 +1111,7 @@ def compare_outputs(self, test_args, golden_args, output_names, params) -> None: @classmethod def compile_chip_callable(cls, platform): """Compile CALLABLE -> ChipCallable (L2). Session-cached.""" - cache_key = (cls.__qualname__, platform, cls._st_runtime, _pto_isa_compile_cache_token()) + cache_key = (cls.__module__, cls.__qualname__, platform, cls._st_runtime, _pto_isa_compile_cache_token()) return _compile_chip_callable_from_spec(cls.CALLABLE, platform, cls._st_runtime, cache_key) @classmethod @@ -1121,7 +1121,7 @@ def _compile_l3_callables(cls, platform): for entry in cls.CALLABLE["callables"]: if "orchestration" in entry: name = entry["name"] - cache_key = l3_compile_cache_key(cls.__qualname__, name, platform, cls._st_runtime) + cache_key = l3_compile_cache_key(cls.__module__, cls.__qualname__, name, platform, cls._st_runtime) chip = _compile_chip_callable_from_spec(entry, platform, cls._st_runtime, cache_key) compiled[name] = chip compiled[f"{name}_sig"] = entry["orchestration"].get("signature", []) @@ -2097,7 +2097,7 @@ def _create_standalone_worker(group, level, args, selected_by_cls): cls_sub_handles[entry["name"]] = handle elif "orchestration" in entry: name = entry["name"] - cache_key = l3_compile_cache_key(cls.__qualname__, name, args.platform, cls._st_runtime) + cache_key = l3_compile_cache_key(cls.__module__, cls.__qualname__, name, args.platform, cls._st_runtime) chip = _compile_chip_callable_from_spec(entry, args.platform, cls._st_runtime, cache_key) handle = worker.register(chip) cls_chip_handles[name] = handle 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 943e4393bb..dc4383eb6d 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..dc4383eb6d --- /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( + "[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("[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..d8df4cec53 --- /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": {"block_dim": 24}, + "params": {"batch": 1, "M": 1, "N": 1, "matmul_batch": 1, "add_batch": 1}, + }, + { + "name": "Case1", + "platforms": ["a5"], + "config": {"block_dim": 24}, + "params": {"batch": 500, "M": 4, "N": 4, "matmul_batch": 4, "add_batch": 4}, + "manual": True, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"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..a377a72ac3 --- /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("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( + "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..3ff3ebd0ab --- /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": {"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": {"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": {"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": {"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": {"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": {"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": {"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": {"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..b369aaa9d9 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -0,0 +1,383 @@ +#!/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() + + # 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(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() + + 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() + + # 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() + + 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..a548a41e2e --- /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": {"block_dim": 24}, + "params": {"producer_count": 64, "consumer_count": 64, "use_real_kernels": 0}, + }, + { + "name": "SwimlaneProducers64Consumers64", + "platforms": ["a5sim", "a5"], + "config": {"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..530375b005 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/heap_empty_ring_rebase/test_heap_empty_ring_rebase.py @@ -0,0 +1,84 @@ +#!/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": { + "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..d783eb580c --- /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": {"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": {"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": {"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": {"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..3e448ca58d --- /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": {"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": {"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": {"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..6c19a8867d --- /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("[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..b99e8495c4 --- /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": {"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/test_l3_dependency.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py new file mode 100644 index 0000000000..a8ca6c8200 --- /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}, + "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_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__) diff --git a/tests/ut/py/test_scene_test_cache.py b/tests/ut/py/test_scene_test_cache.py index f0490c0360..157dfbcbd9 100644 --- a/tests/ut/py/test_scene_test_cache.py +++ b/tests/ut/py/test_scene_test_cache.py @@ -20,12 +20,20 @@ from __future__ import annotations +import importlib + from _task_interface import ArgDirection, ChipCallable # pyright: ignore[reportMissingImports] # ``simpler_setup/__init__.py`` re-exports the ``scene_test`` *decorator*, # which shadows the submodule attribute when accessed via ``simpler_setup``. # Importing the names directly from the submodule avoids that ambiguity. -from simpler_setup.scene_test import _compile_cache, _pto_isa_compile_cache_token, clear_compile_cache +from simpler_setup.scene_test import ( + SceneTestCase, + _compile_cache, + _pto_isa_compile_cache_token, + clear_compile_cache, + l3_compile_cache_key, +) def _build_chip_callable(tag: str) -> ChipCallable: @@ -65,3 +73,25 @@ def test_pto_isa_compile_cache_token_tracks_pin(monkeypatch): assert _pto_isa_compile_cache_token() == pin_a monkeypatch.setattr("simpler_setup.pto_isa.read_pto_isa_pin", lambda: pin_b) assert _pto_isa_compile_cache_token() == pin_b + + +def test_compile_cache_keys_include_module(monkeypatch): + """Same-named scene classes from different test modules must not share binaries.""" + monkeypatch.setattr("simpler_setup.pto_isa.read_pto_isa_pin", lambda: "pin") + captured = [] + scene_test_module = importlib.import_module("simpler_setup.scene_test") + monkeypatch.setattr( + scene_test_module, + "_compile_chip_callable_from_spec", + lambda spec, platform, runtime, cache_key: captured.append(cache_key), + ) + + first = type("TestScene", (SceneTestCase,), {"__module__": "tests.first", "_st_runtime": "a5", "CALLABLE": {}}) + second = type("TestScene", (SceneTestCase,), {"__module__": "tests.second", "_st_runtime": "a5", "CALLABLE": {}}) + first.compile_chip_callable("a5") + second.compile_chip_callable("a5") + + assert captured[0] != captured[1] + assert l3_compile_cache_key("tests.first", "TestScene", "child", "a5", "a5") != l3_compile_cache_key( + "tests.second", "TestScene", "child", "a5", "a5" + )