diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index 2cc97e023..57447633b 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -1,6 +1,8 @@ #pragma once +#include #include +#include #include #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 op); @@ -54,7 +73,21 @@ class Graph { struct DeviceGraph; struct Segment; std::vector> segments_; + std::unordered_map> 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 *lookup_task_update_host_int_array( + const Tensor &tensor); } // namespace infinicore::graph #define INFINICORE_GRAPH_OP_CLASS(__OP_NAME__, ...) \ diff --git a/include/infinirt.h b/include/infinirt.h index 4ada848f5..778d5b3db 100644 --- a/include/infinirt.h +++ b/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__ diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 72f422adc..2bf0c7936 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -6,6 +6,66 @@ #include namespace infinicore::graph { +namespace { + +using HostIntArrayMap = std::unordered_map>; + +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 op; + infinirtGraphTaskGroup_t handle; + }; + infinirtGraph_t graph; infinirtGraphExec_t exec; infinirtGraphNode_t node; std::vector log_buffer; + std::vector 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,8 +153,19 @@ 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(); } } @@ -94,6 +173,80 @@ void Graph::add_operator(std::shared_ptr 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(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(host_values[i]); + } +} + +const std::vector *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( diff --git a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc index ead1ca4e2..5ce5d7803 100644 --- a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc +++ b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc @@ -42,8 +42,64 @@ struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; float scale; + Tensor out_work, q_work, k_work, v_work, block_table_work; }; +static Tensor persistent_contiguous_work_tensor(const Tensor &tensor) { + if (tensor->is_contiguous()) { + return Tensor(tensor); + } + return Tensor::empty( + tensor->shape(), tensor->dtype(), tensor->device()); +} + +static std::vector +get_actual_seq_lengths_k(const PlannedMeta *p, int64_t batch_size) { + if (const auto *bound = graph::lookup_task_update_host_int_array(p->seqlens_k)) { + if (bound->size() != static_cast(batch_size)) { + throw std::runtime_error( + "[mha_kvcache/ascend] bound actualSeqLengthsKv size does not " + "match the captured batch size"); + } + return *bound; + } + + if (graph::is_task_updating()) { + throw std::runtime_error( + "[mha_kvcache/ascend] missing host actualSeqLengthsKv binding " + "during graph task update"); + } + + auto seqlens_k_shape = p->seqlens_k->shape(); + if (seqlens_k_shape.size() != 1 + || seqlens_k_shape[0] != static_cast(batch_size)) { + throw std::runtime_error( + "[mha_kvcache/ascend] seqlens_k must be a 1D tensor whose length " + "matches the batch size"); + } + + std::vector seqlens_k_host(batch_size); + auto copy_ret = aclrtMemcpy( + seqlens_k_host.data(), + batch_size * sizeof(int32_t), + reinterpret_cast(p->seqlens_k->data()), + batch_size * sizeof(int32_t), + ACL_MEMCPY_DEVICE_TO_HOST); + if (copy_ret != ACL_SUCCESS) { + throw std::runtime_error( + std::string( + "[mha_kvcache/ascend] copy seqlens_k to host failed: ") + + std::to_string(copy_ret)); + } + + std::vector actual_seq_k_vec; + actual_seq_k_vec.reserve(batch_size); + for (int64_t i = 0; i < batch_size; ++i) { + actual_seq_k_vec.push_back(seqlens_k_host[i]); + } + return actual_seq_k_vec; +} + void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &seqlens_k, const Tensor &block_table, std::optional alibi_slopes, @@ -57,7 +113,12 @@ void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, alibi_slopes ? std::optional( graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; + scale, + persistent_contiguous_work_tensor(out), + persistent_contiguous_work_tensor(q), + persistent_contiguous_work_tensor(k_cache), + persistent_contiguous_work_tensor(v_cache), + persistent_contiguous_work_tensor(block_table)}; } void run(void *planned_meta) { @@ -98,38 +159,32 @@ void run(void *planned_meta) { "[mha_kvcache/ascend] k_cache and v_cache shapes are incompatible"); } - Tensor q_work = p->q->is_contiguous() ? Tensor(p->q) : p->q->contiguous(); - Tensor k_work = p->k_cache->is_contiguous() ? Tensor(p->k_cache) - : p->k_cache->contiguous(); - Tensor v_work = p->v_cache->is_contiguous() ? Tensor(p->v_cache) - : p->v_cache->contiguous(); - Tensor bt_work = p->block_table->is_contiguous() - ? Tensor(p->block_table) - : p->block_table->contiguous(); - Tensor out_work = p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); + const bool task_updating = graph::is_task_updating(); + if (!task_updating) { + if (!p->q->is_contiguous()) { + p->q_work->copy_from(p->q); + } + if (!p->k_cache->is_contiguous()) { + p->k_work->copy_from(p->k_cache); + } + if (!p->v_cache->is_contiguous()) { + p->v_work->copy_from(p->v_cache); + } + if (!p->block_table->is_contiguous()) { + p->block_table_work->copy_from(p->block_table); + } + } + Tensor q_work = p->q_work; + Tensor k_work = p->k_work; + Tensor v_work = p->v_work; + Tensor bt_work = p->block_table_work; + Tensor out_work = p->out_work; aclDataType q_dtype = to_acl_dtype(q_work->dtype()); - // Read seqlens_k to host - auto seqlens_k_shape = p->seqlens_k->shape(); - int64_t seqlens_k_len = seqlens_k_shape[0]; - std::vector seqlens_k_host(seqlens_k_len); - auto copy_ret = aclrtMemcpy(seqlens_k_host.data(), seqlens_k_len * sizeof(int32_t), - reinterpret_cast(p->seqlens_k->data()), - seqlens_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); - if (copy_ret != ACL_SUCCESS) { - throw std::runtime_error( - std::string("[mha_kvcache/ascend] copy seqlens_k to host failed: ") + std::to_string(copy_ret)); - } - - // Build actual_seq vectors std::vector actual_seq_q_vec(batch_size, 1); // decode: always 1 query token - std::vector actual_seq_k_vec; - actual_seq_k_vec.reserve(batch_size); - for (int64_t i = 0; i < batch_size; ++i) { - actual_seq_k_vec.push_back(seqlens_k_host[i]); - } + auto actual_seq_k_vec = get_actual_seq_lengths_k(p, batch_size); // BNSD [batch, num_heads, 1, head_size], viewed from BSND memory. std::vector q_dims = {batch_size, num_heads, 1, head_size}; @@ -248,8 +303,22 @@ void run(void *planned_meta) { .buf; } + const bool capture_task_group = graph::is_task_group_capturing(); + if (capture_task_group) { + graph::begin_task_group_capture(); + } + const bool update_task_group = graph::is_task_updating(); + if (update_task_group) { + graph::begin_task_update(); + } ret = aclnnFusedInferAttentionScoreV4(workspace, workspace_size, executor, stream); + if (update_task_group) { + graph::end_task_update(); + } + if (capture_task_group) { + graph::end_task_group_capture(); + } // Release aclTensor/aclTensorList/aclIntArray resources aclDestroyTensor(query_acl); @@ -269,7 +338,7 @@ void run(void *planned_meta) { } // Copy back if out was not contiguous - if (!p->out->is_contiguous()) { + if (!task_updating && !p->out->is_contiguous()) { p->out->copy_from(out_work); } } diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc index 0c5b3ae8c..b6d441520 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc @@ -16,6 +16,11 @@ MhaKVCache::MhaKVCache(Tensor out, INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, seqlens_k, block_table); INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); +#if defined(ENABLE_ASCEND_FLASH_ATTN) + if (out->device().getType() == Device::Type::ASCEND) { + enable_task_update(); + } +#endif } void MhaKVCache::execute(Tensor out, diff --git a/src/infinirt/ascend/infinirt_ascend.cc b/src/infinirt/ascend/infinirt_ascend.cc index 252b9b8e4..4217196d2 100644 --- a/src/infinirt/ascend/infinirt_ascend.cc +++ b/src/infinirt/ascend/infinirt_ascend.cc @@ -215,6 +215,34 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_SUCCESS; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + CHECK_ACLRT(aclmdlRICaptureTaskGrpBegin((aclrtStream)stream)); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + aclrtTaskGrp task_group = nullptr; + CHECK_ACLRT(aclmdlRICaptureTaskGrpEnd((aclrtStream)stream, &task_group)); + *handle = reinterpret_cast(task_group); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + CHECK_ACLRT(aclmdlRICaptureTaskUpdateBegin( + (aclrtStream)stream, + reinterpret_cast(handle))); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + CHECK_ACLRT(aclmdlRICaptureTaskUpdateEnd((aclrtStream)stream)); + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/bang/infinirt_bang.cc b/src/infinirt/bang/infinirt_bang.cc index ff40e58cc..14eb91e16 100644 --- a/src/infinirt/bang/infinirt_bang.cc +++ b/src/infinirt/bang/infinirt_bang.cc @@ -182,6 +182,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/cpu/infinirt_cpu.cc b/src/infinirt/cpu/infinirt_cpu.cc index 9036e484f..3a03be06e 100644 --- a/src/infinirt/cpu/infinirt_cpu.cc +++ b/src/infinirt/cpu/infinirt_cpu.cc @@ -196,4 +196,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + } // namespace infinirt::cpu diff --git a/src/infinirt/cuda/infinirt_cuda.cu b/src/infinirt/cuda/infinirt_cuda.cu index f9fc14c42..d33269e7a 100644 --- a/src/infinirt/cuda/infinirt_cuda.cu +++ b/src/infinirt/cuda/infinirt_cuda.cu @@ -567,4 +567,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre CHECK_CUDART(cudaGraphLaunch((cudaGraphExec_t)graph_exec, (cudaStream_t)stream)); return INFINI_STATUS_SUCCESS; } + +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} } diff --git a/src/infinirt/infinirt.cc b/src/infinirt/infinirt.cc index 72614865b..102757ff2 100644 --- a/src/infinirt/infinirt.cc +++ b/src/infinirt/infinirt.cc @@ -249,3 +249,23 @@ __INFINI_C infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exe __INFINI_C infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) { INFINIRT_CALL_DEVICE_API(graphLuanch, (graph_exec, stream)); } + +__INFINI_C infiniStatus_t infinirtGraphTaskGroupBegin(infinirtStream_t stream) { + INFINIRT_CALL_DEVICE_API(graphTaskGroupBegin, (stream)); +} + +__INFINI_C infiniStatus_t infinirtGraphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + INFINIRT_CALL_DEVICE_API(graphTaskGroupEnd, (stream, handle)); +} + +__INFINI_C infiniStatus_t infinirtGraphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + INFINIRT_CALL_DEVICE_API(graphTaskUpdateBegin, (stream, handle)); +} + +__INFINI_C infiniStatus_t infinirtGraphTaskUpdateEnd(infinirtStream_t stream) { + INFINIRT_CALL_DEVICE_API(graphTaskUpdateEnd, (stream)); +} diff --git a/src/infinirt/infinirt_impl.h b/src/infinirt/infinirt_impl.h index 6c26e0b02..ff17801a9 100644 --- a/src/infinirt/infinirt_impl.h +++ b/src/infinirt/infinirt_impl.h @@ -47,7 +47,11 @@ char *log_buffer, \ size_t buffer_size) IMPL; \ INLINE infiniStatus_t graphExecDestroy(infinirtGraphExec_t graph_exec) IMPL; \ - INLINE infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) IMPL; + INLINE infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) IMPL; \ + INLINE infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) IMPL; \ + INLINE infiniStatus_t graphTaskGroupEnd(infinirtStream_t stream, infinirtGraphTaskGroup_t *handle) IMPL; \ + INLINE infiniStatus_t graphTaskUpdateBegin(infinirtStream_t stream, infinirtGraphTaskGroup_t handle) IMPL; \ + INLINE infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) IMPL; #define INFINIRT_DEVICE_API_IMPL INFINIRT_DEVICE_API(, , ) #define INFINIRT_DEVICE_API_NOOP INFINIRT_DEVICE_API( \ diff --git a/src/infinirt/kunlun/infinirt_kunlun.cc b/src/infinirt/kunlun/infinirt_kunlun.cc index 02c46bb33..37a2d655a 100644 --- a/src/infinirt/kunlun/infinirt_kunlun.cc +++ b/src/infinirt/kunlun/infinirt_kunlun.cc @@ -190,6 +190,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/metax/infinirt_metax.cc b/src/infinirt/metax/infinirt_metax.cc index d9ea00b3f..878240b4d 100644 --- a/src/infinirt/metax/infinirt_metax.cc +++ b/src/infinirt/metax/infinirt_metax.cc @@ -269,4 +269,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_SUCCESS; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + } // namespace infinirt::metax diff --git a/src/infinirt/moore/infinirt_moore.cc b/src/infinirt/moore/infinirt_moore.cc index 301cbe77c..54ceee218 100644 --- a/src/infinirt/moore/infinirt_moore.cc +++ b/src/infinirt/moore/infinirt_moore.cc @@ -178,6 +178,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; }