Add manual-scope v0 to tensormap runtime

examples/a2a3/tensormap_and_ringbuffer/paged_attention_partial_manual/golden.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+import sys
+
+_BASE = Path(__file__).resolve().parents[1] / "paged_attention"
+sys.path.insert(0, str(_BASE))
+
+from golden import ALL_CASES, ATOL, DEFAULT_CASE, RTOL, __outputs__, compute_golden, generate_inputs  # noqa: E402,F401

examples/a2a3/tensormap_and_ringbuffer/paged_attention_partial_manual/kernels/kernel_config.py

@@ -0,0 +1,72 @@
+# 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.
+# -----------------------------------------------------------------------------------------------------------
+from pathlib import Path
+
+from task_interface import ArgDirection as D  # pyright: ignore[reportAttributeAccessIssue]
+
+_ROOT = Path(__file__).parent
+_PA_KERNELS = _ROOT.parent.parent / "paged_attention" / "kernels"
+
+ORCHESTRATION = {
+    "source": str(_ROOT / "orchestration" / "paged_attention_orch.cpp"),
+    "function_name": "aicpu_orchestration_entry",
+    "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT],
+}
+
+KERNELS = [
+    {
+        "func_id": 0,
+        "name": "QK",
+        "source": str(_PA_KERNELS / "aic" / "aic_qk_matmul.cpp"),
+        "core_type": "aic",
+        "signature": [D.IN, D.IN, D.OUT],
+    },
+    {
+        "func_id": 2,
+        "name": "PV",
+        "source": str(_PA_KERNELS / "aic" / "aic_pv_matmul.cpp"),
+        "core_type": "aic",
+        "signature": [D.IN, D.IN, D.OUT],
+    },
+    {
+        "func_id": 4,
+        "name": "AIC_HUB",
+        "source": str(_PA_KERNELS / "aic" / "aic_hub.cpp"),
+        "core_type": "aic",
+        "signature": [],
+    },
+    {
+        "func_id": 1,
+        "name": "SF",
+        "source": str(_PA_KERNELS / "aiv" / "aiv_softmax_prepare.cpp"),
+        "core_type": "aiv",
+        "signature": [D.IN, D.OUT, D.OUT, D.OUT],
+    },
+    {
+        "func_id": 3,
+        "name": "UP",
+        "source": str(_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],
+    },
+    {
+        "func_id": 5,
+        "name": "AIV_HUB",
+        "source": str(_PA_KERNELS / "aiv" / "aiv_hub.cpp"),
+        "core_type": "aiv",
+        "signature": [],
+    },
+]
+
+RUNTIME_CONFIG = {
+    "runtime": "tensormap_and_ringbuffer",
+    "aicpu_thread_num": 4,
+    "orch_thread_num": 1,
+    "block_dim": 24,
+}

examples/a2a3/tensormap_and_ringbuffer/paged_attention_partial_manual/kernels/orchestration/paged_attention_orch.cpp

