test(a5/tmr): port a2a3-only tensormap_and_ringbuffer scene/L3 tests to a5

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;
 

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 <cstdint>
+#include <pto/pto-inst.hpp>
+#include <pto/common/constants.hpp>
+#include <pto/common/pto_tile.hpp>
+
+#include "tensor.h"
+
+using namespace pto;
+
+#include "pipe_sync.h"
+
+#ifndef __gm__
+#define __gm__
+#endif
+
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+template <typename T>
+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<int>(total_elems / tile_elems);
+}
+
+template <int TILE>
+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<int>(TILE, 16);
+    constexpr int K = CeilAlign<int>(TILE, blockAlign);
+    constexpr int N = CeilAlign<int>(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<TileType::Mat, float, M, K, BLayout::ColMajor, TILE, TILE, SLayout::RowMajor, 512>;
+    using TileMatB = Tile<TileType::Mat, float, K, N, BLayout::ColMajor, TILE, TILE, SLayout::RowMajor, 512>;
+
+    using LeftTile = TileLeft<float, M, K, TILE, TILE>;
+    using RightTile = TileRight<float, K, N, TILE, TILE>;
+    using AccTile = TileAcc<float, M, N, TILE, TILE>;
+
+    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);
+    }
+}

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 <cstdint>
+#include <pto/pto-inst.hpp>
+
+#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<int>(total_elems / tile_elems);
+}
+
+template <int ROWS, int COLS>
+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<float, DynShapeDim5, DynStridDim5>;
+    using TileData = Tile<TileType::Vec, float, ROWS, COLS, BLayout::RowMajor, -1, -1>;
+
+    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);
+    }
+}

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 <stddef.h>
+#include <stdint.h>
+
+#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<int>(orch_args.scalar(0));
+    int M = static_cast<int>(orch_args.scalar(1));
+    int N = static_cast<int>(orch_args.scalar(2));
+    int matmul_batch = static_cast<int>(orch_args.scalar(3));
+    int add_batch = static_cast<int>(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<uint64_t>(start_task_idx) * MATMUL_ELEMS;
+            uint64_t group_size = static_cast<uint64_t>(matmul_batch) * MATMUL_ELEMS;
+
+            uint32_t matmul_group_shapes[1] = {static_cast<uint32_t>(group_size)};
+            uint32_t view_offsets[1] = {static_cast<uint32_t>(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<uint64_t>(start_task_idx) * ADD_ELEMS;
+            uint64_t group_size = static_cast<uint64_t>(add_batch) * ADD_ELEMS;
+
+            uint32_t add_group_shapes[1] = {static_cast<uint32_t>(group_size)};
+            uint32_t view_offsets[1] = {static_cast<uint32_t>(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"