fix: graph of ascend

include/infinicore/graph/graph.hpp

@@ -1,6 +1,8 @@
 #pragma once
 
+#include <cstdint>
 #include <memory>
+#include <unordered_map>
 #include <vector>
 
 #include "../tensor.hpp"
@@ -20,20 +22,33 @@ class GraphOperator {
     virtual bool is_device_graph_capture_safe() const {
         return true;
     }
+    virtual bool requires_task_update() const {
+        return false;
+    }
     virtual ~GraphOperator() = default;
 };
 
 class DispatchableGraphOperator : public GraphOperator {
 public:
     void run() const override;
+    bool requires_task_update() const override {
+        return requires_task_update_;
+    }
     ~DispatchableGraphOperator() override;
 
 protected:
+    void enable_task_update() {
+        requires_task_update_ = true;
+    }
+
     using run_schema = void (*)(void *);
     using cleanup_schema = void (*)(void **);
-    void *planned_meta_;
-    run_schema runner_;
-    cleanup_schema deleter_;
+    void *planned_meta_ = nullptr;
+    run_schema runner_ = nullptr;
+    cleanup_schema deleter_ = nullptr;
+
+private:
+    bool requires_task_update_ = false;
 };
 
 class Graph {
@@ -42,6 +57,10 @@ class Graph {
     ~Graph();
 
     void run() const;
+    void update_host_int_array(
+        const Tensor &device_tensor,
+        const int32_t *host_values,
+        size_t count);
 
 protected:
     void add_operator(std::shared_ptr<GraphOperator> op);
@@ -54,7 +73,21 @@ class Graph {
     struct DeviceGraph;
     struct Segment;
     std::vector<std::unique_ptr<Segment>> segments_;
+    std::unordered_map<const void *, std::vector<int64_t>> host_int_arrays_;
 };
+
+bool is_task_updating();
+bool is_task_group_capturing();
+void begin_task_group_capture();
+void end_task_group_capture();
+void begin_task_update();
+void end_task_update();
+void stage_task_update_host_int_array(
+    const Tensor &device_tensor,
+    const int32_t *host_values,
+    size_t count);
+const std::vector<int64_t> *lookup_task_update_host_int_array(
+    const Tensor &tensor);
 } // namespace infinicore::graph
 
 #define INFINICORE_GRAPH_OP_CLASS(__OP_NAME__, ...)                        \

include/infinirt.h

@@ -9,6 +9,7 @@ typedef void *infinirtEvent_t;
 typedef void *infinirtGraph_t;
 typedef void *infinirtGraphNode_t;
 typedef void *infinirtGraphExec_t;
+typedef void *infinirtGraphTaskGroup_t;
 
 // Bitmask describing which fields of an infinirtDeviceResourceSnapshot_t
 // have been populated by the backend. Backends without a particular
@@ -130,5 +131,13 @@ __INFINI_C __export infiniStatus_t infinirtGraphInstantiate(
     size_t buffer_size);
 __INFINI_C __export infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exec);
 __INFINI_C __export infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream);
+__INFINI_C __export infiniStatus_t infinirtGraphTaskGroupBegin(infinirtStream_t stream);
+__INFINI_C __export infiniStatus_t infinirtGraphTaskGroupEnd(
+    infinirtStream_t stream,
+    infinirtGraphTaskGroup_t *handle);
+__INFINI_C __export infiniStatus_t infinirtGraphTaskUpdateBegin(
+    infinirtStream_t stream,
+    infinirtGraphTaskGroup_t handle);
+__INFINI_C __export infiniStatus_t infinirtGraphTaskUpdateEnd(infinirtStream_t stream);
 
 #endif // __INFINIRT_API_H__

src/infinicore/graph/graph.cc