@@ -0,0 +1,168 @@
+/*
+ * 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 <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <cinttypes>
+
+#include "pto_orchestration_api.h"  // NOLINT(build/include_subdir)
+
+#define FUNC_QK_MATMUL 0
+#define FUNC_SOFTMAX_PREPARE 1
+#define FUNC_PV_MATMUL 2
+#define FUNC_ONLINE_UPDATE 3
+#define FUNC_AIC_HUB 4
+#define FUNC_AIV_HUB 5
+
+extern "C" {
+
+__attribute__((visibility("default"))) PTO2OrchestrationConfig
+aicpu_orchestration_config(const ChipStorageTaskArgs &orch_args) {
+    (void)orch_args;  // NOLINT(readability/casting)
+    return PTO2OrchestrationConfig{
+        .expected_arg_count = 7,
+    };
+}
+
+__attribute__((visibility("default"))) void
+aicpu_orchestration_entry(const ChipStorageTaskArgs &orch_args, int orch_thread_num, int orch_thread_index) {
+    uint64_t batch = orch_args.tensor(0).shapes[0];
+    uint64_t num_heads = orch_args.tensor(0).shapes[1];
+    uint64_t head_dim = orch_args.tensor(0).shapes[2];
+    DataType data_type = orch_args.tensor(0).dtype;
+    uint64_t block_size = orch_args.tensor(1).shapes[1];
+    uint64_t block_num = orch_args.tensor(3).shapes[1];
+    uint64_t scale_value = orch_args.scalar(0);
+
+    uint64_t q_head_num = num_heads;
+    uint64_t q_tile = 16;
+    uint64_t q_loop = (q_head_num + q_tile - 1) / q_tile;
+
+    uint64_t b_start = batch * orch_thread_index / orch_thread_num;
+    uint64_t b_end = batch * (orch_thread_index + 1) / orch_thread_num;
+
+    void *query_ptr = orch_args.tensor(0).data_as<void>();
+    void *kc_ptr = orch_args.tensor(1).data_as<void>();
+    void *vc_ptr = orch_args.tensor(2).data_as<void>();
+    void *out_ptr = orch_args.tensor(5).data_as<void>();
+
+    uint64_t total_blocks_count = orch_args.tensor(1).shapes[0];
+    uint64_t kv_total_rows = total_blocks_count * block_size;
+
+    uint32_t query_shapes[2] = {static_cast<uint32_t>(batch * num_heads), static_cast<uint32_t>(head_dim)};
+    uint32_t key_cache_shapes[2] = {static_cast<uint32_t>(kv_total_rows), static_cast<uint32_t>(head_dim)};
+    uint32_t value_cache_shapes[2] = {static_cast<uint32_t>(kv_total_rows), static_cast<uint32_t>(head_dim)};
+    uint32_t out_shapes[2] = {static_cast<uint32_t>(batch * num_heads), static_cast<uint32_t>(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);
+
+    int *host_block_table = orch_args.tensor(3).data_as<int>();
+    int *host_context_lens = orch_args.tensor(4).data_as<int>();
+
+    uint32_t tile2d_shapes[2] = {static_cast<uint32_t>(q_tile), static_cast<uint32_t>(head_dim)};
+    uint32_t scalar_shapes[1] = {static_cast<uint32_t>(q_tile)};
+    uint32_t sij_shapes[2] = {static_cast<uint32_t>(q_tile), static_cast<uint32_t>(block_size)};
+    TensorCreateInfo tile2d_ci(tile2d_shapes, 2, DataType::FLOAT32);
+    TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32);
+    TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32);
+    TensorCreateInfo pij_f16_ci(sij_shapes, 2, data_type);
+
+    for (uint64_t b_idx = b_start; b_idx < b_end; b_idx++) {
+        uint64_t cur_seq = host_context_lens[b_idx];
+        uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size;
+        for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) {
+            PTO2_SCOPE() {
+                uint32_t cur_offset = static_cast<uint32_t>(b_idx * q_head_num + q_idx * q_tile);
+                uint32_t qi_offsets[2] = {cur_offset, 0};
+                uint32_t out_view_offsets[2] = {cur_offset, 0};
+                Tensor qi = query.view(tile2d_shapes, qi_offsets);
+                Tensor out_view = out.view(tile2d_shapes, out_view_offsets);
+
+                Arg params_inplace;
+                params_inplace.add_output(tile2d_ci);
+                params_inplace.add_output(scalar_ci);
+                params_inplace.add_output(scalar_ci);
+                TaskOutputTensors hub_outs = pto2_rt_submit_aiv_task(FUNC_AIV_HUB, params_inplace);
+                const Tensor &oi = hub_outs.get_ref(0);
+                const Tensor &li_update = hub_outs.get_ref(1);
+                const Tensor &mi_update = hub_outs.get_ref(2);
+
+                for (uint64_t bn = 0; bn < bn_this_batch; bn++) {
+                    uint64_t cur_block_idx = host_block_table[b_idx * block_num + bn];
+                    uint64_t valid_len = std::min(block_size, cur_seq - bn * block_size);
+                    uint32_t kv_shapes[2] = {static_cast<uint32_t>(block_size), static_cast<uint32_t>(head_dim)};
+                    uint32_t kv_offsets[2] = {static_cast<uint32_t>(cur_block_idx * block_size), 0};
+                    Tensor kj = key_cache.view(kv_shapes, kv_offsets);
+                    Tensor vj = value_cache.view(kv_shapes, kv_offsets);
+
+                    PTO2_SCOPE(PTO2ScopeMode::MANUAL) {
+                        Arg params_qk;
+                        params_qk.add_input(qi);
+                        params_qk.add_input(kj);
+                        params_qk.add_output(sij_ci);
+                        PTO2ManualSubmitResult qk_outs = pto2_rt_submit_aic_task_manual(FUNC_QK_MATMUL, params_qk);
+                        const Tensor &sij = qk_outs.outputs.get_ref(0);
+
+                        uint32_t sij_valid_shapes[2] = {
+                            static_cast<uint32_t>(q_tile), static_cast<uint32_t>(valid_len)
+                        };
+                        uint32_t sij_valid_offsets[2] = {0, 0};
+                        Tensor sij_valid = sij.view(sij_valid_shapes, sij_valid_offsets);
+
+                        Arg params_sf;
+                        params_sf.add_input(sij_valid);
+                        params_sf.add_output(pij_f16_ci);
+                        params_sf.add_output(scalar_ci);
+                        params_sf.add_output(scalar_ci);
+                        params_sf.add_scalar(scale_value);
+                        PTO2ManualSubmitResult sf_outs = pto2_rt_submit_aiv_task_manual_with_deps(
+                            FUNC_SOFTMAX_PREPARE, params_sf, {qk_outs.task_id}
+                        );
+                        const Tensor &pij_f16 = sf_outs.outputs.get_ref(0);
+                        const Tensor &mi = sf_outs.outputs.get_ref(1);
+                        const Tensor &li = sf_outs.outputs.get_ref(2);
+
+                        Arg params_pv;
+                        params_pv.add_input(pij_f16);
+                        params_pv.add_input(vj);
+                        params_pv.add_output(tile2d_ci);
+                        PTO2ManualSubmitResult pv_outs =
+                            pto2_rt_submit_aic_task_manual_with_deps(FUNC_PV_MATMUL, params_pv, {sf_outs.task_id});
+                        const Tensor &oi_tmp = pv_outs.outputs.get_ref(0);
+
+                        uint64_t is_first = (bn == 0) ? 1 : 0;
+                        uint64_t is_last = (bn == bn_this_batch - 1) ? 1 : 0;
+
+                        Arg params_up;
+                        params_up.add_input(mi);
+                        params_up.add_input(li);
+                        params_up.add_input(oi_tmp);
+                        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);
+                        PTO2ManualSubmitResult up_outs = pto2_rt_submit_aiv_task_manual_with_deps(
+                            FUNC_ONLINE_UPDATE, params_up, {sf_outs.task_id, pv_outs.task_id}
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+}  // extern "C"

examples/scripts/run_example.py

@@ -21,12 +21,12 @@
 
 Examples:
     # Run hardware example (requires Ascend device)
-    python examples/scripts/run_example.py -k examples/host_build_graph/vector_example/kernels \
-                                      -g examples/host_build_graph/vector_example/golden.py
+    python examples/scripts/run_example.py -k examples/a2a3/host_build_graph/vector_example/kernels \
+                                      -g examples/a2a3/host_build_graph/vector_example/golden.py
 
     # Run simulation example (no hardware required)
-    python examples/scripts/run_example.py -k examples/host_build_graph/vector_example/kernels \
-                                      -g examples/host_build_graph/vector_example/golden.py \
+    python examples/scripts/run_example.py -k examples/a2a3/host_build_graph/vector_example/kernels \
+                                      -g examples/a2a3/host_build_graph/vector_example/golden.py \
                                       -p a2a3sim
 
     # Run with specific device

src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp

@@ -2364,6 +2364,17 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
         }
     }
 
+    if (rt != nullptr) {
+        void* sm = runtime->get_pto2_gm_sm_ptr();
+        if (sm != nullptr) {
+            int32_t orch_err = static_cast<PTO2SharedMemoryHeader*>(sm)->orch_error_code.load(std::memory_order_acquire);
+            if (orch_err != PTO2_ERROR_NONE) {
+                DEV_ERROR("Thread %d: Exiting with orchestrator error code=%d", thread_idx, orch_err);
+                return -1;
+            }
+        }
+    }
+
     return 0;
 }