diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp index d08f7645bc..49ccd773cd 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp +++ b/tests/st/a2a3/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp @@ -69,6 +69,19 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta int total_matmul_tasks = batch * M; int total_add_tasks = batch * N; + if (matmul_batch <= 0 || add_batch <= 0) { + LOG_ERROR( + "[alternating_orch] batch sizes must be positive (matmul_batch=%d, add_batch=%d)", matmul_batch, add_batch + ); + return; + } + if (total_matmul_tasks % matmul_batch != 0 || total_add_tasks % add_batch != 0) { + LOG_ERROR( + "[alternating_orch] task counts must divide evenly (matmul %d/%d, add %d/%d)", total_matmul_tasks, + matmul_batch, total_add_tasks, add_batch + ); + return; + } int num_matmul_groups = total_matmul_tasks / matmul_batch; int num_add_groups = total_add_tasks / add_batch; diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp new file mode 100644 index 0000000000..607e5d657f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp @@ -0,0 +1,133 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Matrix Multiplication Kernel (Cube Core) + * + * Computes: C = A @ B (TILE x TILE x TILE matmul) + * Uses TMATMUL instruction + * + * Args (Tensor*): + * args[0] = A (INPUT) - TILE x TILE + * args[1] = B (INPUT) - TILE x TILE + * args[2] = C (OUTPUT) - TILE x TILE + */ + +#include +#include +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +AICORE constexpr inline T CeilAlign(T num_1, T num_2) { + if (num_2 == 0) { + return 0; + } + return (num_1 + num_2 - 1) / num_2 * num_2; +} + +static __aicore__ inline int get_num_tiles(__gm__ Tensor *tensor, uint64_t tile_elems) { + uint64_t total_elems = tensor->shapes[0]; + return static_cast(total_elems / tile_elems); +} + +template +static __aicore__ void matmul_impl(__gm__ float *input_a, __gm__ float *input_b, __gm__ float *output) { + constexpr int blockAlign = C0_SIZE_BYTE / sizeof(float); + constexpr int M = CeilAlign(TILE, 16); + constexpr int K = CeilAlign(TILE, blockAlign); + constexpr int N = CeilAlign(TILE, blockAlign); + + using GlobalDataA = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + using GlobalDataB = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + using GlobalDataC = GlobalTensor< + float, Shape<1, 1, 1, TILE, TILE>, pto::Stride<1 * TILE * TILE, 1 * TILE * TILE, TILE * TILE, TILE, 1>>; + + GlobalDataA src0Global(input_a); + GlobalDataB src1Global(input_b); + GlobalDataC dstGlobal(output); + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + TLOAD(aMatTile, src0Global); + TLOAD(bMatTile, src1Global); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(dstGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_a = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *input_b = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *output = reinterpret_cast<__gm__ Tensor *>(args[2]); + + constexpr uint64_t TILE_ELEMS = 128 * 128; + int num_tiles = get_num_tiles(input_a, TILE_ELEMS); + + __gm__ float *base_a = reinterpret_cast<__gm__ float *>(input_a->buffer.addr) + input_a->start_offset; + __gm__ float *base_b = reinterpret_cast<__gm__ float *>(input_b->buffer.addr) + input_b->start_offset; + __gm__ float *base_c = reinterpret_cast<__gm__ float *>(output->buffer.addr) + output->start_offset; + + for (int tile_idx = 0; tile_idx < num_tiles; tile_idx++) { + __gm__ float *a_ptr = base_a + (tile_idx * TILE_ELEMS); + __gm__ float *b_ptr = base_b + (tile_idx * TILE_ELEMS); + __gm__ float *c_ptr = base_c + (tile_idx * TILE_ELEMS); + + matmul_impl<128>(a_ptr, b_ptr, c_ptr); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp new file mode 100644 index 0000000000..ba0cdfa3b4 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp @@ -0,0 +1,93 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Element-wise Tensor Addition Kernel + * + * Implements: out[i] = src0[i] + src1[i] + * Tile size: ROWS x COLS + * + * Args (Tensor*): + * args[0] = src0 (INPUT) - ROWS x COLS + * args[1] = src1 (INPUT) - ROWS x COLS + * args[2] = out (OUTPUT) - ROWS x COLS + */ + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static __aicore__ inline int get_num_tiles(__gm__ Tensor *tensor, uint64_t tile_elems) { + uint64_t total_elems = tensor->shapes[0]; + return static_cast(total_elems / tile_elems); +} + +template +static __aicore__ void add_impl(__gm__ float *src0, __gm__ float *src1, __gm__ float *out) { + using DynShapeDim5 = Shape<1, 1, 1, ROWS, COLS>; + using DynStridDim5 = pto::Stride<1, 1, 1, COLS, 1>; + using GlobalData = GlobalTensor; + using TileData = Tile; + + TileData src0Tile(ROWS, COLS); + TileData src1Tile(ROWS, COLS); + TileData dstTile(ROWS, COLS); + TASSIGN(src0Tile, 0x0); + TASSIGN(src1Tile, 0x10000); + TASSIGN(dstTile, 0x20000); + + GlobalData src0Global(src0); + GlobalData src1Global(src1); + GlobalData dstGlobal(out); + + TLOAD(src0Tile, src0Global); + TLOAD(src1Tile, src1Global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(dstTile, src0Tile, src1Tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(dstGlobal, dstTile); + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *src0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *src1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + constexpr uint64_t TILE_ELEMS = 128 * 128; + int num_tiles = get_num_tiles(src0_tensor, TILE_ELEMS); + + __gm__ float *base_src0 = reinterpret_cast<__gm__ float *>(src0_tensor->buffer.addr) + src0_tensor->start_offset; + __gm__ float *base_src1 = reinterpret_cast<__gm__ float *>(src1_tensor->buffer.addr) + src1_tensor->start_offset; + __gm__ float *base_out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + for (int tile_idx = 0; tile_idx < num_tiles; tile_idx++) { + __gm__ float *src0_ptr = base_src0 + (tile_idx * TILE_ELEMS); + __gm__ float *src1_ptr = base_src1 + (tile_idx * TILE_ELEMS); + __gm__ float *out_ptr = base_out + (tile_idx * TILE_ELEMS); + + add_impl<128, 128>(src0_ptr, src1_ptr, out_ptr); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp new file mode 100644 index 0000000000..49ccd773cd --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp @@ -0,0 +1,139 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Alternating Matmul-Add Orchestration Function (tensormap_and_ringbuffer Runtime) + * + * Submits independent matmul and add tasks per batch. + * + * Configuration read from scalar args: + * - batch: Number of batches + * - M: Number of matmul tasks per batch + * - N: Number of add tasks per batch + * - matmul_batch: Number of matmul tiles per task group + * - add_batch: Number of add tiles per task group + * + * Task pattern: interleaved [matmul_0, add_0, matmul_1, add_1, ...] + * All tasks are completely independent (no dependencies). + * + * Arg layout: [A, B, C, X, Y, Z, batch, M_val, N_val, matmul_batch, add_batch] + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_MATMUL 0 +#define FUNC_ADD 1 + +static constexpr uint64_t MATMUL_ELEMS = 128 * 128; +static constexpr uint64_t ADD_ELEMS = 128 * 128; + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 11, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Tensor args + const Tensor &ext_A = orch_args.tensor(0).ref(); + const Tensor &ext_B = orch_args.tensor(1).ref(); + const Tensor &ext_C = orch_args.tensor(2).ref(); + const Tensor &ext_X = orch_args.tensor(3).ref(); + const Tensor &ext_Y = orch_args.tensor(4).ref(); + const Tensor &ext_Z = orch_args.tensor(5).ref(); + + // Scalar config args + int batch = static_cast(orch_args.scalar(0)); + int M = static_cast(orch_args.scalar(1)); + int N = static_cast(orch_args.scalar(2)); + int matmul_batch = static_cast(orch_args.scalar(3)); + int add_batch = static_cast(orch_args.scalar(4)); + + LOG_INFO_V0( + "[alternating_orch] Batch: %d, M: %d, N: %d, matmul_batch: %d, add_batch: %d", batch, M, N, matmul_batch, + add_batch + ); + + int total_matmul_tasks = batch * M; + int total_add_tasks = batch * N; + if (matmul_batch <= 0 || add_batch <= 0) { + LOG_ERROR( + "[alternating_orch] batch sizes must be positive (matmul_batch=%d, add_batch=%d)", matmul_batch, add_batch + ); + return; + } + if (total_matmul_tasks % matmul_batch != 0 || total_add_tasks % add_batch != 0) { + LOG_ERROR( + "[alternating_orch] task counts must divide evenly (matmul %d/%d, add %d/%d)", total_matmul_tasks, + matmul_batch, total_add_tasks, add_batch + ); + return; + } + int num_matmul_groups = total_matmul_tasks / matmul_batch; + int num_add_groups = total_add_tasks / add_batch; + + int total_matmul = 0; + int total_add = 0; + + int max_groups = num_matmul_groups > num_add_groups ? num_matmul_groups : num_add_groups; + + // Interleaved submit: matmul and add groups alternate + for (int group_idx = 0; group_idx < max_groups; group_idx++) { + if (group_idx < num_matmul_groups) { + int start_task_idx = group_idx * matmul_batch; + uint64_t offset = static_cast(start_task_idx) * MATMUL_ELEMS; + uint64_t group_size = static_cast(matmul_batch) * MATMUL_ELEMS; + + uint32_t matmul_group_shapes[1] = {static_cast(group_size)}; + uint32_t view_offsets[1] = {static_cast(offset)}; + + Tensor A_view = ext_A.view(matmul_group_shapes, view_offsets); + Tensor B_view = ext_B.view(matmul_group_shapes, view_offsets); + Tensor C_view = ext_C.view(matmul_group_shapes, view_offsets); + + L0TaskArgs params_matmul; + params_matmul.add_input(A_view); + params_matmul.add_input(B_view); + params_matmul.add_output(C_view); + rt_submit_aic_task(FUNC_MATMUL, params_matmul); + total_matmul++; + } + + if (group_idx < num_add_groups) { + int start_task_idx = group_idx * add_batch; + uint64_t offset = static_cast(start_task_idx) * ADD_ELEMS; + uint64_t group_size = static_cast(add_batch) * ADD_ELEMS; + + uint32_t add_group_shapes[1] = {static_cast(group_size)}; + uint32_t view_offsets[1] = {static_cast(offset)}; + + Tensor X_view = ext_X.view(add_group_shapes, view_offsets); + Tensor Y_view = ext_Y.view(add_group_shapes, view_offsets); + Tensor Z_view = ext_Z.view(add_group_shapes, view_offsets); + + L0TaskArgs params_add; + params_add.add_input(X_view); + params_add.add_input(Y_view); + params_add.add_output(Z_view); + rt_submit_aiv_task(FUNC_ADD, params_add); + total_add++; + } + } + + LOG_INFO_V9("[alternating_orch] Submitted %d matmul groups and %d add groups", total_matmul, total_add); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py new file mode 100644 index 0000000000..9fc261b4d8 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Alternating matmul + add: interleaved AIC (matmul 128x128) and AIV (add 128x128) tasks. + +Tests AIC+AIV mixed execution with scalar parameters and batched task submission. +C[b,m] = A[b,m] @ B[b,m], Z[b,n] = X[b,n] + Y[b,n]. +""" + +import ctypes + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAlternatingMatmulAdd(SceneTestCase): + """Alternating matmul + add with scalar parameters.""" + + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/alternating_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/kernel_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 1, "M": 1, "N": 1, "matmul_batch": 1, "add_batch": 1}, + }, + { + "name": "Case1", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 500, "M": 4, "N": 4, "matmul_batch": 4, "add_batch": 4}, + "manual": True, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"batch": 512, "M": 2, "N": 5, "matmul_batch": 4, "add_batch": 5}, + "manual": True, + }, + ] + + def generate_args(self, params): + batch = params["batch"] + M = params["M"] + N = params["N"] + matmul_batch = params.get("matmul_batch", 1) + add_batch = params.get("add_batch", 1) + matmul_size = 128 + add_rows = 128 + add_cols = 128 + + torch.manual_seed(42) + A = torch.randn(batch, M, matmul_size, matmul_size, dtype=torch.float32) * 0.01 + B = torch.randn(batch, M, matmul_size, matmul_size, dtype=torch.float32) * 0.01 + C = torch.zeros(batch, M, matmul_size, matmul_size, dtype=torch.float32) + X = torch.randn(batch, N, add_rows, add_cols, dtype=torch.float32) * 0.01 + Y = torch.randn(batch, N, add_rows, add_cols, dtype=torch.float32) * 0.01 + Z = torch.zeros(batch, N, add_rows, add_cols, dtype=torch.float32) + + return TaskArgsBuilder( + Tensor("A", A.flatten()), + Tensor("B", B.flatten()), + Tensor("C", C.flatten()), + Tensor("X", X.flatten()), + Tensor("Y", Y.flatten()), + Tensor("Z", Z.flatten()), + Scalar("batch", ctypes.c_int64(batch)), + Scalar("M_val", ctypes.c_int64(M)), + Scalar("N_val", ctypes.c_int64(N)), + Scalar("matmul_batch", ctypes.c_int64(matmul_batch)), + Scalar("add_batch", ctypes.c_int64(add_batch)), + ) + + def compute_golden(self, args, params): + batch = params["batch"] + M = params["M"] + N = params["N"] + matmul_size = 128 + add_rows = 128 + add_cols = 128 + + A = args.A.reshape(batch, M, matmul_size, matmul_size) + B = args.B.reshape(batch, M, matmul_size, matmul_size) + C = args.C.reshape(batch, M, matmul_size, matmul_size) + X = args.X.reshape(batch, N, add_rows, add_cols) + Y = args.Y.reshape(batch, N, add_rows, add_cols) + Z = args.Z.reshape(batch, N, add_rows, add_cols) + + for b in range(batch): + for m in range(M): + C[b, m] = torch.matmul(A[b, m], B[b, m]) + for b in range(batch): + for n in range(N): + Z[b, n] = X[b, n] + Y[b, n] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..5dcd577ea6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,137 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched PV Matmul Kernel: for each batch b, pij(M, K) @ vj(K, N) -> oi_new(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// Template: M=q_tile, K=block_size, N=head_dim + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void pv_matmul_batch_impl( + __gm__ Tensor *pij_batch, __gm__ Tensor *value_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *oi_new_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t block_num, uint64_t batch_start +) { + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ bfloat16_t *val_base = reinterpret_cast<__gm__ bfloat16_t *>(value_cache->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + using GlobalOut = GlobalTensor, pto::Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *pij_addr = pij_base + b * M * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *vj_addr = val_base + static_cast(phys_block) * K * N; + __gm__ float *oi_addr = oi_base + b * M * N; + + GlobalA pijGlobal(pij_addr); + GlobalB vjGlobal(vj_addr); + GlobalOut oiGlobal(oi_addr); + + TLOAD(aMatTile, pijGlobal); + TLOAD(bMatTile, vjGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(oiGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *value_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t block_num = static_cast(args[6]); + uint64_t batch_start = static_cast(args[7]); + + uint64_t q_tile_size = static_cast(pij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + pv_matmul_batch_impl<16, 16, 16>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else if (q_tile_size == 16) { + pv_matmul_batch_impl<16, 128, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else { + pv_matmul_batch_impl<64, 64, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..64f4147cc3 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,144 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched QK Matmul Kernel: for each batch b, qi(M, K) @ kj.T(K, N) -> sij(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// Template: M=q_tile, K=head_dim, N=block_size + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void qk_matmul_batch_impl( + __gm__ Tensor *query, __gm__ Tensor *key_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *sij_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t q_offset, uint64_t block_num, uint64_t num_heads, + uint64_t batch_start +) { + __gm__ bfloat16_t *query_base = reinterpret_cast<__gm__ bfloat16_t *>(query->buffer.addr); + __gm__ bfloat16_t *key_base = reinterpret_cast<__gm__ bfloat16_t *>(key_cache->buffer.addr); + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride, Layout::DN>; + using GlobalOut = GlobalTensor, pto::Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *qi_addr = query_base + ((batch_start + b) * num_heads + q_offset) * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *kj_addr = key_base + static_cast(phys_block) * N * K; + __gm__ float *sij_addr = sij_base + b * M * N; + + GlobalA qiGlobal(qi_addr); + GlobalB kjGlobal(kj_addr); + GlobalOut sijGlobal(sij_addr); + + TLOAD(aMatTile, qiGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TLOAD(bMatTile, kjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *query = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t q_offset = static_cast(args[6]); + uint64_t block_num = static_cast(args[7]); + uint64_t num_heads = static_cast(args[8]); + uint64_t batch_start = static_cast(args[9]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(sij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + qk_matmul_batch_impl<16, 16, 16>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else if (q_tile_size == 16) { + qk_matmul_batch_impl<16, 128, 128>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else { + qk_matmul_batch_impl<64, 128, 64>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..f55220b1cc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,221 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +// Batched Online Softmax Update + Normalize Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b, updates accumulators mi/li/oi with new block's mij/lij/oi_new. +// On is_last, normalizes and writes to the output tensor at the correct batch offset. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, head_dim=128 +// Case2: (64, 128) -- q_tile=64, head_dim=128 +// +// Scalar layout strategy: +// M scalar floats stored contiguously in GM can be loaded as either: +// - ND (kScalarRows, kScalarCols) RowMajor for element-wise ops +// - DN (kAlignedRows, 1) ColMajor for row-broadcast ops +// Conversion between layouts uses TRESHAPE (UB-internal, zero GM access). + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_batch_impl( + __gm__ Tensor *mij_batch, __gm__ Tensor *lij_batch, __gm__ Tensor *oi_new_batch, __gm__ Tensor *mi_batch, + __gm__ Tensor *li_batch, __gm__ Tensor *oi_batch, __gm__ Tensor *out, uint64_t is_first, uint64_t is_last, + uint64_t batch_count, uint64_t q_offset, uint64_t num_heads, uint64_t batch_start +) { + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ float *oi_new_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ float *mi_base = reinterpret_cast<__gm__ float *>(mi_batch->buffer.addr); + __gm__ float *li_base = reinterpret_cast<__gm__ float *>(li_batch->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_batch->buffer.addr); + __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out->buffer.addr); + + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + using TileDataMxN = Tile; + using TileScalarND = + Tile; + using TileScalarDN = Tile; + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarNDBytes = kScalarRows * kScalarCols * sizeof(float); + + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + TileScalarND mijND, lijND, miND, liND; + TileScalarND miNewND, alphaND, betaND, tmpND; + + TileScalarDN alphaDN, betaDN, liDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijND, 2 * kDataBytes); + TASSIGN(lijND, 2 * kDataBytes + kScalarNDBytes); + TASSIGN(miND, 2 * kDataBytes + 2 * kScalarNDBytes); + TASSIGN(liND, 2 * kDataBytes + 3 * kScalarNDBytes); + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarNDBytes); + TASSIGN(alphaND, 2 * kDataBytes + 5 * kScalarNDBytes); + TASSIGN(betaND, 2 * kDataBytes + 6 * kScalarNDBytes); + TASSIGN(tmpND, 2 * kDataBytes + 7 * kScalarNDBytes); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ float *mij_ptr = mij_base + b * M; + __gm__ float *lij_ptr = lij_base + b * M; + __gm__ float *oi_new_ptr = oi_new_base + b * M * N; + __gm__ float *mi_ptr = mi_base + b * M; + __gm__ float *li_ptr = li_base + b * M; + __gm__ float *oi_ptr = oi_base + b * M * N; + __gm__ float *dst_ptr = out_base + ((batch_start + b) * num_heads + q_offset) * N; + + GlobalDataMxN oiNewGlobal(oi_new_ptr); + GlobalDataMxN oiGlobal(oi_ptr); + GlobalDataMxN dstGlobal(dst_ptr); + + GlobalScalarND mijGlobalND(mij_ptr); + GlobalScalarND lijGlobalND(lij_ptr); + GlobalScalarND miGlobalND(mi_ptr); + GlobalScalarND liGlobalND(li_ptr); + + if (is_first) { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); + TSTORE(liGlobalND, lijND); + TSTORE(oiGlobal, oiNewTile); + + if (is_last) { + TRESHAPE(liDN, lijND); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TROWEXPANDDIV(oiNewTile, oiNewTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + TLOAD(miND, miGlobalND); + TLOAD(liND, liGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TMAX(miNewND, miND, mijND); + TSUB(alphaND, miND, miNewND); + TSUB(betaND, mijND, miNewND); + TEXP(alphaND, alphaND); + TEXP(betaND, betaND); + TMUL(liND, alphaND, liND); + TMUL(tmpND, betaND, lijND); + TADD(liND, liND, tmpND); + + TRESHAPE(alphaDN, alphaND); + TRESHAPE(betaDN, betaND); + if (is_last) { + TRESHAPE(liDN, liND); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); + TSTORE(liGlobalND, liND); + + TROWEXPANDMUL(oiTile, oiTile, alphaDN); + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); + TADD(oiTile, oiTile, oiNewTile); + + if (is_last) { + TROWEXPANDDIV(oiTile, oiTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiTile); + } else { + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(oiGlobal, oiTile); + } + } + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi_batch = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *out = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t batch_count = static_cast(args[9]); + uint64_t q_offset = static_cast(args[10]); + uint64_t num_heads = static_cast(args[11]); + uint64_t batch_start = static_cast(args[12]); + + uint64_t q_tile_size = static_cast(mij_batch->shapes[0] / batch_count); + uint64_t head_dim = static_cast(oi_new_batch->shapes[1]); + + if (q_tile_size == 16 && head_dim <= 16) { + online_update_batch_impl<16, 16>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else if (q_tile_size == 16) { + online_update_batch_impl<16, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else { + online_update_batch_impl<64, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..0dd0b24256 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,191 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched Softmax Preparation Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b at block_idx bn: +// valid_len = min(N, context_lens[b] - bn * N) +// sij_masked = pad(sij[b], valid_len, -inf) +// sij_scale = sij_masked * scale +// mij[b] = row_max(sij_scale) +// pij[b] = exp(sij_scale - mij[b]) (truncated to bf16 then back) +// lij[b] = row_sum(pij[b]) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, block_size=128 +// Case2: (64, 64) -- q_tile=64, block_size=64 + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void softmax_prepare_batch_impl( + __gm__ Tensor *sij_batch, __gm__ Tensor *context_lens_t, __gm__ Tensor *pij_batch, __gm__ Tensor *mij_batch, + __gm__ Tensor *lij_batch, float scale_value, uint64_t batch_count, uint64_t block_idx, uint64_t batch_start +) { + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ int32_t *ctx_lens = reinterpret_cast<__gm__ int32_t *>(context_lens_t->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, pto::Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + + using TileSijDyn = Tile; + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + TASSIGN(sijTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + for (uint64_t b = 0; b < batch_count; b++) { + int32_t cur_seq = ctx_lens[batch_start + b]; + uint64_t start = block_idx * N; + uint64_t valid_len = 0; + if (start < static_cast(cur_seq)) { + uint64_t remaining = static_cast(cur_seq) - start; + valid_len = (remaining < N) ? remaining : N; + } + + __gm__ float *sij_addr = sij_base + b * M * N; + __gm__ bfloat16_t *pij_addr = pij_base + b * M * N; + __gm__ float *mij_addr = mij_base + b * M; + __gm__ float *lij_addr = lij_base + b * M; + + GlobalDataMxN sijGlobal(sij_addr); + GlobalDataMxN_bf16 pijGlobal(pij_addr); + GlobalScalarDN mijGlobal(mij_addr); + GlobalScalarDN lijGlobal(lij_addr); + + if (valid_len == 0) { + // Block entirely beyond sequence: write mij=-1e30, lij=0, pij=0 + // Use -1e30 instead of -inf to avoid NaN in online_update (exp(-inf - (-inf)) = NaN) + constexpr float NEG_LARGE = -1e30f; + for (int i = 0; i < kAlignedRows; i++) { + maxTile.SetValue(i, NEG_LARGE); + sumTile.SetValue(i, 0.0f); + } + for (int i = 0; i < M * N; i++) { + pijBf16Tile.SetValue(i, static_cast(0.0f)); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(mijGlobal, maxTile); + TSTORE(lijGlobal, sumTile); + TSTORE(pijGlobal, pijBf16Tile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + continue; + } + + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TileSijDyn sijDynTile(static_cast(valid_len)); + TASSIGN(sijDynTile, 0x0); + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + + TMULS(sijTile, sijTile, scale_value); + TROWMAX(maxTile, sijTile, tmpTile); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first, then compute lij from truncated values (matches golden) + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + TROWSUM(sumTile, pijTile, tmpTile); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + TSTORE(mijGlobal, maxTile); + TSTORE(lijGlobal, sumTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *context_lens_t = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[5]); + float scale_value = scale_conv.f; + uint64_t batch_count = static_cast(args[6]); + uint64_t block_idx = static_cast(args[7]); + uint64_t batch_start = static_cast(args[8]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + softmax_prepare_batch_impl<16, 16>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else if (q_tile_size == 16) { + softmax_prepare_batch_impl<16, 128>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else { + softmax_prepare_batch_impl<64, 64>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..1717ebc48a --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,215 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Batch Paged Attention Orchestration Function - Production Scale + * + * Chunked batched architecture: the full batch is split into chunks of + * IN_CORE_BATCH size. Each chunk's QK/SF/PV/UP tasks are independent + * and can be scheduled to different cores in parallel. + * + * Task count = num_chunks * (1 + max_bn * 4), where + * num_chunks = ceil(batch / IN_CORE_BATCH) + * + * For batch <= IN_CORE_BATCH, behavior is identical to the non-chunked version. + * + * Memory Layout: + * Query: (batch * num_heads, head_dim) bf16 + * Key: (total_blocks, block_size, head_dim) bf16 (stored as K^T for QK) + * Value: (total_blocks, block_size, head_dim) bf16 + * + * Per-chunk intermediate tensors (contiguous across chunk_bc dimension): + * sij: (chunk_bc * q_tile, block_size) fp32 + * pij: (chunk_bc * q_tile, block_size) bf16 + * mij/lij: (chunk_bc * q_tile) fp32 + * oi_new: (chunk_bc * q_tile, head_dim) fp32 + * oi: (chunk_bc * q_tile, head_dim) fp32 accumulator + * mi/li: (chunk_bc * q_tile) fp32 accumulator + * + * Kernels receive global tensors + scalar metadata (including batch_start) + * and compute per-batch addresses internally. + */ + +#include +#include + +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + uint64_t elem_size = get_element_size(data_type); + + LOG_INFO_V0("batch_paged_attention: batch=%" PRIu64 ", num_heads=%" PRIu64, batch, num_heads); + + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + uint64_t max_bn = 0; + for (uint64_t b = 0; b < batch; b++) { + uint32_t cl_idx[1] = {static_cast(b)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_b = (cur_seq + block_size - 1) / block_size; + if (bn_b > max_bn) max_bn = bn_b; + } + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + uint64_t kv_total_rows = total_blocks_count * block_size; + uint32_t key_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t value_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32, true); + + constexpr uint64_t IN_CORE_BATCH = 16; + uint64_t num_chunks = (batch + IN_CORE_BATCH - 1) / IN_CORE_BATCH; + + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + uint64_t q_offset = q_idx * q_tile; + + for (uint64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + uint64_t chunk_bc = batch - chunk_idx * IN_CORE_BATCH; + if (chunk_bc > IN_CORE_BATCH) chunk_bc = IN_CORE_BATCH; + uint64_t batch_start = chunk_idx * IN_CORE_BATCH; + + PTO2_SCOPE() { + uint32_t oi_acc_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(head_dim)}; + uint32_t scalar_acc_shapes[1] = {static_cast(chunk_bc * q_tile)}; + TensorCreateInfo oi_batch_ci(oi_acc_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_acc_ci(scalar_acc_shapes, 1, DataType::FLOAT32); + TaskOutputTensors alloc_outs = alloc_tensors(oi_batch_ci, scalar_acc_ci, scalar_acc_ci); + const Tensor &oi_batch = alloc_outs.get_ref(0); + const Tensor &li_batch = alloc_outs.get_ref(1); + const Tensor &mi_batch = alloc_outs.get_ref(2); + + // Inner-loop create infos: shapes are loop-invariant, hoist out of bn loop + uint32_t sij_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(block_size)}; + uint32_t vec_shapes[1] = {static_cast(chunk_bc * q_tile)}; + uint32_t oi_new_shapes[2] = {static_cast(chunk_bc * q_tile), static_cast(head_dim)}; + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_ci(sij_shapes, 2, data_type); + TensorCreateInfo vec_ci(vec_shapes, 1, DataType::FLOAT32); + TensorCreateInfo oi_new_ci(oi_new_shapes, 2, DataType::FLOAT32); + + for (uint64_t bn = 0; bn < max_bn; bn++) { + PTO2_SCOPE() { + L0TaskArgs params_qk; + params_qk.add_input(query); + params_qk.add_input(key_cache); + params_qk.add_input(block_table); + params_qk.add_output(sij_ci); + params_qk.add_scalar(chunk_bc); + params_qk.add_scalar(bn); + params_qk.add_scalar(q_offset); + params_qk.add_scalar(block_num); + params_qk.add_scalar(num_heads); + params_qk.add_scalar(batch_start); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij_b = qk_outs.get_ref(0); + + L0TaskArgs params_sf; + params_sf.add_input(sij_b); + params_sf.add_input(context_lens); + params_sf.add_output(pij_ci); + params_sf.add_output(vec_ci); + params_sf.add_output(vec_ci); + params_sf.add_scalar(scale_value); + params_sf.add_scalar(chunk_bc); + params_sf.add_scalar(bn); + params_sf.add_scalar(batch_start); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_b = sf_outs.get_ref(0); + const Tensor &mij_b = sf_outs.get_ref(1); + const Tensor &lij_b = sf_outs.get_ref(2); + + L0TaskArgs params_pv; + params_pv.add_input(pij_b); + params_pv.add_input(value_cache); + params_pv.add_input(block_table); + params_pv.add_output(oi_new_ci); + params_pv.add_scalar(chunk_bc); + params_pv.add_scalar(bn); + params_pv.add_scalar(block_num); + params_pv.add_scalar(batch_start); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_new_b = pv_outs.get_ref(0); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == max_bn - 1) ? 1 : 0; + L0TaskArgs params_up; + params_up.add_input(mij_b); + params_up.add_input(lij_b); + params_up.add_input(oi_new_b); + params_up.add_inout(mi_batch); + params_up.add_inout(li_batch); + params_up.add_inout(oi_batch); + params_up.add_inout(out); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + params_up.add_scalar(chunk_bc); + params_up.add_scalar(q_offset); + params_up.add_scalar(num_heads); + params_up.add_scalar(batch_start); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + } + } + } + } + } + + LOG_INFO_V0( + "batch_paged_attention: %" PRIu64 " tasks (batch=%" PRIu64 ", max_bn=%" PRIu64 ", chunks=%" PRIu64 + ", IN_CORE_BATCH=%" PRIu64 ")", + static_cast(num_chunks * (1 + max_bn * 4)), batch, max_bn, num_chunks, IN_CORE_BATCH + ); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py new file mode 100644 index 0000000000..e7392531cb --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Batch paged attention: batched online softmax with AIC/AIV subgraph splitting (bfloat16).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestBatchPagedAttention(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case3", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 256, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 31, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall3", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py b/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py new file mode 100644 index 0000000000..e5ba8404e7 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""End-to-end ST for post-start Worker.register(ChipCallable) at L3. + +Exercises the _CTRL_REGISTER IPC path end-to-end: parent stages a +ChipCallable in shared memory after child startup, broadcasts CTRL_REGISTER to +every chip child, the child mmaps + prepares, and the resulting +CallableHandle is indistinguishable from a pre-start preparation when used +in run(). + +The UT suite (tests/ut/py/test_worker/test_host_worker.py) already covers +the facade-level paths (lock guard, capacity overflow, lambda rejection, run +race detection, shm name generator). This file's job is to prove the +bytes actually traverse shm to the chip child and prepare succeeds — +which only a real (sim or device) chip child can confirm. +""" + +import os + +import pytest +import torch +from _task_interface import MAX_REGISTERED_CALLABLE_IDS # pyright: ignore[reportMissingImports] +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import CallConfig, ChipCallable +from simpler.worker import Worker + +from simpler_setup import TaskArgsBuilder, Tensor +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.scene_test import _build_l3_task_args + +_RUNTIME = "tensormap_and_ringbuffer" +_HERE = os.path.dirname(os.path.abspath(__file__)) +_KERNELS = os.path.join( + _HERE, + "..", + "..", + "..", + "..", + "..", + "examples", + "a5", + "tensormap_and_ringbuffer", + "vector_example", + "kernels", +) +_ORCH_SRC = os.path.join(_KERNELS, "orchestration", "example_orchestration.cpp") +_AIV_ADD = os.path.join(_KERNELS, "aiv", "kernel_add.cpp") +_AIV_ADD_SCALAR = os.path.join(_KERNELS, "aiv", "kernel_add_scalar.cpp") +_AIV_MUL = os.path.join(_KERNELS, "aiv", "kernel_mul.cpp") + +_ORCH_SIG = [D.IN, D.IN, D.OUT] + + +def _build_vector_callable(platform: str, *, extra_unused_child: bool = False) -> ChipCallable: + """Compile the vector_example orchestration + 3 AIV kernels. + + Mirrors how SceneTestCase._compile_chip_callable_from_spec assembles + a ChipCallable, but inline so the test can call register_callable() on it both + before and after init(). + """ + from simpler.task_interface import CoreCallable # noqa: PLC0415 + + from simpler_setup.elf_parser import extract_text_section # noqa: PLC0415 + from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: PLC0415 + + kc = KernelCompiler(platform=platform) + pto_isa_root = ensure_pto_isa_root() + inc_dirs = kc.get_orchestration_include_dirs(_RUNTIME) + + orch_bytes = kc.compile_orchestration(runtime_name=_RUNTIME, source_path=_ORCH_SRC) + + def _aiv(path: str) -> bytes: + raw = kc.compile_incore(path, core_type="aiv", pto_isa_root=pto_isa_root, extra_include_dirs=inc_dirs) + return raw if platform.endswith("sim") else extract_text_section(raw) + + add = CoreCallable.build(signature=[D.IN, D.IN, D.OUT], binary=_aiv(_AIV_ADD)) + add_scalar = CoreCallable.build(signature=[D.IN, D.OUT], binary=_aiv(_AIV_ADD_SCALAR)) + mul = CoreCallable.build(signature=[D.IN, D.IN, D.OUT], binary=_aiv(_AIV_MUL)) + + children = [(0, add), (1, add_scalar), (2, mul)] + if extra_unused_child: + children.append((99, add)) + + return ChipCallable.build( + signature=_ORCH_SIG, + func_name="aicpu_orchestration_entry", + binary=orch_bytes, + children=children, + ) + + +def _unique_py_callable(index: int): + def fn(args, _index=index): + return _index + + return fn + + +def _make_args(a: float, b: float) -> TaskArgsBuilder: + size = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((size,), a, dtype=torch.float32).share_memory_()), + Tensor("b", torch.full((size,), b, dtype=torch.float32).share_memory_()), + Tensor("f", torch.zeros(size, dtype=torch.float32).share_memory_()), + ) + + +def _golden(a: float, b: float) -> float: + # Matches the orchestration: f = (a+b+1) * (a+b+2) + (a+b) + s = a + b + return (s + 1) * (s + 2) + s + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_new_identity_after_start_then_run(st_platform, st_device_ids): + """Happy path: prepare one identity pre-start and another post-start. + + Proves the post-start control path delivers a usable handle for a + previously unseen hashid. Both identities execute equivalent kernels and + must produce numerically identical outputs. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + # Pre-allocate both runs' tensors BEFORE Worker.init() so the + # share_memory_() mappings are inherited by the forked chip child. + # share_memory_ regions created after fork in the parent are not visible + # to the chip child, so dispatch on those would segfault. + a, b = 2.0, 3.0 + expected = _golden(a, b) + args_pre = _make_args(a, b) + args_post = _make_args(a, b) + chip_args_pre, output_names_pre = _build_l3_task_args(args_pre, _ORCH_SIG) + chip_args_post, output_names_post = _build_l3_task_args(args_post, _ORCH_SIG) + assert output_names_pre == ["f"] and output_names_post == ["f"] + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + # 1. Run pre_handle once against the snapshot-prepared chip callable. + # init() already forked the chip children and prepared the startup + # snapshot, so they are in _run_chip_main_loop — the only state in + # which the CTRL_REGISTER broadcast in step 2 can be ACKed. + def orch_pre(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + got_pre = args_pre.f + assert torch.allclose(got_pre, torch.full_like(got_pre, expected), rtol=1e-5, atol=1e-5), ( + f"pre_handle={pre_handle.hashid}: expected {expected}, got {got_pre[:4].tolist()}..." + ) + + # 2. Now do the post-start dynamic prepare. The parent stages bytes + # in shm and broadcasts CTRL_REGISTER; the child mmaps and calls + # prepare_callable_from_blob. post_handle is unknown to the + # CoW-inherited registry on the child side — only the IPC path + # can deliver it. + post_handle = worker.register(post_callable) + assert post_handle.hashid != pre_handle.hashid + + # 3. Run with post_handle. If CTRL_REGISTER delivered correctly, the + # child has the identity prepared; otherwise dispatch will fail. + def orch_post(o, _args, _cfg): + o.submit_next_level(post_handle, chip_args_post, config, worker=0) + + worker.run(orch_post) + got_post = args_post.f + assert torch.allclose(got_post, torch.full_like(got_post, expected), rtol=1e-5, atol=1e-5), ( + f"post_handle={post_handle.hashid}: expected {expected}, got {got_post[:4].tolist()}..." + ) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(2) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_new_identity_after_start_parallel_broadcast(st_platform, st_device_ids): + """Two chip children, post-start prepare broadcasts to both in parallel. + + Asserts that the prepared handle runs successfully on each chip — proving + the C++ broadcast (one std::thread per WorkerThread) delivers the bytes + to every chip's mailbox and each prepare_callable_from_blob runs without + racing against the others. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + device_ids = [int(d) for d in st_device_ids[:2]] + worker = Worker( + level=3, + device_ids=device_ids, + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + a, b = 2.0, 3.0 + expected = _golden(a, b) + # Use one first-run bundle to start the control-capable child loop, then one + # post-start bundle per chip so the dynamically prepared handle is executed + # by both broadcast targets. + args_pre = _make_args(a, b) + args_post = [_make_args(a, b) for _ in device_ids] + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) + chip_args_post = [_build_l3_task_args(args, _ORCH_SIG)[0] for args in args_post] + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_pre(o, _a, _c): + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + assert torch.allclose(args_pre.f, torch.full_like(args_pre.f, expected), rtol=1e-5, atol=1e-5) + + # Now broadcast CTRL_REGISTER to BOTH chip mailboxes in parallel. + post_handle = worker.register(post_callable) + + def orch_post(o, _a, _c): + o.submit_next_level_group(post_handle, chip_args_post, config, workers=[0, 1]) + + worker.run(orch_post) + for args in args_post: + assert torch.allclose(args.f, torch.full_like(args.f, expected), rtol=1e-5, atol=1e-5) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_prepare_capacity_overflow_post_start(st_platform, st_device_ids): + """Saturate callable capacity pre-start, then verify post-start prepare hits + the same ``MAX_REGISTERED_CALLABLE_IDS`` ceiling for a new hashid. + + Confirms the public capacity guard is shared between pre-start preparation + and the post-start control path (and that the error message is + protocol-aware so the operator sees the same diagnostic in both paths). + """ + chip_callable = _build_vector_callable(st_platform) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + # A sub worker gives the LOCAL_PYTHON fillers below a valid resolver, so + # they exercise the shared capacity ceiling as eligible registrations + # (not inert ones) alongside the chip callable. + num_sub_workers=1, + platform=st_platform, + runtime=_RUNTIME, + ) + # Fill the registry pre-start with distinct sub fn identities (cheap, no + # device cost). + for i in range(MAX_REGISTERED_CALLABLE_IDS - 1): + worker.register(_unique_py_callable(i)) + chip_handle = worker.register(chip_callable) # final capacity entry + + a, b = 2.0, 3.0 + args_pre = _make_args(a, b) + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_pre(o, _a, _c): + o.submit_next_level(chip_handle, chip_args_pre, config, worker=0) + + worker.run(orch_pre) + + # The very next dynamic prepare of a new identity hits the capacity + # ceiling. Re-preparing ``chip_callable`` itself would only create + # another handle to the existing identity. + with pytest.raises(RuntimeError, match="MAX_REGISTERED_CALLABLE_IDS"): + worker.register(_build_vector_callable(st_platform, extra_unused_child=True)) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_duplicate_prepare_same_hashid_survives_one_unregister(st_platform, st_device_ids): + """prepare same hashid twice, unregister one handle, run the other. + + This is the hashid-specific post-start path: the second + ``register_callable(same_chip_callable)`` must return a distinct handle for + the same hashid. Unregistering one handle must only drop that public + handle; the remaining handle must still dispatch successfully. + """ + chip_callable = _build_vector_callable(st_platform) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + a, b = 2.0, 3.0 + expected = _golden(a, b) + # Two runs total — preallocate both args bundles BEFORE init() so + # the share_memory_ mappings are inherited by the forked chip child. + args_one = _make_args(a, b) + args_two = _make_args(a, b) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + # 1. Trigger fork via pre_handle to put the chip child into the main loop. + def orch_one(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) + + worker.run(orch_one) + assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) + + # 2. Prepare the same callable after start. This returns another + # public handle for the same hashid, not a new identity. + duplicate_handle = worker.register(chip_callable) + assert duplicate_handle.hashid == pre_handle.hashid + assert duplicate_handle.digest == pre_handle.digest + assert duplicate_handle._handle_id != pre_handle._handle_id + + # 3. Drop the first handle. The child must keep the prepared identity + # alive for duplicate_handle. + worker.unregister(pre_handle) + + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(pre_handle, chip_args_one, config, worker=0)) + + def orch_two(o, _args, _cfg): + o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0) + + worker.run(orch_two) + assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) + + # 4. Dropping the final handle invalidates it through the public API. + worker.unregister(duplicate_handle) + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0)) + finally: + worker.close() + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_unregister_last_handle_allows_reprepare_same_hashid(st_platform, st_device_ids): + """prepare → run → unregister final handle → prepare same identity again. + + Proves the IPC unregister path works end-to-end: after CTRL_UNREGISTER + propagates to the chip child, the old handle is invalid and a subsequent + post-start prepare of that identity materializes a usable handle again. + """ + chip_callable = _build_vector_callable(st_platform) + post_callable = _build_vector_callable(st_platform, extra_unused_child=True) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + pre_handle = worker.register(chip_callable) + + a, b = 2.0, 3.0 + expected = _golden(a, b) + args_one = _make_args(a, b) + args_two = _make_args(a, b) + args_three = _make_args(a, b) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) + chip_args_three, _ = _build_l3_task_args(args_three, _ORCH_SIG) + + worker.init() + try: + config = CallConfig() + config.block_dim = 3 + config.aicpu_thread_num = 4 + + def orch_one(o, _args, _cfg): + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) + + worker.run(orch_one) + assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) + + dyn_handle = worker.register(post_callable) + + def orch_two(o, _args, _cfg): + o.submit_next_level(dyn_handle, chip_args_two, config, worker=0) + + worker.run(orch_two) + assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) + + worker.unregister(dyn_handle) + with pytest.raises(KeyError, match="not live"): + worker.run(lambda o, _args, _cfg: o.submit_next_level(dyn_handle, chip_args_two, config, worker=0)) + + # Re-prepare the same hashid after its final handle was dropped. + again_handle = worker.register(post_callable) + assert again_handle.hashid == dyn_handle.hashid + assert again_handle.digest == dyn_handle.digest + assert again_handle._handle_id != dyn_handle._handle_id + + def orch_three(o, _args, _cfg): + o.submit_next_level(again_handle, chip_args_three, config, worker=0) + + worker.run(orch_three) + assert torch.allclose(args_three.f, torch.full_like(args_three.f, expected), rtol=1e-5, atol=1e-5) + finally: + worker.close() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp new file mode 100644 index 0000000000..543cb83618 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp @@ -0,0 +1,55 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + // Keep the swimlane bars visible. Lookup-only timing uses dummy tasks, so + // this spin does not affect the fanin lookup cost measurement. + volatile uint32_t spin = 0; + for (uint32_t i = 0; i < 4096; i++) { + spin += i; + } + + out[0] = 42.0f; + if (spin == 0xffffffffu) { + out[0] = 43.0f; + } + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp new file mode 100644 index 0000000000..566e213cde --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp @@ -0,0 +1,93 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Explicit 64x64 fanin DAG validation scene. + * + * Args layout: + * tensor[0]: producer outputs, split into disjoint producer slices + * tensor[1]: consumer outputs, split into disjoint consumer slices + * scalar[0]: producer_count + * scalar[1]: consumer_count + * scalar[2]: use_real_kernels + * + * The scene submits producer_count independent producers, then consumer_count + * independent consumers where every consumer explicitly depends on every + * producer. Real-kernel mode writes disjoint tensor slices so tensormap + * auto-deps do not add producer chains or consumer chains. + * + * When use_real_kernels is false, the same dependency shape is submitted with + * dummy tasks to isolate orchestrator fanin lookup cost. + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +static constexpr int32_t MAX_PRODUCERS = 64; +static constexpr int32_t MAX_CONSUMERS = 64; +static constexpr uint32_t SLOT_ELEMS = 16; + +#define FUNC_WRITE_CONST 0 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 5, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &producer_outputs = orch_args.tensor(0).ref(); + const Tensor &consumer_outputs = orch_args.tensor(1).ref(); + int32_t producer_count = static_cast(orch_args.scalar(0)); + int32_t consumer_count = static_cast(orch_args.scalar(1)); + bool use_real_kernels = orch_args.scalar(2) != 0; + if (producer_count < 1 || producer_count > MAX_PRODUCERS || consumer_count < 1 || consumer_count > MAX_CONSUMERS) { + rt_report_fatal( + PTO2_ERROR_INVALID_ARGS, + "producer_count=%d consumer_count=%d exceed supported range producers=[1, %d] consumers=[1, %d]", + producer_count, consumer_count, MAX_PRODUCERS, MAX_CONSUMERS + ); + return; + } + + PTO2TaskId producer_ids[MAX_PRODUCERS]; + uint32_t slot_shape[1] = {SLOT_ELEMS}; + for (int32_t i = 0; i < producer_count; i++) { + L0TaskArgs args; + if (use_real_kernels) { + uint32_t offset[1] = {static_cast(i) * SLOT_ELEMS}; + Tensor producer_out = producer_outputs.view(slot_shape, offset); + args.add_inout(producer_out); + producer_ids[i] = rt_submit_aic_task(FUNC_WRITE_CONST, args).task_id(); + } else { + producer_ids[i] = rt_submit_dummy_task(args).task_id(); + } + } + + for (int32_t c = 0; c < consumer_count; c++) { + L0TaskArgs args; + args.set_dependencies(producer_ids, static_cast(producer_count)); + if (use_real_kernels) { + uint32_t offset[1] = {static_cast(c) * SLOT_ELEMS}; + Tensor consumer_out = consumer_outputs.view(slot_shape, offset); + args.add_inout(consumer_out); + rt_submit_aic_task(FUNC_WRITE_CONST, args); + } else { + rt_submit_dummy_task(args); + } + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py new file mode 100644 index 0000000000..155894272a --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""fanin_lookup_perf: validate a 64x64 explicit fanin DAG. + +The orchestration submits 64 independent producers and 64 independent +consumers. Each consumer explicitly depends on all 64 producers. The real +kernel case uses disjoint tensor slices so tensormap auto-deps do not add +producer chains or consumer chains. +""" + +import ctypes + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestFaninLookupPerf(SceneTestCase): + """Validate a wide 64-producer/64-consumer explicit fanin DAG.""" + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/fanin_lookup_perf_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "WRITE_CONST", + "source": "kernels/aic/kernel_write_const_visible.cpp", + "core_type": "aic", + # Single-AIC task with one INOUT tensor at payload slot 0. + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "LookupOnlyProducers64Consumers64", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"producer_count": 64, "consumer_count": 64, "use_real_kernels": 0}, + }, + { + "name": "SwimlaneProducers64Consumers64", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"producer_count": 64, "consumer_count": 64, "use_real_kernels": 1}, + }, + ] + + def generate_args(self, params): + slot_elems = 16 + producer_count = int(params["producer_count"]) + consumer_count = int(params["consumer_count"]) + return TaskArgsBuilder( + Tensor("producer_out", torch.full((producer_count * slot_elems,), -1.0, dtype=torch.float32)), + Tensor("consumer_out", torch.full((consumer_count * slot_elems,), -1.0, dtype=torch.float32)), + Scalar("producer_count", ctypes.c_int64(producer_count)), + Scalar("consumer_count", ctypes.c_int64(consumer_count)), + Scalar("use_real_kernels", ctypes.c_int64(int(params["use_real_kernels"]))), + ) + + def compute_golden(self, args, params): + if not params["use_real_kernels"]: + return + slot_elems = 16 + for i in range(int(params["producer_count"])): + args.producer_out[i * slot_elems] = 42.0 + for c in range(int(params["consumer_count"])): + args.consumer_out[c * slot_elems] = 42.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py b/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py new file mode 100644 index 0000000000..28a54467fc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Multi-round paged attention: benchmark multi-round execution (default 10 rounds). + +Run with --rounds 10 --skip-golden for benchmarking. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + +_PA_KERNELS = "../../../../../examples/a5/tensormap_and_ringbuffer/paged_attention/kernels" + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestMultiRoundPagedAttention(SceneTestCase): + RTOL = 1e-2 + ATOL = 1e-2 + + CALLABLE = { + "orchestration": { + "source": f"{_PA_KERNELS}/orchestration/paged_attention_orch.cpp", + "function_name": "build_paged_attention_graph", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": f"{_PA_KERNELS}/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": f"{_PA_KERNELS}/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": f"{_PA_KERNELS}/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": f"{_PA_KERNELS}/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a5sim", "a5"], + "manual": True, + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..ed14f900b6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,171 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// SplitK PV Matmul Kernel: Accumulated P @ V across n_blocks +// +// Processes n_blocks blocks using SplitK accumulation pattern: +// Block 0: TMATMUL(C, A, B) — initialize accumulator +// Block i: TMATMUL_ACC(C, C, A, B) — accumulate into same C +// +// Per-block pij addresses: contiguous slices of pij_buf (n_blocks * M * K) +// Per-block vj addresses: value_cache base + block_indices lookup +// Single output: oi_new (M, N) fp32 = sum of P_i @ V_i across all blocks +// +// Optimizations: +// - Double-buffered L1 tiles (ping/pong for A and B via MTE2) +// - Double-buffered L0 tiles (ping/pong for L0A and L0B via MTE1) +// - TLOAD(next) overlaps with TMATMUL(current) via MTE2/M-pipe parallelism +// - Canonical 3-stage pipeline: TLOAD(MTE2) → TMOV(MTE1) → TMATMUL(M) +// - Reverse-dependency events ensure buffer safety across iterations +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// pij is bfloat16 (from softmax_prepare TCVT). +// vj is stored as (K, N) = (block_size, head_dim) in row-major (ND) layout. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void pv_matmul_n_impl( + __gm__ Tensor *pij_buf, __gm__ Tensor *value_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *oi_new, + uint64_t n_blocks, uint64_t bt_offset +) { + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_buf->buffer.addr) + pij_buf->start_offset; + __gm__ bfloat16_t *val_base = reinterpret_cast<__gm__ bfloat16_t *>(value_cache->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr) + oi_new->start_offset; + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride<1, M * K, M * K, K, 1>>; + using GlobalB = GlobalTensor, pto::Stride>; + using GlobalOut = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + // L1 memory layout: double-buffered A and B tiles (tightly packed) + constexpr int kATileBytes = M * K * static_cast(sizeof(bfloat16_t)); + constexpr int kBTileBytes = K * N * static_cast(sizeof(bfloat16_t)); + + TileMatA aMatTile[2]; + TileMatB bMatTile[2]; + TASSIGN(aMatTile[0], 0x0); + TASSIGN(aMatTile[1], kATileBytes); + TASSIGN(bMatTile[0], 2 * kATileBytes); + TASSIGN(bMatTile[1], 2 * kATileBytes + kBTileBytes); + + // L0 memory layout: double-buffered L0A and L0B, single accumulator L0C + LeftTile aTile[2]; + RightTile bTile[2]; + AccTile cTile; + TASSIGN(aTile[0], 0x0); + TASSIGN(aTile[1], kATileBytes); + TASSIGN(bTile[0], 0x0); + TASSIGN(bTile[1], kBTileBytes); + TASSIGN(cTile, 0x0); + + GlobalOut oiGlobal(oi_base); + + // Seed reverse-dependency flags: all ping/pong buffers initially free + // PIPE_MTE1 → PIPE_MTE2: L1 buffer [0/1] safe for TLOAD to overwrite + // PIPE_M → PIPE_MTE1: L0 buffer [0/1] safe for TMOV to overwrite + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + + for (uint64_t i = 0; i < n_blocks; i++) { + int cur = static_cast(i % 2); + GlobalA pijGlobal(pij_base + i * M * K); + GlobalB vjGlobal(val_base + bt[bt_offset + i] * K * N); + + // Stage 1: TLOAD (MTE2: GM → L1[cur]) + // Wait for MTE1 to release L1[cur] (reverse dep from previous iteration) + wait_flag(PIPE_MTE1, PIPE_MTE2, static_cast<::event_t>(cur)); + TLOAD(aMatTile[cur], pijGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // forward: A in L1 ready + TLOAD(bMatTile[cur], vjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // forward: B in L1 ready + + // Stage 2: TMOV (MTE1: L1[cur] → L0[cur]) + // Wait for M-pipe to release L0[cur] (reverse dep from previous iteration) + wait_flag(PIPE_M, PIPE_MTE1, static_cast<::event_t>(cur)); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // forward: wait A loaded + TMOV(aTile[cur], aMatTile[cur]); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // forward: wait B loaded + TMOV(bTile[cur], bMatTile[cur]); + set_flag(PIPE_MTE1, PIPE_MTE2, static_cast<::event_t>(cur)); // reverse: release L1[cur] + + // Stage 3: TMATMUL (M-pipe: L0A[cur] × L0B[cur] → L0C) + set_flag(PIPE_MTE1, PIPE_M, static_cast<::event_t>(cur)); // forward: L0[cur] ready + wait_flag(PIPE_MTE1, PIPE_M, static_cast<::event_t>(cur)); + if (i == 0) { + TMATMUL(cTile, aTile[cur], bTile[cur]); + } else { + TMATMUL_ACC(cTile, cTile, aTile[cur], bTile[cur]); + } + set_flag(PIPE_M, PIPE_MTE1, static_cast<::event_t>(cur)); // reverse: release L0[cur] + } + + // Drain outstanding reverse-dependency flags + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE(oiGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij_buf = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *value_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t n_blocks = static_cast(args[4]); + uint64_t bt_offset = static_cast(args[5]); + + // pij_buf is 4D (1, 1, q_tile, n_blocks*block_size) to match qk's 4D output. + uint64_t q_tile_size = static_cast(pij_buf->shapes[2]); + + if (q_tile_size == 16) { + pv_matmul_n_impl<16, 128, 128>(pij_buf, value_cache, block_table_t, oi_new, n_blocks, bt_offset); + } else { + pv_matmul_n_impl<64, 64, 128>(pij_buf, value_cache, block_table_t, oi_new, n_blocks, bt_offset); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..02ccd5987f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,135 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Multi-block QK Matmul Kernel: qi(M, K) @ kj.T(K, N) -> sij(M, N) for each block +// +// Processes n_blocks blocks in a single kernel invocation. +// Per-block kj addresses computed from key_cache base + block_indices lookup. +// qi is shared across all blocks (same query head against different key blocks). +// +// Output layout: n_blocks contiguous (M, N) tiles stacked vertically. +// Block i occupies sij[i*M : (i+1)*M, 0:N]. +// +// Optimizations: +// - qi TLOAD hoisted before the loop (constant across all iterations) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// Template: M=q_tile, K=head_dim, N=block_size + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void qk_matmul_n_impl( + __gm__ Tensor *qi, __gm__ Tensor *key_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *sij_buf, + uint64_t n_blocks, uint64_t bt_offset +) { + // Decode 4D query view: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + __gm__ bfloat16_t *qi_base = reinterpret_cast<__gm__ bfloat16_t *>(qi->buffer.addr) + qi->start_offset; + __gm__ bfloat16_t *key_base = reinterpret_cast<__gm__ bfloat16_t *>(key_cache->buffer.addr); + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_buf->buffer.addr) + sij_buf->start_offset; + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, pto::Stride<1, M * K, M * K, K, 1>>; + using GlobalB = GlobalTensor, pto::Stride, Layout::DN>; + using GlobalOut = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Hoist qi TLOAD before the loop (qi is constant across all blocks) + GlobalA qiGlobal(qi_base); + TLOAD(aMatTile, qiGlobal); + + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalB kjGlobal(key_base + bt[bt_offset + i] * N * K); + GlobalOut sijGlobal(sij_base + i * M * N); + + // Load only B each iteration (qi already in L1 from hoist) + TLOAD(bMatTile, kjGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + // TMOV qi from L1→L0A (re-copy since TMATMUL consumed L0A) and kj from L1→L0B + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + if (i + 1 < n_blocks) { + pipe_barrier(PIPE_ALL); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *qi = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *sij_buf = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t n_blocks = static_cast(args[4]); + uint64_t bt_offset = static_cast(args[5]); + + // qi is a 4D view (batch, q_len, num_heads_tile, head_dim); decode fixes batch=q_len=1. + uint64_t q_tile_size = static_cast(qi->shapes[2]); + + if (q_tile_size == 16) { + qk_matmul_n_impl<16, 128, 128>(qi, key_cache, block_table_t, sij_buf, n_blocks, bt_offset); + } else { + qk_matmul_n_impl<64, 128, 64>(qi, key_cache, block_table_t, sij_buf, n_blocks, bt_offset); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..b3f51fc11c --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,251 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + // 4D data views (1, 1, q_tile, head_dim) — oi, oi_new, dst. + using GlobalData4D = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalData4D oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalData4D oiGlobal(oi_ptr + oi->start_offset); + GlobalData4D dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to same UB as DN tiles for ND-format store + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs as DN (ColMajor) + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + TMUL(tmpRow, alphaRow, liRow); // alpha * li + TMUL(liNewRow, betaRow, lijRow); // beta * lij + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + // mij is 3D (1, 1, q_tile) to match softmax's 3D scalar output. + uint64_t q_tile_size = static_cast(mij->shapes[2]); + + if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8c05ea8a00 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,265 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Two-Pass Softmax Kernel (AIV) for n_blocks tiles +// +// Input: sij_buf (n_blocks * M, N) fp32 — QK results stacked vertically +// Output: pij_buf (n_blocks * M, N) bf16 — attention weights per block +// mij (M,) fp32 — global row max across all blocks +// lij (M,) fp32 — total row sum across all blocks +// +// Pass 1: Iterate over n_blocks tiles, apply scale, mask last block, +// find global m = max over all blocks of rowmax(S_i * scale) +// Uses TRESHAPE for DN↔Row conversion to keep globalMax in UB +// (eliminates 63 × 4 GM round-trip operations). +// Pass 2: Iterate again, compute P_i = exp(S_i * scale - m) -> bf16, +// accumulate l = sum over all blocks of rowsum(P_i) +// Uses double-buffered sij tiles to overlap TLOAD with computation. +// +// Two-pass ensures all P_i tiles share the same scale (global max), +// enabling direct TMATMUL_ACC accumulation in the PV kernel. +// +// Supports two tile configurations via runtime dispatch: +// Case1: M=16, N=128 (q_tile=16, block_size=128) +// Case2: M=64, N=64 (q_tile=64, block_size=64) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_n_impl( + __gm__ Tensor *sij_buf, __gm__ Tensor *pij_buf, __gm__ Tensor *mij, __gm__ Tensor *lij, float scale_value, + uint64_t n_blocks, uint64_t valid_len_last +) { + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_buf->buffer.addr) + sij_buf->start_offset; + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_buf->buffer.addr) + pij_buf->start_offset; + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr) + mij->start_offset; + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr) + lij->start_offset; + + // Decode 4D semantic: batch/q_len are constexpr 1. + static constexpr int BATCH = 1; + static constexpr int Q_LEN = 1; + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + + // --- GlobalTensor types --- + // 4D data views (1, 1, q_tile, n_blocks*block_size) for sij/pij. + using GlobalDataMxN = GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + using GlobalDataMxN_bf16 = + GlobalTensor, pto::Stride<1, M * N, M * N, N, 1>>; + // DN/ND scalar globals stay 2D: scalar vectors only need per-element layout. + using GlobalScalarDN = GlobalTensor, pto::Stride<1, 1, 1, 1, 1>, Layout::DN>; + using GlobalScalarND = + GlobalTensor, pto::Stride<1, 1, 1, kScalarCols, 1>>; + + // --- Tile types --- + using TileSijDyn = Tile; + using TileSijPad = Tile; + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + using TileScalarND = + Tile; + // RowMajor (1, M) tile for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // --- UB memory layout (double-buffered sij) --- + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Double-buffered sij tiles + TileVecMxN sijTile_A; + TileSijPad sijPadTile_A; + TileVecMxN sijTile_B; + TileSijPad sijPadTile_B; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileVecMxN sumAccTile; + TileScalarDN localMaxDN; + TileScalarDN globalMaxDN; + TileScalarDN sumDN; + TileVecMxN_bf16 pijBf16Tile; + + // TRESHAPE aliases (same UB address as their DN counterparts) + TileScalarRow localMaxRow; + TileScalarRow globalMaxRow; + + // ND alias for storing globalMax to GM + TileScalarND globalMaxND; + + TASSIGN(sijTile_A, 0x0); + TASSIGN(sijPadTile_A, 0x0); + TASSIGN(sijTile_B, kDataBytes); + TASSIGN(sijPadTile_B, kDataBytes); + TASSIGN(pijTile, 2 * kDataBytes); + TASSIGN(tmpTile, 3 * kDataBytes); + TASSIGN(sumAccTile, 4 * kDataBytes); + int scalarBase = 5 * kDataBytes; + TASSIGN(localMaxDN, scalarBase); + TASSIGN(localMaxRow, scalarBase); // alias: same UB as localMaxDN + TASSIGN(globalMaxDN, scalarBase + kScalarDNBytes); + TASSIGN(globalMaxRow, scalarBase + kScalarDNBytes); // alias: same UB as globalMaxDN + TASSIGN(globalMaxND, scalarBase + kScalarDNBytes); // alias: same UB as globalMaxDN + TASSIGN(sumDN, scalarBase + 2 * kScalarDNBytes); + TASSIGN(pijBf16Tile, scalarBase + 3 * kScalarDNBytes); + + // GM aliases (mij/lij output buffers) + GlobalScalarND mijGlobalND(mij_addr); + GlobalScalarDN lijGlobalDN(lij_addr); + + // ======== Pass 1: Find global row max via TRESHAPE (no GM round-trip) ======== + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalDataMxN sijGlobal(sij_base + i * M * N); + TLOAD(sijTile_A, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if (i == n_blocks - 1 && valid_len_last < static_cast(N)) { + TileSijDyn sijDynTile(static_cast(valid_len_last)); + TASSIGN(sijDynTile, 0x0); + TFILLPAD_INPLACE(sijPadTile_A, sijDynTile); + } + + TMULS(sijTile_A, sijTile_A, scale_value); + TROWMAX(localMaxDN, sijTile_A, tmpTile); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise TMAX + TRESHAPE(localMaxRow, localMaxDN); + if (i == 0) { + TMAX(globalMaxRow, localMaxRow, localMaxRow); + } else { + TMAX(globalMaxRow, globalMaxRow, localMaxRow); + } + } + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for Pass 2's TROWEXPANDSUB + TRESHAPE(globalMaxDN, globalMaxRow); + + // Store final global max to mij for online_update to consume + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(mijGlobalND, globalMaxND); + + // ======== Pass 2: Compute softmax with double-buffered sij ======== + // globalMaxDN is already in UB from TRESHAPE — no reload needed. + // Sync MTE3→MTE2 to ensure the mij TSTORE completed before first sij TLOAD. + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + + // Pre-load first sij tile into buffer A + GlobalDataMxN sijGlobal_0(sij_base); + TLOAD(sijTile_A, sijGlobal_0); + + for (uint64_t i = 0; i < n_blocks; i++) { + GlobalDataMxN_bf16 pijGlobal(pij_base + i * M * N); + + // Wait for current tile's TLOAD to complete + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TFILLPAD on current buffer if last block with partial valid length + if (i == n_blocks - 1 && valid_len_last < static_cast(N)) { + TileSijDyn curSijDyn(static_cast(valid_len_last)); + if (i % 2 == 0) { + TASSIGN(curSijDyn, 0x0); + TFILLPAD_INPLACE(sijPadTile_A, curSijDyn); + } else { + TASSIGN(curSijDyn, static_cast(kDataBytes)); + TFILLPAD_INPLACE(sijPadTile_B, curSijDyn); + } + } + + // Compute on current buffer (select A or B based on iteration parity) + if (i % 2 == 0) { + TMULS(sijTile_A, sijTile_A, scale_value); + TROWEXPANDSUB(pijTile, sijTile_A, globalMaxDN); + } else { + TMULS(sijTile_B, sijTile_B, scale_value); + TROWEXPANDSUB(pijTile, sijTile_B, globalMaxDN); + } + TEXP(pijTile, pijTile); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + + if (i == 0) { + TMULS(sumAccTile, pijTile, 1.0f); + } else { + TADD(sumAccTile, sumAccTile, pijTile); + } + + // Store pij (must complete before next iteration's TCVT overwrites pijBf16Tile) + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Prefetch next sij into alternate buffer (after TSTORE to avoid UB race) + if (i + 1 < n_blocks) { + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + GlobalDataMxN sijGlobal_next(sij_base + (i + 1) * M * N); + if (i % 2 == 0) { + TLOAD(sijTile_B, sijGlobal_next); + } else { + TLOAD(sijTile_A, sijGlobal_next); + } + } + } + + // Compute final row sum from accumulated pij values + TROWSUM(sumDN, sumAccTile, tmpTile); + + // Store lij (total sum). mij already stored after Pass 1. + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(lijGlobalDN, sumDN); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij_buf = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij_buf = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t n_blocks = static_cast(args[5]); + uint64_t valid_len_last = static_cast(args[6]); + + // sij_buf is 4D (1, 1, q_tile, n_blocks*block_size) to match qk's 4D output semantic. + uint64_t q_tile_size = static_cast(sij_buf->shapes[2]); + + if (q_tile_size == 16) { + softmax_prepare_n_impl<16, 128>(sij_buf, pij_buf, mij, lij, scale_value, n_blocks, valid_len_last); + } else { + softmax_prepare_n_impl<64, 64>(sij_buf, pij_buf, mij, lij, scale_value, n_blocks, valid_len_last); + } +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..fb5b2015ed --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,186 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Paged Attention Orchestration - 4D input shapes, N_UNROLL=64, 4 Tasks Per Group + * + * Batches up to N_UNROLL blocks per group. Each group submits exactly 4 tasks: + * 1. QK matmul: qi @ K^T for n_blocks → sij_buf (1, 1, q_tile, n_blocks * block_size) + * 2. Softmax: two-pass over sij_buf → pij_buf, mi, li + * 3. PV matmul: SplitK accumulated P @ V → oi_new (1, 1, q_tile, head_dim) + * 4. Update: online softmax accumulation with group-level mi, li, oi_new + * + * Memory Layout (4D throughout): + * Query: (batch, seq_len=1, num_heads, head_dim) bf16 + * Key: (total_blocks, block_size, kv_head_num, head_dim) bf16 + * Value: (total_blocks, block_size, kv_head_num, head_dim) bf16 + * Out: (batch, seq_len=1, num_heads, head_dim) fp32 + */ + +#include +#include + +#include "pto_orchestration_api.h" + +#define N_UNROLL 64 + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 + +extern "C" { +/** + * Orchestration config — the executor reads these values to set up + * shared memory and runtime before calling aicpu_orchestration_entry. + */ +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // Read dimensions from tensor metadata + // query: shape=[batch, seq_len, num_heads, head_dim] + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[2]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[3]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + // key_cache: shape=[total_blocks, block_size, kv_head_num, head_dim] + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + + // block_table: shape=[batch, max_num_blocks_per_req] + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + // scale from scalar arg + uint64_t scale_value = orch_args.scalar(0); + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + + // External 4D tensors inherit shape/dtype from TaskArg (golden provides 4D). + const Tensor &query = orch_args.tensor(0).ref(); + const Tensor &key_cache = orch_args.tensor(1).ref(); + const Tensor &value_cache = orch_args.tensor(2).ref(); + const Tensor &block_table = orch_args.tensor(3).ref(); + const Tensor &out = orch_args.tensor(5).ref(); + + int *host_context_lens = orch_args.tensor(4).ref().data_as(); + + // Loop-invariant shape descriptors: 4D data tiles (1, 1, q_tile, head_dim), + // 3D scalar vectors (1, 1, q_tile). + uint32_t tile4d_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)head_dim}; + uint32_t scalar_shapes[3] = {1, 1, (uint32_t)q_tile}; + TensorCreateInfo tile4d_ci(tile4d_shapes, 4, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 3, DataType::FLOAT32); + + // Prefetch first block host_context_lens data into cache + __builtin_prefetch(&host_context_lens[0], 0, 3); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint64_t cur_seq = host_context_lens[b_idx]; + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + + // Prefetch next block host_context_lens data while processing current batch + if (b_idx + 1 < batch) { + __builtin_prefetch(&host_context_lens[b_idx + 1], 0, 3); + } + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE() { + // 4D views into query/out, matching (1, 1, q_tile, head_dim). + uint32_t view_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)head_dim}; + uint32_t view_offsets[4] = {(uint32_t)b_idx, 0, (uint32_t)(q_idx * q_tile), 0}; + Tensor qi = query.view(view_shapes, view_offsets); + Tensor out_view = out.view(view_shapes, view_offsets, true); + + // Per-group accumulators: oi (4D data), mi_update/li_update (3D scalars). + TaskOutputTensors alloc_outs = alloc_tensors(tile4d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + + // Reusable Arg objects — reset() before each use avoids + // repeated stack-frame construction in the inner loop. + L0TaskArgs params_qk, params_sf, params_pv, params_up; + + for (uint64_t bn = 0; bn < bn_this_batch; bn += N_UNROLL) { + uint64_t n_blocks = std::min((uint64_t)N_UNROLL, bn_this_batch - bn); + + // Valid length for last block in this group + uint64_t last_block_seq_start = (bn + n_blocks - 1) * block_size; + uint64_t valid_len_last = std::min(block_size, cur_seq - last_block_seq_start); + + // === Task 1: Batched QK matmul — produces 4D sij_buf === + uint32_t sij_buf_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)(n_blocks * block_size)}; + TensorCreateInfo sij_buf_ci(sij_buf_shapes, 4, DataType::FLOAT32); + + params_qk.reset(); + params_qk.add_input(qi); + params_qk.add_input(key_cache); + params_qk.add_input(block_table); + params_qk.add_output(sij_buf_ci); + params_qk.add_scalar(n_blocks); + params_qk.add_scalar(b_idx * block_num + bn); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij_buf = qk_outs.get_ref(0); + + // === Task 2: Two-pass softmax — produces 4D pij_buf, 3D mi, li === + uint32_t pij_buf_shapes[4] = {1, 1, (uint32_t)q_tile, (uint32_t)(n_blocks * block_size)}; + TensorCreateInfo pij_buf_ci(pij_buf_shapes, 4, data_type); + + params_sf.reset(); + params_sf.add_input(sij_buf); + params_sf.add_output(pij_buf_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + params_sf.add_scalar(scale_value); + params_sf.add_scalar(n_blocks); + params_sf.add_scalar(valid_len_last); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_buf = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + + // === Task 3: SplitK PV matmul — produces 4D oi_new === + params_pv.reset(); + params_pv.add_input(pij_buf); + params_pv.add_input(value_cache); + params_pv.add_input(block_table); + params_pv.add_output(tile4d_ci); + params_pv.add_scalar(n_blocks); + params_pv.add_scalar(b_idx * block_num + bn); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_new = pv_outs.get_ref(0); + + // === Task 4: Online update (per-group) === + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn + n_blocks >= bn_this_batch) ? 1 : 0; + + params_up.reset(); + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_new); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + } + } + } + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py new file mode 100644 index 0000000000..3eb8d49691 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention unroll with 4D input shapes (batch, seq_len, num_heads, head_dim). + +Query and output tensors use 4D format instead of the standard 3D. +6 kernels: QK/PV matmul (AIC), softmax_prepare/online_update (AIV). +Orchestration with N_UNROLL=64, 4 tasks per group, online softmax accumulation. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPagedAttentionUnroll4dims(SceneTestCase): + """Paged attention unroll with 4D query/out shapes.""" + + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "Case2", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + "manual": True, + }, + { + "name": "Case3", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 256, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + "manual": True, + }, + ] + + def generate_args(self, params): + inputs = _pa_generate_inputs(params) + batch = params["batch"] + num_heads = params["num_heads"] + head_dim = params["head_dim"] + specs = [] + for name, val in inputs: + if isinstance(val, torch.Tensor): + if name in ("query", "out"): + val = val.reshape(batch, 1, num_heads, head_dim) + specs.append(Tensor(name, val)) + else: + specs.append(Scalar(name, val)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + batch = params["batch"] + num_heads = params["num_heads"] + head_dim = params["head_dim"] + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + # Reshape 4D out to 3D for shared golden, then restore + out_4d = tensors["out"] + tensors["out"] = out_4d.reshape(batch, num_heads, head_dim) + _pa_compute_golden(tensors, params) + tensors["out"] = out_4d + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp new file mode 100644 index 0000000000..e9fb327157 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp @@ -0,0 +1,61 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * AIC kernel: writes float(block_idx) at cache line (base_cl + block_idx*3 + 0). + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 0) * FLOATS_PER_CACHE_LINE; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp new file mode 100644 index 0000000000..1a7fe065bf --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp @@ -0,0 +1,63 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * AIV kernel: writes float(block_idx) at cache line + * (base_cl + block_idx*3 + 1 + sub_block_id). + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t block_idx = get_block_idx(args); + int32_t sub_block_id = get_sub_block_id(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 1 + sub_block_id) * FLOATS_PER_CACHE_LINE; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp new file mode 100644 index 0000000000..366dabbdd1 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp @@ -0,0 +1,64 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Regression test for batch dispatch OOB (issue #565). + * + * Submits two MIX tasks with block_num=48 back-to-back so they are both + * in the ready queue when the scheduler runs pop_ready_tasks_batch. + * + * Args layout: [output] + */ + +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_AIC 0 +#define FUNC_AIV0 1 +#define FUNC_AIV1 2 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +static void submit_spmd_mix(const Tensor &out, int16_t block_num, int64_t base_cl) { + MixedKernels mk; + mk.aic_kernel_id = FUNC_AIC; + mk.aiv0_kernel_id = FUNC_AIV0; + mk.aiv1_kernel_id = FUNC_AIV1; + + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.launch_spec.set_core_num(block_num); + rt_submit_task(mk, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + + // Two back-to-back tasks with block_num=48 (2x cluster count). + // Both land in the ready queue simultaneously, triggering got=2 in + // pop_ready_tasks_batch — the scenario that causes OOB without the fix. + submit_spmd_mix(ext_output, 48, 0); + submit_spmd_mix(ext_output, 48, 144); + + LOG_INFO_V9("[spmd_batch_dispatch_oob] Submitted 2 MIX tasks: block_num=48,48"); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py new file mode 100644 index 0000000000..0e8dadeac9 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Regression test for batch dispatch OOB (issue #565). + +Submits two back-to-back MIX tasks each with block_num=48 (>> 24 clusters). +When both tasks enter the ready queue simultaneously, pop_ready_tasks_batch +returns got=2. Without the fix, the first task's do-while drains all idle +clusters, and the second task's do-while calls pop_first() on an empty mask, +returning -1 as cluster_offset — an out-of-bounds array index. + +Each block writes float(block_idx) at 3 cache lines (AIC, AIV0, AIV1). +Output tensor: 2 * 48 * 3 = 288 cache lines = 4608 float32. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +SLOTS_PER_BLOCK = 3 +TASKS = [(48, 0), (48, 144)] +TOTAL_CL = sum(bn * SLOTS_PER_BLOCK for bn, _ in TASKS) + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdBatchDispatchOob(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aic/kernel_write.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/kernel_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "source": "kernels/aiv/kernel_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {}, + } + ] + + def generate_args(self, params): + return TaskArgsBuilder(Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32))) + + def compute_golden(self, args, params): + out = args.output + for block_num, base_cl in TASKS: + for block_idx in range(block_num): + for slot in range(SLOTS_PER_BLOCK): + cl = base_cl + block_idx * SLOTS_PER_BLOCK + slot + out[cl * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py new file mode 100644 index 0000000000..d27e9ba3a5 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""L3 ChipTask → SubTask dependency via TensorMap. + +Worker(level=3) submits a ChipTask then a SubTask that depends on it. +Verifies: TensorMap dependency inference, cross-fork data visibility, +SubWorker reads result produced by ChipWorker. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test +from simpler_setup.scene_test import _build_l3_task_args + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + + +def verify(args): + """SubCallable — dependency target, runs after ChipTask completes.""" + + +def run_dag(orch, callables, task_args, config): + """L3 orchestration: ChipTask → SubTask dependency.""" + # ChipTask: tags inside chip_args drive deps (INPUT → lookup; OUTPUT_EXISTING → insert). + chip_args, _ = _build_l3_task_args(task_args, callables.vector_kernel_sig) + callables.keep(chip_args) # prevent GC before drain + + orch.submit_next_level(callables.vector_kernel, chip_args, config, worker=0) + + # SubTask: tag the chip output as INPUT — Orchestrator wires the dep via TensorMap. + sub_args = TaskArgs() + sub_args.add_tensor(make_tensor_arg(task_args.f), TensorArgType.INPUT) + orch.submit_sub(callables.verify, sub_args) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3Dependency(SceneTestCase): + """L3: ChipTask produces output, SubTask depends on it via TensorMap.""" + + CALLABLE = { + "orchestration": run_dag, + "callables": [ + { + "name": "vector_kernel", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "verify", "callable": verify}, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"device_count": 1, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + SIZE = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + ) + + def compute_golden(self, args, params): + args.f[:] = (args.a + args.b + 1) * (args.a + args.b + 2) + (args.a + args.b) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py new file mode 100644 index 0000000000..f628b0a4c6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""L3 group task — 2 ChipWorkers (process-isolated) on 1 DAG node. + +Each chip runs the same kernel with its own args (different tensors). +A downstream SubTask depends on the group output. +Verifies: fork+shm process isolation, 2-chip concurrent execution, +group completion aggregation, downstream dependency waits for group. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + + +def verify(args): + """SubCallable — runs after group completes.""" + + +def _chip_args(in_a, in_b, out_f) -> TaskArgs: + """Build per-chip TaskArgs with INPUT/INPUT/OUTPUT_EXISTING tags.""" + a = TaskArgs() + a.add_tensor(make_tensor_arg(in_a), TensorArgType.INPUT) + a.add_tensor(make_tensor_arg(in_b), TensorArgType.INPUT) + a.add_tensor(make_tensor_arg(out_f), TensorArgType.OUTPUT_EXISTING) + return a + + +def run_dag(orch, callables, task_args, config): + """L3 orchestration: group of 2 chips → SubTask dependency.""" + args0 = _chip_args(task_args.a0, task_args.b0, task_args.f0) + args1 = _chip_args(task_args.a1, task_args.b1, task_args.f1) + callables.keep(args0, args1) # prevent GC before drain + + orch.submit_next_level_group(callables.vector_kernel, [args0, args1], config, workers=[0, 1]) + + # SubTask depends on both group outputs (f0, f1) — tag both as INPUT. + sub_args = TaskArgs() + sub_args.add_tensor(make_tensor_arg(task_args.f0), TensorArgType.INPUT) + sub_args.add_tensor(make_tensor_arg(task_args.f1), TensorArgType.INPUT) + orch.submit_sub(callables.verify, sub_args) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3Group(SceneTestCase): + """L3: Group of 2 ChipWorkers as 1 DAG node, SubTask depends on group.""" + + CALLABLE = { + "orchestration": run_dag, + "callables": [ + { + "name": "vector_kernel", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "verify", "callable": verify}, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5sim", "a5"], + "config": {"device_count": 2, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + SIZE = 128 * 128 + return TaskArgsBuilder( + Tensor("a0", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b0", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f0", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + Tensor("a1", torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_()), + Tensor("b1", torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_()), + Tensor("f1", torch.zeros(SIZE, dtype=torch.float32).share_memory_()), + ) + + def compute_golden(self, args, params): + args.f0[:] = (args.a0 + args.b0 + 1) * (args.a0 + args.b0 + 2) + (args.a0 + args.b0) + args.f1[:] = (args.a1 + args.b1 + 1) * (args.a1 + args.b1 + 2) + (args.a1 + args.b1) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py b/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py new file mode 100644 index 0000000000..4e88784efc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""L3 post-fork zero-copy host buffers (issue #1027 #1). + +A host tensor created *after* the chip children are forked (lazily on the +first ``Worker.run()``) is not visible to those children: the orch fn runs in +the parent and carries a raw parent VA that is unmapped (or stale) in the child. +``Worker.create_host_buffer`` hands back born-shared memory already attached into +every chip child, so a tensor built over it with ``torch.frombuffer`` round-trips +with **no per-run copy** — the child reads and writes the same physical pages the +parent sees. + +Covers the mechanism end-to-end (allocate a post-fork buffer, fill it in place, +run, read the result back). The host-side staging (born-shared bytes need no +copy; an in-range view is validated to fit) is unit-tested in +``tests/ut/py/test_worker/test_host_buffer_registration.py``. + +a5sim: ``create_host_buffer`` is pure host-side (POSIX shm + a control +broadcast to the forked chip children) with no platform branching, so the sim +backend exercises the full mechanism without needing a device. The +vector_example orchestration kernels exist only for a5. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import CallConfig, TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, make_tensor_arg, scene_test + +KERNELS_BASE = "../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" + +SIZE = 128 * 128 +DTYPE = torch.float32 + + +def _golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + s = a + b + return (s + 1) * (s + 2) + s + + +def _one_task_orch(chip_handle, a, b, out): + def orch_fn(orch, _args, cfg): + ta = TaskArgs() + ta.add_tensor(make_tensor_arg(a), TensorArgType.INPUT) + ta.add_tensor(make_tensor_arg(b), TensorArgType.INPUT) + ta.add_tensor(make_tensor_arg(out), TensorArgType.OUTPUT_EXISTING) + orch.submit_next_level(chip_handle, ta, cfg, worker=0) + + return orch_fn + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestPostForkHostBufferZeroCopy(SceneTestCase): + """Post-fork zero-copy host buffers on a single L3 worker (issue #1027 #1).""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + ], + } + + CASES = [ + {"name": "post_fork_zero_copy", "platforms": ["a5sim"]}, + ] + + def test_run(self, st_worker): + """Zero-copy: buffers allocated AFTER the fork via ``create_host_buffer``, + filled in place, run, and read back — all without a per-run copy. + + ``Worker.init()`` is eager, so the chip child is already forked when the + buffers are created; a born-shared ``create_host_buffer`` is the mapped + path a post-init host tensor takes to reach that child. + """ + worker = st_worker + chip_handle = type(self)._st_chip_handles["vector"] + + nbytes = SIZE * DTYPE.itemsize # element count × dtype size, not a magic 4 + ba = worker.create_host_buffer(nbytes) + bb = worker.create_host_buffer(nbytes) + bout = worker.create_host_buffer(nbytes) + a = b = out = None + result = False + try: + a = torch.frombuffer(ba.buffer, dtype=DTYPE, count=SIZE) + b = torch.frombuffer(bb.buffer, dtype=DTYPE, count=SIZE) + out = torch.frombuffer(bout.buffer, dtype=DTYPE, count=SIZE) + a.fill_(5.0) # in place → lands directly in the child-visible shm + b.fill_(7.0) + out.zero_() + worker.run(_one_task_orch(chip_handle, a, b, out), args=None, config=CallConfig()) + result = torch.allclose(out, _golden(a, b), rtol=self.RTOL, atol=self.ATOL) + finally: + # Drop the views before freeing so the shm releases promptly, and do it + # in finally so a run()/assert failure above still cleans up all three + # buffers instead of leaving live views warning on the first free. + del a, b, out + worker.free_host_buffer(ba) + worker.free_host_buffer(bb) + worker.free_host_buffer(bout) + assert result + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py b/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py new file mode 100644 index 0000000000..2cea0d9cd9 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py @@ -0,0 +1,35 @@ +#!/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 runtime-arena prewarm during L2 worker initialization.""" + +import pytest +from simpler.task_interface import CallConfig +from simpler.worker import Worker + +_RUNTIME = "tensormap_and_ringbuffer" + + +@pytest.mark.platforms(["a5sim"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_l2_init_with_prewarm_config(st_platform, st_device_ids): + config = CallConfig() + config.runtime_env.ring_task_window = 64 + + worker = Worker( + level=2, + device_id=int(st_device_ids[0]), + platform=st_platform, + runtime=_RUNTIME, + ) + try: + worker.init(prewarm_config=config) + finally: + worker.close()