已合并
feat: Eager自定义算子RT2路径辅流申请功能 #5139
Chang-an-HW创建于 9月23日
feat: Eager自定义算子RT2路径辅流申请功能 #5139
已合并
共 20 个文件变更+851-29
| @@ -13,8 +13,8 @@ | |||
| 13 | 13 | ||
| 14 | 按 key 申请框架托管的物理辅流，供 Eager 类自定义算子在[`Execute`](../../ge/EagerExecuteOp/Execute.md)回调中申请辅流，并在辅流上下发 kernel。 | 14 | 按 key 申请框架托管的物理辅流，供 Eager 类自定义算子在[`Execute`](../../ge/EagerExecuteOp/Execute.md)回调中申请辅流，并在辅流上下发 kernel。 |
| 15 | 15 | ||
| 16 | -- 首次申请时，框架创建物理流并绑定到当前模型实例；同一模型实例内使用相同 key 的后续申请直接返回已创建的流句柄，支持跨节点、跨算子类型复用。 | 16 | +- 首次申请时，框架创建物理流并托管其生命周期；同一模型执行实例内使用相同 key 的后续申请直接返回已创建的流句柄，支持跨节点、跨算子类型复用。 |
| 17 | -- 不同 key 返回不同的流句柄；不同模型实例之间的辅流相互隔离，即使 key 相同也不共享。 | 17 | +- 不同 key 返回不同的流句柄；不同模型执行实例之间的辅流相互隔离，即使 key 相同也不共享。 |
| 18 | 18 | ||
| 19 | ## 函数原型 | 19 | ## 函数原型 |
| 20 | 20 | ||
| @@ -26,19 +26,21 @@ rtStream RequestAttachedStream(const ge::AscendString &key) | |||
| 26 | 26 | ||
| 27 | | 参数名 | 输入/输出 | 描述 | | 27 | | 参数名 | 输入/输出 | 描述 | |
| 28 | | --- | --- | --- | | 28 | | --- | --- | --- | |
| 29 | -| key | 输入 | 辅流复用 key。不能为空字符串，且不能以`CANN-FMK-`保留前缀开头。 | | 29 | +| key | 输入 | 辅流复用 key。不能为空字符串，为空时返回nullptr。 | |
| 30 | 30 | ||
| 31 | ## 返回值说明 | 31 | ## 返回值说明 |
| 32 | 32 | ||
| 33 | | 参数名 | 类型 | 说明 | | 33 | | 参数名 | 类型 | 说明 | |
| 34 | | --- | --- | --- | | 34 | | --- | --- | --- | |
| 35 | -| - | rtStream | 成功返回框架托管的辅流句柄；key无效、当前执行链路不支持辅流申请、或物理流创建/绑定失败时返回nullptr。 | | 35 | +| - | rtStream | 成功返回框架托管的辅流句柄；key为空、当前执行链路不支持辅流申请、或物理流创建/绑定失败时返回nullptr。 | |
| 36 | 36 | ||
| 37 | ## 约束说明 | 37 | ## 约束说明 |
| 38 | 38 | ||
| 39 | - 主辅流之间的执行顺序由算子自行管理，如：通过 event（`aclrtRecordEvent`/`aclrtStreamWaitEvent`）自行同步，框架不保证执行顺序。 | 39 | - 主辅流之间的执行顺序由算子自行管理，如：通过 event（`aclrtRecordEvent`/`aclrtStreamWaitEvent`）自行同步，框架不保证执行顺序。 |
| 40 | +- 算子必须在[`Execute`](../../ge/EagerExecuteOp/Execute.md)返回前把辅流上的工作 join 回`GetStream()`返回的执行流（辅流 record event、执行流 wait event）。RT2 链路的内存回收只感知框架逻辑流，若不 join 就返回，本轮分配的 args 与输出内存可能在辅流 kernel 执行完成前被回收复用。 | ||
| 40 | - 流的生命周期由框架管理，算子不得销毁返回的辅流（禁止调用`aclrtDestroyStream`）。 | 41 | - 流的生命周期由框架管理，算子不得销毁返回的辅流（禁止调用`aclrtDestroyStream`）。 |
| 41 | - 此接口不支持并发调用。 | 42 | - 此接口不支持并发调用。 |
| 43 | +- `CANN-FMK-`是框架内部使用的辅流key前缀。用户自定义算子应避免以该前缀命名key，以免与框架内部申请的辅流相互覆盖。 | ||
| 42 | 44 | ||
| 43 | ## 调用示例 | 45 | ## 调用示例 |
| 44 | 46 | ||
| @@ -15,19 +15,10 @@ | |||
| 15 | #include "graph/ascend_string.h" | 15 | #include "graph/ascend_string.h" |
| 16 | 16 | ||
| 17 | namespace ge { | 17 | namespace ge { |
| 18 | -namespace { | ||
| 19 | -constexpr const char *kReservedPrefix = "CANN-FMK-"; | ||
| 20 | -} | ||
| 21 | - | ||
| 22 | bool AttachedStreamCollection::IsValidKey(const AscendString &key) const { | 18 | bool AttachedStreamCollection::IsValidKey(const AscendString &key) const { |
H | |||
| 23 | const char_t *const key_str = key.GetString(); | 19 | const char_t *const key_str = key.GetString(); |
| 24 | - if (key_str == nullptr) { | 20 | + if ((key_str == nullptr) || (key_str[0] == '\0')) { |
| 25 | - GELOGE(PARAM_INVALID, "Invalid eager attached stream key: key string is null."); | 21 | + GELOGE(PARAM_INVALID, "Invalid eager attached stream key: key is null or empty."); |
| 26 | - return false; | ||
| 27 | - } | ||
| 28 | - const std::string value(key_str); | ||
| 29 | - if (value.empty() || value.rfind(kReservedPrefix, 0U) == 0U) { | ||
| 30 | - GELOGE(PARAM_INVALID, "Invalid eager attached stream key: empty or reserved prefix, key=%s.", value.c_str()); | ||
| 31 | return false; | 22 | return false; |
| 32 | } | 23 | } |
| 33 | return true; | 24 | return true; |
| @@ -57,6 +48,7 @@ gert::rtStream AttachedStreamCollection::RequestAttachedStream(const AscendStrin | |||
| 57 | const auto bind_ret = aclmdlRIBindStream(model_, stream, ACL_MODEL_STREAM_FLAG_HEAD); | 48 | const auto bind_ret = aclmdlRIBindStream(model_, stream, ACL_MODEL_STREAM_FLAG_HEAD); |
| 58 | if (bind_ret != ACL_SUCCESS) { | 49 | if (bind_ret != ACL_SUCCESS) { |
| 59 | GELOGE(FAILED, "Failed to bind eager attached stream, key=%s, ret=%d.", value.c_str(), bind_ret); | 50 | GELOGE(FAILED, "Failed to bind eager attached stream, key=%s, ret=%d.", value.c_str(), bind_ret); |
| 51 | + // 刚建好的流上没有任何任务，无需先同步即可销毁 | ||
| 60 | (void)aclrtDestroyStream(stream); | 52 | (void)aclrtDestroyStream(stream); |
| 61 | return nullptr; | 53 | return nullptr; |
| 62 | } | 54 | } |
| @@ -66,6 +58,12 @@ gert::rtStream AttachedStreamCollection::RequestAttachedStream(const AscendStrin | |||
| 66 | 58 | ||
| 67 | void AttachedStreamCollection::UnbindAndDestroy() { | 59 | void AttachedStreamCollection::UnbindAndDestroy() { |
| 68 | for (const auto &item : streams_) { | 60 | for (const auto &item : streams_) { |
| 61 | + // 销毁前必须同步：aclrtDestroyStream 的前置约束要求流上任务已执行完，且解绑也会改动 rtModel 的流表。 | ||
| 62 | + // 同步不设超时：超时后仍会继续解绑销毁、并由 DestroyResources 释放模型内存，等于把阻塞换成设备侧写已释放内存。 | ||
| 63 | + // 时序与 Rt2AttachedStreamCollection::Destroy 一致（RT2 无 rtModel，故少一步解绑）。 | ||
| 64 | + GELOGI("Synchronizing eager attached stream before destroy, key=%s.", item.first.c_str()); | ||
| 65 | + GE_LOGW_IF(aclrtSynchronizeStream(item.second) != ACL_SUCCESS, | ||
| 66 | + "Failed to synchronize eager attached stream, key=%s.", item.first.c_str()); | ||
| 69 | if (model_ != nullptr) { | 67 | if (model_ != nullptr) { |
| 70 | GE_LOGW_IF(aclmdlRIUnbindStream(model_, item.second) != ACL_SUCCESS, | 68 | GE_LOGW_IF(aclmdlRIUnbindStream(model_, item.second) != ACL_SUCCESS, |
| 71 | "Failed to unbind eager attached stream, key=%s.", item.first.c_str()); | 69 | "Failed to unbind eager attached stream, key=%s.", item.first.c_str()); |
| @@ -380,7 +380,7 @@ DavinciModel::~DavinciModel() noexcept { | |||
| 380 | // clear exception dump info before stream release | 380 | // clear exception dump info before stream release |
| 381 | exception_dumper_.Clear(); | 381 | exception_dumper_.Clear(); |
| 382 | UnbindTaskSinkStream(); | 382 | UnbindTaskSinkStream(); |
| 383 | - attached_stream_collection_.UnbindAndDestroy(); | 383 | + UnbindAndDestroyAttachedStreams(); |
| 384 | DestroyStream(); | 384 | DestroyStream(); |
| 385 | DestroyResources(); | 385 | DestroyResources(); |
| 386 | GELOGI("Npu model: %u success to finalize.", model_id_); | 386 | GELOGI("Npu model: %u success to finalize.", model_id_); |
| @@ -5874,6 +5874,17 @@ void DavinciModel::UnbindTaskSinkStream() { | |||
| 5874 | GELOGD("Npu model: %u success to unbind streams.", model_id_); | 5874 | GELOGD("Npu model: %u success to unbind streams.", model_id_); |
| 5875 | } | 5875 | } |
| 5876 | 5876 | ||
| 5877 | +void DavinciModel::UnbindAndDestroyAttachedStreams() { | ||
| 5878 | + // check rt ctx is exist. rt api call will cause error log when ctx does not exist | ||
| 5879 | + aclrtContext current_ctx = nullptr; | ||
| 5880 | + if (aclrtGetCurrentContext(¤t_ctx) != ACL_SUCCESS) { | ||
| 5881 | + return; | ||
| 5882 | + } | ||
| 5883 | + // 内部已清空缓存并置空 model 句柄，重复调用无副作用 | ||
| 5884 | + attached_stream_collection_.UnbindAndDestroy(); | ||
T 这个新方法带了 rt ctx 检查，但 另外同步语义也不一致：RT2 版   | |||
| 5885 | + GELOGD("Npu model: %u success to unbind and destroy attached streams.", model_id_); | ||
| 5886 | +} | ||
| 5887 | + | ||
| 5877 | void DavinciModel::DestroyStream() { | 5888 | void DavinciModel::DestroyStream() { |
| 5878 | GELOGD("Npu model: %u start to destroy stream.", model_id_); | 5889 | GELOGD("Npu model: %u start to destroy stream.", model_id_); |
| 5879 | // check rt ctx is exist. rt api call will cause error log when ctx does not exist | 5890 | // check rt ctx is exist. rt api call will cause error log when ctx does not exist |
| @@ -163,6 +163,10 @@ class DavinciModel { | |||
| 163 | /// being destroyed | 163 | /// being destroyed |
| 164 | void UnbindTaskSinkStream(); | 164 | void UnbindTaskSinkStream(); |
| 165 | 165 | ||
| 166 | + /// @ingroup ge | ||
| 167 | + /// @brief 解绑并销毁 Eager 自定义算子申请的辅流，必须在 DestroyStream 之前调用；重复调用无副作用 | ||
| 168 | + void UnbindAndDestroyAttachedStreams(); | ||
| 169 | + | ||
| 166 | void DestroyStream(); | 170 | void DestroyStream(); |
| 167 | 171 | ||
| 168 | void DestroyResources(); | 172 | void DestroyResources(); |
| @@ -12,6 +12,7 @@ | |||
| 12 | #include "graph_builder/converter_checker.h" | 12 | #include "graph_builder/converter_checker.h" |
| 13 | #include "graph_builder/bg_tensor.h" | 13 | #include "graph_builder/bg_tensor.h" |
| 14 | #include "graph_builder/bg_infer_shape.h" | 14 | #include "graph_builder/bg_infer_shape.h" |
| 15 | +#include "graph_builder/bg_attached_stream_provider.h" | ||
| 15 | #include "common/checker.h" | 16 | #include "common/checker.h" |
| 16 | #include "exe_graph/lowering/frame_selector.h" | 17 | #include "exe_graph/lowering/frame_selector.h" |
| 17 | #include "kernel/common_kernel_impl/build_tensor.h" | 18 | #include "kernel/common_kernel_impl/build_tensor.h" |
| @@ -146,6 +147,11 @@ LowerResult LoweringCustomNode(const ge::NodePtr &node, const LowerInput &lower_ | |||
| 146 | input_holders.emplace_back(allocator_holder); | 147 | input_holders.emplace_back(allocator_holder); |
| 147 | input_holders.emplace_back(lower_input.global_data->GetStream()); | 148 | input_holders.emplace_back(lower_input.global_data->GetStream()); |
| 148 | input_holders.emplace_back(custom_executor_func); | 149 | input_holders.emplace_back(custom_executor_func); |
| 150 | + // 辅流容器：Init 图创建、执行器析构时释放，per-executor 隔离。 | ||
| 151 | + // 必须插在 custom_executor_func 之后、infer 模板 tensor 之前，与 CustomOpInput 枚举顺序保持一致。 | ||
| 152 | + const auto attached_stream_provider = bg::GetAttachedStreamProvider(*lower_input.global_data); | ||
| 153 | + LOWER_REQUIRE_NOTNULL(attached_stream_provider); | ||
| 154 | + input_holders.emplace_back(attached_stream_provider); | ||
| 149 | // Check inference_rule | 155 | // Check inference_rule |
| 150 | const auto op_desc = node->GetOpDesc(); | 156 | const auto op_desc = node->GetOpDesc(); |
| 151 | LOWER_REQUIRE_NOTNULL(op_desc); | 157 | LOWER_REQUIRE_NOTNULL(op_desc); |
| @@ -0,0 +1,40 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#include "common/checker.h" | ||
| 12 | +#include "common/plugin/ge_make_unique_util.h" | ||
| 13 | +#include "engine/custom/kernel/rt2_attached_stream_collection.h" | ||
| 14 | +#include "register/kernel_registry.h" | ||
| 15 | + | ||
| 16 | +namespace gert { | ||
| 17 | +namespace kernel { | ||
| 18 | +// Init 图节点：辅流容器对象由 OutputsCreator 创建，RunFunc 无需做任何事 | ||
| 19 | +ge::graphStatus CreateAttachedStreamProvider(KernelContext *context) { | ||
| 20 | + (void)context; | ||
| 21 | + return ge::GRAPH_SUCCESS; | ||
| 22 | +} | ||
| 23 | + | ||
| 24 | +ge::graphStatus BuildAttachedStreamProviderOutputs(const ge::FastNode *node, KernelContext *context) { | ||
| 25 | + (void)node; | ||
| 26 | + auto provider_chain = context->GetOutput(0); | ||
| 27 | + GE_ASSERT_NOTNULL(provider_chain); | ||
| 28 | + auto provider = ge::MakeUnique<Rt2AttachedStreamCollection>(); | ||
| 29 | + GE_ASSERT_NOTNULL(provider); | ||
| 30 | + // 容器归 Chain 所有，执行器析构时同步并销毁全部辅流（ACL 与单算子路径均在 UnLoad 后立即销毁执行器） | ||
| 31 | + provider_chain->SetWithDefaultDeleter(provider.release()); | ||
| 32 | + GELOGD("Create eager attached stream provider for current rt2 executor."); | ||
| 33 | + return ge::GRAPH_SUCCESS; | ||
| 34 | +} | ||
| 35 | + | ||
| 36 | +REGISTER_KERNEL(CreateAttachedStreamProvider) | ||
| 37 | + .RunFunc(CreateAttachedStreamProvider) | ||
| 38 | + .OutputsCreator(BuildAttachedStreamProviderOutputs); | ||
| 39 | +} // namespace kernel | ||
| 40 | +} // namespace gert | ||
| @@ -24,12 +24,17 @@ | |||
| 24 | #include "rt_external_kernel.h" | 24 | #include "rt_external_kernel.h" |
| 25 | #include "core/executor/multi_thread_topological/executor/schedule/producer/producers/kernel_tags/critical_section_config.h" | 25 | #include "core/executor/multi_thread_topological/executor/schedule/producer/producers/kernel_tags/critical_section_config.h" |
| 26 | #include "runtime/v2/engine/custom/kernel/eager_args_handler.h" | 26 | #include "runtime/v2/engine/custom/kernel/eager_args_handler.h" |
| 27 | +#include "runtime/v2/engine/custom/kernel/rt2_attached_stream_collection.h" | ||
| 27 | 28 | ||
| 28 | namespace gert { | 29 | namespace gert { |
| 29 | namespace kernel { | 30 | namespace kernel { |
| 30 | namespace { | 31 | namespace { |
| 31 | // 自定义算子特有的输入，从 AdditionalInputIndex::kNum 开始 | 32 | // 自定义算子特有的输入，从 AdditionalInputIndex::kNum 开始 |
| 32 | -enum class CustomOpInput { kFunc = static_cast<uint32_t>(EagerOpExecutionContext::AdditionalInputIndex::kNum), kEnd }; | 33 | +enum class CustomOpInput { |
| 34 | + kFunc = static_cast<uint32_t>(EagerOpExecutionContext::AdditionalInputIndex::kNum), | ||
| 35 | + kAttachedStreamProvider, | ||
| 36 | + kEnd | ||
| 37 | +}; | ||
| 33 | enum class HostCustomOpInput { | 38 | enum class HostCustomOpInput { |
| 34 | kFunc = static_cast<uint32_t>(HostCpuOpExecutionContext::AdditionalInputIndex::kNum), | 39 | kFunc = static_cast<uint32_t>(HostCpuOpExecutionContext::AdditionalInputIndex::kNum), |
| 35 | kEnd | 40 | kEnd |
| @@ -201,7 +206,12 @@ static ge::graphStatus ExecuteCustomOpImpl(KernelContext *context) { | |||
| 201 | GE_ASSERT_NOTNULL(allocator); | 206 | GE_ASSERT_NOTNULL(allocator); |
| 202 | auto stream_id = allocator->GetStreamId(); | 207 | auto stream_id = allocator->GetStreamId(); |
| 203 | args_handler->Initialize(allocator, stream_id); | 208 | args_handler->Initialize(allocator, stream_id); |
| 204 | - GELOGD("EagerArgsHandler initialized in RunFunc with allocator %p", allocator); | 209 | + // 辅流容器由 Init 图创建；取不到时为 nullptr，此时 RequestAttachedStream 按既有语义降级返回 nullptr |
| 210 | + auto *attached_streams = context->GetInputValue<Rt2AttachedStreamCollection *>( | ||
| 211 | + node_input_num + static_cast<size_t>(CustomOpInput::kAttachedStreamProvider)); | ||
| 212 | + args_handler->SetAttachedStreamProvider(attached_streams); | ||
| 213 | + GELOGD("EagerArgsHandler initialized in RunFunc with allocator %p, attached stream provider %p", allocator, | ||
| 214 | + attached_streams); | ||
| 205 | } | 215 | } |
| 206 | 216 | ||
| 207 | auto custom_op_ptr = | 217 | auto custom_op_ptr = |
| @@ -25,18 +25,32 @@ class EagerArgsHandler : public ArgsHandler { | |||
| 25 | public: | 25 | public: |
| 26 | EagerArgsHandler() : allocator_(nullptr), stream_id_(-1) {} | 26 | EagerArgsHandler() : allocator_(nullptr), stream_id_(-1) {} |
| 27 | EagerArgsHandler(GertAllocator *allocator, int64_t stream_id); | 27 | EagerArgsHandler(GertAllocator *allocator, int64_t stream_id); |
| 28 | - ~EagerArgsHandler() override { Release(); } | 28 | + ~EagerArgsHandler() override { |
| 29 | + Release(); | ||
| 30 | + } | ||
| 29 | 31 | ||
| 30 | void Initialize(GertAllocator *allocator, int64_t stream_id); | 32 | void Initialize(GertAllocator *allocator, int64_t stream_id); |
| 31 | - bool IsInitialized() const { return allocator_ != nullptr; } | 33 | + bool IsInitialized() const { |
| 32 | - const KernelArgs* MallocReadOnlyDevArgs(void *host_args, size_t args_size) override; | 34 | + return allocator_ != nullptr; |
| 33 | - const std::deque<KernelArgs>& GetKernelArgs(Placement placement) const override; | 35 | + } |
| 36 | + const KernelArgs *MallocReadOnlyDevArgs(void *host_args, size_t args_size) override; | ||
| 37 | + const std::deque<KernelArgs> &GetKernelArgs(Placement placement) const override; | ||
| 38 | + /// 注入 Eager 自定义算子辅流容器；容器由 Init 图创建、执行器析构时销毁，本类不拥有其生命周期。 | ||
| 39 | + /// 注意：provider_ 是裸指针，析构顺序不保证容器晚于本对象释放，因此 Release() 与本类析构函数中 | ||
| 40 | + /// 都不得使用 provider_（Release 每轮执行都会调用，也不能在此销毁容器）。 | ||
| 41 | + void SetAttachedStreamProvider(AttachedStreamProvider *provider) { | ||
| 42 | + provider_ = provider; | ||
| 43 | + } | ||
| 44 | + AttachedStreamProvider *GetAttachedStreamProvider() override { | ||
| 45 | + return provider_; | ||
| 46 | + } | ||
| 34 | 47 | ||
| 35 | void Release(); | 48 | void Release(); |
| 36 | 49 | ||
| 37 | private: | 50 | private: |
| 38 | GertAllocator *allocator_; // 生命周期由框架保证，覆盖整个执行过程 | 51 | GertAllocator *allocator_; // 生命周期由框架保证，覆盖整个执行过程 |
| 39 | int64_t stream_id_; | 52 | int64_t stream_id_; |
| 53 | + AttachedStreamProvider *provider_{nullptr}; | ||
| 40 | std::deque<KernelArgs> device_args_; | 54 | std::deque<KernelArgs> device_args_; |
| 41 | std::vector<GertMemBlock *> allocated_blocks_; | 55 | std::vector<GertMemBlock *> allocated_blocks_; |
| 42 | }; | 56 | }; |
| @@ -0,0 +1,56 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#include "rt2_attached_stream_collection.h" | ||
| 12 | + | ||
| 13 | +#include <string_view> | ||
| 14 | + | ||
| 15 | +#include "acl/acl_rt.h" | ||
| 16 | +#include "framework/common/debug/ge_log.h" | ||
| 17 | +#include "graph/ascend_string.h" | ||
| 18 | + | ||
| 19 | +namespace gert { | ||
| 20 | +rtStream Rt2AttachedStreamCollection::RequestAttachedStream(const ge::AscendString &key) { | ||
| 21 | + // 用 GetString 判空而不用 GetLength：stub 版 libgraph.so 的 GetLength 恒返回 0 | ||
| 22 | + const char *const key_str = key.GetString(); | ||
| 23 | + if ((key_str == nullptr) || (key_str[0] == '\0')) { | ||
| 24 | + GELOGE(ge::PARAM_INVALID, "Invalid eager attached stream key: key is null or empty."); | ||
| 25 | + return nullptr; | ||
| 26 | + } | ||
| 27 | + const auto found = streams_.find(std::string_view(key_str)); | ||
| 28 | + if (found != streams_.end()) { | ||
| 29 | + return found->second; | ||
| 30 | + } | ||
| 31 | + rtStream_t stream = nullptr; | ||
| 32 | + if (rtStreamCreateWithFlags(&stream, RT_STREAM_PRIORITY_DEFAULT, RT_STREAM_DEFAULT) != RT_ERROR_NONE) { | ||
| 33 | + GELOGE(ge::FAILED, "Failed to create eager attached stream, key=%s.", key_str); | ||
| 34 | + return nullptr; | ||
| 35 | + } | ||
| 36 | + (void)streams_.emplace(key_str, stream); | ||
| 37 | + GELOGI("Created eager attached stream, key=%s, total=%zu.", key_str, streams_.size()); | ||
| 38 | + return stream; | ||
| 39 | +} | ||
| 40 | + | ||
| 41 | +void Rt2AttachedStreamCollection::Destroy() { | ||
| 42 | + for (const auto &item : streams_) { | ||
| 43 | + // 先流同步再销毁：aclrtDestroyStream 的前置约束要求流上任务已执行完。 | ||
| 44 | + // 同步不设超时：超时后仍会继续销毁并释放执行器内存，等于把阻塞换成设备侧写已释放内存。 | ||
| 45 | + // 时序与 V1 的 ge::AttachedStreamCollection::UnbindAndDestroy 一致（V1 多一步 rtModel 解绑）。 | ||
| 46 | + GELOGI("Synchronizing eager attached stream before destroy, key=%s.", item.first.c_str()); | ||
| 47 | + if (aclrtSynchronizeStream(item.second) != ACL_SUCCESS) { | ||
| 48 | + GELOGW("Failed to synchronize eager attached stream, key=%s.", item.first.c_str()); | ||
| 49 | + } | ||
| 50 | + if (aclrtDestroyStream(item.second) != ACL_SUCCESS) { | ||
| 51 | + GELOGW("Failed to destroy eager attached stream, key=%s.", item.first.c_str()); | ||
| 52 | + } | ||
| 53 | + } | ||
| 54 | + streams_.clear(); | ||
| 55 | +} | ||
| 56 | +} // namespace gert | ||
| @@ -0,0 +1,51 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#ifndef AIR_CXX_RUNTIME_V2_ENGINE_CUSTOM_KERNEL_RT2_ATTACHED_STREAM_COLLECTION_H_ | ||
| 12 | +#define AIR_CXX_RUNTIME_V2_ENGINE_CUSTOM_KERNEL_RT2_ATTACHED_STREAM_COLLECTION_H_ | ||
| 13 | + | ||
| 14 | +#include <functional> | ||
| 15 | +#include <map> | ||
| 16 | +#include <string> | ||
| 17 | + | ||
| 18 | +#include "framework/runtime/attached_stream_provider.h" | ||
| 19 | +#include "rt_external_stream.h" | ||
| 20 | + | ||
| 21 | +namespace gert { | ||
| 22 | +/** | ||
| 23 | + * @brief RT2 链路的 Eager 自定义算子辅流容器 | ||
| 24 | + * | ||
| 25 | + * 生命周期归属单个 ModelV2Executor：Init 图节点的 OutputsCreator 创建对象并交给输出 Chain 持有 | ||
| 26 | + * （SetWithDefaultDeleter），Chain 释放对象时经析构函数进入 Destroy()。无 DeInit 侧接线。 | ||
| 27 | + * 与 V1 的 ge::AttachedStreamCollection 的差异：RT2 没有 rtModel_t，不需要 aclmdlRIBindStream 绑定， | ||
| 28 | + * 也没有模型流 WAIT_ACTIVE 激活语义，因此按默认 flag 建流即可。 | ||
| 29 | + * 注意：本类与编译期的 _attached_stream_num（HCCL/Mix Vector Core 使用的附着流）无关，两者同名不同义。 | ||
| 30 | + */ | ||
| 31 | +class Rt2AttachedStreamCollection final : public AttachedStreamProvider { | ||
| 32 | + public: | ||
| 33 | + Rt2AttachedStreamCollection() = default; | ||
| 34 | + ~Rt2AttachedStreamCollection() override { | ||
| 35 | + Destroy(); | ||
| 36 | + } | ||
| 37 | + Rt2AttachedStreamCollection(const Rt2AttachedStreamCollection &) = delete; | ||
| 38 | + Rt2AttachedStreamCollection &operator=(const Rt2AttachedStreamCollection &) = delete; | ||
| 39 | + | ||
| 40 | + rtStream RequestAttachedStream(const ge::AscendString &key) override; | ||
| 41 | + | ||
| 42 | + /// 同步并销毁全部辅流；幂等。生产路径仅由析构函数调用，公开是为便于单测直接校验销毁时序 | ||
| 43 | + void Destroy(); | ||
| 44 | + | ||
| 45 | + private: | ||
| 46 | + // 透明比较器：命中路径以 string_view 查找，避免每轮执行都构造 std::string | ||
| 47 | + std::map<std::string, rtStream_t, std::less<>> streams_; | ||
| 48 | +}; | ||
| 49 | +} // namespace gert | ||
| 50 | + | ||
| 51 | +#endif // AIR_CXX_RUNTIME_V2_ENGINE_CUSTOM_KERNEL_RT2_ATTACHED_STREAM_COLLECTION_H_ | ||
| @@ -0,0 +1,32 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#include "bg_attached_stream_provider.h" | ||
| 12 | + | ||
| 13 | +#include "exe_graph/lowering/frame_selector.h" | ||
| 14 | + | ||
| 15 | +namespace gert { | ||
| 16 | +namespace bg { | ||
| 17 | +namespace { | ||
| 18 | +const std::string kCreateAttachedStreamProvider = "CreateAttachedStreamProvider"; | ||
| 19 | +} // namespace | ||
| 20 | + | ||
| 21 | +ValueHolderPtr GetAttachedStreamProvider(LoweringGlobalData &global_data) { | ||
| 22 | + auto builder = []() -> std::vector<ValueHolderPtr> { | ||
| 23 | + return FrameSelector::OnInitRoot([]() -> std::vector<ValueHolderPtr> { | ||
| 24 | + return ValueHolder::CreateDataOutput(kCreateAttachedStreamProvider.c_str(), {}, 1U); | ||
| 25 | + }); | ||
| 26 | + }; | ||
| 27 | + const auto &providers = global_data.GetOrCreateUniqueValueHolder(kCreateAttachedStreamProvider, builder); | ||
| 28 | + GE_ASSERT_TRUE(!providers.empty()); | ||
| 29 | + return providers[0]; | ||
| 30 | +} | ||
| 31 | +} // namespace bg | ||
| 32 | +} // namespace gert | ||
| @@ -0,0 +1,27 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#ifndef AIR_CXX_RUNTIME_V2_GRAPH_BUILDER_BG_ATTACHED_STREAM_PROVIDER_H_ | ||
| 12 | +#define AIR_CXX_RUNTIME_V2_GRAPH_BUILDER_BG_ATTACHED_STREAM_PROVIDER_H_ | ||
| 13 | +#include "exe_graph/lowering/lowering_global_data.h" | ||
| 14 | + | ||
| 15 | +namespace gert { | ||
| 16 | +namespace bg { | ||
| 17 | +/** | ||
| 18 | + * @brief 获取当前执行器独有的 Eager 自定义算子辅流容器，首次调用时在 Init 图创建 | ||
| 19 | + * | ||
| 20 | + * 容器对象由 Init 图节点的输出 Chain 持有（SetWithDefaultDeleter），执行器析构时随之释放： | ||
| 21 | + * ACL 模型卸载（UnloadRt2Model -> DeleteExecutor）与单算子流卸载（StreamExecutor::Erase） | ||
| 22 | + * 都在 UnLoad 之后立即销毁执行器，因此辅流不会跨 UnLoad 长期占用设备流配额。 | ||
| 23 | + */ | ||
| 24 | +ValueHolderPtr GetAttachedStreamProvider(LoweringGlobalData &global_data); | ||
| 25 | +} // namespace bg | ||
| 26 | +} // namespace gert | ||
| 27 | +#endif // AIR_CXX_RUNTIME_V2_GRAPH_BUILDER_BG_ATTACHED_STREAM_PROVIDER_H_ | ||
| @@ -28,6 +28,10 @@ ge::graphStatus DavinciModelFinalizer(KernelContext *context) { | |||
| 28 | for (const auto davinci_model : davinci_models) { | 28 | for (const auto davinci_model : davinci_models) { |
| 29 | davinci_model->UnbindTaskSinkStream(); | 29 | davinci_model->UnbindTaskSinkStream(); |
| 30 | } | 30 | } |
| 31 | + // Eager 自定义算子申请的辅流也绑定在 rt model 上，必须先解绑销毁，再销毁模型流 | ||
| 32 | + for (const auto davinci_model : davinci_models) { | ||
| 33 | + davinci_model->UnbindAndDestroyAttachedStreams(); | ||
| 34 | + } | ||
| 31 | for (const auto davinci_model : davinci_models) { | 35 | for (const auto davinci_model : davinci_models) { |
| 32 | davinci_model->DestroyStream(); | 36 | davinci_model->DestroyStream(); |
| 33 | } | 37 | } |
| @@ -109,6 +109,7 @@ set(DISTINCT_GRAPH_LOAD_TEST_FILES | |||
| 109 | "graph/load/hccl_task_info_unittest.cc" | 109 | "graph/load/hccl_task_info_unittest.cc" |
| 110 | "graph/load/custom_task_info_unittest.cc" | 110 | "graph/load/custom_task_info_unittest.cc" |
| 111 | "graph/load/attached_stream_collection_unittest.cc" | 111 | "graph/load/attached_stream_collection_unittest.cc" |
| 112 | + "graph/load/davinci_model_attached_stream_unittest.cc" | ||
| 112 | "graph/load/kernel_ex_task_info_unittest.cc" | 113 | "graph/load/kernel_ex_task_info_unittest.cc" |
| 113 | "graph/load/kernel_task_info_unittest.cc" | 114 | "graph/load/kernel_task_info_unittest.cc" |
| 114 | "graph/load/kernel_manager_unittest.cc" | 115 | "graph/load/kernel_manager_unittest.cc" |
| @@ -46,6 +46,11 @@ class AclForAttachedStreamTest final : public AclRuntimeStub { | |||
| 46 | events.push_back("bind"); | 46 | events.push_back("bind"); |
| 47 | return bind_success ? ACL_SUCCESS : ACL_ERROR_FAILURE; | 47 | return bind_success ? ACL_SUCCESS : ACL_ERROR_FAILURE; |
| 48 | } | 48 | } |
| 49 | + aclError aclrtSynchronizeStream(aclrtStream stream) override { | ||
| 50 | + synced_streams.push_back(stream); | ||
| 51 | + events.push_back("sync"); | ||
| 52 | + return sync_success ? ACL_SUCCESS : ACL_ERROR_RT_FAILURE; | ||
| 53 | + } | ||
| 49 | aclError aclmdlRIUnbindStream(aclmdlRI, aclrtStream stream) override { | 54 | aclError aclmdlRIUnbindStream(aclmdlRI, aclrtStream stream) override { |
| 50 | unbound_streams.push_back(stream); | 55 | unbound_streams.push_back(stream); |
| 51 | events.push_back("unbind"); | 56 | events.push_back("unbind"); |
| @@ -59,17 +64,35 @@ class AclForAttachedStreamTest final : public AclRuntimeStub { | |||
| 59 | 64 | ||
| 60 | std::vector<aclrtStream> bound_streams; | 65 | std::vector<aclrtStream> bound_streams; |
| 61 | std::vector<uint32_t> bind_flags; | 66 | std::vector<uint32_t> bind_flags; |
| 67 | + std::vector<aclrtStream> synced_streams; | ||
| 62 | std::vector<aclrtStream> unbound_streams; | 68 | std::vector<aclrtStream> unbound_streams; |
| 63 | std::vector<aclrtStream> destroyed_streams; | 69 | std::vector<aclrtStream> destroyed_streams; |
| 64 | std::vector<std::string> events; | 70 | std::vector<std::string> events; |
| 65 | bool bind_success{true}; | 71 | bool bind_success{true}; |
| 72 | + bool sync_success{true}; | ||
| 66 | }; | 73 | }; |
| 67 | 74 | ||
| 68 | -TEST(AttachedStreamCollectionTest, RejectsEmptyAndReservedKeysBeforeRuntimeCall) { | 75 | +TEST(AttachedStreamCollectionTest, RejectsEmptyKeysBeforeRuntimeCall) { |
| 69 | AttachedStreamCollection collection(nullptr, RT_STREAM_PRIORITY_DEFAULT, RT_STREAM_DEFAULT); | 76 | AttachedStreamCollection collection(nullptr, RT_STREAM_PRIORITY_DEFAULT, RT_STREAM_DEFAULT); |
| 70 | EXPECT_EQ(collection.RequestAttachedStream(AscendString()), nullptr); | 77 | EXPECT_EQ(collection.RequestAttachedStream(AscendString()), nullptr); |
| 71 | EXPECT_EQ(collection.RequestAttachedStream(AscendString("")), nullptr); | 78 | EXPECT_EQ(collection.RequestAttachedStream(AscendString("")), nullptr); |
| 72 | - EXPECT_EQ(collection.RequestAttachedStream(AscendString("CANN-FMK-internal")), nullptr); | 79 | +} |
| 80 | + | ||
| 81 | +TEST(AttachedStreamCollectionTest, AcceptsFrameworkReservedPrefixKey) { | ||
| 82 | + auto runtime = std::make_shared<RuntimeForAttachedStreamTest>(); | ||
| 83 | + auto acl = std::make_shared<AclForAttachedStreamTest>(); | ||
| 84 | + RuntimeStub::SetInstance(runtime); | ||
| 85 | + AclRuntimeStub::SetInstance(acl); | ||
| 86 | + | ||
| 87 | + AttachedStreamCollection collection(reinterpret_cast<rtModel_t>(0x55U), RT_STREAM_PRIORITY_DEFAULT, | ||
| 88 | + RT_STREAM_DEFAULT); | ||
| 89 | + // CANN-FMK- 前缀保留给 HCCL 等框架组件，属 API 文档软约束，代码不拦截 | ||
| 90 | + EXPECT_NE(collection.RequestAttachedStream(AscendString("CANN-FMK-hccl")), nullptr); | ||
| 91 | + EXPECT_EQ(runtime->create_count_, 1); | ||
| 92 | + | ||
| 93 | + collection.UnbindAndDestroy(); | ||
| 94 | + RuntimeStub::Reset(); | ||
| 95 | + AclRuntimeStub::Reset(); | ||
| 73 | } | 96 | } |
| 74 | 97 | ||
| 75 | TEST(AttachedStreamCollectionTest, ReusesKeyAndCleansUpAfterUnbind) { | 98 | TEST(AttachedStreamCollectionTest, ReusesKeyAndCleansUpAfterUnbind) { |
| @@ -91,12 +114,16 @@ TEST(AttachedStreamCollectionTest, ReusesKeyAndCleansUpAfterUnbind) { | |||
| 91 | EXPECT_EQ(acl->bind_flags[0], static_cast<uint32_t>(ACL_MODEL_STREAM_FLAG_HEAD)); | 114 | EXPECT_EQ(acl->bind_flags[0], static_cast<uint32_t>(ACL_MODEL_STREAM_FLAG_HEAD)); |
| 92 | 115 | ||
| 93 | collection.UnbindAndDestroy(); | 116 | collection.UnbindAndDestroy(); |
| 117 | + EXPECT_EQ(acl->synced_streams.size(), 1U); | ||
| 94 | EXPECT_EQ(acl->unbound_streams.size(), 1U); | 118 | EXPECT_EQ(acl->unbound_streams.size(), 1U); |
| 95 | EXPECT_EQ(acl->destroyed_streams.size(), 1U); | 119 | EXPECT_EQ(acl->destroyed_streams.size(), 1U); |
| 96 | - ASSERT_EQ(acl->events.size(), 3U); | 120 | + // 顺序必须是 bind -> sync -> unbind -> destroy：aclrtDestroyStream 要求流上任务已执行完， |
| 121 | + // 且解绑会改动 rtModel 流表，故同步先于解绑 | ||
| 122 | + ASSERT_EQ(acl->events.size(), 4U); | ||
| 97 | EXPECT_EQ(acl->events[0], "bind"); | 123 | EXPECT_EQ(acl->events[0], "bind"); |
| 98 | - EXPECT_EQ(acl->events[1], "unbind"); | 124 | + EXPECT_EQ(acl->events[1], "sync"); |
| 99 | - EXPECT_EQ(acl->events[2], "destroy"); | 125 | + EXPECT_EQ(acl->events[2], "unbind"); |
| 126 | + EXPECT_EQ(acl->events[3], "destroy"); | ||
| 100 | RuntimeStub::Reset(); | 127 | RuntimeStub::Reset(); |
| 101 | AclRuntimeStub::Reset(); | 128 | AclRuntimeStub::Reset(); |
| 102 | } | 129 | } |
| @@ -117,12 +144,35 @@ TEST(AttachedStreamCollectionTest, BindFailureRollsBackTemporaryStream) { | |||
| 117 | RT_STREAM_DEFAULT); | 144 | RT_STREAM_DEFAULT); |
| 118 | EXPECT_EQ(collection.RequestAttachedStream(AscendString("rollback")), nullptr); | 145 | EXPECT_EQ(collection.RequestAttachedStream(AscendString("rollback")), nullptr); |
| 119 | EXPECT_EQ(runtime->create_count_, 1); | 146 | EXPECT_EQ(runtime->create_count_, 1); |
| 147 | + // 回滚路径销毁的是刚建好、未下发任何任务的流，不需要同步 | ||
| 148 | + EXPECT_TRUE(acl->synced_streams.empty()); | ||
| 120 | EXPECT_EQ(acl->unbound_streams.size(), 0U); | 149 | EXPECT_EQ(acl->unbound_streams.size(), 0U); |
| 121 | EXPECT_EQ(acl->destroyed_streams.size(), 1U); | 150 | EXPECT_EQ(acl->destroyed_streams.size(), 1U); |
| 122 | RuntimeStub::Reset(); | 151 | RuntimeStub::Reset(); |
| 123 | AclRuntimeStub::Reset(); | 152 | AclRuntimeStub::Reset(); |
| 124 | } | 153 | } |
| 125 | 154 | ||
| 155 | +// 同步失败仍须继续解绑销毁：跳过会让 streams_ 清空后句柄彻底丢失，变成确定性泄漏 | ||
| 156 | +TEST(AttachedStreamCollectionTest, SyncFailureStillUnbindsAndDestroys) { | ||
| 157 | + auto runtime = std::make_shared<RuntimeForAttachedStreamTest>(); | ||
| 158 | + auto acl = std::make_shared<AclForAttachedStreamTest>(); | ||
| 159 | + acl->sync_success = false; | ||
| 160 | + RuntimeStub::SetInstance(runtime); | ||
| 161 | + AclRuntimeStub::SetInstance(acl); | ||
| 162 | + | ||
| 163 | + AttachedStreamCollection collection(reinterpret_cast<rtModel_t>(0x55U), RT_STREAM_PRIORITY_DEFAULT, | ||
| 164 | + RT_STREAM_DEFAULT); | ||
| 165 | + ASSERT_NE(collection.RequestAttachedStream(AscendString("sync_fail")), nullptr); | ||
| 166 | + | ||
| 167 | + collection.UnbindAndDestroy(); | ||
| 168 | + ASSERT_EQ(acl->events.size(), 4U); | ||
| 169 | + EXPECT_EQ(acl->events[1], "sync"); | ||
| 170 | + EXPECT_EQ(acl->unbound_streams.size(), 1U); | ||
| 171 | + EXPECT_EQ(acl->destroyed_streams.size(), 1U); | ||
| 172 | + RuntimeStub::Reset(); | ||
| 173 | + AclRuntimeStub::Reset(); | ||
| 174 | +} | ||
| 175 | + | ||
| 126 | TEST(AttachedStreamCollectionTest, CreateFailureReturnsNullptrWithoutBindOrDestroy) { | 176 | TEST(AttachedStreamCollectionTest, CreateFailureReturnsNullptrWithoutBindOrDestroy) { |
| 127 | auto runtime = std::make_shared<RuntimeForAttachedStreamTest>(); | 177 | auto runtime = std::make_shared<RuntimeForAttachedStreamTest>(); |
| 128 | auto acl = std::make_shared<AclForAttachedStreamTest>(); | 178 | auto acl = std::make_shared<AclForAttachedStreamTest>(); |
| @@ -0,0 +1,127 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#include <gtest/gtest.h> | ||
| 12 | +#include <memory> | ||
| 13 | +#include <string> | ||
| 14 | +#include <vector> | ||
| 15 | + | ||
| 16 | +#include "graph/ascend_string.h" | ||
| 17 | +#include "graph/load/model_manager/davinci_model.h" | ||
| 18 | +#include "depends/runtime/src/runtime_stub.h" | ||
| 19 | +#include "depends/ascendcl/src/ascendcl_stub.h" | ||
| 20 | + | ||
| 21 | +namespace ge { | ||
| 22 | +namespace { | ||
| 23 | +class RuntimeForModelAttachedStreamTest final : public RuntimeStub { | ||
| 24 | + public: | ||
| 25 | + rtError_t rtStreamCreateWithFlags(rtStream_t *stream, int32_t, uint32_t) override { | ||
| 26 | + ++create_count_; | ||
| 27 | + *stream = reinterpret_cast<rtStream_t>(++next_stream_); | ||
| 28 | + return RT_ERROR_NONE; | ||
| 29 | + } | ||
| 30 | + | ||
| 31 | + int create_count_{0}; | ||
| 32 | + | ||
| 33 | + private: | ||
| 34 | + uintptr_t next_stream_{0x2000U}; | ||
| 35 | +}; | ||
| 36 | + | ||
| 37 | +class AclForModelAttachedStreamTest : public AclRuntimeStub { | ||
| 38 | + public: | ||
| 39 | + aclError aclrtSynchronizeStream(aclrtStream stream) override { | ||
| 40 | + synced_streams.push_back(stream); | ||
| 41 | + events.push_back("sync"); | ||
| 42 | + return ACL_SUCCESS; | ||
| 43 | + } | ||
| 44 | + aclError aclmdlRIUnbindStream(aclmdlRI, aclrtStream stream) override { | ||
| 45 | + unbound_streams.push_back(stream); | ||
| 46 | + events.push_back("unbind"); | ||
| 47 | + return ACL_SUCCESS; | ||
| 48 | + } | ||
| 49 | + aclError aclrtDestroyStream(aclrtStream stream) override { | ||
| 50 | + destroyed_streams.push_back(stream); | ||
| 51 | + events.push_back("destroy"); | ||
| 52 | + return ACL_SUCCESS; | ||
| 53 | + } | ||
| 54 | + | ||
| 55 | + std::vector<aclrtStream> synced_streams; | ||
| 56 | + std::vector<aclrtStream> unbound_streams; | ||
| 57 | + std::vector<aclrtStream> destroyed_streams; | ||
| 58 | + std::vector<std::string> events; | ||
| 59 | +}; | ||
| 60 | + | ||
| 61 | +class AclWithoutCtxForModelAttachedStreamTest final : public AclForModelAttachedStreamTest { | ||
| 62 | + public: | ||
| 63 | + aclError aclrtGetCurrentContext(aclrtContext *context) override { | ||
| 64 | + (void)context; | ||
| 65 | + return ACL_ERROR_RT_FAILURE; | ||
| 66 | + } | ||
| 67 | +}; | ||
| 68 | + | ||
| 69 | +// RT2 内嵌 V1 静态子图场景：DavinciModelFinalizer 调用本接口清理 Eager 算子申请的辅流， | ||
| 70 | +// 且 DestroyResources 会置 has_finalized_，析构不再兜底，因此本接口必须真正释放流且可重复调用。 | ||
| 71 | +TEST(DavinciModelAttachedStreamTest, UnbindAndDestroyAttachedStreamsIsIdempotent) { | ||
| 72 | + auto runtime = std::make_shared<RuntimeForModelAttachedStreamTest>(); | ||
| 73 | + auto acl = std::make_shared<AclForModelAttachedStreamTest>(); | ||
| 74 | + RuntimeStub::SetInstance(runtime); | ||
| 75 | + AclRuntimeStub::SetInstance(acl); | ||
| 76 | + | ||
| 77 | + DavinciModel model(0, nullptr); | ||
| 78 | + // UT 目标以 -fno-access-control 编译，可直接初始化私有成员 | ||
| 79 | + model.attached_stream_collection_.Initialize(reinterpret_cast<rtModel_t>(0x55U), RT_STREAM_PRIORITY_DEFAULT, | ||
| 80 | + RT_STREAM_DEFAULT); | ||
| 81 | + ASSERT_NE(model.GetAttachedStreamProvider()->RequestAttachedStream(AscendString("eager_aux")), nullptr); | ||
| 82 | + ASSERT_EQ(runtime->create_count_, 1); | ||
| 83 | + | ||
| 84 | + model.UnbindAndDestroyAttachedStreams(); | ||
| 85 | + // 顺序必须是 sync -> unbind -> destroy，与 RT2 的 Rt2AttachedStreamCollection::Destroy 一致 | ||
| 86 | + ASSERT_EQ(acl->events.size(), 3U); | ||
| 87 | + EXPECT_EQ(acl->events[0], "sync"); | ||
| 88 | + EXPECT_EQ(acl->events[1], "unbind"); | ||
| 89 | + EXPECT_EQ(acl->events[2], "destroy"); | ||
| 90 | + EXPECT_EQ(acl->synced_streams.size(), 1U); | ||
| 91 | + EXPECT_EQ(acl->unbound_streams.size(), 1U); | ||
| 92 | + EXPECT_EQ(acl->destroyed_streams.size(), 1U); | ||
| 93 | + | ||
| 94 | + // 幂等：Finalizer 与析构路径都可能调用，重复调用不得产生额外的同步/解绑/销毁 | ||
| 95 | + model.UnbindAndDestroyAttachedStreams(); | ||
| 96 | + EXPECT_EQ(acl->events.size(), 3U); | ||
| 97 | + EXPECT_EQ(acl->unbound_streams.size(), 1U); | ||
| 98 | + EXPECT_EQ(acl->destroyed_streams.size(), 1U); | ||
| 99 | + | ||
| 100 | + RuntimeStub::Reset(); | ||
| 101 | + AclRuntimeStub::Reset(); | ||
| 102 | +} | ||
| 103 | + | ||
| 104 | +// ~DavinciModel 兜底清理必须与 Finalizer 路径同一守卫：ctx 已销毁时不得再调 RT 接口， | ||
| 105 | +// 否则每条辅流都会刷 sync/unbind/destroy 的 RT 错误日志（对齐 UnbindTaskSinkStream/DestroyStream/DestroyResources） | ||
| 106 | +TEST(DavinciModelAttachedStreamTest, DestructorSkipsAttachedStreamCleanupWhenContextMissing) { | ||
| 107 | + auto runtime = std::make_shared<RuntimeForModelAttachedStreamTest>(); | ||
| 108 | + auto acl = std::make_shared<AclWithoutCtxForModelAttachedStreamTest>(); | ||
| 109 | + RuntimeStub::SetInstance(runtime); | ||
| 110 | + AclRuntimeStub::SetInstance(acl); | ||
| 111 | + { | ||
| 112 | + DavinciModel model(0, nullptr); | ||
| 113 | + model.attached_stream_collection_.Initialize(reinterpret_cast<rtModel_t>(0x55U), RT_STREAM_PRIORITY_DEFAULT, | ||
| 114 | + RT_STREAM_DEFAULT); | ||
| 115 | + ASSERT_NE(model.GetAttachedStreamProvider()->RequestAttachedStream(AscendString("eager_aux")), nullptr); | ||
| 116 | + ASSERT_EQ(runtime->create_count_, 1); | ||
| 117 | + } | ||
| 118 | + EXPECT_TRUE(acl->events.empty()); | ||
| 119 | + EXPECT_TRUE(acl->synced_streams.empty()); | ||
| 120 | + EXPECT_TRUE(acl->unbound_streams.empty()); | ||
| 121 | + EXPECT_TRUE(acl->destroyed_streams.empty()); | ||
| 122 | + | ||
| 123 | + RuntimeStub::Reset(); | ||
| 124 | + AclRuntimeStub::Reset(); | ||
| 125 | +} | ||
| 126 | +} // namespace | ||
| 127 | +} // namespace ge | ||
| @@ -121,10 +121,15 @@ TEST_F(CustomNodeConverterUT, custom_op_convert_test) { | |||
| 121 | {"FreeMemory", 4}, | 121 | {"FreeMemory", 4}, |
| 122 | {"SelectL2Allocator", 1}, | 122 | {"SelectL2Allocator", 1}, |
| 123 | {"SplitRtStreams", 1}, | 123 | {"SplitRtStreams", 1}, |
| 124 | - {"InnerData", 4}, | 124 | + {"InnerData", 5}, |
| 125 | {"MakeSureTensorAtDevice", 3}, | 125 | {"MakeSureTensorAtDevice", 3}, |
| 126 | {"SplitDataTensor", 4}}), | 126 | {"SplitDataTensor", 4}}), |
| 127 | "success"); | 127 | "success"); |
| 128 | + // 输入顺序：3 个 tensor + allocator + stream + custom_op_func + attached_stream_provider， | ||
| 129 | + // 与 CustomOpInput 枚举保持一致，新增附加输入必须插在 func 之后 | ||
| 130 | + const auto *execute_node = ge::ExecuteGraphUtils::FindFirstNodeMatchType(exe_graph, "ExecuteCustomOp"); | ||
| 131 | + ASSERT_NE(execute_node, nullptr); | ||
| 132 | + ASSERT_EQ(execute_node->GetDataInNum(), 7); | ||
| 128 | } | 133 | } |
| 129 | 134 | ||
| 130 | TEST_F(CustomNodeConverterUT, host_cpu_custom_op_convert_test) { | 135 | TEST_F(CustomNodeConverterUT, host_cpu_custom_op_convert_test) { |
| @@ -38,6 +38,7 @@ | |||
| 38 | #include "register/kernel_registry_impl.h" | 38 | #include "register/kernel_registry_impl.h" |
| 39 | #include "exe_graph/runtime/extended_kernel_context.h" | 39 | #include "exe_graph/runtime/extended_kernel_context.h" |
| 40 | #include "exe_graph/runtime/eager_op_execution_context.h" | 40 | #include "exe_graph/runtime/eager_op_execution_context.h" |
| 41 | +#include "depends/ascendcl/src/ascendcl_stub.h" | ||
| 41 | #include "exe_graph/runtime/host_cpu_op_execution_context.h" | 42 | #include "exe_graph/runtime/host_cpu_op_execution_context.h" |
| 42 | #include "exe_graph/runtime/storage_shape.h" | 43 | #include "exe_graph/runtime/storage_shape.h" |
| 43 | #include "exe_graph/runtime/gert_tensor_data.h" | 44 | #include "exe_graph/runtime/gert_tensor_data.h" |
| @@ -82,6 +83,77 @@ class TestBaseCustomOp : public EagerExecuteOp { | |||
| 82 | } | 83 | } |
| 83 | }; | 84 | }; |
| 84 | 85 | ||
| 86 | +namespace { | ||
| 87 | +// 记录 Eager 算子在 Execute 中申请辅流的结果，用于验证 provider 注入链路是否打通 | ||
| 88 | +constexpr size_t kAttachedStreamRoundNum = 2U; | ||
| 89 | +size_t g_attached_execute_count = 0U; | ||
| 90 | +rtStream g_attached_aux[kAttachedStreamRoundNum] = {}; | ||
| 91 | +rtStream g_attached_reuse[kAttachedStreamRoundNum] = {}; | ||
| 92 | +rtStream g_attached_other[kAttachedStreamRoundNum] = {}; | ||
| 93 | +rtStream g_attached_empty_key[kAttachedStreamRoundNum] = {}; | ||
| 94 | + | ||
| 95 | +void ResetAttachedStreamRecord() { | ||
| 96 | + g_attached_execute_count = 0U; | ||
| 97 | + for (size_t i = 0U; i < kAttachedStreamRoundNum; ++i) { | ||
| 98 | + g_attached_aux[i] = nullptr; | ||
| 99 | + g_attached_reuse[i] = nullptr; | ||
| 100 | + g_attached_other[i] = nullptr; | ||
| 101 | + g_attached_empty_key[i] = nullptr; | ||
| 102 | + } | ||
| 103 | +} | ||
| 104 | + | ||
| 105 | +// 记录被销毁的流句柄（其余行为委托基类），用于验证执行器析构时辅流被释放 | ||
| 106 | +class DestroyRecordingAclStub final : public ge::AclRuntimeStub { | ||
| 107 | + public: | ||
| 108 | + aclError aclrtDestroyStream(aclrtStream stream) override { | ||
| 109 | + destroyed.push_back(stream); | ||
| 110 | + return ge::AclRuntimeStub::aclrtDestroyStream(stream); | ||
| 111 | + } | ||
| 112 | + | ||
| 113 | + bool IsDestroyed(aclrtStream stream) const { | ||
| 114 | + for (const auto item : destroyed) { | ||
| 115 | + if (item == stream) { | ||
| 116 | + return true; | ||
| 117 | + } | ||
| 118 | + } | ||
| 119 | + return false; | ||
| 120 | + } | ||
| 121 | + | ||
| 122 | + std::vector<aclrtStream> destroyed; | ||
| 123 | +}; | ||
| 124 | + | ||
| 125 | +// 保证即使用例中途 ASSERT 失败也会恢复默认 stub，避免污染后续用例 | ||
| 126 | +class AclStubGuard { | ||
| 127 | + public: | ||
| 128 | + explicit AclStubGuard(const std::shared_ptr<ge::AclRuntimeStub> &stub) { | ||
| 129 | + ge::AclRuntimeStub::SetInstance(stub); | ||
| 130 | + } | ||
| 131 | + ~AclStubGuard() { | ||
| 132 | + ge::AclRuntimeStub::Reset(); | ||
| 133 | + } | ||
| 134 | +}; | ||
| 135 | +} // namespace | ||
| 136 | + | ||
| 137 | +class TestAttachedStreamCustomOp : public EagerExecuteOp { | ||
| 138 | + public: | ||
| 139 | + graphStatus Execute(gert::EagerOpExecutionContext *ctx) override { | ||
| 140 | + auto workspaces = ctx->MallocWorkSpace(1024); | ||
| 141 | + GE_ASSERT_NOTNULL(workspaces); | ||
| 142 | + auto output_tensor = ctx->MallocOutputTensor(0, StorageShape({2048}, {2048}), | ||
| 143 | + StorageFormat(FORMAT_ND, FORMAT_ND, ExpandDimsType()), DT_FLOAT); | ||
| 144 | + GE_ASSERT_NOTNULL(output_tensor); | ||
| 145 | + | ||
| 146 | + const size_t idx = | ||
| 147 | + (g_attached_execute_count < kAttachedStreamRoundNum) ? g_attached_execute_count : kAttachedStreamRoundNum - 1U; | ||
| 148 | + ++g_attached_execute_count; | ||
| 149 | + g_attached_aux[idx] = ctx->RequestAttachedStream("ut_aux"); | ||
| 150 | + g_attached_reuse[idx] = ctx->RequestAttachedStream("ut_aux"); | ||
| 151 | + g_attached_other[idx] = ctx->RequestAttachedStream("ut_other"); | ||
| 152 | + g_attached_empty_key[idx] = ctx->RequestAttachedStream(""); | ||
| 153 | + return SUCCESS; | ||
| 154 | + } | ||
| 155 | +}; | ||
| 156 | + | ||
| 85 | class TestBaseCustomOpWithInferShape : public EagerExecuteOp { | 157 | class TestBaseCustomOpWithInferShape : public EagerExecuteOp { |
| 86 | public: | 158 | public: |
| 87 | graphStatus Execute(gert::EagerOpExecutionContext *ctx) override { | 159 | graphStatus Execute(gert::EagerOpExecutionContext *ctx) override { |
| @@ -332,6 +404,123 @@ TEST_F(CustomNodeKernelUT, custom_op_kernel_execute_test) { | |||
| 332 | aclrtDestroyStream(stream); | 404 | aclrtDestroyStream(stream); |
| 333 | } | 405 | } |
| 334 | 406 | ||
| 407 | +/** | ||
| 408 | + * 用例描述：验证 RT2 链路下 Eager 自定义算子可申请框架托管辅流 | ||
| 409 | + * 预置条件：注册一个在 Execute 中调用 RequestAttachedStream 的 EagerExecuteOp | ||
| 410 | + * 测试步骤：构建执行器并连续执行两轮 | ||
| 411 | + * 预期结果：provider 注入生效，同 key 复用、异 key 隔离、空 key 返回 nullptr，且句柄跨执行稳定 | ||
| 412 | + */ | ||
| 413 | +TEST_F(CustomNodeKernelUT, custom_op_attached_stream_execute_test) { | ||
| 414 | + ResetAttachedStreamRecord(); | ||
| 415 | + auto graph = ShareGraph::BuildCustomOpGraph(); | ||
| 416 | + graph->TopologicalSorting(); | ||
| 417 | + auto custom_op_registry = BuildCustomOpRegistryForRt2( | ||
| 418 | + []() -> std::unique_ptr<BaseCustomOp> { return std::make_unique<TestAttachedStreamCustomOp>(); }); | ||
| 419 | + GeModelBuilder builder(graph); | ||
| 420 | + auto ge_root_model = builder.BuildGeRootModel(); | ||
| 421 | + ge_root_model->SetCustomOpRegistry(custom_op_registry); | ||
| 422 | + bg::ValueHolder::PopGraphFrame(); | ||
| 423 | + auto exe_graph = ModelConverter().ConvertGeModelToExecuteGraph(ge_root_model, {}); | ||
| 424 | + ASSERT_NE(exe_graph, nullptr); | ||
| 425 | + TaskProducerFactory::GetInstance().SetProducerType(TaskProducerType::KERNEL); | ||
| 426 | + auto model_executor = | ||
| 427 | + ModelV2Executor::Create(exe_graph, ExecutorOption(ExecutorType::kTopologicalPriority), ge_root_model); | ||
| 428 | + ASSERT_NE(model_executor, nullptr); | ||
| 429 | + ASSERT_EQ(model_executor->Load(), GRAPH_SUCCESS); | ||
| 430 | + | ||
| 431 | + auto outputs = FakeTensors({2048}, 1); | ||
| 432 | + auto inputs = FakeTensors({2048}, 3); | ||
| 433 | + rtStream_t stream; | ||
| 434 | + ASSERT_EQ(aclrtCreateStreamWithConfig(&stream, static_cast<uint32_t>(RT_STREAM_PRIORITY_DEFAULT), 0U), RT_ERROR_NONE); | ||
| 435 | + auto i3 = FakeValue<uint64_t>(reinterpret_cast<uint64_t>(stream)); | ||
| 436 | + | ||
| 437 | + for (size_t round = 0U; round < kAttachedStreamRoundNum; ++round) { | ||
| 438 | + ASSERT_EQ(model_executor->Execute({i3.value}, inputs.GetTensorList(), inputs.size(), | ||
| 439 | + reinterpret_cast<Tensor **>(outputs.GetAddrList()), outputs.size()), | ||
| 440 | + GRAPH_SUCCESS); | ||
| 441 | + } | ||
| 442 | + | ||
| 443 | + ASSERT_EQ(g_attached_execute_count, kAttachedStreamRoundNum); | ||
| 444 | + // provider 注入生效：RT2 下不再恒返回 nullptr | ||
| 445 | + ASSERT_NE(g_attached_aux[0], nullptr); | ||
| 446 | + EXPECT_EQ(g_attached_reuse[0], g_attached_aux[0]); // 同 key 复用 | ||
| 447 | + ASSERT_NE(g_attached_other[0], nullptr); | ||
| 448 | + EXPECT_NE(g_attached_other[0], g_attached_aux[0]); // 异 key 隔离 | ||
| 449 | + EXPECT_EQ(g_attached_empty_key[0], nullptr); // 空 key 降级 | ||
| 450 | + EXPECT_EQ(g_attached_aux[1], g_attached_aux[0]); // 跨执行句柄稳定，未重复建流 | ||
| 451 | + EXPECT_EQ(g_attached_other[1], g_attached_other[0]); | ||
| 452 | + aclrtDestroyStream(stream); | ||
| 453 | + // 执行器析构时 Chain deleter 释放辅流容器，容器析构同步并销毁物理流 | ||
| 454 | + model_executor.reset(); | ||
| 455 | + ResetAttachedStreamRecord(); | ||
| 456 | +} | ||
| 457 | + | ||
| 458 | +/** | ||
| 459 | + * 用例描述：验证辅流按执行器隔离（决策 D2），且执行器析构时释放辅流（决策 D3） | ||
| 460 | + * 预置条件：同一份 exe_graph 与 root_model 上创建两个 ModelV2Executor， | ||
| 461 | + * 模拟 StreamExecutor 为每条 aclrtStream 各建一个执行器的场景 | ||
| 462 | + * 测试步骤：两个执行器各执行一轮，相同 key 申请辅流；再依次析构 | ||
| 463 | + * 预期结果：两个执行器拿到不同句柄（隔离）；析构 A 只销毁 A 的辅流，析构 B 后 B 的辅流也被销毁 | ||
| 464 | + */ | ||
| 465 | +TEST_F(CustomNodeKernelUT, custom_op_attached_stream_isolated_per_executor_test) { | ||
| 466 | + ResetAttachedStreamRecord(); | ||
| 467 | + auto acl_stub = std::make_shared<DestroyRecordingAclStub>(); | ||
| 468 | + AclStubGuard acl_stub_guard(acl_stub); | ||
| 469 | + | ||
| 470 | + auto graph = ShareGraph::BuildCustomOpGraph(); | ||
| 471 | + graph->TopologicalSorting(); | ||
| 472 | + auto custom_op_registry = BuildCustomOpRegistryForRt2( | ||
| 473 | + []() -> std::unique_ptr<BaseCustomOp> { return std::make_unique<TestAttachedStreamCustomOp>(); }); | ||
| 474 | + GeModelBuilder builder(graph); | ||
| 475 | + auto ge_root_model = builder.BuildGeRootModel(); | ||
| 476 | + ge_root_model->SetCustomOpRegistry(custom_op_registry); | ||
| 477 | + bg::ValueHolder::PopGraphFrame(); | ||
| 478 | + auto exe_graph = ModelConverter().ConvertGeModelToExecuteGraph(ge_root_model, {}); | ||
| 479 | + ASSERT_NE(exe_graph, nullptr); | ||
| 480 | + TaskProducerFactory::GetInstance().SetProducerType(TaskProducerType::KERNEL); | ||
| 481 | + | ||
| 482 | + auto outputs = FakeTensors({2048}, 1); | ||
| 483 | + auto inputs = FakeTensors({2048}, 3); | ||
| 484 | + rtStream_t stream; | ||
| 485 | + ASSERT_EQ(aclrtCreateStreamWithConfig(&stream, static_cast<uint32_t>(RT_STREAM_PRIORITY_DEFAULT), 0U), RT_ERROR_NONE); | ||
| 486 | + auto i3 = FakeValue<uint64_t>(reinterpret_cast<uint64_t>(stream)); | ||
| 487 | + | ||
| 488 | + auto executor_a = | ||
| 489 | + ModelV2Executor::Create(exe_graph, ExecutorOption(ExecutorType::kTopologicalPriority), ge_root_model); | ||
| 490 | + ASSERT_NE(executor_a, nullptr); | ||
| 491 | + ASSERT_EQ(executor_a->Load(), GRAPH_SUCCESS); | ||
| 492 | + ASSERT_EQ(executor_a->Execute({i3.value}, inputs.GetTensorList(), inputs.size(), | ||
| 493 | + reinterpret_cast<Tensor **>(outputs.GetAddrList()), outputs.size()), | ||
| 494 | + GRAPH_SUCCESS); | ||
| 495 | + const auto aux_a = g_attached_aux[0]; | ||
| 496 | + const auto other_a = g_attached_other[0]; | ||
| 497 | + ASSERT_NE(aux_a, nullptr); | ||
| 498 | + | ||
| 499 | + auto executor_b = | ||
| 500 | + ModelV2Executor::Create(exe_graph, ExecutorOption(ExecutorType::kTopologicalPriority), ge_root_model); | ||
| 501 | + ASSERT_NE(executor_b, nullptr); | ||
| 502 | + ASSERT_EQ(executor_b->Load(), GRAPH_SUCCESS); | ||
| 503 | + ASSERT_EQ(executor_b->Execute({i3.value}, inputs.GetTensorList(), inputs.size(), | ||
| 504 | + reinterpret_cast<Tensor **>(outputs.GetAddrList()), outputs.size()), | ||
| 505 | + GRAPH_SUCCESS); | ||
| 506 | + const auto aux_b = g_attached_aux[1]; | ||
| 507 | + ASSERT_NE(aux_b, nullptr); | ||
| 508 | + | ||
| 509 | + // D2：同一算子实例被两个执行器共享，但辅流按执行器隔离，相同 key 也不共享 | ||
| 510 | + EXPECT_NE(aux_a, aux_b); | ||
| 511 | + | ||
| 512 | + // D3：执行器析构 → Chain deleter → 容器 Destroy()，只影响本执行器的辅流 | ||
| 513 | + executor_a.reset(); | ||
| 514 | + EXPECT_TRUE(acl_stub->IsDestroyed(aux_a)); | ||
| 515 | + EXPECT_TRUE(acl_stub->IsDestroyed(other_a)); | ||
| 516 | + EXPECT_FALSE(acl_stub->IsDestroyed(aux_b)); | ||
| 517 | + executor_b.reset(); | ||
| 518 | + EXPECT_TRUE(acl_stub->IsDestroyed(aux_b)); | ||
| 519 | + | ||
| 520 | + aclrtDestroyStream(stream); | ||
| 521 | + ResetAttachedStreamRecord(); | ||
| 522 | +} | ||
| 523 | + | ||
| 335 | TEST_F(CustomNodeKernelUT, find_custom_op_uses_model_registry) { | 524 | TEST_F(CustomNodeKernelUT, find_custom_op_uses_model_registry) { |
| 336 | const std::string node_type = "RegistryOnlyCustomOpForRt2"; | 525 | const std::string node_type = "RegistryOnlyCustomOpForRt2"; |
| 337 | auto custom_op_registry = std::make_shared<CustomOpRegistry>(); | 526 | auto custom_op_registry = std::make_shared<CustomOpRegistry>(); |
| @@ -30,6 +30,7 @@ class TrackingAllocatorFaker : public AllocatorFaker { | |||
| 30 | } | 30 | } |
| 31 | GertMemBlockFaker::Free(stream_id); | 31 | GertMemBlockFaker::Free(stream_id); |
| 32 | } | 32 | } |
| 33 | + | ||
| 33 | private: | 34 | private: |
| 34 | size_t *free_count_; | 35 | size_t *free_count_; |
| 35 | }; | 36 | }; |
| @@ -40,7 +41,9 @@ class TrackingAllocatorFaker : public AllocatorFaker { | |||
| 40 | delete block; | 41 | delete block; |
| 41 | } | 42 | } |
| 42 | } | 43 | } |
| 43 | - size_t GetFreeCount() const { return free_count_; } | 44 | + size_t GetFreeCount() const { |
| 45 | + return free_count_; | ||
| 46 | + } | ||
| 44 | 47 | ||
| 45 | GertMemBlock *Malloc(size_t size) override { | 48 | GertMemBlock *Malloc(size_t size) override { |
| 46 | void *addr = malloc(size); | 49 | void *addr = malloc(size); |
| @@ -51,11 +54,23 @@ class TrackingAllocatorFaker : public AllocatorFaker { | |||
| 51 | blocks_.push_back(block); | 54 | blocks_.push_back(block); |
| 52 | return block; | 55 | return block; |
| 53 | } | 56 | } |
| 57 | + | ||
| 54 | private: | 58 | private: |
| 55 | size_t free_count_ = 0; | 59 | size_t free_count_ = 0; |
| 56 | std::vector<GertMemBlock *> blocks_; | 60 | std::vector<GertMemBlock *> blocks_; |
| 57 | }; | 61 | }; |
| 58 | 62 | ||
| 63 | +class FakeAttachedStreamProvider final : public AttachedStreamProvider { | ||
| 64 | + public: | ||
| 65 | + rtStream RequestAttachedStream(const ge::AscendString &key) override { | ||
| 66 | + (void)key; | ||
| 67 | + return reinterpret_cast<rtStream>(&dummy_); | ||
| 68 | + } | ||
| 69 | + | ||
| 70 | + private: | ||
| 71 | + int32_t dummy_{0}; | ||
| 72 | +}; | ||
| 73 | + | ||
| 59 | class EagerArgsHandlerTest : public ::testing::Test { | 74 | class EagerArgsHandlerTest : public ::testing::Test { |
| 60 | protected: | 75 | protected: |
| 61 | void SetUp() override { | 76 | void SetUp() override { |
| @@ -194,5 +209,29 @@ TEST_F(EagerArgsHandlerTest, MallocReadOnlyDevArgs_WithoutInitialize_ReturnsNull | |||
| 194 | EXPECT_EQ(result, nullptr); | 209 | EXPECT_EQ(result, nullptr); |
| 195 | } | 210 | } |
| 196 | 211 | ||
| 212 | +TEST_F(EagerArgsHandlerTest, GetAttachedStreamProvider_DefaultNullptr) { | ||
| 213 | + handler_ = std::make_unique<EagerArgsHandler>(); | ||
| 214 | + EXPECT_EQ(handler_->GetAttachedStreamProvider(), nullptr); | ||
| 215 | +} | ||
| 216 | + | ||
| 217 | +TEST_F(EagerArgsHandlerTest, SetAttachedStreamProvider_ReturnsSamePointer) { | ||
| 218 | + handler_ = std::make_unique<EagerArgsHandler>(allocator_.get(), stream_id_); | ||
| 219 | + FakeAttachedStreamProvider provider; | ||
| 220 | + handler_->SetAttachedStreamProvider(&provider); | ||
| 221 | + EXPECT_EQ(handler_->GetAttachedStreamProvider(), &provider); | ||
| 222 | +} | ||
| 223 | + | ||
| 224 | +TEST_F(EagerArgsHandlerTest, Release_KeepsAttachedStreamProvider) { | ||
| 225 | + handler_ = std::make_unique<EagerArgsHandler>(allocator_.get(), stream_id_); | ||
| 226 | + FakeAttachedStreamProvider provider; | ||
| 227 | + handler_->SetAttachedStreamProvider(&provider); | ||
| 228 | + uint8_t host_data[256] = {}; | ||
| 229 | + handler_->MallocReadOnlyDevArgs(host_data, sizeof(host_data)); | ||
| 230 | + | ||
| 231 | + // Release 每轮执行末尾都会跑，辅流容器跨执行复用，不能被清掉 | ||
| 232 | + handler_->Release(); | ||
| 233 | + EXPECT_EQ(handler_->GetAttachedStreamProvider(), &provider); | ||
| 234 | +} | ||
| 235 | + | ||
| 197 | } // namespace | 236 | } // namespace |
| 198 | } // namespace gert | 237 | } // namespace gert |
| @@ -0,0 +1,156 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#include <gtest/gtest.h> | ||
| 12 | +#include <memory> | ||
| 13 | +#include <string> | ||
| 14 | +#include <vector> | ||
| 15 | + | ||
| 16 | +#include "engine/custom/kernel/rt2_attached_stream_collection.h" | ||
| 17 | +#include "graph/ascend_string.h" | ||
| 18 | +#include "depends/runtime/src/runtime_stub.h" | ||
| 19 | +#include "depends/ascendcl/src/ascendcl_stub.h" | ||
| 20 | + | ||
| 21 | +namespace gert { | ||
| 22 | +namespace { | ||
| 23 | +constexpr rtError_t kRtErrorForTest = 1; | ||
| 24 | +// 辅流数量不设上限，取一个较大的数量验证不会被截断 | ||
| 25 | +constexpr size_t kLargeStreamNum = 16U; | ||
| 26 | + | ||
| 27 | +class RuntimeForRt2AttachedStreamTest final : public ge::RuntimeStub { | ||
| 28 | + public: | ||
| 29 | + rtError_t rtStreamCreateWithFlags(rtStream_t *stream, int32_t, uint32_t) override { | ||
| 30 | + ++create_count_; | ||
| 31 | + if (create_fail) { | ||
| 32 | + return kRtErrorForTest; | ||
| 33 | + } | ||
| 34 | + *stream = reinterpret_cast<rtStream_t>(++next_stream_); | ||
| 35 | + return RT_ERROR_NONE; | ||
| 36 | + } | ||
| 37 | + | ||
| 38 | + int create_count_{0}; | ||
| 39 | + bool create_fail{false}; | ||
| 40 | + | ||
| 41 | + private: | ||
| 42 | + uintptr_t next_stream_{0x1000U}; | ||
| 43 | +}; | ||
| 44 | + | ||
| 45 | +class AclForRt2AttachedStreamTest final : public ge::AclRuntimeStub { | ||
| 46 | + public: | ||
| 47 | + aclError aclrtSynchronizeStream(aclrtStream stream) override { | ||
| 48 | + synced_streams.push_back(stream); | ||
| 49 | + events.push_back("sync"); | ||
| 50 | + return ACL_SUCCESS; | ||
| 51 | + } | ||
| 52 | + aclError aclrtDestroyStream(aclrtStream stream) override { | ||
| 53 | + destroyed_streams.push_back(stream); | ||
| 54 | + events.push_back("destroy"); | ||
| 55 | + return ACL_SUCCESS; | ||
| 56 | + } | ||
| 57 | + | ||
| 58 | + std::vector<aclrtStream> synced_streams; | ||
| 59 | + std::vector<aclrtStream> destroyed_streams; | ||
| 60 | + std::vector<std::string> events; | ||
| 61 | +}; | ||
| 62 | + | ||
| 63 | +class Rt2AttachedStreamCollectionTest : public testing::Test { | ||
| 64 | + protected: | ||
| 65 | + void SetUp() override { | ||
| 66 | + runtime_ = std::make_shared<RuntimeForRt2AttachedStreamTest>(); | ||
| 67 | + acl_ = std::make_shared<AclForRt2AttachedStreamTest>(); | ||
| 68 | + ge::RuntimeStub::SetInstance(runtime_); | ||
| 69 | + ge::AclRuntimeStub::SetInstance(acl_); | ||
| 70 | + } | ||
| 71 | + void TearDown() override { | ||
| 72 | + ge::RuntimeStub::Reset(); | ||
| 73 | + ge::AclRuntimeStub::Reset(); | ||
| 74 | + } | ||
| 75 | + | ||
| 76 | + std::shared_ptr<RuntimeForRt2AttachedStreamTest> runtime_; | ||
| 77 | + std::shared_ptr<AclForRt2AttachedStreamTest> acl_; | ||
| 78 | +}; | ||
| 79 | + | ||
| 80 | +TEST_F(Rt2AttachedStreamCollectionTest, RejectsEmptyKeyWithoutRuntimeCall) { | ||
| 81 | + Rt2AttachedStreamCollection collection; | ||
| 82 | + EXPECT_EQ(collection.RequestAttachedStream(ge::AscendString()), nullptr); | ||
| 83 | + EXPECT_EQ(collection.RequestAttachedStream(ge::AscendString("")), nullptr); | ||
| 84 | + EXPECT_EQ(runtime_->create_count_, 0); | ||
| 85 | +} | ||
| 86 | + | ||
| 87 | +TEST_F(Rt2AttachedStreamCollectionTest, ReusesSameKeyAndIsolatesDistinctKeys) { | ||
| 88 | + Rt2AttachedStreamCollection collection; | ||
| 89 | + const auto first = collection.RequestAttachedStream(ge::AscendString("rt2_aux")); | ||
| 90 | + const auto reuse = collection.RequestAttachedStream(ge::AscendString("rt2_aux")); | ||
| 91 | + const auto other = collection.RequestAttachedStream(ge::AscendString("rt2_other")); | ||
| 92 | + ASSERT_NE(first, nullptr); | ||
| 93 | + EXPECT_EQ(reuse, first); | ||
| 94 | + ASSERT_NE(other, nullptr); | ||
| 95 | + EXPECT_NE(other, first); | ||
| 96 | + EXPECT_EQ(runtime_->create_count_, 2); | ||
| 97 | +} | ||
| 98 | + | ||
| 99 | +TEST_F(Rt2AttachedStreamCollectionTest, AcceptsFrameworkReservedPrefixKey) { | ||
| 100 | + Rt2AttachedStreamCollection collection; | ||
| 101 | + // CANN-FMK- 前缀保留给 HCCL 等框架组件，属 API 文档软约束，代码不拦截 | ||
| 102 | + EXPECT_NE(collection.RequestAttachedStream(ge::AscendString("CANN-FMK-hccl")), nullptr); | ||
| 103 | + EXPECT_EQ(runtime_->create_count_, 1); | ||
| 104 | +} | ||
| 105 | + | ||
| 106 | +TEST_F(Rt2AttachedStreamCollectionTest, CreateFailureReturnsNullptrWithoutDestroy) { | ||
| 107 | + runtime_->create_fail = true; | ||
| 108 | + Rt2AttachedStreamCollection collection; | ||
| 109 | + EXPECT_EQ(collection.RequestAttachedStream(ge::AscendString("rt2_aux")), nullptr); | ||
| 110 | + EXPECT_EQ(runtime_->create_count_, 1); | ||
| 111 | + EXPECT_TRUE(acl_->destroyed_streams.empty()); | ||
| 112 | +} | ||
| 113 | + | ||
| 114 | +TEST_F(Rt2AttachedStreamCollectionTest, HasNoStreamNumLimit) { | ||
| 115 | + Rt2AttachedStreamCollection collection; | ||
| 116 | + for (size_t i = 0U; i < kLargeStreamNum; ++i) { | ||
| 117 | + ASSERT_NE(collection.RequestAttachedStream(ge::AscendString(("rt2_key_" + std::to_string(i)).c_str())), nullptr); | ||
| 118 | + } | ||
| 119 | + EXPECT_EQ(static_cast<size_t>(runtime_->create_count_), kLargeStreamNum); | ||
| 120 | + // 数量不受限后，复用与隔离语义仍须成立 | ||
| 121 | + EXPECT_EQ(collection.RequestAttachedStream(ge::AscendString("rt2_key_0")), | ||
| 122 | + collection.RequestAttachedStream(ge::AscendString("rt2_key_0"))); | ||
| 123 | + EXPECT_EQ(static_cast<size_t>(runtime_->create_count_), kLargeStreamNum); | ||
| 124 | +} | ||
| 125 | + | ||
| 126 | +TEST_F(Rt2AttachedStreamCollectionTest, DestroySyncsBeforeDestroyAndIsIdempotent) { | ||
| 127 | + Rt2AttachedStreamCollection collection; | ||
| 128 | + const auto first = collection.RequestAttachedStream(ge::AscendString("rt2_aux")); | ||
| 129 | + ASSERT_NE(first, nullptr); | ||
| 130 | + | ||
| 131 | + collection.Destroy(); | ||
| 132 | + ASSERT_EQ(acl_->events.size(), 2U); | ||
| 133 | + EXPECT_EQ(acl_->events[0], "sync"); | ||
| 134 | + EXPECT_EQ(acl_->events[1], "destroy"); | ||
| 135 | + EXPECT_EQ(acl_->synced_streams[0], first); | ||
| 136 | + EXPECT_EQ(acl_->destroyed_streams[0], first); | ||
| 137 | + | ||
| 138 | + // 幂等：重复 Destroy 不再产生任何 RTS/ACL 调用 | ||
| 139 | + collection.Destroy(); | ||
| 140 | + EXPECT_EQ(acl_->events.size(), 2U); | ||
| 141 | +} | ||
| 142 | + | ||
| 143 | +TEST_F(Rt2AttachedStreamCollectionTest, RecreatesStreamAfterDestroy) { | ||
| 144 | + Rt2AttachedStreamCollection collection; | ||
| 145 | + const auto first = collection.RequestAttachedStream(ge::AscendString("rt2_aux")); | ||
| 146 | + ASSERT_NE(first, nullptr); | ||
| 147 | + collection.Destroy(); | ||
| 148 | + | ||
| 149 | + // 防御性契约：Destroy 后对象仍可继续使用（生产路径 Destroy 只由析构调用，不会走到这里） | ||
| 150 | + const auto second = collection.RequestAttachedStream(ge::AscendString("rt2_aux")); | ||
| 151 | + ASSERT_NE(second, nullptr); | ||
| 152 | + EXPECT_NE(second, first); | ||
| 153 | + EXPECT_EQ(runtime_->create_count_, 2); | ||
| 154 | +} | ||
| 155 | +} // namespace | ||
| 156 | +} // namespace gert | ||
类名为什么是个名词