@@ -6,6 +6,66 @@
 #include <infinirt.h>
 
 namespace infinicore::graph {
+namespace {
+
+using HostIntArrayMap = std::unordered_map<const void *, std::vector<int64_t>>;
+
+thread_local const HostIntArrayMap *current_host_int_arrays = nullptr;
+thread_local bool task_update_active = false;
+thread_local HostIntArrayMap staged_host_int_arrays;
+thread_local infinirtGraphTaskGroup_t *capture_task_group_handle = nullptr;
+thread_local infinirtGraphTaskGroup_t task_update_handle = nullptr;
+
+class HostIntArrayScope {
+public:
+    HostIntArrayScope(
+        const HostIntArrayMap *host_int_arrays,
+        bool is_task_update) {
+        current_host_int_arrays = host_int_arrays;
+        task_update_active = is_task_update;
+    }
+
+    ~HostIntArrayScope() {
+        task_update_active = false;
+        current_host_int_arrays = nullptr;
+    }
+
+    HostIntArrayScope(const HostIntArrayScope &) = delete;
+    HostIntArrayScope &operator=(const HostIntArrayScope &) = delete;
+};
+
+class TaskGroupCaptureScope {
+public:
+    explicit TaskGroupCaptureScope(infinirtGraphTaskGroup_t *handle) {
+        INFINICORE_ASSERT(capture_task_group_handle == nullptr);
+        capture_task_group_handle = handle;
+    }
+
+    ~TaskGroupCaptureScope() {
+        capture_task_group_handle = nullptr;
+    }
+
+    TaskGroupCaptureScope(const TaskGroupCaptureScope &) = delete;
+    TaskGroupCaptureScope &operator=(const TaskGroupCaptureScope &) = delete;
+};
+
+class TaskUpdateHandleScope {
+public:
+    explicit TaskUpdateHandleScope(infinirtGraphTaskGroup_t handle) {
+        INFINICORE_ASSERT(task_update_handle == nullptr);
+        INFINICORE_ASSERT(handle != nullptr);
+        task_update_handle = handle;
+    }
+
+    ~TaskUpdateHandleScope() {
+        task_update_handle = nullptr;
+    }
+
+    TaskUpdateHandleScope(const TaskUpdateHandleScope &) = delete;
+    TaskUpdateHandleScope &operator=(const TaskUpdateHandleScope &) = delete;
+};
+
+} // namespace
 
 /* =========================
  * GraphTensor
@@ -33,10 +93,16 @@ DispatchableGraphOperator::~DispatchableGraphOperator() {
  * ========================= */
 
 struct Graph::DeviceGraph {
+    struct UpdatableTask {
+        std::shared_ptr<GraphOperator> op;
+        infinirtGraphTaskGroup_t handle;
+    };
+
     infinirtGraph_t graph;
     infinirtGraphExec_t exec;
     infinirtGraphNode_t node;
     std::vector<char> log_buffer;
+    std::vector<UpdatableTask> updatable_tasks;
 
     DeviceGraph() : graph(nullptr), exec(nullptr), node(nullptr) {
         log_buffer.resize(4 * 1024);
@@ -75,7 +141,9 @@ struct Graph::Segment {
     }
 };
 
-Graph::Graph() {
+Graph::Graph()
+    : host_int_arrays_(std::move(staged_host_int_arrays)) {
+    staged_host_int_arrays.clear();
 }
 
 void Graph::run() const {
@@ -85,15 +153,100 @@ void Graph::run() const {
         }
         return;
     }
+
     for (const auto &segment : segments_) {
-        segment->run();
+        if (!segment->device_graph) {
+            segment->run();
+            continue;
+        }
+
+        HostIntArrayScope update_scope(&host_int_arrays_, true);
+        for (const auto &task : segment->device_graph->updatable_tasks) {
+            TaskUpdateHandleScope task_update_scope(task.handle);
+            task.op->run();
+        }
+        segment->device_graph->launch();
     }
 }
 
 void Graph::add_operator(std::shared_ptr<GraphOperator> op) {
     op_list_.push_back(op);
 }
 
+void Graph::update_host_int_array(
+    const Tensor &device_tensor,
+    const int32_t *host_values,
+    size_t count) {
+    INFINICORE_ASSERT(device_tensor);
+    INFINICORE_ASSERT(host_values != nullptr || count == 0);
+
+    auto &values = host_int_arrays_[device_tensor->data()];
+    values.resize(count);
+    for (size_t i = 0; i < count; ++i) {
+        values[i] = static_cast<int64_t>(host_values[i]);
+    }
+}
+
+bool is_task_updating() {
+    return task_update_active;
+}
+
+bool is_task_group_capturing() {
+    return capture_task_group_handle != nullptr;
+}
+
+void begin_task_group_capture() {
+    INFINICORE_ASSERT(capture_task_group_handle != nullptr);
+    INFINICORE_ASSERT(*capture_task_group_handle == nullptr);
+    INFINICORE_CHECK_ERROR(
+        infinirtGraphTaskGroupBegin(context::getStream()));
+}
+
+void end_task_group_capture() {
+    INFINICORE_ASSERT(capture_task_group_handle != nullptr);
+    INFINICORE_ASSERT(*capture_task_group_handle == nullptr);
+    INFINICORE_CHECK_ERROR(infinirtGraphTaskGroupEnd(
+        context::getStream(), capture_task_group_handle));
+}
+
+void begin_task_update() {
+    INFINICORE_ASSERT(task_update_handle != nullptr);
+    INFINICORE_CHECK_ERROR(infinirtGraphTaskUpdateBegin(
+        context::getStream(), task_update_handle));
+}
+
+void end_task_update() {
+    INFINICORE_ASSERT(task_update_handle != nullptr);
+    INFINICORE_CHECK_ERROR(
+        infinirtGraphTaskUpdateEnd(context::getStream()));
+}
+
+void stage_task_update_host_int_array(
+    const Tensor &device_tensor,
+    const int32_t *host_values,
+    size_t count) {
+    INFINICORE_ASSERT(device_tensor);
+    INFINICORE_ASSERT(host_values != nullptr || count == 0);
+
+    auto &values = staged_host_int_arrays[device_tensor->data()];
+    values.resize(count);
+    for (size_t i = 0; i < count; ++i) {
+        values[i] = static_cast<int64_t>(host_values[i]);
+    }
+}
+
+const std::vector<int64_t> *lookup_task_update_host_int_array(
+    const Tensor &tensor) {
+    if (current_host_int_arrays == nullptr || !tensor) {
+        return nullptr;
+    }
+    auto it = current_host_int_arrays->find(tensor->data());
+    if (it == current_host_int_arrays->end()) {
+        return nullptr;
+    }
+    return &it->second;
+}
+
 void Graph::instantiate() {
     segments_.clear();
 
@@ -135,8 +288,21 @@ void Graph::instantiate() {
             throw std::runtime_error("failed to begin device graph capture");
         }
 
+        device_graph.updatable_tasks.clear();
+        HostIntArrayScope capture_scope(&host_int_arrays_, false);
         for (const auto &op : segment->ops) {
-            op->run();
+            if (!op->requires_task_update()) {
+                op->run();
+                continue;
+            }
+
+            infinirtGraphTaskGroup_t handle = nullptr;
+            {
+                TaskGroupCaptureScope task_group_scope(&handle);
+                op->run();
+            }
+            INFINICORE_ASSERT(handle != nullptr);
+            device_graph.updatable_tasks.push_back({op, handle});
         }
 
         if (infinirtStreamEndCapture(