Add: fully_distributed_within_core runtime — SPMD on-core orchestration

examples/a2a3/fully_distributed_within_core/async_notify_demo/kernels/aiv/kernel_consumer.cpp

@@ -0,0 +1,59 @@
+/*
+ * 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 <cstdint>
+
+#ifndef __gm__
+#define __gm__
+#endif
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+#include <pto/pto-inst.hpp>
+#include "pto/common/pto_tile.hpp"
+
+#include "tensor.h"
+
+using namespace pto;
+
+extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) {
+    __gm__ Tensor *src_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]);
+    __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]);
+    __gm__ int32_t *notify_counter = reinterpret_cast<__gm__ int32_t *>(args[3]);
+
+    __gm__ float *src = reinterpret_cast<__gm__ float *>(src_tensor->buffer.addr) + src_tensor->start_offset;
+    __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset;
+
+    constexpr int kRows = 128;
+    constexpr int kCols = 128;
+    using DynShapeDim5 = Shape<1, 1, 1, kRows, kCols>;
+    using DynStridDim5 = Stride<1, 1, 1, kCols, 1>;
+    using GlobalData = GlobalTensor<float, DynShapeDim5, DynStridDim5>;
+    using TileData = Tile<TileType::Vec, float, kRows, kCols, BLayout::RowMajor, -1, -1>;
+
+    TileData src_tile(kRows, kCols);
+    TileData dst_tile(kRows, kCols);
+    TASSIGN(src_tile, 0x0);
+    TASSIGN(dst_tile, 0x10000);
+
+    GlobalData src_global(src);
+    GlobalData dst_global(result);
+    TLOAD(src_tile, src_global);
+    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+    TADDS(dst_tile, src_tile, static_cast<float>(*notify_counter));
+    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    TSTORE(dst_global, dst_tile);
+    set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
+    wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
+}

examples/a2a3/fully_distributed_within_core/async_notify_demo/kernels/aiv/kernel_notify_wait.cpp

@@ -0,0 +1,31 @@
+/*
+ * 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 <cstdint>
+
+#ifndef __gm__
+#define __gm__
+#endif
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+#include <pto/pto-inst.hpp>
+#include "pto_async_kernel_api.h"
+
+extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) {
+    uint64_t notify_counter_addr = static_cast<uint64_t>(args[1]);
+    uint32_t expected_value = static_cast<uint32_t>(args[2]);
+    AsyncCtx async_ctx = get_async_ctx(args);
+    save_expected_notification_counter(
+        async_ctx, reinterpret_cast<volatile __gm__ void *>(notify_counter_addr), expected_value
+    );
+}

examples/a2a3/fully_distributed_within_core/async_notify_demo/kernels/aiv/kernel_producer_notify.cpp

@@ -0,0 +1,79 @@
+/*
+ * 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 <cstdint>
+
+#ifndef __gm__
+#define __gm__
+#endif
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+#include <pto/pto-inst.hpp>
+
+#include "platform_comm/comm_context.h"
+#include "pto_async_kernel_api.h"
+#include "tensor.h"
+
+using namespace pto;
+
+template <typename T>
+static inline __aicore__ __gm__ T *comm_remote_ptr(__gm__ CommContext *ctx, __gm__ T *local_ptr, int peer_rank) {
+    uint64_t local_base = ctx->windowsIn[ctx->rankId];
+    uint64_t offset = reinterpret_cast<uint64_t>(local_ptr) - local_base;
+    return reinterpret_cast<__gm__ T *>(ctx->windowsIn[peer_rank] + offset);
+}
+
+extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) {
+    __gm__ Tensor *in_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]);
+    __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]);
+    __gm__ int32_t *local_counter = reinterpret_cast<__gm__ int32_t *>(args[2]);
+    __gm__ CommContext *comm_ctx = reinterpret_cast<__gm__ CommContext *>(args[3]);
+
+    __gm__ float *in_data = reinterpret_cast<__gm__ float *>(in_tensor->buffer.addr) + in_tensor->start_offset;
+    __gm__ float *out_data = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset;
+
+    int my_rank = static_cast<int>(comm_ctx->rankId);
+    int peer_rank = 1 - my_rank;
+
+    constexpr int kRows = 128;
+    constexpr int kCols = 128;
+    using DynShapeDim5 = Shape<1, 1, 1, kRows, kCols>;
+    using DynStridDim5 = Stride<1, 1, 1, kCols, 1>;
+    using GlobalData = GlobalTensor<float, DynShapeDim5, DynStridDim5>;
+    using TileData = Tile<TileType::Vec, float, kRows, kCols, BLayout::RowMajor, -1, -1>;
+
+    TileData in_tile(kRows, kCols);
+    TileData out_tile(kRows, kCols);
+    TASSIGN(in_tile, 0x0);
+    TASSIGN(out_tile, 0x10000);
+
+    GlobalData in_global(in_data);
+    GlobalData out_global(out_data);
+    TLOAD(in_tile, in_global);
+    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+    TADD(out_tile, in_tile, in_tile);
+    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
+    TSTORE(out_global, out_tile);
+    set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
+    wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
+
+    if (my_rank == 1) {
+        for (volatile int i = 0; i < 2000000; ++i) {}
+    }
+
+    __gm__ int32_t *remote_counter = comm_remote_ptr(comm_ctx, local_counter, peer_rank);
+    send_notification(remote_counter, 1, pto::comm::NotifyOp::AtomicAdd);
+    pipe_barrier(PIPE_ALL);
+}

examples/a2a3/fully_distributed_within_core/async_notify_demo/kernels/orchestration/async_notify_orchestration.cpp

@@ -0,0 +1,65 @@
+/*
+ * 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 <stdint.h>
+
+#include "platform_comm/comm_context.h"
+#include "pto_orchestration_api.h"
+
+extern "C" {
+
+__attribute__((visibility("default"))) PTO2OrchestrationConfig
+async_notify_orchestration_config(const L2TaskArgs &orch_args) {
+    (void)orch_args;
+    return PTO2OrchestrationConfig{.expected_arg_count = 5};
+}
+
+__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) {
+    return async_notify_orchestration_config(orch_args);
+}
+
+__attribute__((visibility("default"))) void async_notify_orchestration(const L2TaskArgs &orch_args) {
+    if (orch_args.tensor_count() + orch_args.scalar_count() != 5) {
+        LOG_ERROR("async_notify_demo: expected 5 args");
+        return;
+    }
+
+    const Tensor &input = orch_args.tensor(0).ref();
+    const Tensor &output = orch_args.tensor(1).ref();
+    const Tensor &result = orch_args.tensor(2).ref();
+    const Tensor &notify_counter = orch_args.tensor(3).ref();
+    auto *comm_ctx = reinterpret_cast<CommContext *>(static_cast<uintptr_t>(orch_args.scalar(0)));
+
+    L0TaskArgs params_producer;
+    params_producer.add_input(input);
+    params_producer.add_output(output);
+    params_producer.add_scalar(notify_counter.buffer.addr);
+    params_producer.add_scalar(reinterpret_cast<uint64_t>(comm_ctx));
+    rt_submit_aiv_task(0, params_producer);
+
+    uint32_t notify_token_shape[1] = {1};
+    TensorCreateInfo notify_token_info(notify_token_shape, 1, DataType::INT32);
+    L0TaskArgs params_notify;
+    params_notify.add_output(notify_token_info);
+    params_notify.add_scalar(notify_counter.buffer.addr);
+    params_notify.add_scalar(static_cast<uint64_t>(1));
+    TaskOutputTensors notify_outputs = rt_submit_aiv_task(2, params_notify);
+    Tensor notify_token = notify_outputs.get_ref(0);
+
+    L0TaskArgs params_consumer;
+    params_consumer.add_input(notify_token);
+    params_consumer.add_input(output);
+    params_consumer.add_output(result);
+    params_consumer.add_scalar(notify_counter.buffer.addr);
+    rt_submit_aiv_task(1, params_consumer);
+}
+
+}  // extern "C"