feat：声明式自定义算子的辅流申请功能-ge-AtomGit

Mcompiler/CMakeLists.txt+2-0

Mcompiler/engines/custom_engine/custom_ops_kernel_builder.cc+210-82

	@@ -10,6 +10,7 @@
10		10
11	#include "engines/custom_engine/custom_ops_kernel_builder.h"	11	#include "engines/custom_engine/custom_ops_kernel_builder.h"
12	#include <cinttypes>	12	#include <cinttypes>
		13	+#include <algorithm>
13	#include <limits>	14	#include <limits>
14	#include <map>	15	#include <map>
15	#include <memory>	16	#include <memory>
	@@ -35,6 +36,9 @@
35	#include "exe_graph/runtime/gert_tensor_data.h"	36	#include "exe_graph/runtime/gert_tensor_data.h"
36	#include "exe_graph/runtime/annotated_args_context.h"	37	#include "exe_graph/runtime/annotated_args_context.h"
37	#include "graph_metadef/exe_graph/runtime/annotated_args_handler.h"	38	#include "graph_metadef/exe_graph/runtime/annotated_args_handler.h"
		39	+#include "graph/build/stream/declarative_stream_registry.h"
		40	+#include "graph/build/stream/annotated_args_task_plan.h"
		41	+#include "graph/build/task_generator_utils.h"
38	#include "exe_graph/runtime/storage_shape.h"	42	#include "exe_graph/runtime/storage_shape.h"
39	#include "graph/buffer.h"	43	#include "graph/buffer.h"
40	#include "graph/custom_op.h"	44	#include "graph/custom_op.h"
	@@ -56,16 +60,11 @@ constexpr uint32_t kMaxCustomOpSqeNum = 5;
56	constexpr uint32_t kKernelArgSlotSize = sizeof(uint64_t);	60	constexpr uint32_t kKernelArgSlotSize = sizeof(uint64_t);
57	constexpr const char_t *kCustomOmcAppendWs = "_custom_omc_append_ws";	61	constexpr const char_t *kCustomOmcAppendWs = "_custom_omc_append_ws";
58	constexpr const char_t *kAppendWs = "_append_ws";	62	constexpr const char_t *kAppendWs = "_append_ws";
59	-constexpr const char_t *kAnnotatedArgsTaskPlan = "_custom_annotated_args_task_plan";	63	+constexpr const char_t *kAnnotatedArgsTaskPlanBytes = "_custom_annotated_args_task_plan_bytes";
60	constexpr const char_t *kCustomLaunchPrefix = "_custom_launch_";	64	constexpr const char_t *kCustomLaunchPrefix = "_custom_launch_";
61	constexpr const char_t *kDefaultCustomKernelMagic = "RT_DEV_BINARY_MAGIC_ELF_AIVEC";	65	constexpr const char_t *kDefaultCustomKernelMagic = "RT_DEV_BINARY_MAGIC_ELF_AIVEC";
62	constexpr size_t kWorkspaceAlignment = 512U;	66	constexpr size_t kWorkspaceAlignment = 512U;
63		67
64	-struct AnnotatedArgsTaskPlan {
65	- std::vector<domi::TaskDef> task_templates;
66	-};
67	-using AnnotatedArgsTaskPlanPtr = std::shared_ptr<const AnnotatedArgsTaskPlan>;
68	-
69	bool IsMobileSocVersion(const std::string &soc_version) {	68	bool IsMobileSocVersion(const std::string &soc_version) {
70	return (soc_version == "KirinX90") \|\| (soc_version == "Kirin9030");	69	return (soc_version == "KirinX90") \|\| (soc_version == "Kirin9030");
71	}	70	}
	@@ -390,52 +389,66 @@ Status PrepareAnnotatedArgsTaskForMaterialization(const OpDescPtr &op_desc, cons
390	return SUCCESS;	389	return SUCCESS;
391	}	390	}
392		391
		392	+Status MaterializeAnnotatedTaskKernelArgs(const OpDescPtr &op_desc, const RunContext &context, domi::TaskDef &task,
		393	+ const std::vector<ArgDesc> &arg_descs,
		394	+ const std::vector<std::unique_ptr<gert::Tensor>> &input_tensor_holders,
		395	+ const std::vector<std::unique_ptr<gert::Tensor>> &output_tensor_holders) {
		396	+ auto *const kernel = task.mutable_kernel();
		397	+ std::string args = kernel->args();
		398	+ for (size_t i = 0U; i < arg_descs.size(); ++i) {
		399	+ uint64_t addr = 0U;
		400	+ switch (arg_descs[i].addr_type) {
		401	+ case AddrType::INPUT_INSTANCE:
		T T tang-haojie9月23日只有这几个argsformat吗？ C Chang-an-HW 9月23日评论： 9月23日评论：
		402	+ GE_ASSERT_SUCCESS(ResolveTensorPlanAddr(arg_descs[i].ir_idx, input_tensor_holders, addr));
		403	+ break;
		404	+ case AddrType::OUTPUT_INSTANCE:
		405	+ GE_ASSERT_SUCCESS(ResolveTensorPlanAddr(arg_descs[i].ir_idx, output_tensor_holders, addr));
		406	+ break;
		407	+ case AddrType::INPUT:
		408	+ case AddrType::OUTPUT:
		409	+ GELOGE(INTERNAL_ERROR, "Custom op %s(%s) planned arg[%zu] uses unsupported legacy addr_type %d.",
		410	+ op_desc->GetNamePtr(), op_desc->GetTypePtr(), i, static_cast<int32_t>(arg_descs[i].addr_type));
		411	+ return INTERNAL_ERROR;
		412	+ case AddrType::WORKSPACE:
		413	+ GE_ASSERT_SUCCESS(ResolveWorkspacePlanAddr(op_desc, context, arg_descs[i].ir_idx, addr));
		414	+ break;
		415	+ case AddrType::CUSTOM_VALUE:
		416	+ continue;
		417	+ default:
		418	+ GELOGE(INTERNAL_ERROR, "Custom op %s(%s) planned arg[%zu] type %d is not materializable.",
		419	+ op_desc->GetNamePtr(), op_desc->GetTypePtr(), i, static_cast<int32_t>(arg_descs[i].addr_type));
		420	+ return INTERNAL_ERROR;
		421	+ }
		422	+ GE_ASSERT_EOK(memcpy_s(&args[i * kKernelArgSlotSize], kKernelArgSlotSize, &addr, sizeof(addr)));
		423	+ }
		424	+ kernel->set_args(args);
		425	+ return SUCCESS;
		426	+}
		427	+
393	Status MaterializeAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const RunContext &context,	428	Status MaterializeAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const RunContext &context,
394	const AnnotatedArgsTaskPlan &plan, std::vector<domi::TaskDef> &tasks) {	429	const AnnotatedArgsTaskPlan &plan, std::vector<domi::TaskDef> &tasks) {
395	std::vector<std::unique_ptr<gert::Tensor>> input_tensor_holders;	430	std::vector<std::unique_ptr<gert::Tensor>> input_tensor_holders;
396	std::vector<std::unique_ptr<gert::Tensor>> output_tensor_holders;	431	std::vector<std::unique_ptr<gert::Tensor>> output_tensor_holders;
397	GE_ASSERT_SUCCESS(ConstructInputTensors(op_desc, context, input_tensor_holders));	432	GE_ASSERT_SUCCESS(ConstructInputTensors(op_desc, context, input_tensor_holders));
398	GE_ASSERT_SUCCESS(ConstructOutputTensors(op_desc, context, output_tensor_holders));	433	GE_ASSERT_SUCCESS(ConstructOutputTensors(op_desc, context, output_tensor_holders));
399	-
400	GE_ASSERT_TRUE((op_desc->GetStreamId() >= 0) &&	434	GE_ASSERT_TRUE((op_desc->GetStreamId() >= 0) &&
401	(op_desc->GetStreamId() <= static_cast<int64_t>(std::numeric_limits<uint32_t>::max())),	435	(op_desc->GetStreamId() <= static_cast<int64_t>(std::numeric_limits<uint32_t>::max())),
402	"Custom op stream id %ld is invalid, op_name:%s, op_type:%s", op_desc->GetStreamId(),	436	"Custom op stream id %ld is invalid, op_name:%s, op_type:%s", op_desc->GetStreamId(),
403	op_desc->GetNamePtr(), op_desc->GetTypePtr());	437	op_desc->GetNamePtr(), op_desc->GetTypePtr());
404	const auto stream_id = static_cast<uint32_t>(op_desc->GetStreamId());	438	const auto stream_id = static_cast<uint32_t>(op_desc->GetStreamId());
405	- for (const auto &task_template : plan.task_templates) {	439	+ GE_ASSERT_TRUE(plan.launch_stream_ids.size() == plan.task_templates.size(),
		440	+ "Custom op %s(%s) launch stream count does not match task count.", op_desc->GetNamePtr(),
		441	+ op_desc->GetTypePtr());
		442	+ // 无依赖且无附着流的legacy计划在物化时刷新为节点当前主流，支持缓存计划跨流刷新复用
		443	+ const bool refresh_legacy_stream = plan.dependencies.empty() && plan.attached_stream_ids.empty();
		444	+ for (size_t task_index = 0U; task_index < plan.task_templates.size(); ++task_index) {
406	domi::TaskDef task;	445	domi::TaskDef task;
407	std::vector<ArgDesc> arg_descs;	446	std::vector<ArgDesc> arg_descs;
408	- GE_CHK_STATUS_RET_NOLOG(	447	+ GE_CHK_STATUS_RET_NOLOG(PrepareAnnotatedArgsTaskForMaterialization(
409	- PrepareAnnotatedArgsTaskForMaterialization(op_desc, stream_id, task_template, task, arg_descs));	448	+ op_desc, refresh_legacy_stream ? stream_id : plan.launch_stream_ids[task_index],
410	- auto *const kernel = task.mutable_kernel();	449	+ plan.task_templates[task_index], task, arg_descs));
411	- std::string args = kernel->args();	450	+ GE_ASSERT_SUCCESS(MaterializeAnnotatedTaskKernelArgs(op_desc, context, task, arg_descs, input_tensor_holders,
412	- for (size_t i = 0U; i < arg_descs.size(); ++i) {	451	+ output_tensor_holders));
413	- uint64_t addr = 0U;
414	- switch (arg_descs[i].addr_type) {
415	- case AddrType::INPUT_INSTANCE:
416	- GE_ASSERT_SUCCESS(ResolveTensorPlanAddr(arg_descs[i].ir_idx, input_tensor_holders, addr));
417	- break;
418	- case AddrType::OUTPUT_INSTANCE:
419	- GE_ASSERT_SUCCESS(ResolveTensorPlanAddr(arg_descs[i].ir_idx, output_tensor_holders, addr));
420	- break;
421	- case AddrType::INPUT:
422	- case AddrType::OUTPUT:
423	- GELOGE(INTERNAL_ERROR, "Custom op %s(%s) planned arg[%zu] uses unsupported legacy addr_type %d.",
424	- op_desc->GetNamePtr(), op_desc->GetTypePtr(), i, static_cast<int32_t>(arg_descs[i].addr_type));
425	- return INTERNAL_ERROR;
426	- case AddrType::WORKSPACE:
427	- GE_ASSERT_SUCCESS(ResolveWorkspacePlanAddr(op_desc, context, arg_descs[i].ir_idx, addr));
428	- break;
429	- case AddrType::CUSTOM_VALUE:
430	- continue;
431	- default:
432	- GELOGE(INTERNAL_ERROR, "Custom op %s(%s) planned arg[%zu] type %d is not materializable.",
433	- op_desc->GetNamePtr(), op_desc->GetTypePtr(), i, static_cast<int32_t>(arg_descs[i].addr_type));
434	- return INTERNAL_ERROR;
435	- }
436	- GE_ASSERT_EOK(memcpy_s(&args[i * kKernelArgSlotSize], kKernelArgSlotSize, &addr, sizeof(addr)));
437	- }
438	- kernel->set_args(args);
439	tasks.emplace_back(std::move(task));	452	tasks.emplace_back(std::move(task));
440	}	453	}
441	return SUCCESS;	454	return SUCCESS;
	@@ -447,7 +460,7 @@ enum class OfflineLaunchGenMode {
447	};	460	};
448		461
449	Status CheckAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const gert::AnnotatedArgsHandler &args_handler,	462	Status CheckAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const gert::AnnotatedArgsHandler &args_handler,
450	- const uint32_t stream_id) {	463	+ const uint32_t stream_id, const DeclarativeStreamRegistryPtr &registry) {
451	GE_ASSERT_NOTNULL(op_desc);	464	GE_ASSERT_NOTNULL(op_desc);
452	const size_t launch_count = args_handler.GetLaunchCount();	465	const size_t launch_count = args_handler.GetLaunchCount();
453	if (launch_count == 0U) {	466	if (launch_count == 0U) {
	@@ -455,22 +468,36 @@ Status CheckAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const gert::Annotate
455	op_desc->GetTypePtr());	468	op_desc->GetTypePtr());
456	return INTERNAL_ERROR;	469	return INTERNAL_ERROR;
457	}	470	}
		471	+ bool has_main_stream_launch = false;
458	for (size_t i = 0U; i < launch_count; ++i) {	472	for (size_t i = 0U; i < launch_count; ++i) {
459	const auto *launch = args_handler.GetLaunch(i);	473	const auto *launch = args_handler.GetLaunch(i);
460	GE_ASSERT_NOTNULL(launch);	474	GE_ASSERT_NOTNULL(launch);
461	- if (launch->GetStreamId() != stream_id) {	475	+ has_main_stream_launch = has_main_stream_launch \|\| (launch->GetStreamId() == stream_id);
		476	+ if ((launch->GetStreamId() != stream_id) && ((registry == nullptr) \|\| !registry->Contains(launch->GetStreamId()))) {
462	GELOGE(INTERNAL_ERROR,	477	GELOGE(INTERNAL_ERROR,
463	"Custom op launch[%zu] stream id %u does not match node stream id %u, op_name:%s, op_type:%s", i,	478	"Custom op launch[%zu] stream id %u does not match node stream id %u, op_name:%s, op_type:%s", i,
464	launch->GetStreamId(), stream_id, op_desc->GetNamePtr(), op_desc->GetTypePtr());	479	launch->GetStreamId(), stream_id, op_desc->GetNamePtr(), op_desc->GetTypePtr());
465	return INTERNAL_ERROR;	480	return INTERNAL_ERROR;
466	}	481	}
		482	+ GE_ASSERT_TRUE((launch->GetDependencyCount() == 0U) \|\| (launch->GetDependencies() != nullptr),
		483	+ "Custom op launch[%zu] dependencies are missing.", i);
		484	+ for (size_t j = 0U; j < launch->GetDependencyCount(); ++j) {
		485	+ const auto token = launch->GetDependencies()[j];
		486	+ GE_ASSERT_TRUE(token < i, "Custom op launch[%zu] dependency token %u is invalid.", i, token);
		487	+ }
		488	+ }
		489	+ if (!has_main_stream_launch) {
		490	+ GELOGE(INTERNAL_ERROR, "Custom op annotated args plan has no main stream launch, op_name:%s, op_type:%s",
		491	+ op_desc->GetNamePtr(), op_desc->GetTypePtr());
		492	+ return INTERNAL_ERROR;
467	}	493	}
468	return SUCCESS;	494	return SUCCESS;
469	}	495	}
470		496
471	Status CheckAnnotatedArgsTaskPlanForMode(const OpDescPtr &op_desc, const gert::AnnotatedArgsHandler &args_handler,	497	Status CheckAnnotatedArgsTaskPlanForMode(const OpDescPtr &op_desc, const gert::AnnotatedArgsHandler &args_handler,
472	- const uint32_t stream_id, const OfflineLaunchGenMode mode) {	498	+ const uint32_t stream_id, const OfflineLaunchGenMode mode,
473	- const auto check_ret = CheckAnnotatedArgsTaskPlan(op_desc, args_handler, stream_id);	499	+ const DeclarativeStreamRegistryPtr &registry) {
		500	+ const auto check_ret = CheckAnnotatedArgsTaskPlan(op_desc, args_handler, stream_id, registry);
474	if (check_ret != SUCCESS) {	501	if (check_ret != SUCCESS) {
475	return check_ret;	502	return check_ret;
476	}	503	}
	@@ -480,6 +507,13 @@ Status CheckAnnotatedArgsTaskPlanForMode(const OpDescPtr &op_desc, const gert::A
480	args_handler.GetLaunchCount(), op_desc->GetNamePtr(), op_desc->GetTypePtr());	507	args_handler.GetLaunchCount(), op_desc->GetNamePtr(), op_desc->GetTypePtr());
481	return INTERNAL_ERROR;	508	return INTERNAL_ERROR;
482	}	509	}
		510	+ if (mode == OfflineLaunchGenMode::kMobileLegacySingleTask &&
		511	+ (args_handler.GetLaunch(0U)->GetDependencyCount() != 0U \|\| !args_handler.GetAttachedStreamIds().empty())) {
		512	+ GELOGE(INTERNAL_ERROR,
		513	+ "Custom op mobile OMC does not support annotated dependencies or attached streams, op_name:%s",
		514	+ op_desc->GetNamePtr());
		515	+ return INTERNAL_ERROR;
		516	+ }
483	return SUCCESS;	517	return SUCCESS;
484	}	518	}
485		519
	@@ -623,7 +657,8 @@ Status ValidateKernelBinNames(const OpDescPtr &op_desc, const gert::AnnotatedArg
623		657
624	Status DeclareLaunchArgsTaskPlan(const OpDescPtr &op_desc, AnnotatedArgsOp &annotated_args_op,	658	Status DeclareLaunchArgsTaskPlan(const OpDescPtr &op_desc, AnnotatedArgsOp &annotated_args_op,
625	WorkspaceAllocator &workspace_allocator, const uint32_t stream_id,	659	WorkspaceAllocator &workspace_allocator, const uint32_t stream_id,
626	- const OfflineLaunchGenMode mode, AnnotatedArgsCapture &capture) {	660	+ const OfflineLaunchGenMode mode, const DeclarativeStreamRegistryPtr &registry,
		661	+ AnnotatedArgsCapture &capture) {
627	std::vector<void *> inputs = GetHoldersRawPtr(capture.input_tensor_holders);	662	std::vector<void *> inputs = GetHoldersRawPtr(capture.input_tensor_holders);
628	inputs.push_back(&workspace_allocator);	663	inputs.push_back(&workspace_allocator);
629		664
	@@ -657,7 +692,7 @@ Status DeclareLaunchArgsTaskPlan(const OpDescPtr &op_desc, AnnotatedArgsOp &anno
657	op_desc->GetTypePtr());	692	op_desc->GetTypePtr());
658	return GRAPH_FAILED;	693	return GRAPH_FAILED;
659	}	694	}
660	- return CheckAnnotatedArgsTaskPlanForMode(op_desc, capture.args_handler, stream_id, mode);	695	+ return CheckAnnotatedArgsTaskPlanForMode(op_desc, capture.args_handler, stream_id, mode, registry);
661	}	696	}
662		697
663	Status UpdateAnnotatedArgsWorkspaceAttrs(const OpDescPtr &op_desc, const WorkspaceAllocator &workspace_allocator) {	698	Status UpdateAnnotatedArgsWorkspaceAttrs(const OpDescPtr &op_desc, const WorkspaceAllocator &workspace_allocator) {
	@@ -673,16 +708,39 @@ Status UpdateAnnotatedArgsWorkspaceAttrs(const OpDescPtr &op_desc, const Workspa
673	}	708	}
674		709
675	Status CacheAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const std::vector<domi::TaskDef> &tasks,	710	Status CacheAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const std::vector<domi::TaskDef> &tasks,
676	- const size_t task_start) {	711	+ const size_t task_start, const gert::AnnotatedArgsHandler &args_handler) {
677	GE_ASSERT_TRUE(task_start < tasks.size(), "Custom op %s(%s) generated no task for annotated args plan.",	712	GE_ASSERT_TRUE(task_start < tasks.size(), "Custom op %s(%s) generated no task for annotated args plan.",
678	op_desc->GetNamePtr(), op_desc->GetTypePtr());	713	op_desc->GetNamePtr(), op_desc->GetTypePtr());
679	auto mutable_plan = ge::ComGraphMakeShared<AnnotatedArgsTaskPlan>();	714	auto mutable_plan = ge::ComGraphMakeShared<AnnotatedArgsTaskPlan>();
680	GE_ASSERT_NOTNULL(mutable_plan);	715	GE_ASSERT_NOTNULL(mutable_plan);
681	mutable_plan->task_templates.assign(tasks.cbegin() + static_cast<std::ptrdiff_t>(task_start), tasks.cend());	716	mutable_plan->task_templates.assign(tasks.cbegin() + static_cast<std::ptrdiff_t>(task_start), tasks.cend());
		717	+ for (size_t i = 0U; i < args_handler.GetLaunchCount(); ++i) {
		718	+ const auto *launch = args_handler.GetLaunch(i);
		719	+ GE_ASSERT_NOTNULL(launch);
		720	+ mutable_plan->launch_stream_ids.emplace_back(launch->GetStreamId());
		721	+ for (size_t j = 0U; j < launch->GetDependencyCount(); ++j) {
		722	+ mutable_plan->dependencies.push_back({launch->GetDependencies()[j], static_cast<uint32_t>(i),
		723	+ static_cast<uint32_t>(mutable_plan->dependencies.size() + 1U)});
		724	+ }
		725	+ }
		726	+ for (const auto id : args_handler.GetAttachedStreamIds()) {
		727	+ if (std::find(mutable_plan->attached_stream_ids.cbegin(), mutable_plan->attached_stream_ids.cend(), id) ==
		728	+ mutable_plan->attached_stream_ids.cend()) {
		729	+ mutable_plan->attached_stream_ids.emplace_back(id);
		730	+ }
		731	+ }
682	AnnotatedArgsTaskPlanPtr plan = std::move(mutable_plan);	732	AnnotatedArgsTaskPlanPtr plan = std::move(mutable_plan);
683	- GE_ASSERT_TRUE(op_desc->SetExtAttr(kAnnotatedArgsTaskPlan, plan),	733	+ Buffer encoded;
684	- "Cache annotated args task plan failed for custom op %s(%s).", op_desc->GetNamePtr(),	734	+ GE_ASSERT_SUCCESS(SerializeAnnotatedArgsTaskPlan(*plan, encoded));
		735	+ GE_ASSERT_TRUE(AttrUtils::SetBytes(op_desc, kAnnotatedArgsTaskPlanBytes, encoded),
		T T tang-haojie9月23日计划缓存把整个 TaskDef 模板（带 kernel args 二进制）序列化进这个属性跟着 OM 走，而 task 列表里已经有一份物化结果了，kernel args 部分在 OM 里存了两份，加载时属性还会整个读进内存。可以考虑模板部分不进 OM（比如只存 launch 布局和依赖关系），或者加载侧把属性清掉。顺带：上面第 712 行的 `registry` 参数传进来又 `(void)` 掉了，是死参数，一起清理下。 C Chang-an-HW 9月23日评论： 9月23日评论：
		736	+ "Cache serialized annotated args task plan failed for custom op %s(%s).", op_desc->GetNamePtr(),
685	op_desc->GetTypePtr());	737	op_desc->GetTypePtr());
		738	+ // 带依赖或附着流的计划：拆流阶段会往本节点任务列表插入事件任务、并把流号改写为真实流号，
		739	+ // 而二次生成只能按本缓存重物化（会丢事件任务、写回拆流前的逻辑流号），因此打标让 TaskGenerator 跳过重生成
		740	+ if ((!plan->dependencies.empty()) \|\| (!plan->attached_stream_ids.empty())) {
		741	+ GE_ASSERT_TRUE(AttrUtils::SetBool(op_desc, kKeepGeneratedTasksAttr, true), "Set %s failed for custom op %s(%s).",
		742	+ kKeepGeneratedTasksAttr, op_desc->GetNamePtr(), op_desc->GetTypePtr());
		743	+ }
686	return SUCCESS;	744	return SUCCESS;
687	}	745	}
688		746
	@@ -694,46 +752,67 @@ Status GenerateMobileAnnotatedArgsTask(const Node &node, const OpDescPtr &op_des
694	GE_ASSERT_NOTNULL(launch);	752	GE_ASSERT_NOTNULL(launch);
695	GE_ASSERT_SUCCESS(FillKernelTask(node, *launch, "", task_def));	753	GE_ASSERT_SUCCESS(FillKernelTask(node, *launch, "", task_def));
696	tasks.emplace_back(std::move(task_def));	754	tasks.emplace_back(std::move(task_def));
697	- GE_ASSERT_SUCCESS(CacheAnnotatedArgsTaskPlan(op_desc, tasks, task_start));	755	+ GE_ASSERT_SUCCESS(CacheAnnotatedArgsTaskPlan(op_desc, tasks, task_start, args_handler));
698	return SUCCESS;	756	return SUCCESS;
699	}	757	}
700		758
701	-Status GenerateAnnotatedArgsTask(const Node &node, const OpDescPtr &op_desc, RunContext &context,	759	+Status ValidateAnnotatedArgsTaskPlanStreams(const AnnotatedArgsTaskPlan &cached_plan, const uint32_t main_stream_id,
702	- const OfflineLaunchGenMode mode, std::vector<domi::TaskDef> &tasks) {	760	+ const DeclarativeStreamRegistryPtr &plan_registry) {
703	- GE_ASSERT_SUCCESS(CheckStaticGraph(node));	761	+ // 无依赖且无附着流的legacy缓存计划：物化时统一改写为节点当前主流（与refresh_legacy_stream条件一致），
704	- const auto plan = op_desc->TryGetExtAttr(kAnnotatedArgsTaskPlan, AnnotatedArgsTaskPlanPtr());	762	+ // 且ReGetTaskInfo等二次生成时节点流ID可能已被拆流改写，因此跳过缓存launch流ID与当前主流的一致性校验。
705	- if (plan != nullptr) {	763	+ const bool is_legacy_plan = cached_plan.dependencies.empty() && cached_plan.attached_stream_ids.empty();
706	- if ((mode == OfflineLaunchGenMode::kMobileLegacySingleTask) && (plan->task_templates.size() != 1U)) {	764	+ if (!is_legacy_plan) {
707	- GELOGE(INTERNAL_ERROR,	765	+ bool has_main_stream_launch = false;
708	- "Custom op planned task count %zu is not supported in mobile OMC scenario, op_name:%s, op_type:%s",	766	+ for (const auto id : cached_plan.launch_stream_ids) {
709	- plan->task_templates.size(), op_desc->GetNamePtr(), op_desc->GetTypePtr());	767	+ has_main_stream_launch = has_main_stream_launch \|\| (id == main_stream_id);
710	- return INTERNAL_ERROR;	768	+ GE_ASSERT_TRUE((id == main_stream_id) \|\| (plan_registry != nullptr && plan_registry->Contains(id)),
		769	+ "Custom op cached launch stream id %u is not declared.", id);
		T T tang-haojie9月23日这里有个循环依赖问题想确认下：带辅流的缓存计划恢复时要用 `plan_registry->Contains(id)` 校验，而 registry 是 `BuildModelForGetTask` 时新建的空表，靠算子回调时才填；但走缓存恢复这条路（RestoreCachedAnnotatedArgsTaskPlan）恰恰不会回调算子。推演下来：二次 GetTaskInfo 或者重打包一个已经带计划的 OM 时，registry 是空的，计划里的辅流 ID 一个都查不到，直接在这里报 "not declared" 编译失败。另外注释里只给 legacy 计划（无依赖无辅流）做了主流 ID 刷新，新格式的计划没这个待遇——节点流被拆流改写过的话，缓存里的主流 launch ID 和当前主流也对不上，同样过不了这个校验。这个场景会不会真的走到（ReGetTaskInfo、已有 OM 重打包）需要确认下，会走到的话就是功能缺陷。 C Chang-an-HW 9月23日评论： 9月23日评论：
711	}	770	}
712	- return MaterializeAnnotatedArgsTaskPlan(op_desc, context, *plan, tasks);	771	+ GE_ASSERT_TRUE(has_main_stream_launch, "Custom op cached annotated args plan has no main stream launch.");
713	}	772	}
714	-	773	+ if (!cached_plan.attached_stream_ids.empty()) {
715	- AnnotatedArgsOp *annotated_args_op = nullptr;	774	+ GE_ASSERT_NOTNULL(plan_registry, "Custom op cached attached streams require a declarative stream registry.");
716	- GE_ASSERT_SUCCESS(GetAnnotatedArgsOp(op_desc, annotated_args_op));
717	-
718	- const size_t task_start = tasks.size();
719	- AnnotatedArgsCapture capture;
720	- GE_ASSERT_SUCCESS(ConstructInputTensors(op_desc, context, capture.input_tensor_holders));
721	- GE_ASSERT_SUCCESS(ConstructOutputTensors(op_desc, context, capture.output_tensor_holders));
722	- WorkspaceAllocator workspace_allocator(op_desc, context);
723	- uint32_t stream_id = 0U;
724	- GE_ASSERT_SUCCESS(GetCustomOpStreamId(op_desc, stream_id));
725	- const auto launch_task_plan_ret =
726	- DeclareLaunchArgsTaskPlan(op_desc, *annotated_args_op, workspace_allocator, stream_id, mode, capture);
727	- if (launch_task_plan_ret != SUCCESS) {
728	- return launch_task_plan_ret;
729	}	775	}
730	- GE_ASSERT_SUCCESS(UpdateAnnotatedArgsWorkspaceAttrs(op_desc, workspace_allocator));	776	+ for (const auto id : cached_plan.attached_stream_ids) {
731	-	777	+ GE_ASSERT_TRUE(plan_registry->Contains(id), "Custom op cached attached stream id %u is not declared.", id);
732	- GE_ASSERT_SUCCESS(ValidateKernelBinNames(op_desc, capture.args_handler));
733	- if (mode == OfflineLaunchGenMode::kMobileLegacySingleTask) {
734	- return GenerateMobileAnnotatedArgsTask(node, op_desc, capture.args_handler, task_start, tasks);
735	}	778	}
		779	+ return SUCCESS;
		780	+}
736		781
		782	+DeclarativeStreamRegistryPtr FindDeclarativeStreamRegistry(const Node &node) {
		783	+ for (auto graph = node.GetOwnerComputeGraph(); graph != nullptr; graph = graph->GetParentGraph()) {
		784	+ const auto registry = graph->TryGetExtAttr(kDeclarativeStreamRegistryAttr, DeclarativeStreamRegistryPtr());
		785	+ if (registry != nullptr) {
		786	+ return registry;
		787	+ }
		788	+ }
		789	+ return nullptr;
		790	+}
		791	+
		792	+Status RestoreCachedAnnotatedArgsTaskPlan(const OpDescPtr &op_desc, const uint32_t main_stream_id,
		793	+ const DeclarativeStreamRegistryPtr &plan_registry,
		794	+ const OfflineLaunchGenMode mode, const RunContext &context,
		795	+ std::vector<domi::TaskDef> &tasks, bool &restored) {
		796	+ restored = false;
		797	+ Buffer encoded_plan;
		798	+ if (!AttrUtils::GetBytes(op_desc, kAnnotatedArgsTaskPlanBytes, encoded_plan)) {
		799	+ return SUCCESS;
		800	+ }
		801	+ restored = true;
		802	+ auto cached_plan = ge::ComGraphMakeShared<AnnotatedArgsTaskPlan>();
		803	+ GE_ASSERT_NOTNULL(cached_plan);
		804	+ GE_ASSERT_SUCCESS(DeserializeAnnotatedArgsTaskPlan(encoded_plan, *cached_plan));
		805	+ GE_ASSERT_SUCCESS(ValidateAnnotatedArgsTaskPlanStreams(*cached_plan, main_stream_id, plan_registry));
		806	+ if ((mode == OfflineLaunchGenMode::kMobileLegacySingleTask) &&
		807	+ (cached_plan->task_templates.size() != 1U \|\| !cached_plan->dependencies.empty() \|\|
		808	+ !cached_plan->attached_stream_ids.empty())) {
		809	+ return INTERNAL_ERROR;
		810	+ }
		811	+ return MaterializeAnnotatedArgsTaskPlan(op_desc, context, *cached_plan, tasks);
		812	+}
		813	+
		814	+Status EmitAnnotatedArgsLaunchTasks(const Node &node, const OpDescPtr &op_desc, const AnnotatedArgsCapture &capture,
		815	+ const size_t task_start, std::vector<domi::TaskDef> &tasks) {
737	std::vector<std::string> prefixes;	816	std::vector<std::string> prefixes;
738	prefixes.reserve(capture.args_handler.GetLaunchCount());	817	prefixes.reserve(capture.args_handler.GetLaunchCount());
739	for (size_t i = 0U; i < capture.args_handler.GetLaunchCount(); ++i) {	818	for (size_t i = 0U; i < capture.args_handler.GetLaunchCount(); ++i) {
	@@ -748,10 +827,58 @@ Status GenerateAnnotatedArgsTask(const Node &node, const OpDescPtr &op_desc, Run
748	GE_ASSERT_TRUE(AttrUtils::SetListStr(op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, prefixes),	827	GE_ASSERT_TRUE(AttrUtils::SetListStr(op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, prefixes),
749	"Set %s failed for custom op %s(%s).", ATTR_NAME_KERNEL_NAMES_PREFIX.c_str(), op_desc->GetNamePtr(),	828	"Set %s failed for custom op %s(%s).", ATTR_NAME_KERNEL_NAMES_PREFIX.c_str(), op_desc->GetNamePtr(),
750	op_desc->GetTypePtr());	829	op_desc->GetTypePtr());
751	- GE_ASSERT_SUCCESS(CacheAnnotatedArgsTaskPlan(op_desc, tasks, task_start));	830	+ GE_ASSERT_SUCCESS(CacheAnnotatedArgsTaskPlan(op_desc, tasks, task_start, capture.args_handler));
752	return SUCCESS;	831	return SUCCESS;
753	}	832	}
754		833
		834	+Status GenerateAnnotatedArgsTask(const Node &node, const OpDescPtr &op_desc, RunContext &context,
		835	+ const OfflineLaunchGenMode mode, std::vector<domi::TaskDef> &tasks) {
		836	+ GE_ASSERT_SUCCESS(CheckStaticGraph(node));
		837	+ uint32_t main_stream_id = 0U;
		838	+ GE_ASSERT_SUCCESS(GetCustomOpStreamId(op_desc, main_stream_id));
		839	+ const auto plan_registry = FindDeclarativeStreamRegistry(node);
		840	+ bool plan_restored = false;
		841	+ const auto restore_ret =
		842	+ RestoreCachedAnnotatedArgsTaskPlan(op_desc, main_stream_id, plan_registry, mode, context, tasks, plan_restored);
		843	+ if (restore_ret != SUCCESS) {
		844	+ GELOGE(restore_ret, "Custom op %s(%s) failed to restore annotated args task plan from cache.",
		845	+ op_desc->GetNamePtr(), op_desc->GetTypePtr());
		846	+ return restore_ret;
		T T tang-haojie9月23日异常分支需要加日志吗？下面也有类似的分支 C Chang-an-HW 9月23日评论： 9月23日评论：
		847	+ }
		848	+ if (plan_restored) {
		849	+ GELOGI("Custom op %s(%s) annotated args task plan restored from cache.", op_desc->GetNamePtr(),
		850	+ op_desc->GetTypePtr());
		851	+ return SUCCESS;
		852	+ }
		853	+
		854	+ AnnotatedArgsOp *annotated_args_op = nullptr;
		855	+ GE_ASSERT_SUCCESS(GetAnnotatedArgsOp(op_desc, annotated_args_op));
		856	+
		857	+ const size_t task_start = tasks.size();
		858	+ AnnotatedArgsCapture capture;
		859	+ GE_ASSERT_SUCCESS(ConstructInputTensors(op_desc, context, capture.input_tensor_holders));
		860	+ GE_ASSERT_SUCCESS(ConstructOutputTensors(op_desc, context, capture.output_tensor_holders));
		861	+ WorkspaceAllocator workspace_allocator(op_desc, context);
		862	+ const uint32_t stream_id = main_stream_id;
		863	+ DeclarativeStreamRegistryPtr registry = plan_registry;
		864	+ if (registry != nullptr) {
		865	+ capture.args_handler.SetAttachedStreamRequestFunc(
		866	+ [registry](const ge::AscendString &key) { return registry->RequestAttachedStream(key); });
		867	+ }
		868	+ const auto launch_task_plan_ret =
		869	+ DeclareLaunchArgsTaskPlan(op_desc, *annotated_args_op, workspace_allocator, stream_id, mode, registry, capture);
		870	+ if (launch_task_plan_ret != SUCCESS) {
		871	+ return launch_task_plan_ret;
		872	+ }
		873	+ GE_ASSERT_SUCCESS(UpdateAnnotatedArgsWorkspaceAttrs(op_desc, workspace_allocator));
		874	+
		875	+ GE_ASSERT_SUCCESS(ValidateKernelBinNames(op_desc, capture.args_handler));
		876	+ if (mode == OfflineLaunchGenMode::kMobileLegacySingleTask) {
		877	+ return GenerateMobileAnnotatedArgsTask(node, op_desc, capture.args_handler, task_start, tasks);
		878	+ }
		879	+ return EmitAnnotatedArgsLaunchTasks(node, op_desc, capture, task_start, tasks);
		880	+}
		881	+
755	Status FillBasicCustomKernelTask(const Node &node, domi::TaskDef &task_def) {	882	Status FillBasicCustomKernelTask(const Node &node, domi::TaskDef &task_def) {
756	const auto op_desc = node.GetOpDesc();	883	const auto op_desc = node.GetOpDesc();
757	GE_ASSERT_NOTNULL(op_desc);	884	GE_ASSERT_NOTNULL(op_desc);
	@@ -797,9 +924,10 @@ Status CustomOpsKernelBuilder::CalcOpRunningParam(Node &node) {
797	GE_ASSERT_NOTNULL(op_desc);	924	GE_ASSERT_NOTNULL(op_desc);
798	const auto args_refresh_strategy = CustomOpFactory::GetArgsRefreshStrategy(AscendString(op_desc->GetTypePtr()));	925	const auto args_refresh_strategy = CustomOpFactory::GetArgsRefreshStrategy(AscendString(op_desc->GetTypePtr()));
799	if (args_refresh_strategy == ArgsRefreshStrategy::kAnnotatedArgs) {	926	if (args_refresh_strategy == ArgsRefreshStrategy::kAnnotatedArgs) {
800	- (void)op_desc->DelExtAttr(kAnnotatedArgsTaskPlan);	927	+ (void)op_desc->DelAttr(kAnnotatedArgsTaskPlanBytes);
801	(void)op_desc->DelAttr(kAppendWs);	928	(void)op_desc->DelAttr(kAppendWs);
802	(void)op_desc->DelAttr(kCustomOmcAppendWs);	929	(void)op_desc->DelAttr(kCustomOmcAppendWs);
		930	+ (void)op_desc->DelAttr(kKeepGeneratedTasksAttr);
803	}	931	}
804	for (size_t i = 0; i < op_desc->GetOutputsSize(); i++) {	932	for (size_t i = 0; i < op_desc->GetOutputsSize(); i++) {
805	if (op_desc->GetOutputDesc(i).GetShape().IsUnknownShape()) {	933	if (op_desc->GetOutputDesc(i).GetShape().IsUnknownShape()) {

Mcompiler/graph/build/graph_builder.cc+7-0

	@@ -380,6 +380,13 @@ Status GraphBuilder::BuildForKnownShapeGraph(ComputeGraphPtr &comp_graph, GeMode
380	"[Get][FirstTaskInfo] fail, Graph[%s].", comp_graph->GetName().c_str());	380	"[Get][FirstTaskInfo] fail, Graph[%s].", comp_graph->GetName().c_str());
381	GE_COMPILE_TRACE_TIMESTAMP_END(GetTaskInfo, "GraphBuilder::GetTaskInfo");	381	GE_COMPILE_TRACE_TIMESTAMP_END(GetTaskInfo, "GraphBuilder::GetTaskInfo");
382		382
		383	+ GE_ASSERT_NOTNULL(graph_2_task_generator_[comp_graph]);
		384	+ GE_ASSERT_SUCCESS(builder.FinalizeDeclarativeAttachedStreams(), "FinalizeDeclarativeAttachedStreams fail, Graph[%s].",
		T T tang-haojie9月23日流分配这块找国正看下
		385	+ comp_graph->GetName().c_str());
		386	+ GE_ASSERT_SUCCESS(
		387	+ builder.MaterializeAnnotatedArgsTaskDependencies(graph_2_task_generator_[comp_graph]->MutableNodeId2TaskDefs()),
		388	+ "MaterializeAnnotatedArgsTaskDependencies fail, Graph[%s].", comp_graph->GetName().c_str());
		389	+
383	GE_TRACE_START(RefreshRealStream);	390	GE_TRACE_START(RefreshRealStream);
384	GE_ASSERT_NOTNULL(graph_2_task_generator_[comp_graph]);	391	GE_ASSERT_NOTNULL(graph_2_task_generator_[comp_graph]);
385	GE_ASSERT_SUCCESS(builder.RefreshRealStream(graph_2_task_generator_[comp_graph]->MutableNodeId2TaskDefs()),	392	GE_ASSERT_SUCCESS(builder.RefreshRealStream(graph_2_task_generator_[comp_graph]->MutableNodeId2TaskDefs()),

Mcompiler/graph/build/model_builder.cc+76-0

	@@ -22,7 +22,9 @@
22	#include "graph/build/input_h2d_overlap_planner.h"	22	#include "graph/build/input_h2d_overlap_planner.h"
23	#include "common/ge_common/ge_types.h"	23	#include "common/ge_common/ge_types.h"
24	#include "graph/build/stream/dynamic_stream_allocator.h"	24	#include "graph/build/stream/dynamic_stream_allocator.h"
		25	+#include "graph/build/stream/annotated_args_task_plan.h"
25	#include "graph/build/stream_graph_optimizer.h"	26	#include "graph/build/stream_graph_optimizer.h"
		27	+#include "graph/build/task_generator_utils.h"
26	#include "common/omg_util/omg_util.h"	28	#include "common/omg_util/omg_util.h"
27	#include "common/compile_profiling/ge_trace_wrapper.h"	29	#include "common/compile_profiling/ge_trace_wrapper.h"
28	#include "graph/ge_context.h"	30	#include "graph/ge_context.h"
	@@ -1225,6 +1227,11 @@ void ModelBuilder::DelNodeRepeatSaveAttr() {
1225	GE_IF_BOOL_EXEC(node_op_desc == nullptr, continue);	1227	GE_IF_BOOL_EXEC(node_op_desc == nullptr, continue);
1226	(void)node_op_desc->DelAttr(ATTR_NAME_TBE_KERNEL_BUFFER);	1228	(void)node_op_desc->DelAttr(ATTR_NAME_TBE_KERNEL_BUFFER);
1227	(void)node_op_desc->DelAttr(ATTR_NAME_TBE_KERNEL_NAME);	1229	(void)node_op_desc->DelAttr(ATTR_NAME_TBE_KERNEL_NAME);
		1230	+ // 声明式自定义算子的任务计划缓存仅供编译期二次生成任务消费，OM加载后无读取方，
		1231	+ // 保存前删除，避免kernel args模板在OM中与任务列表双份存储
		1232	+ (void)node_op_desc->DelAttr("_custom_annotated_args_task_plan_bytes");
		1233	+ // 同上，禁止二次生成任务的标记也只在编译期有效
		1234	+ (void)node_op_desc->DelAttr(kKeepGeneratedTasksAttr);
1228	// Load need kernel_name attr as key, so just remove kernel_buffer, which is more larger	1235	// Load need kernel_name attr as key, so just remove kernel_buffer, which is more larger
1229	std::vector<std::string> names_prefix;	1236	std::vector<std::string> names_prefix;
1230	(void)AttrUtils::GetListStr(node_op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, names_prefix);	1237	(void)AttrUtils::GetListStr(node_op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, names_prefix);
	@@ -1283,6 +1290,58 @@ Status ModelBuilder::RefreshRealStream(std::unordered_map<int64_t, std::vector<d
1283	return SUCCESS;	1290	return SUCCESS;
1284	}	1291	}
1285		1292
		1293	+Status ModelBuilder::FinalizeDeclarativeAttachedStreams() {
		C C Chang-an-HW9月14日把 AnnotatedArgs 计划中记录的声明式附加流，正式同步到图和流分配器中让流分配器和模型最终的 stream 数量覆盖这些在任务生成阶段动态申请出来的附加流
		1294	+ GE_CHECK_NOTNULL(compute_graph_);
		1295	+ auto registry = compute_graph_->TryGetExtAttr(kDeclarativeStreamRegistryAttr, DeclarativeStreamRegistryPtr());
		1296	+ uint32_t max_stream_id = 0U;
		1297	+ bool has_attached_stream = false;
		1298	+ if (registry != nullptr) {
		1299	+ for (const auto id : registry->GetAllocatedStreamIds()) {
		1300	+ max_stream_id = std::max(max_stream_id, id);
		1301	+ has_attached_stream = true;
		1302	+ }
		1303	+ }
		1304	+ for (const auto &node : compute_graph_->GetAllNodes()) {
		1305	+ GE_CHECK_NOTNULL(node);
		1306	+ const auto op_desc = node->GetOpDesc();
		1307	+ GE_CHECK_NOTNULL(op_desc);
		1308	+ Buffer encoded_plan;
		1309	+ if (AttrUtils::GetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan)) {
		1310	+ AnnotatedArgsTaskPlan plan;
		1311	+ GE_ASSERT_SUCCESS(DeserializeAnnotatedArgsTaskPlan(encoded_plan, plan));
		1312	+ if (plan.attached_stream_ids.empty()) {
		1313	+ continue;
		1314	+ }
		1315	+ GE_ASSERT_NOTNULL(registry, "Declarative stream registry is missing for custom op %s(%s).", op_desc->GetNamePtr(),
		1316	+ op_desc->GetTypePtr());
		1317	+ has_attached_stream = true;
		1318	+ std::vector<int64_t> attached_stream_ids;
		1319	+ attached_stream_ids.reserve(plan.attached_stream_ids.size());
		1320	+ for (const auto id : plan.attached_stream_ids) {
		1321	+ GE_ASSERT_TRUE(registry->Contains(id),
		1322	+ "Declarative stream registry does not contain stream id %u for custom op %s(%s).", id,
		1323	+ op_desc->GetNamePtr(), op_desc->GetTypePtr());
		1324	+ max_stream_id = std::max(max_stream_id, id);
		1325	+ attached_stream_ids.emplace_back(static_cast<int64_t>(id));
		1326	+ }
		1327	+ op_desc->SetAttachedStreamIds(attached_stream_ids);
		1328	+ }
		1329	+ }
		1330	+ if (has_attached_stream) {
		1331	+ GE_ASSERT_TRUE(max_stream_id < std::numeric_limits<uint32_t>::max(), "Max stream id %u exceeds uint32_t max value.",
		1332	+ max_stream_id);
		1333	+ const int64_t required_stream_num = static_cast<int64_t>(max_stream_id) + 1L;
		1334	+ stream_num_ = std::max(stream_num_, required_stream_num);
		1335	+ stream_allocator_.ExtendLogicalStreamNum(required_stream_num);
		1336	+ }
		1337	+ return SUCCESS;
		1338	+}
		1339	+
		1340	+Status ModelBuilder::MaterializeAnnotatedArgsTaskDependencies(
		C C Chang-an-HW9月14日把 AnnotatedArgs 计划中声明的 launch 依赖，物化成真正的 runtime event task
		1341	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks) {
		1342	+ return stream_allocator_.InsertAnnotatedArgsEventTasks(node_id_2_node_tasks, event_num_);
		1343	+}
		1344	+
1286	Status ModelBuilder::BuildModelForGetTask(ge::Model &model) {	1345	Status ModelBuilder::BuildModelForGetTask(ge::Model &model) {
1287	GE_CHK_STATUS_RET(AdjustInputTensorFlag(), "[Adjust][InputTensorFlag] failed! graph:%s",	1346	GE_CHK_STATUS_RET(AdjustInputTensorFlag(), "[Adjust][InputTensorFlag] failed! graph:%s",
1288	compute_graph_->GetName().c_str());	1347	compute_graph_->GetName().c_str());
	@@ -1293,6 +1352,8 @@ Status ModelBuilder::BuildModelForGetTask(ge::Model &model) {
1293	"[Assign][LogicalStreams] failed. graph:%s", compute_graph_->GetName().c_str());	1352	"[Assign][LogicalStreams] failed. graph:%s", compute_graph_->GetName().c_str());
1294	GE_COMPILE_TRACE_TIMESTAMP_END(AssignLogicalStreams, "GraphBuilder::AssignLogicalStreams");	1353	GE_COMPILE_TRACE_TIMESTAMP_END(AssignLogicalStreams, "GraphBuilder::AssignLogicalStreams");
1295	GE_DUMP(compute_graph_, "AfterAssignLogicalStreams");	1354	GE_DUMP(compute_graph_, "AfterAssignLogicalStreams");
		1355	+ GE_CHK_STATUS_RET(PrepareDeclarativeAttachedStreamRegistry(),
		1356	+ "[Prepare][DeclarativeAttachedStreamRegistry] failed. graph:%s", compute_graph_->GetName().c_str());
1296		1357
1297	// Assign functional op labels.	1358	// Assign functional op labels.
1298	auto root_graph = GraphUtils::FindRootGraph(compute_graph_);	1359	auto root_graph = GraphUtils::FindRootGraph(compute_graph_);
	@@ -1352,6 +1413,21 @@ Status ModelBuilder::BuildModelForGetTask(ge::Model &model) {
1352	return SUCCESS;	1413	return SUCCESS;
1353	}	1414	}
1354		1415
		1416	+Status ModelBuilder::PrepareDeclarativeAttachedStreamRegistry() {
		1417	+ GE_CHECK_NOTNULL(compute_graph_);
		1418	+ // registry 必须挂在当前编译子图上，起始 ID 取本子图的逻辑流数：运行期物理流池按子模型独立
		1419	+ // （V1 每个 GeModel 对应一个 DavinciModel 的 stream_list_，V2 静态子模型各自成 rt model），
		1420	+ // 因此声明式辅流 ID 的作用域就是子图。若统一挂到根图，同一 key 在各子图拿到相同 ID 也不会共享物理流，
		1421	+ // 反而会让每个子模型按全局最大 ID 虚增流数，且根图在下沉切图路径下不经过 ModelBuilder 构建，无起始流数可用。
		1422	+ const auto logic_stream_num = stream_allocator_.GetStreamNum();
		1423	+ const auto first_stream_id = (logic_stream_num < 0) ? 0U : static_cast<uint32_t>(logic_stream_num);
		1424	+ auto registry = std::make_shared<DeclarativeStreamRegistry>(first_stream_id);
		1425	+ GE_CHECK_NOTNULL(registry);
		1426	+ GE_ASSERT_TRUE(compute_graph_->SetExtAttr(kDeclarativeStreamRegistryAttr, registry),
		1427	+ "Set declarative stream registry failed, graph:%s", compute_graph_->GetName().c_str());
		1428	+ return SUCCESS;
		1429	+}
		1430	+
1355	Status ModelBuilder::BuildModelDefForStream(ge::Model &model) {	1431	Status ModelBuilder::BuildModelDefForStream(ge::Model &model) {
1356	GE_ASSERT_SUCCESS(CollectEagerCustomOpStreamNum(model), "[Collect][EagerCustomOpStreamNum] failed! graph:%s",	1432	GE_ASSERT_SUCCESS(CollectEagerCustomOpStreamNum(model), "[Collect][EagerCustomOpStreamNum] failed! graph:%s",
1357	compute_graph_->GetName().c_str());	1433	compute_graph_->GetName().c_str());

Mcompiler/graph/build/model_builder.h+6-0

	@@ -27,6 +27,7 @@
27	#include "graph/manager/graph_manager_utils.h"	27	#include "graph/manager/graph_manager_utils.h"
28	#include "graph/manager/graph_var_manager.h"	28	#include "graph/manager/graph_var_manager.h"
29	#include "graph/build/stream/graph_stream_allocator.h"	29	#include "graph/build/stream/graph_stream_allocator.h"
		30	+#include "graph/build/stream/declarative_stream_registry.h"
30	#include "graph/model.h"	31	#include "graph/model.h"
31	#include "graph/node.h"	32	#include "graph/node.h"
32	#include "common/model/ge_model.h"	33	#include "common/model/ge_model.h"
	@@ -54,6 +55,9 @@ class ModelBuilder {
54	Status BuildModelDefForStream(ge::Model &model);	55	Status BuildModelDefForStream(ge::Model &model);
55	Status SaveInputH2DOverlapPlan(ge::Model &model);	56	Status SaveInputH2DOverlapPlan(ge::Model &model);
56	Status RefreshRealStream(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks);	57	Status RefreshRealStream(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks);
		58	+ Status FinalizeDeclarativeAttachedStreams();
		59	+ Status MaterializeAnnotatedArgsTaskDependencies(
		60	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks);
57	ge::Status BuildModelForGetDynShapeTask(ge::Model &model_def);	61	ge::Status BuildModelForGetDynShapeTask(ge::Model &model_def);
58	Status AssignStreamForDynamicShapeGraph(ComputeGraphPtr &compute_graph);	62	Status AssignStreamForDynamicShapeGraph(ComputeGraphPtr &compute_graph);
59		63
	@@ -105,6 +109,8 @@ class ModelBuilder {
105		109
106	Status BuildModelDef(ge::Model &model);	110	Status BuildModelDef(ge::Model &model);
107		111
		112	+ Status PrepareDeclarativeAttachedStreamRegistry();
		113	+
108	Status BuildModelDefForMem(ge::Model &model);	114	Status BuildModelDefForMem(ge::Model &model);
109		115
110	Status InitL1FusionOption();	116	Status InitL1FusionOption();

Acompiler/graph/build/stream/annotated_args_task_plan.cc+290-0

@@ -0,0 +1,290 @@
	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 "annotated_args_task_plan.h"
	12	+
	13	+#include <climits>
	14	+#include <cstddef>
	15	+#include <limits>
	16	+#include <string>
	17	+
	18	+#include "framework/omg/omg_inner_types.h"
	19	+#include "framework/common/debug/ge_log.h"
	20	+
	21	+namespace ge {
	22	+namespace {
	23	+constexpr uint32_t kVersion = 1U;
	24	+constexpr size_t kU32EntrySize = 4U;
	25	+constexpr size_t kDependencyEntrySize = 12U;
	26	+
	27	+bool CanAppend(const std::vector<uint8_t> &out, const size_t count) {
	28	+ return count <= (std::numeric_limits<size_t>::max() - out.size());
	29	+}
	30	+
	31	+bool PutU32(std::vector<uint8_t> &out, const uint32_t value) {
	32	+ if (!CanAppend(out, sizeof(value))) {
	33	+ return false;
	34	+ }
	35	+ out.push_back(static_cast<uint8_t>(value));
	36	+ out.push_back(static_cast<uint8_t>(value >> 8U));
	37	+ out.push_back(static_cast<uint8_t>(value >> 16U));
	38	+ out.push_back(static_cast<uint8_t>(value >> 24U));
	39	+ return true;
	40	+}
	41	+
	42	+bool PutCount(std::vector<uint8_t> &out, const size_t count) {
	43	+ if (count > std::numeric_limits<uint32_t>::max()) {
	44	+ return false;
	45	+ }
	46	+ return PutU32(out, static_cast<uint32_t>(count));
	47	+}
	48	+
	49	+bool PutIdList(std::vector<uint8_t> &out, const std::vector<uint32_t> &ids, const char *const section) {
	50	+ if (!PutCount(out, ids.size())) {
	51	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] %s count cannot be encoded, count=%zu.", section, ids.size());
	52	+ return false;
	53	+ }
	54	+ for (size_t i = 0U; i < ids.size(); ++i) {
	55	+ if (!PutU32(out, ids[i])) {
	56	+ GELOGE(INTERNAL_ERROR, "[AnnotatedArgsTaskPlan] failed to serialize %s id, index=%zu, id=%u.", section, i,
	57	+ ids[i]);
	58	+ return false;
	59	+ }
	60	+ }
	61	+ return true;
	62	+}
	63	+
	64	+bool IsValidDependencyEdge(const AnnotatedArgsLaunchDependency &dep, const size_t task_count) {
	65	+ return (dep.predecessor_launch_index < task_count) && (dep.successor_launch_index < task_count) &&
	66	+ (dep.predecessor_launch_index < dep.successor_launch_index);
	67	+}
	68	+
	69	+bool SerializeTaskTemplates(const AnnotatedArgsTaskPlan &plan, std::vector<uint8_t> &bytes) {
	70	+ if (!PutCount(bytes, plan.task_templates.size())) {
	71	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] task template count cannot be encoded, count=%zu.",
	72	+ plan.task_templates.size());
	73	+ return false;
	74	+ }
	75	+ for (size_t i = 0U; i < plan.task_templates.size(); ++i) {
	76	+ std::string serialized;
	77	+ if (!plan.task_templates[i].SerializeToString(&serialized) \|\| !PutCount(bytes, serialized.size()) \|\|
	78	+ !CanAppend(bytes, serialized.size())) {
	79	+ GELOGE(INTERNAL_ERROR, "[AnnotatedArgsTaskPlan] failed to serialize kernel task template, index=%zu, size=%zu.",
	80	+ i, serialized.size());
	81	+ return false;
	82	+ }
	83	+ bytes.insert(bytes.end(), serialized.begin(), serialized.end());
	84	+ }
	85	+ return true;
	86	+}
	87	+
	88	+bool SerializeDependencies(const AnnotatedArgsTaskPlan &plan, std::vector<uint8_t> &bytes) {
	89	+ if (!PutCount(bytes, plan.dependencies.size())) {
	90	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] dependency count cannot be encoded, count=%zu.",
	91	+ plan.dependencies.size());
	92	+ return false;
	93	+ }
	94	+ for (size_t i = 0U; i < plan.dependencies.size(); ++i) {
	95	+ const auto &dep = plan.dependencies[i];
	96	+ if (!PutU32(bytes, dep.predecessor_launch_index) \|\| !PutU32(bytes, dep.successor_launch_index) \|\|
	97	+ !PutU32(bytes, dep.event_id)) {
	98	+ GELOGE(INTERNAL_ERROR,
	99	+ "[AnnotatedArgsTaskPlan] failed to serialize dependency, index=%zu, predecessor=%u, successor=%u.", i,
	100	+ dep.predecessor_launch_index, dep.successor_launch_index);
	101	+ return false;
	102	+ }
	103	+ }
	104	+ return true;
	105	+}
	106	+
	107	+class Reader {
	108	+ public:
	109	+ explicit Reader(const Buffer &buffer) : data_(buffer.GetData()), size_(buffer.GetSize()) {}
	110	+ bool U32(uint32_t &value) {
	111	+ if (pos_ > size_ \|\| size_ - pos_ < 4U) {
	112	+ return false;
	113	+ }
	114	+ value = static_cast<uint32_t>(data_[pos_]) \| (static_cast<uint32_t>(data_[pos_ + 1U]) << 8U) \|
	115	+ (static_cast<uint32_t>(data_[pos_ + 2U]) << 16U) \| (static_cast<uint32_t>(data_[pos_ + 3U]) << 24U);
	116	+ pos_ += 4U;
	117	+ return true;
	118	+ }
	119	+ bool Bytes(const uint8_t *&bytes, size_t &length) {
	120	+ uint32_t len = 0U;
	121	+ if (!U32(len) \|\| pos_ > size_ \|\| static_cast<size_t>(len) > size_ - pos_) {
	122	+ return false;
	123	+ }
	124	+ bytes = data_ + pos_;
	125	+ length = len;
	126	+ pos_ += len;
	127	+ return true;
	128	+ }
	129	+ bool Done() const {
	130	+ return pos_ == size_;
	131	+ }
	132	+
	133	+ private:
	134	+ const uint8_t *data_ = nullptr;
	135	+ size_t size_ = 0U;
	136	+ size_t pos_ = 0U;
	137	+};
	138	+
	139	+bool ReadCount(Reader &reader, const size_t buffer_size, const size_t entry_size, const char *const section,
	140	+ uint32_t &count) {
	141	+ if (!reader.U32(count) \|\| static_cast<size_t>(count) > buffer_size / entry_size) {
	142	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] invalid %s count in serialized plan, count=%u, size=%zu.", section,
	143	+ count, buffer_size);
	144	+ return false;
	145	+ }
	146	+ return true;
	147	+}
	148	+
	149	+bool ReadTaskTemplates(Reader &reader, const size_t buffer_size, AnnotatedArgsTaskPlan &result) {
	150	+ uint32_t count = 0U;
	151	+ if (!ReadCount(reader, buffer_size, kU32EntrySize, "task template", count)) {
	152	+ return false;
	153	+ }
	154	+ result.task_templates.reserve(count);
	155	+ for (uint32_t i = 0U; i < count; ++i) {
	156	+ const uint8_t *data = nullptr;
	157	+ size_t size = 0U;
	158	+ if (!reader.Bytes(data, size)) {
	159	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] truncated task template bytes, index=%u.", i);
	160	+ return false;
	161	+ }
	162	+ domi::TaskDef task;
	163	+ if (size > static_cast<size_t>(INT_MAX) \|\| !task.ParseFromArray(data, static_cast<int>(size)) \|\|
	164	+ task.type() != static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL)) {
	165	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] invalid task template, index=%u, size=%zu, type=%u.", i, size,
	166	+ task.type());
	167	+ return false;
	168	+ }
	169	+ result.task_templates.emplace_back(std::move(task));
	170	+ }
	171	+ return true;
	172	+}
	173	+
	174	+bool ReadIdList(Reader &reader, const size_t buffer_size, const char *const section, std::vector<uint32_t> &ids) {
	175	+ uint32_t count = 0U;
	176	+ if (!ReadCount(reader, buffer_size, kU32EntrySize, section, count)) {
	177	+ return false;
	178	+ }
	179	+ ids.resize(count);
	180	+ for (size_t i = 0U; i < ids.size(); ++i) {
	181	+ if (!reader.U32(ids[i])) {
	182	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] truncated %s id, index=%zu.", section, i);
	183	+ return false;
	184	+ }
	185	+ }
	186	+ return true;
	187	+}
	188	+
	189	+bool ReadDependencies(Reader &reader, const size_t buffer_size, AnnotatedArgsTaskPlan &result) {
	190	+ uint32_t count = 0U;
	191	+ if (!ReadCount(reader, buffer_size, kDependencyEntrySize, "dependency", count)) {
	192	+ return false;
	193	+ }
	194	+ result.dependencies.resize(count);
	195	+ for (size_t i = 0U; i < result.dependencies.size(); ++i) {
	196	+ auto &dep = result.dependencies[i];
	197	+ if (!reader.U32(dep.predecessor_launch_index) \|\| !reader.U32(dep.successor_launch_index) \|\|
	198	+ !reader.U32(dep.event_id)) {
	199	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] truncated dependency, index=%zu.", i);
	200	+ return false;
	201	+ }
	202	+ if (!IsValidDependencyEdge(dep, result.task_templates.size())) {
	203	+ GELOGE(PARAM_INVALID,
	204	+ "[AnnotatedArgsTaskPlan] invalid dependency, index=%zu, predecessor=%u, successor=%u, tasks=%zu.", i,
	205	+ dep.predecessor_launch_index, dep.successor_launch_index, result.task_templates.size());
	206	+ return false;
	207	+ }
	208	+ }
	209	+ return true;
	210	+}
	211	+
	212	+bool ValidatePlanIntegrity(const AnnotatedArgsTaskPlan &plan) {
	213	+ if (plan.launch_stream_ids.size() != plan.task_templates.size()) {
	214	+ GELOGE(PARAM_INVALID,
	215	+ "[AnnotatedArgsTaskPlan] launch stream count does not match task count, tasks=%zu, streams=%zu.",
	216	+ plan.task_templates.size(), plan.launch_stream_ids.size());
	217	+ return false;
	218	+ }
	219	+ if (plan.task_templates.size() > std::numeric_limits<uint32_t>::max() \|\|
	220	+ plan.launch_stream_ids.size() > std::numeric_limits<uint32_t>::max() \|\|
	221	+ plan.dependencies.size() > std::numeric_limits<uint32_t>::max() \|\|
	222	+ plan.attached_stream_ids.size() > std::numeric_limits<uint32_t>::max()) {
	223	+ GELOGE(PARAM_INVALID,
	224	+ "[AnnotatedArgsTaskPlan] plan count exceeds uint32 limit, tasks=%zu, streams=%zu, dependencies=%zu, "
	225	+ "attached_streams=%zu.",
	226	+ plan.task_templates.size(), plan.launch_stream_ids.size(), plan.dependencies.size(),
	227	+ plan.attached_stream_ids.size());
	228	+ return false;
	229	+ }
	230	+ for (const auto &dep : plan.dependencies) {
	231	+ if (!IsValidDependencyEdge(dep, plan.task_templates.size())) {
	232	+ GELOGE(PARAM_INVALID,
	233	+ "[AnnotatedArgsTaskPlan] invalid dependency edge, predecessor=%u, successor=%u, task_count=%zu; "
	234	+ "indices must be in range and predecessor must be before successor.",
	235	+ dep.predecessor_launch_index, dep.successor_launch_index, plan.task_templates.size());
	236	+ return false;
	237	+ }
	238	+ }
	239	+ for (const auto &task : plan.task_templates) {
	240	+ if (task.type() != static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL)) {
	241	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] task template type=%u is not custom kernel.", task.type());
	242	+ return false;
	243	+ }
	244	+ }
	245	+ return true;
	246	+}
	247	+} // namespace
	248	+
	249	+Status SerializeAnnotatedArgsTaskPlan(const AnnotatedArgsTaskPlan &plan, Buffer &encoded) {
	250	+ if (!ValidatePlanIntegrity(plan)) {
	251	+ return GRAPH_FAILED;
	252	+ }
	253	+ std::vector<uint8_t> bytes;
	254	+ if (!PutU32(bytes, kVersion)) {
	255	+ GELOGE(INTERNAL_ERROR, "[AnnotatedArgsTaskPlan] failed to serialize plan version.");
	256	+ return GRAPH_FAILED;
	257	+ }
	258	+ if (!SerializeTaskTemplates(plan, bytes) \|\| !PutIdList(bytes, plan.launch_stream_ids, "launch stream") \|\|
	259	+ !SerializeDependencies(plan, bytes) \|\| !PutIdList(bytes, plan.attached_stream_ids, "attached stream")) {
	260	+ return GRAPH_FAILED;
	261	+ }
	262	+ encoded = Buffer::CopyFrom(bytes.data(), bytes.size());
	263	+ return SUCCESS;
	264	+}
	265	+
	266	+Status DeserializeAnnotatedArgsTaskPlan(const Buffer &encoded, AnnotatedArgsTaskPlan &plan) {
	267	+ Reader reader(encoded);
	268	+ uint32_t version = 0U;
	269	+ if (!reader.U32(version) \|\| version != kVersion) {
	270	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsTaskPlan] unsupported or truncated plan version, version=%u.", version);
	271	+ return GRAPH_FAILED;
	272	+ }
	273	+ AnnotatedArgsTaskPlan result;
	274	+ const size_t buffer_size = encoded.GetSize();
	275	+ if (!ReadTaskTemplates(reader, buffer_size, result) \|\|
	276	+ !ReadIdList(reader, buffer_size, "launch stream", result.launch_stream_ids) \|\|
	277	+ !ReadDependencies(reader, buffer_size, result) \|\|
	278	+ !ReadIdList(reader, buffer_size, "attached stream", result.attached_stream_ids)) {
	279	+ return GRAPH_FAILED;
	280	+ }
	281	+ if (!reader.Done() \|\| result.launch_stream_ids.size() != result.task_templates.size()) {
	282	+ GELOGE(PARAM_INVALID,
	283	+ "[AnnotatedArgsTaskPlan] malformed serialized plan, remaining_bytes=%zu, tasks=%zu, streams=%zu.",
	284	+ buffer_size, result.task_templates.size(), result.launch_stream_ids.size());
	285	+ return GRAPH_FAILED;
	286	+ }
	287	+ plan = std::move(result);
	288	+ return SUCCESS;
	289	+}
	290	+} // namespace ge

Acompiler/graph/build/stream/annotated_args_task_plan.h+45-0

@@ -0,0 +1,45 @@
	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 GE_COMPILER_GRAPH_BUILD_STREAM_ANNOTATED_ARGS_TASK_PLAN_H_
	12	+#define GE_COMPILER_GRAPH_BUILD_STREAM_ANNOTATED_ARGS_TASK_PLAN_H_
	13	+
	14	+#include <cstdint>
	15	+#include <memory>
	16	+#include <vector>
	17	+
	18	+#include "graph/buffer.h"
	19	+#include "graph/op_desc.h"
	20	+#include "ge/ge_api_error_codes.h"
	21	+#include "framework/omg/omg_inner_types.h"
	22	+#include "common/opskernel/ops_kernel_info_types.h"
	23	+#include "proto/task.pb.h"
	24	+
	25	+namespace ge {
	26	+struct AnnotatedArgsLaunchDependency {
	27	+ uint32_t predecessor_launch_index;
	28	+ uint32_t successor_launch_index;
	29	+ uint32_t event_id;
	30	+};
	31	+
	32	+struct AnnotatedArgsTaskPlan {
	33	+ std::vector<domi::TaskDef> task_templates;
	34	+ std::vector<uint32_t> launch_stream_ids;
	35	+ std::vector<AnnotatedArgsLaunchDependency> dependencies;
	36	+ std::vector<uint32_t> attached_stream_ids;
	37	+};
	38	+
	39	+using AnnotatedArgsTaskPlanPtr = std::shared_ptr<const AnnotatedArgsTaskPlan>;
	40	+
	41	+Status SerializeAnnotatedArgsTaskPlan(const AnnotatedArgsTaskPlan &plan, Buffer &encoded);
	42	+Status DeserializeAnnotatedArgsTaskPlan(const Buffer &encoded, AnnotatedArgsTaskPlan &plan);
	43	+} // namespace ge
	44	+
	45	+#endif

Acompiler/graph/build/stream/declarative_stream_registry.cc+51-0

@@ -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	+#include "declarative_stream_registry.h"
	12	+
	13	+#include <limits>
	14	+#include <memory>
	15	+#include <algorithm>
	16	+
	17	+#include "framework/common/debug/ge_log.h"
	18	+
	19	+namespace ge {
	20	+DeclarativeStreamRegistry::DeclarativeStreamRegistry(const uint32_t first_stream_id)
	21	+ : next_stream_id_(first_stream_id) {}
	22	+
	23	+uint32_t DeclarativeStreamRegistry::RequestAttachedStream(const AscendString &key) {
	24	+ const std::string key_string(key.GetString());
	25	+ if (key_string.empty()) {
	26	+ GELOGE(PARAM_INVALID, "[DeclarativeStreamRegistry] empty attached stream key.");
	27	+ return std::numeric_limits<uint32_t>::max();
	T T tang-haojie9月23日此条代码评论区间+25至+27 这里只拦了空 key，`CANN-FMK-` 保留前缀没有查，但 RequestAttachedStream 的 API 文档写了"不能使用保留前缀"。而且 UT（stream_allocator_unittest.cc 的 DeduplicatesKeysAndRejectsEmptyKey）里专门断言了 `CANN-FMK-internal` 能正常拿到流 ID，测试和文档是反着的。5035 那个 PR 的运行时版本（AttachedStreamCollection::IsValidKey）是拦的。要么把前缀拦截补上，要么把文档那句删掉。 C Chang-an-HW 9月23日评论： 9月23日评论：
	28	+ }
	29	+ const auto found = key_to_stream_id_.find(key_string);
	30	+ if (found != key_to_stream_id_.end()) {
	31	+ return found->second;
	32	+ }
	33	+ if (next_stream_id_ == std::numeric_limits<uint32_t>::max()) {
	34	+ GELOGE(INTERNAL_ERROR, "[DeclarativeStreamRegistry] attached stream id overflow, next_id=%u.", next_stream_id_);
	35	+ return std::numeric_limits<uint32_t>::max();
	36	+ }
	37	+ const auto stream_id = next_stream_id_++;
	38	+ key_to_stream_id_.emplace(key_string, stream_id);
	39	+ allocated_stream_ids_.emplace_back(stream_id);
	40	+ return stream_id;
	41	+}
	42	+
	43	+bool DeclarativeStreamRegistry::Contains(const uint32_t stream_id) const {
	44	+ return std::find(allocated_stream_ids_.cbegin(), allocated_stream_ids_.cend(), stream_id) !=
	45	+ allocated_stream_ids_.cend();
	46	+}
	47	+
	48	+const std::vector<uint32_t> &DeclarativeStreamRegistry::GetAllocatedStreamIds() const {
	49	+ return allocated_stream_ids_;
	50	+}
	51	+} // namespace ge

Acompiler/graph/build/stream/declarative_stream_registry.h+39-0

@@ -0,0 +1,39 @@
	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 GE_GRAPH_BUILD_STREAM_DECLARATIVE_STREAM_REGISTRY_H_
	12	+#define GE_GRAPH_BUILD_STREAM_DECLARATIVE_STREAM_REGISTRY_H_
	13	+
	14	+#include <cstdint>
	15	+#include <memory>
	16	+#include <map>
	17	+#include <string>
	18	+#include <vector>
	19	+
	20	+#include "graph/ascend_string.h"
	21	+
	22	+namespace ge {
	23	+constexpr char kDeclarativeStreamRegistryAttr[] = "declarative_stream_registry";
	24	+class DeclarativeStreamRegistry {
	25	+ public:
	26	+ explicit DeclarativeStreamRegistry(uint32_t first_stream_id);
	27	+ uint32_t RequestAttachedStream(const AscendString &key);
	28	+ bool Contains(uint32_t stream_id) const;
	29	+ const std::vector<uint32_t> &GetAllocatedStreamIds() const;
	30	+
	31	+ private:
	32	+ std::map<std::string, uint32_t> key_to_stream_id_;
	33	+ std::vector<uint32_t> allocated_stream_ids_;
	34	+ uint32_t next_stream_id_;
	35	+};
	36	+using DeclarativeStreamRegistryPtr = std::shared_ptr<DeclarativeStreamRegistry>;
	37	+} // namespace ge
	38	+
	39	+#endif // GE_GRAPH_BUILD_STREAM_DECLARATIVE_STREAM_REGISTRY_H_

Mcompiler/graph/build/stream/graph_stream_allocator.h+25-0

	@@ -17,6 +17,8 @@
17	#include <string>	17	#include <string>
18	#include <utility>	18	#include <utility>
19	#include <vector>	19	#include <vector>
		20	+#include <unordered_map>
		21	+#include "proto/task.pb.h"
20		22
21	#include "framework/common/ge_inner_error_codes.h"	23	#include "framework/common/ge_inner_error_codes.h"
22	#include "graph/compute_graph.h"	24	#include "graph/compute_graph.h"
	@@ -24,6 +26,7 @@
24	#include "graph/utils/node_utils.h"	26	#include "graph/utils/node_utils.h"
25		27
26	namespace ge {	28	namespace ge {
		29	+struct AnnotatedArgsTaskPlan;
27	enum class EventType : std::uint32_t {	30	enum class EventType : std::uint32_t {
28	kEvent,	31	kEvent,
29	kNotify,	32	kNotify,
	@@ -105,6 +108,17 @@ class StreamAllocator {
105	Status SplitStreamAndRefreshTaskDef(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,	108	Status SplitStreamAndRefreshTaskDef(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
106	int64_t &stream_num, int64_t &event_num, int64_t &notify_num);	109	int64_t &stream_num, int64_t &event_num, int64_t &notify_num);
107		110
		111	+ Status InsertAnnotatedArgsEventTasks(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		112	+ int64_t &event_num);
		113	+ Status ReconcileAnnotatedArgsEventTasks(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		114	+ int64_t &event_num);
		115	+
		116	+ void ExtendLogicalStreamNum(const int64_t stream_num) {
		117	+ if (stream_num > stream_num_) {
		118	+ stream_num_ = stream_num;
		119	+ }
		120	+ }
		121	+
108	const std::vector<int64_t> &GetHugeStreams() const {	122	const std::vector<int64_t> &GetHugeStreams() const {
109	return huge_streams_;	123	return huge_streams_;
110	}	124	}
	@@ -299,6 +313,17 @@ class StreamAllocator {
299	Node2AttachedStreamId2EventId attached_node_to_stream_id_to_recv_event_id_;	313	Node2AttachedStreamId2EventId attached_node_to_stream_id_to_recv_event_id_;
300		314
301	std::vector<StandaloneWaitEvent> external_wait_events_;	315	std::vector<StandaloneWaitEvent> external_wait_events_;
		316	+ Status ExpandAnnotatedArgsNodeEvents(const OpDescPtr &op_desc, const AnnotatedArgsTaskPlan &plan,
		317	+ std::vector<domi::TaskDef> &tasks, int64_t &event_num);
		318	+ Status CollectAnnotatedEventPairs(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		319	+ std::map<uint32_t, std::pair<domi::TaskDef , domi::TaskDef >> &pairs);
		320	+
		321	+ struct AnnotatedEventInfo {
		322	+ int64_t node_id;
		323	+ uint32_t predecessor_launch_index;
		324	+ uint32_t successor_launch_index;
		325	+ };
		326	+ std::map<uint32_t, AnnotatedEventInfo> annotated_event_infos_;
302	};	327	};
303	} // namespace ge	328	} // namespace ge
304	#endif // GE_GRAPH_BUILD_STREAM_GRAPH_STREAM_ALLOCATOR_H_	329	#endif // GE_GRAPH_BUILD_STREAM_GRAPH_STREAM_ALLOCATOR_H_

Mcompiler/graph/build/stream/stream_allocator.cc+226-0

	@@ -26,6 +26,7 @@
26	#include "api/gelib/gelib.h"	26	#include "api/gelib/gelib.h"
27	#include "assign_attached_notify_pass.h"	27	#include "assign_attached_notify_pass.h"
28	#include "assign_attached_event_pass.h"	28	#include "assign_attached_event_pass.h"
		29	+#include "annotated_args_task_plan.h"
29	#include "common/util.h"	30	#include "common/util.h"
30	#include "common/ge_rts_decl.h"	31	#include "common/ge_rts_decl.h"
31		32
	@@ -306,6 +307,102 @@ void CollectSubgraphStreams(const OpDescPtr &op_desc, std::set<int64_t> &subgrap
306	bool IsBelongToSameGraph(const NodePtr &node1, const NodePtr &node2) {	307	bool IsBelongToSameGraph(const NodePtr &node1, const NodePtr &node2) {
307	return node1->GetOwnerComputeGraphBarePtr() == node2->GetOwnerComputeGraphBarePtr();	308	return node1->GetOwnerComputeGraphBarePtr() == node2->GetOwnerComputeGraphBarePtr();
308	}	309	}
		310	+bool IsAnnotatedEventTask(const domi::TaskDef &task) {
		311	+ return task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD) \|\|
		312	+ task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT);
		313	+}
		314	+domi::TaskDef MakeAnnotatedEventTask(const ModelTaskType type, const uint32_t stream_id, const uint32_t event_id,
		315	+ const uint32_t op_index) {
		316	+ domi::TaskDef task;
		317	+ task.set_type(static_cast<uint32_t>(type));
		318	+ task.set_stream_id(stream_id);
		319	+ task.set_event_id(event_id);
		320	+ task.set_sqe_num(1U);
		321	+ auto *event_ex = task.mutable_event_ex();
		322	+ event_ex->set_event_type(static_cast<uint32_t>(type));
		323	+ event_ex->set_op_index(op_index);
		324	+ return task;
		325	+}
		326	+void CollectAnnotatedKernelPositions(const std::vector<domi::TaskDef> &tasks, std::vector<size_t> &kernel_positions) {
		327	+ for (size_t i = 0U; i < tasks.size(); ++i) {
		328	+ if (tasks[i].type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL)) {
		329	+ kernel_positions.emplace_back(i);
		330	+ }
		331	+ }
		332	+}
		333	+void MergeAnnotatedEventTasks(std::vector<domi::TaskDef> &tasks,
		334	+ const std::map<size_t, std::vector<domi::TaskDef>> &records,
		335	+ const std::map<size_t, std::vector<domi::TaskDef>> &waits,
		336	+ const size_t dependency_count) {
		337	+ std::vector<domi::TaskDef> expanded;
		338	+ expanded.reserve(tasks.size() + dependency_count * 2U);
		339	+ for (size_t i = 0U; i < tasks.size(); ++i) {
		340	+ const auto wait_iter = waits.find(i);
		341	+ if (wait_iter != waits.end()) {
		342	+ expanded.insert(expanded.end(), wait_iter->second.begin(), wait_iter->second.end());
		343	+ }
		344	+ expanded.emplace_back(std::move(tasks[i]));
		345	+ const auto record_iter = records.find(i);
		346	+ if (record_iter != records.end()) {
		347	+ expanded.insert(expanded.end(), record_iter->second.begin(), record_iter->second.end());
		348	+ }
		349	+ }
		350	+ tasks = std::move(expanded);
		351	+}
		352	+void LogAnnotatedEventInsertStats(const std::string &graph_name,
		353	+ const std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		354	+ const int64_t event_num) {
		355	+ size_t total_task_count = 0U;
		356	+ size_t total_event_task_count = 0U;
		357	+ for (const auto &entry : node_id_2_node_tasks) {
		358	+ total_task_count += entry.second.size();
		359	+ for (const auto &task : entry.second) {
		360	+ total_event_task_count += IsAnnotatedEventTask(task);
		361	+ }
		362	+ }
		363	+ GELOGD(
		364	+ "[AnnotatedArgsEvents] inserted tasks, graph=%s, node_count=%zu, task_count=%zu, event_task_count=%zu, "
		365	+ "event_num=%ld.",
		366	+ graph_name.c_str(), node_id_2_node_tasks.size(), total_task_count, total_event_task_count, event_num);
		367	+}
		368	+Status BuildAnnotatedEventRemap(const std::map<uint32_t, std::pair<domi::TaskDef , domi::TaskDef >> &pairs,
		369	+ const int64_t event_num, std::map<uint32_t, uint32_t> &remap, uint32_t &next_id) {
		370	+ next_id = static_cast<uint32_t>(event_num - static_cast<int64_t>(pairs.size()));
		371	+ for (const auto &entry : pairs) {
		372	+ GE_ASSERT_NOTNULL(entry.second.first);
		373	+ GE_ASSERT_NOTNULL(entry.second.second);
		374	+ if (entry.second.first->stream_id() == entry.second.second->stream_id()) {
		375	+ remap.emplace(entry.first, std::numeric_limits<uint32_t>::max());
		376	+ } else {
		377	+ remap.emplace(entry.first, next_id++);
		378	+ }
		379	+ }
		380	+ return SUCCESS;
		381	+}
		382	+void ApplyAnnotatedEventRemap(std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		383	+ const std::map<uint32_t, uint32_t> &remap) {
		384	+ for (auto &entry : node_id_2_node_tasks) {
		385	+ auto &tasks = entry.second;
		386	+ tasks.erase(std::remove_if(tasks.begin(), tasks.end(),
		387	+ [&remap](const domi::TaskDef &task) {
		388	+ if (!IsAnnotatedEventTask(task)) {
		389	+ return false;
		390	+ }
		391	+ const auto iter = remap.find(task.event_id());
		392	+ return iter != remap.end() && iter->second == std::numeric_limits<uint32_t>::max();
		393	+ }),
		394	+ tasks.end());
		395	+ for (auto &task : tasks) {
		396	+ if (!IsAnnotatedEventTask(task)) {
		397	+ continue;
		398	+ }
		399	+ const auto iter = remap.find(task.event_id());
		400	+ if (iter != remap.end() && iter->second != std::numeric_limits<uint32_t>::max()) {
		401	+ task.set_event_id(iter->second);
		402	+ }
		403	+ }
		404	+ }
		405	+}
309	} // namespace	406	} // namespace
310	StreamAllocator::StreamAllocator(ComputeGraphPtr whole_graph, const Graph2SubGraphInfoList &subgraphs)	407	StreamAllocator::StreamAllocator(ComputeGraphPtr whole_graph, const Graph2SubGraphInfoList &subgraphs)
311	: whole_graph_(std::move(whole_graph)), subgraphs_(subgraphs) {	408	: whole_graph_(std::move(whole_graph)), subgraphs_(subgraphs) {
	@@ -341,6 +438,127 @@ Status StreamAllocator::AssignStandaloneWaitEvents() {
341	return SUCCESS;	438	return SUCCESS;
342	}	439	}
343		440
		441	+Status StreamAllocator::InsertAnnotatedArgsEventTasks(
		C C Chang-an-HW9月14日插入EVENT record/wait
		442	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks, int64_t &event_num) {
		443	+ if ((event_num < 0) \|\| (event_num > static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) \|\|
		444	+ (whole_graph_ == nullptr)) {
		445	+ GELOGE(PARAM_INVALID, "[AnnotatedArgsEvents] invalid input, event_num=%ld, graph=%p.", event_num,
		446	+ whole_graph_.get());
		447	+ return FAILED;
		448	+ }
		449	+ event_num_ = static_cast<uint32_t>(event_num);
		450	+ annotated_event_infos_.clear();
		451	+ for (const auto &node : whole_graph_->GetAllNodes()) {
		452	+ GE_CHECK_NOTNULL(node);
		453	+ const auto op_desc = node->GetOpDesc();
		454	+ GE_CHECK_NOTNULL(op_desc);
		455	+ AnnotatedArgsTaskPlan restored_plan;
		456	+ Buffer encoded_plan;
		457	+ if (!AttrUtils::GetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan)) {
		458	+ continue;
		459	+ }
		460	+ GE_ASSERT_SUCCESS(DeserializeAnnotatedArgsTaskPlan(encoded_plan, restored_plan),
		461	+ "[AnnotatedArgsEvents] failed to deserialize task plan, node=%s(%s).", op_desc->GetNamePtr(),
		462	+ op_desc->GetTypePtr());
		463	+ if (restored_plan.dependencies.empty()) {
		464	+ continue;
		465	+ }
		466	+ auto task_iter = node_id_2_node_tasks.find(op_desc->GetId());
		467	+ GE_ASSERT_TRUE(task_iter != node_id_2_node_tasks.end(),
		468	+ "[AnnotatedArgsEvents] task list missing for node=%s, node_id=%ld.", op_desc->GetNamePtr(),
		469	+ op_desc->GetId());
		470	+ GE_ASSERT_SUCCESS(ExpandAnnotatedArgsNodeEvents(op_desc, restored_plan, task_iter->second, event_num));
		471	+ }
		472	+ event_num_ = static_cast<uint32_t>(event_num);
		473	+ LogAnnotatedEventInsertStats(whole_graph_->GetName(), node_id_2_node_tasks, event_num);
		474	+ return SUCCESS;
		475	+}
		476	+
		477	+Status StreamAllocator::ExpandAnnotatedArgsNodeEvents(const OpDescPtr &op_desc, const AnnotatedArgsTaskPlan &plan,
		478	+ std::vector<domi::TaskDef> &tasks, int64_t &event_num) {
		479	+ std::vector<size_t> kernel_positions;
		480	+ CollectAnnotatedKernelPositions(tasks, kernel_positions);
		481	+ if (kernel_positions.size() != plan.task_templates.size()) {
		482	+ GELOGE(GRAPH_FAILED, "[AnnotatedArgsEvents] kernel task count mismatch, node=%s, actual=%zu, planned=%zu.",
		483	+ op_desc->GetNamePtr(), kernel_positions.size(), plan.task_templates.size());
		484	+ return FAILED;
		485	+ }
		486	+ std::map<size_t, std::vector<domi::TaskDef>> records;
		487	+ std::map<size_t, std::vector<domi::TaskDef>> waits;
		488	+ for (const auto &dep : plan.dependencies) {
		489	+ GE_ASSERT_TRUE(dep.predecessor_launch_index < kernel_positions.size() &&
		490	+ dep.successor_launch_index < kernel_positions.size() &&
		491	+ dep.predecessor_launch_index < dep.successor_launch_index,
		492	+ "[AnnotatedArgsEvents] invalid dependency, node=%s, predecessor=%u, successor=%u, "
		493	+ "kernel_count=%zu; indices must be in range and predecessor must be before successor.",
		494	+ op_desc->GetNamePtr(), dep.predecessor_launch_index, dep.successor_launch_index,
		495	+ kernel_positions.size());
		496	+ GE_ASSERT_TRUE(event_num < static_cast<int64_t>(std::numeric_limits<uint32_t>::max()),
		497	+ "[AnnotatedArgsEvents] event id overflow, node=%s, event_num=%ld.", op_desc->GetNamePtr(),
		498	+ event_num);
		499	+ const auto event_id = static_cast<uint32_t>(event_num++);
		500	+ const auto op_index = static_cast<uint32_t>(op_desc->GetId());
		501	+ records[kernel_positions[dep.predecessor_launch_index]].emplace_back(
		502	+ MakeAnnotatedEventTask(ModelTaskType::MODEL_TASK_EVENT_RECORD,
		503	+ tasks[kernel_positions[dep.predecessor_launch_index]].stream_id(), event_id, op_index));
		504	+ waits[kernel_positions[dep.successor_launch_index]].emplace_back(
		505	+ MakeAnnotatedEventTask(ModelTaskType::MODEL_TASK_EVENT_WAIT,
		506	+ tasks[kernel_positions[dep.successor_launch_index]].stream_id(), event_id, op_index));
		507	+ annotated_event_infos_[event_id] = {op_desc->GetId(), dep.predecessor_launch_index, dep.successor_launch_index};
		508	+ }
		509	+ MergeAnnotatedEventTasks(tasks, records, waits, plan.dependencies.size());
		510	+ return SUCCESS;
		511	+}
		512	+
		513	+Status StreamAllocator::ReconcileAnnotatedArgsEventTasks(
		C C Chang-an-HW9月15日在流分配或物理流拆分完成后，重新校正声明式 AnnotatedArgs 算子产生的 Event Record/Wait 任务收集 AnnotatedArgs 事件的 Record/Wait 配对 InsertAnnotatedArgsEventTasks 插入事件时，会把事件 ID 记录在 annotated_event_infos_ 中。该函数遍历所有 TaskDef，只处理这些已登记的 AnnotatedArgs 事件。删除已经没有跨流必要的事件如果某个事件的 Record 和 Wait 最终位于同一条流上，说明经过流重分配后两次 launch 已经天然按序执行，该 Event 不再需要，于是同时删除 Record 和 Wait 任务。重新压缩保留事件的 ID 对仍然跨流的事件重新分配连续的 Event ID，从原有事件区间的起始位置开始，更新对应的 Record/Wait 任务，并同步更新 event_num。这样可以避免删除事件后产生 ID 空洞。
		514	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks, int64_t &event_num) {
		515	+ if (event_num < 0) {
		516	+ return FAILED;
		517	+ }
		518	+ // 无声明式标注事件时直接返回：拆流阶段AddEventPair基于成员event_num_累加事件ID，
		519	+ // 此处若继续执行会用陈旧的入参event_num回写event_num_，导致事件总数回退。
		520	+ if (annotated_event_infos_.empty()) {
		521	+ return SUCCESS;
		522	+ }
		523	+ std::map<uint32_t, std::pair<domi::TaskDef , domi::TaskDef >> pairs;
		524	+ GE_ASSERT_SUCCESS(CollectAnnotatedEventPairs(node_id_2_node_tasks, pairs));
		525	+ GE_ASSERT_EQ(pairs.size(), annotated_event_infos_.size());
		526	+ GE_ASSERT_TRUE(event_num >= static_cast<int64_t>(pairs.size()));
		527	+ std::map<uint32_t, uint32_t> remap;
		528	+ uint32_t next_id = 0U;
		529	+ GE_ASSERT_SUCCESS(BuildAnnotatedEventRemap(pairs, event_num, remap, next_id));
		530	+ ApplyAnnotatedEventRemap(node_id_2_node_tasks, remap);
		531	+ annotated_event_infos_.clear();
		532	+ event_num_ = next_id;
		533	+ event_num = static_cast<int64_t>(next_id);
		534	+ return SUCCESS;
		535	+}
		536	+
		537	+Status StreamAllocator::CollectAnnotatedEventPairs(
		538	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> &node_id_2_node_tasks,
		539	+ std::map<uint32_t, std::pair<domi::TaskDef , domi::TaskDef >> &pairs) {
		540	+ for (auto &entry : node_id_2_node_tasks) {
		541	+ for (auto &task : entry.second) {
		542	+ if (!IsAnnotatedEventTask(task)) {
		543	+ continue;
		544	+ }
		545	+ const auto id = task.event_id();
		546	+ if (annotated_event_infos_.find(id) == annotated_event_infos_.end()) {
		547	+ continue;
		548	+ }
		549	+ auto &pair = pairs[id];
		550	+ if (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD)) {
		551	+ GE_ASSERT_TRUE(pair.first == nullptr);
		552	+ pair.first = &task;
		553	+ } else {
		554	+ GE_ASSERT_TRUE(pair.second == nullptr);
		555	+ pair.second = &task;
		556	+ }
		557	+ }
		558	+ }
		559	+ return SUCCESS;
		560	+}
		561	+
344	void StreamAllocator::BuildEventReuseMap(const EventType event_type, const std::vector<uint32_t> &events,	562	void StreamAllocator::BuildEventReuseMap(const EventType event_type, const std::vector<uint32_t> &events,
345	std::map<uint32_t, uint32_t> &event_seen, uint32_t &event_id) const {	563	std::map<uint32_t, uint32_t> &event_seen, uint32_t &event_id) const {
346	for (size_t i = 0U; i < events.size(); ++i) {	564	for (size_t i = 0U; i < events.size(); ++i) {
	@@ -714,7 +932,12 @@ Status StreamAllocator::SplitStreamAndRefreshTaskDef(
714	static_cast<uint32_t>(rt_ret));	932	static_cast<uint32_t>(rt_ret));
715	if (value == RT_CAPABILITY_SUPPORT) {	933	if (value == RT_CAPABILITY_SUPPORT) {
716	GELOGI("Move split stream from ge to rts");	934	GELOGI("Move split stream from ge to rts");
		935	+ GE_ASSERT_SUCCESS(ReconcileAnnotatedArgsEventTasks(node_id_2_node_tasks, event_num),
		936	+ "[Reconcile][AnnotatedArgsEvents] failed! graph:%s", whole_graph_->GetName().c_str());
717	GE_ASSERT_SUCCESS(PostProcessOfSplitStreams());	937	GE_ASSERT_SUCCESS(PostProcessOfSplitStreams());
		938	+ stream_num = stream_num_;
		939	+ event_num = static_cast<int64_t>(event_num_);
		940	+ notify_num = static_cast<int64_t>(notify_num_);
718	return SUCCESS;	941	return SUCCESS;
719	}	942	}
720		943
	@@ -727,6 +950,9 @@ Status StreamAllocator::SplitStreamAndRefreshTaskDef(
727	GE_ASSERT_SUCCESS(SplitStreams(node_id_2_node_tasks, split_streams), "[Split][Streams] failed! graph:%s",	950	GE_ASSERT_SUCCESS(SplitStreams(node_id_2_node_tasks, split_streams), "[Split][Streams] failed! graph:%s",
728	whole_graph_->GetName().c_str());	951	whole_graph_->GetName().c_str());
729		952
		953	+ GE_ASSERT_SUCCESS(ReconcileAnnotatedArgsEventTasks(node_id_2_node_tasks, event_num),
		954	+ "[Reconcile][AnnotatedArgsEvents] failed! graph:%s", whole_graph_->GetName().c_str());
		955	+
730	GE_ASSERT_SUCCESS(UpdateActiveStreams(split_streams), "[Update][ActiveStreams] failed! graph:%s",	956	GE_ASSERT_SUCCESS(UpdateActiveStreams(split_streams), "[Update][ActiveStreams] failed! graph:%s",
731	whole_graph_->GetName().c_str());	957	whole_graph_->GetName().c_str());
732		958

Mcompiler/graph/build/task_generator.cc+40-10

	@@ -675,7 +675,8 @@ Status TaskGenerator::SaveFusionNodes(std::map<int64_t, std::vector<NodePtr>> &f
675		675
676	Status TaskGenerator::GenerateTaskForFusionNode(Node *const node,	676	Status TaskGenerator::GenerateTaskForFusionNode(Node *const node,
677	const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,	677	const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,
678	- std::unordered_set<Node *> &fusion_nodes_seen) {	678	+ std::unordered_set<Node *> &fusion_nodes_seen,
		679	+ const bool keep_marked_tasks) {
679	int64_t group_key;	680	int64_t group_key;
680	const auto &fusion_op_desc = node->GetOpDesc();	681	const auto &fusion_op_desc = node->GetOpDesc();
681	// without attr or has been seen, skip directly	682	// without attr or has been seen, skip directly
	@@ -686,9 +687,17 @@ Status TaskGenerator::GenerateTaskForFusionNode(Node *const node,
686	GELOGI("Fusion: start fusion group index[%ld], nodes size[%zu].", group_key, fusion_nodes.at(group_key).size());	687	GELOGI("Fusion: start fusion group index[%ld], nodes size[%zu].", group_key, fusion_nodes.at(group_key).size());
687	// If op_desc have this attr, call nodes with same group key in a stream together	688	// If op_desc have this attr, call nodes with same group key in a stream together
688	for (const auto &fusion_node : fusion_nodes.at(group_key)) {	689	for (const auto &fusion_node : fusion_nodes.at(group_key)) {
689	- auto &task_defs = node_id_2_node_tasks_[fusion_node->GetOpDesc()->GetId()];	690	+ const auto node_id = fusion_node->GetOpDesc()->GetId();
		691	+ // 打标节点保留拆流前生成的task，只补齐fusion场景的task顺序记录
		692	+ if (keep_marked_tasks && NeedKeepGeneratedTasks(fusion_node->GetOpDesc())) {
		693	+ fusion_ordered_node_list_.emplace_back(node_id);
		694	+ fusion_task_node_name_list_.emplace_back(fusion_op_desc->GetName());
		695	+ fusion_nodes_seen.insert(fusion_node.get());
		696	+ continue;
		697	+ }
		698	+ auto &task_defs = node_id_2_node_tasks_[node_id];
690	GE_ASSERT_SUCCESS(GenTaskForNormalNode(fusion_node.get(), "Fusion inner", task_defs));	699	GE_ASSERT_SUCCESS(GenTaskForNormalNode(fusion_node.get(), "Fusion inner", task_defs));
691	- fusion_ordered_node_list_.emplace_back(fusion_node->GetOpDesc()->GetId());	700	+ fusion_ordered_node_list_.emplace_back(node_id);
692	fusion_task_node_name_list_.emplace_back(fusion_op_desc->GetName());	701	fusion_task_node_name_list_.emplace_back(fusion_op_desc->GetName());
693	// record nodes which have call generate task successfully	702	// record nodes which have call generate task successfully
694	fusion_nodes_seen.insert(fusion_node.get());	703	fusion_nodes_seen.insert(fusion_node.get());
	@@ -696,7 +705,8 @@ Status TaskGenerator::GenerateTaskForFusionNode(Node *const node,
696	return SUCCESS;	705	return SUCCESS;
697	}	706	}
698		707
699	-Status TaskGenerator::GenTaskForFusionNodes(const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes) {	708	+Status TaskGenerator::GenTaskForFusionNodes(const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,
		709	+ const bool keep_marked_tasks) {
700	fusion_nodes_seen_.clear();	710	fusion_nodes_seen_.clear();
701	fusion_ordered_node_list_.clear();	711	fusion_ordered_node_list_.clear();
702	fusion_task_node_name_list_.clear();	712	fusion_task_node_name_list_.clear();
	@@ -707,7 +717,7 @@ Status TaskGenerator::GenTaskForFusionNodes(const std::map<int64_t, std::vector<
707		717
708	const auto &name = op_desc->GetName();	718	const auto &name = op_desc->GetName();
709	const auto &type = op_desc->GetType();	719	const auto &type = op_desc->GetType();
710	- GE_ASSERT_SUCCESS(GenerateTaskForFusionNode(node, fusion_nodes, fusion_nodes_seen_),	720	+ GE_ASSERT_SUCCESS(GenerateTaskForFusionNode(node, fusion_nodes, fusion_nodes_seen_, keep_marked_tasks),
711	"[Call][GenerateTaskForFusionNode] node:%s(%s) failed", name.c_str(), type.c_str());	721	"[Call][GenerateTaskForFusionNode] node:%s(%s) failed", name.c_str(), type.c_str());
712	// continue directly	722	// continue directly
713	if (ge::AttrUtils::GetInt(op_desc, ATTR_NAME_FUSION_GROUP_KEY, group_key)) {	723	if (ge::AttrUtils::GetInt(op_desc, ATTR_NAME_FUSION_GROUP_KEY, group_key)) {
	@@ -715,6 +725,12 @@ Status TaskGenerator::GenTaskForFusionNodes(const std::map<int64_t, std::vector<
715	continue;	725	continue;
716	}	726	}
717	const auto node_id = op_desc->GetId();	727	const auto node_id = op_desc->GetId();
		728	+ if (keep_marked_tasks && NeedKeepGeneratedTasks(op_desc)) {
		729	+ GELOGI("Node[name:%s, type:%s] keeps generated tasks in fusion mode.", name.c_str(), type.c_str());
		730	+ fusion_ordered_node_list_.emplace_back(node_id);
		731	+ fusion_task_node_name_list_.emplace_back(name);
		732	+ continue;
		733	+ }
718	auto &task_defs = node_id_2_node_tasks_[node_id];	734	auto &task_defs = node_id_2_node_tasks_[node_id];
719	GE_ASSERT_SUCCESS(GenTaskForNormalNode(node, "Fusion outer", task_defs));	735	GE_ASSERT_SUCCESS(GenTaskForNormalNode(node, "Fusion outer", task_defs));
720	fusion_ordered_node_list_.emplace_back(node_id);	736	fusion_ordered_node_list_.emplace_back(node_id);
	@@ -723,7 +739,7 @@ Status TaskGenerator::GenTaskForFusionNodes(const std::map<int64_t, std::vector<
723	return SUCCESS;	739	return SUCCESS;
724	}	740	}
725		741
726	-Status TaskGenerator::GenerateTaskForNodes(const std::vector<Node *> nodes) {	742	+Status TaskGenerator::GenerateTaskForNodes(const std::vector<Node *> nodes, const bool keep_marked_tasks) {
727	if (NeedDoFusionTask()) {	743	if (NeedDoFusionTask()) {
728	// fusion node场景所有node都重新做gen task，因为fusion node不只是对自己GenTask，还会对同一个group key的node GenTask	744	// fusion node场景所有node都重新做gen task，因为fusion node不只是对自己GenTask，还会对同一个group key的node GenTask
729	// todo 可以改成第二次GenTask时就只对部分node GenTaskForNormalNode	745	// todo 可以改成第二次GenTask时就只对部分node GenTaskForNormalNode
	@@ -731,13 +747,19 @@ Status TaskGenerator::GenerateTaskForNodes(const std::vector<Node *> nodes) {
731	GE_TRACE_RUN(TaskGenerator, SaveFusionNodes, fusion_nodes, nodes_);	747	GE_TRACE_RUN(TaskGenerator, SaveFusionNodes, fusion_nodes, nodes_);
732	if (!fusion_nodes.empty()) {	748	if (!fusion_nodes.empty()) {
733	for (const auto &node : nodes_) {	749	for (const auto &node : nodes_) {
		750	+ if (keep_marked_tasks && NeedKeepGeneratedTasks(node->GetOpDesc())) {
		751	+ continue;
		752	+ }
734	// 个别节点会二次GenTask，所以需要将第一次的结果清空	753	// 个别节点会二次GenTask，所以需要将第一次的结果清空
735	node_id_2_node_tasks_[node->GetOpDesc()->GetId()].clear();	754	node_id_2_node_tasks_[node->GetOpDesc()->GetId()].clear();
736	}	755	}
737	- return GenTaskForFusionNodes(fusion_nodes);	756	+ return GenTaskForFusionNodes(fusion_nodes, keep_marked_tasks);
738	}	757	}
739	}	758	}
740	for (const auto &node : nodes) {	759	for (const auto &node : nodes) {
		760	+ if (keep_marked_tasks && NeedKeepGeneratedTasks(node->GetOpDesc())) {
		761	+ continue;
		762	+ }
741	// 个别节点会二次GenTask，所以需要将第一次的结果清空	763	// 个别节点会二次GenTask，所以需要将第一次的结果清空
742	node_id_2_node_tasks_[node->GetOpDesc()->GetId()].clear();	764	node_id_2_node_tasks_[node->GetOpDesc()->GetId()].clear();
743	}	765	}
	@@ -754,6 +776,11 @@ Status TaskGenerator::GenerateTaskForNodes(const std::vector<Node *> nodes) {
754	(void)aclrtGetDevice(&device_id);	776	(void)aclrtGetDevice(&device_id);
755	GELOGI("Get device id %d", device_id);	777	GELOGI("Get device id %d", device_id);
756	for (const auto node : nodes) {	778	for (const auto node : nodes) {
		779	+ if (keep_marked_tasks && NeedKeepGeneratedTasks(node->GetOpDesc())) {
		780	+ GELOGI("Node[name:%s, type:%s] keeps generated tasks, skip regenerating.", node->GetNamePtr(),
		781	+ node->GetTypePtr());
		782	+ continue;
		783	+ }
757	const auto key = GetKey(node);	784	const auto key = GetKey(node);
758	auto &task_defs = node_id_2_node_tasks_[node->GetOpDesc()->GetId()];	785	auto &task_defs = node_id_2_node_tasks_[node->GetOpDesc()->GetId()];
759	// key must be valid	786	// key must be valid
	@@ -1031,13 +1058,16 @@ Status TaskGenerator::ReGetTaskInfo(const ComputeGraphPtr &comp_graph) {
1031	for (auto &node : need_to_gen_task_nodes) {	1058	for (auto &node : need_to_gen_task_nodes) {
1032	second_gen_task_nodes.emplace_back(node);	1059	second_gen_task_nodes.emplace_back(node);
1033	}	1060	}
1034	- GE_ASSERT_SUCCESS(GenerateTaskForNodes(second_gen_task_nodes));	1061	+ GE_ASSERT_SUCCESS(GenerateTaskForNodes(second_gen_task_nodes, true));
1035	for (const auto &node : need_to_gen_task_nodes) {	1062	for (const auto &node : need_to_gen_task_nodes) {
1036	auto stream_id = node->GetOpDesc()->GetStreamId();	1063	auto stream_id = node->GetOpDesc()->GetStreamId();
1037	const auto &iter = node_id_2_node_tasks_.find(node->GetOpDesc()->GetId());	1064	const auto &iter = node_id_2_node_tasks_.find(node->GetOpDesc()->GetId());
1038	GE_ASSERT_TRUE(iter != node_id_2_node_tasks_.end(), "node: %s doesn't have taskdef", node->GetNamePtr());	1065	GE_ASSERT_TRUE(iter != node_id_2_node_tasks_.end(), "node: %s doesn't have taskdef", node->GetNamePtr());
1039	bool has_attached_stream = node->GetOpDesc()->HasValidAttachedStreamId();	1066	bool has_attached_stream = node->GetOpDesc()->HasValidAttachedStreamId();
1040	- if (!has_attached_stream) {	1067	+ // 打标节点的task在拆流阶段已完成流号重映射，且可能包含非主流上的事件任务，这里不再统一改写。
		1068	+ // 前提：拆流之后不再有改动内存大小/offset/workspace的阶段（最后一次地址刷新是pass1内的ProcessAppendWs），
		1069	+ // 若将来新增此类阶段，必须同步处理打标节点，否则其保留的task里是失效地址
		1070	+ if ((!has_attached_stream) && (!NeedKeepGeneratedTasks(node->GetOpDesc()))) {
1041	RefreshTaskDefStreamId(has_attached_stream, stream_id, stream_id, iter->second);	1071	RefreshTaskDefStreamId(has_attached_stream, stream_id, stream_id, iter->second);
1042	}	1072	}
1043	}	1073	}
	@@ -1076,7 +1106,7 @@ Status TaskGenerator::GenModelTaskDef(const ComputeGraphPtr &graph, uint64_t ses
1076	} else {	1106	} else {
1077	for (const auto &node : graph->GetNodes(graph->GetGraphUnknownFlag(), nullptr, ffts_filter)) {	1107	for (const auto &node : graph->GetNodes(graph->GetGraphUnknownFlag(), nullptr, ffts_filter)) {
1078	if (!NoNeedGenTask(node->GetOpDesc())) {	1108	if (!NoNeedGenTask(node->GetOpDesc())) {
1079	- const auto &iter = node_id_2_node_tasks_.find(node->GetOpDesc()->GetId());	1109	+ const auto iter = node_id_2_node_tasks_.find(node->GetOpDesc()->GetId());
1080	GE_ASSERT_TRUE(iter != node_id_2_node_tasks_.end(), "node %s does not gen task", node->GetNamePtr());	1110	GE_ASSERT_TRUE(iter != node_id_2_node_tasks_.end(), "node %s does not gen task", node->GetNamePtr());
1081	op_names_.insert(op_names_.end(), iter->second.size(), node->GetOpDesc()->GetName());	1111	op_names_.insert(op_names_.end(), iter->second.size(), node->GetOpDesc()->GetName());
1082	task_def_list.insert(task_def_list.end(), iter->second.begin(), iter->second.end());	1112	task_def_list.insert(task_def_list.end(), iter->second.begin(), iter->second.end());

Mcompiler/graph/build/task_generator.h+11-3

	@@ -51,7 +51,14 @@ class TaskGenerator {
51	*/	51	*/
52	Status FindProfilingNodeIndex(const ComputeGraphPtr &graph, ProfilingPoint &profiling_point);	52	Status FindProfilingNodeIndex(const ComputeGraphPtr &graph, ProfilingPoint &profiling_point);
53	Status GenModelTaskDef(const ComputeGraphPtr &graph, uint64_t session_id, Model &model);	53	Status GenModelTaskDef(const ComputeGraphPtr &graph, uint64_t session_id, Model &model);
54	- Status GenerateTaskForNodes(const std::vector<Node *> nodes);	54	+ /**
		55	+ * 对指定节点重新生成task
		56	+ * @param nodes 待生成task的节点
		57	+ * @param keep_marked_tasks 拆流后的二次生成需置true：打了kKeepGeneratedTasksAttr标记的节点保留已有task，
		58	+ * 不清空也不重新生成（其task已包含拆流阶段插入的事件任务与真实流号）
		59	+ * @return
		60	+ */
		61	+ Status GenerateTaskForNodes(const std::vector<Node *> nodes, const bool keep_marked_tasks = false);
55	Status ReGetTaskInfo(const ComputeGraphPtr &comp_graph);	62	Status ReGetTaskInfo(const ComputeGraphPtr &comp_graph);
56	std::unordered_map<int64_t, std::vector<domi::TaskDef>> &MutableNodeId2TaskDefs();	63	std::unordered_map<int64_t, std::vector<domi::TaskDef>> &MutableNodeId2TaskDefs();
57		64
	@@ -65,9 +72,10 @@ class TaskGenerator {
65	Status GenerateTask(const ComputeGraphPtr &graph, Model &model);	72	Status GenerateTask(const ComputeGraphPtr &graph, Model &model);
66	Status UpdateTaskDef();	73	Status UpdateTaskDef();
67	Status SaveFusionNodes(std::map<int64_t, std::vector<NodePtr>> &fusion_nodes, const std::vector<Node *> nodes) const;	74	Status SaveFusionNodes(std::map<int64_t, std::vector<NodePtr>> &fusion_nodes, const std::vector<Node *> nodes) const;
68	- Status GenTaskForFusionNodes(const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes);	75	+ Status GenTaskForFusionNodes(const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,
		76	+ const bool keep_marked_tasks);
69	Status GenerateTaskForFusionNode(Node *const node, const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,	77	Status GenerateTaskForFusionNode(Node *const node, const std::map<int64_t, std::vector<NodePtr>> &fusion_nodes,
70	- std::unordered_set<Node *> &fusion_nodes_seen);	78	+ std::unordered_set<Node *> &fusion_nodes_seen, const bool keep_marked_tasks);
71	Status PrepareForGenerateTask(const ComputeGraphPtr &graph);	79	Status PrepareForGenerateTask(const ComputeGraphPtr &graph);
72	Status GenerateTaskForFftsNode(Node *ffts_node, const std::string &tag,	80	Status GenerateTaskForFftsNode(Node *ffts_node, const std::string &tag,
73	std::vector<domi::TaskDef> &task_def_list_per_node,	81	std::vector<domi::TaskDef> &task_def_list_per_node,

Mcompiler/graph/build/task_generator_utils.cc+9-0

	@@ -30,6 +30,15 @@ bool NoNeedGenTask(const OpDescPtr &op_desc) {
30	return false;	30	return false;
31	}	31	}
32		32
		33	+bool NeedKeepGeneratedTasks(const OpDescPtr &op_desc) {
		34	+ if (op_desc == nullptr) {
		35	+ return false;
		36	+ }
		37	+ bool keep_generated_tasks = false;
		38	+ (void)ge::AttrUtils::GetBool(op_desc, kKeepGeneratedTasksAttr, keep_generated_tasks);
		39	+ return keep_generated_tasks;
		40	+}
		41	+
33	void RefreshTaskDefStreamId(bool has_attached_stream, int64_t logical_stream_id, int64_t real_stream_id,	42	void RefreshTaskDefStreamId(bool has_attached_stream, int64_t logical_stream_id, int64_t real_stream_id,
34	std::vector<domi::TaskDef> &task_defs) {	43	std::vector<domi::TaskDef> &task_defs) {
35	if (has_attached_stream) {	44	if (has_attached_stream) {

Mcompiler/graph/build/task_generator_utils.h+5-0

	@@ -14,7 +14,12 @@
14	#include "graph/op_desc.h"	14	#include "graph/op_desc.h"
15	#include "proto/task.pb.h"	15	#include "proto/task.pb.h"
16	namespace ge {	16	namespace ge {
		17	+// 标记节点的任务已生成完毕且不可再次生成：拆流阶段会往任务列表插入事件任务并把流号改写为真实流号，
		18	+// 二次生成只能从算子侧的缓存计划重物化，会丢掉这些结果。目前仅声明式自定义算子（AnnotatedArgsOp）
		19	+// 的带依赖/附着流计划会打标，见 custom_ops_kernel_builder.cc 的 CacheAnnotatedArgsTaskPlan。
		20	+constexpr char kKeepGeneratedTasksAttr[] = "_keep_generated_tasks";
17	bool NoNeedGenTask(const OpDescPtr &op_desc);	21	bool NoNeedGenTask(const OpDescPtr &op_desc);
		22	+bool NeedKeepGeneratedTasks(const OpDescPtr &op_desc);
18	void RefreshTaskDefStreamId(bool has_attached_stream, int64_t logical_stream_id, int64_t real_stream_id,	23	void RefreshTaskDefStreamId(bool has_attached_stream, int64_t logical_stream_id, int64_t real_stream_id,
19	std::vector<domi::TaskDef> &task_defs);	24	std::vector<domi::TaskDef> &task_defs);
20	} // namespace ge	25	} // namespace ge

Mdocs/zh/api/graph_engine_api/README.md+1-0

	@@ -281,6 +281,7 @@
281	- [MallocWorkSpace](cpp/gert/AnnotatedArgsContext/MallocWorkSpace.md)	281	- [MallocWorkSpace](cpp/gert/AnnotatedArgsContext/MallocWorkSpace.md)
282	- [GetStreamId](cpp/gert/AnnotatedArgsContext/GetStreamId.md)	282	- [GetStreamId](cpp/gert/AnnotatedArgsContext/GetStreamId.md)
283	- [AddLaunch](cpp/gert/AnnotatedArgsContext/AddLaunch.md)	283	- [AddLaunch](cpp/gert/AnnotatedArgsContext/AddLaunch.md)
		284	+ - [RequestAttachedStream](cpp/gert/AnnotatedArgsContext/RequestAttachedStream.md)
284	- [GetInputTensor](cpp/gert/AnnotatedArgsContext/GetInputTensor.md)	285	- [GetInputTensor](cpp/gert/AnnotatedArgsContext/GetInputTensor.md)
285	- [GetOutputTensor](cpp/gert/AnnotatedArgsContext/GetOutputTensor.md)	286	- [GetOutputTensor](cpp/gert/AnnotatedArgsContext/GetOutputTensor.md)
286	- [GetRequiredInputTensor](cpp/gert/AnnotatedArgsContext/GetRequiredInputTensor.md)	287	- [GetRequiredInputTensor](cpp/gert/AnnotatedArgsContext/GetRequiredInputTensor.md)

Mdocs/zh/api/graph_engine_api/cpp/ge/AnnotatedArgsOp/DeclareLaunchArgs.md+2-2

	@@ -33,8 +33,8 @@ virtual graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) = 0
33		33
34	## 约束说明	34	## 约束说明
35		35
36	-- 回调成功返回前，必须至少调用一次`AddLaunch`。	36	+- 回调成功返回前，必须至少在主流上下发一个任务。
37	-- 调用方必须检查每次`AddLaunch`的返回值；若返回非`GRAPH_SUCCESS`，应立即将该状态返回，使本接口返回失败。	37	+- 调用方必须检查每次`AddLaunch`的返回值：两参数原型返回非`GRAPH_SUCCESS`、带`predecessors`参数的重载返回UINT32\_MAX均表示失败，应立即返回失败状态，使本接口返回失败。
38	<!-- npu="x90,9030" id1 -->	38	<!-- npu="x90,9030" id1 -->
39	- 端侧场景：只允许调用一次`AddLaunch`。	39	- 端侧场景：只允许调用一次`AddLaunch`。
40	<!-- end id1 -->	40	<!-- end id1 -->

Mdocs/zh/api/graph_engine_api/cpp/gert/AnnotatedArgsContext/AddLaunch.md+47-1

	@@ -13,6 +13,8 @@
13		13
14	添加一个声明式kernel launch。接口提取[`AnnotatedKernelArgs`](../AnnotatedKernelArgs/overview.md)的参数，并记录kernel名称、二进制、block dim和逻辑stream ID。	14	添加一个声明式kernel launch。接口提取[`AnnotatedKernelArgs`](../AnnotatedKernelArgs/overview.md)的参数，并记录kernel名称、二进制、block dim和逻辑stream ID。
15		15
		16	+第二个原型在添加launch的同时返回其`AnnotatedLaunchToken`，并可声明依赖的前置launch token，用于表达多次launch之间的先后关系。依赖关系通过`AnnotatedLaunchToken`表达：每个经第二个原型成功提交的launch获得一个token，后续launch可引用这些token作为前驱。
		17	+
16	<!-- npu="x90,9030" id1 -->	18	<!-- npu="x90,9030" id1 -->
17	端侧场景：要求恰好调用一次。	19	端侧场景：要求恰好调用一次。
18	<!-- end id1 -->	20	<!-- end id1 -->
	@@ -21,21 +23,35 @@
21		23
22	```c++	24	```c++
23	ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args)	25	ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args)
		26	+AnnotatedLaunchToken AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args,
		27	+ const std::vector<AnnotatedLaunchToken> &predecessors)
24	```	28	```
25		29
		30	+- 第一个原型：添加一个kernel launch，不产生token，也不能作为其他launch的依赖源。
		31	+- 第二个原型：添加一个kernel launch并返回其token；`predecessors`声明本launch依赖的前驱token列表，无依赖时传空vector（`{}`）。
		32	+
26	## 参数说明	33	## 参数说明
27		34
28	\| 参数名 \| 输入/输出 \| 说明 \|	35	\| 参数名 \| 输入/输出 \| 说明 \|
29	\| --- \| --- \| --- \|	36	\| --- \| --- \| --- \|
30	-\| launch_info \| 输入 \| kernel launch信息。`kernel_name`和`kernel_bin`必须非空，`kernel_bin_size`和`block_dim`必须大于0；`stream_id`应使用[`GetStreamId`](GetStreamId.md)的返回值。详见[`AnnotatedKernelLaunchInfo`](../AnnotatedKernelLaunchInfo.md)。接口在调用期间GE会复制并保存`kernel_name`和`kernel_bin`数据；调用方可在接口返回后结束这些临时数据的生命周期。 \|	37	+\| launch_info \| 输入 \| kernel launch信息。`kernel_name`和`kernel_bin`必须非空，`kernel_bin_size`和`block_dim`必须大于0；`stream_id`应使用[`GetStreamId`](GetStreamId.md)或[`RequestAttachedStream`](RequestAttachedStream.md)返回的合法逻辑辅流ID。详见[`AnnotatedKernelLaunchInfo`](../AnnotatedKernelLaunchInfo.md)。接口在调用期间GE会复制并保存`kernel_name`和`kernel_bin`数据；调用方可在接口返回后结束这些临时数据的生命周期。 \|
31	\| args \| 输入 \| kernel launch参数构建器，以右值引用移交。参数必须非空，且此前的参数追加操作均成功。调用期间GE会复制并保存args数据；调用方可在接口返回后结束args的生命周期。调用方以`std::move`移交后，不应依赖该对象的后续状态。 \|	38	\| args \| 输入 \| kernel launch参数构建器，以右值引用移交。参数必须非空，且此前的参数追加操作均成功。调用期间GE会复制并保存args数据；调用方可在接口返回后结束args的生命周期。调用方以`std::move`移交后，不应依赖该对象的后续状态。 \|
		39	+\| predecessors \| 输入 \| 前驱launch的token列表，可为空。仅在第二个原型中使用。元素必须是本Context内此前第二个原型返回的有效token，不可重复、不可前向依赖；列表内容在调用期间复制，调用方可在接口返回后结束其生命周期。 \|
32		40
33	## 返回值说明	41	## 返回值说明
34		42
		43	+- 第一个原型：
		44	+
35	\| 参数名 \| 类型 \| 说明 \|	45	\| 参数名 \| 类型 \| 说明 \|
36	\| --- \| --- \| --- \|	46	\| --- \| --- \| --- \|
37	\| - \| ge::graphStatus \| `GRAPH_SUCCESS(0)`：添加成功；其他值：参数或Context状态异常，添加失败。 \|	47	\| - \| ge::graphStatus \| `GRAPH_SUCCESS(0)`：添加成功；其他值：参数或Context状态异常，添加失败。 \|
38		48
		49	+- 第二个原型：
		50	+
		51	+\| 参数名 \| 类型 \| 说明 \|
		52	+\| --- \| --- \| --- \|
		53	+\| - \| AnnotatedLaunchToken \| 成功时返回本次launch的新token（`uint32_t`，从0递增），供后续launch声明依赖；失败时返回UINT32\_MAX，且不发布launch。 \|
		54	+
39	## 约束说明	55	## 约束说明
40		56
41	- 调用方必须检查每次调用的返回值，并在失败时立即返回该错误状态。	57	- 调用方必须检查每次调用的返回值，并在失败时立即返回该错误状态。
	@@ -44,6 +60,8 @@ ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, Annotate
44		60
45	以下`kKernelBin`和`kKernelBinSize`仅表示用户实际编译得到的kernel二进制及其大小，并非可执行的示例二进制。	61	以下`kKernelBin`和`kKernelBinSize`仅表示用户实际编译得到的kernel二进制及其大小，并非可执行的示例二进制。
46		62
		63	+- 第一个原型：
		64	+
47	```c++	65	```c++
48	extern const uint8_t kKernelBin[];	66	extern const uint8_t kKernelBin[];
49	extern const size_t kKernelBinSize;	67	extern const size_t kKernelBinSize;
	@@ -67,3 +85,31 @@ if (ret != ge::GRAPH_SUCCESS) {
67	}	85	}
68	return ge::GRAPH_SUCCESS;	86	return ge::GRAPH_SUCCESS;
69	```	87	```
		88	+
		89	+- 第二个原型（`args0`和`args1`为已构造完成的[`AnnotatedKernelArgs`](../AnnotatedKernelArgs/overview.md)）：
		90	+
		91	+```c++
		92	+extern const uint8_t kKernelBin[];
		93	+extern const size_t kKernelBinSize;
		94	+
		95	+// 在主流上提交首个launch，无前驱，获取token。
		96	+const auto token0 = ctx.AddLaunch(
		97	+ gert::AnnotatedKernelLaunchInfo{"main_kernel", kKernelBin, kKernelBinSize, 32U, ctx.GetStreamId()},
		98	+ std::move(args0), {});
		99	+if (token0 == UINT32_MAX) {
		100	+ return ge::GRAPH_FAILED;
		101	+}
		102	+
		103	+// 申请逻辑辅流，并声明辅流launch依赖主流launch。
		104	+const auto aux_stream_id = ctx.RequestAttachedStream(ge::AscendString("aux"));
		105	+if (aux_stream_id == UINT32_MAX) {
		106	+ return ge::GRAPH_FAILED;
		107	+}
		108	+const auto token1 = ctx.AddLaunch(
		109	+ gert::AnnotatedKernelLaunchInfo{"aux_kernel", kKernelBin, kKernelBinSize, 32U, aux_stream_id},
		110	+ std::move(args1), {token0});
		111	+if (token1 == UINT32_MAX) {
		112	+ return ge::GRAPH_FAILED;
		113	+}
		114	+return ge::GRAPH_SUCCESS;
		115	+```

Adocs/zh/api/graph_engine_api/cpp/gert/AnnotatedArgsContext/RequestAttachedStream.md+47-0

@@ -0,0 +1,47 @@
	1	+# RequestAttachedStream
	2	+
	3	+## 产品支持情况
	4	+
	5	+全量芯片支持。
	6	+
	7	+## 头文件/库文件
	8	+
	9	+- 头文件：\#include <exe\_graph/runtime/annotated\_args\_context.h\>
	10	+- 库文件：liblowering.so
	11	+
	12	+## 功能说明
	13	+
	14	+按key申请逻辑辅流。key的作用域为算子所在的子图（对应一个编译子模型）：同一子图内相同key复用同一逻辑stream ID；不同子图之间key相互独立，即使key相同也各自分配stream ID，运行时对应各自独立的物理流。返回的逻辑stream ID可用于[`AnnotatedKernelLaunchInfo`](../AnnotatedKernelLaunchInfo.md)的`stream_id`字段，在辅流上提交kernel launch。
	15	+
	16	+## 函数原型
	17	+
	18	+```c++
	19	+uint32_t RequestAttachedStream(const ge::AscendString &key)
	20	+```
	21	+
	22	+## 参数说明
	23	+
	24	+\| 参数名 \| 输入/输出 \| 说明 \|
	25	+\| --- \| --- \| --- \|
	26	+\| key \| 输入 \| 辅流标识，作用域为算子所在的子图。不能为空，且不能使用`CANN-FMK-`保留前缀。 \|
	27	+
	28	+## 返回值说明
	29	+
	30	+正常时返回key对应的逻辑辅流ID；key为空或申请失败时返回UINT32_MAX。
	31	+
	32	+## 约束说明
	33	+
	34	+- 返回的编译期逻辑ID只能用于当前Context的`AnnotatedKernelLaunchInfo::stream_id`。
	35	+- Context销毁后该ID不可作为运行时句柄使用。
	36	+- key不跨子图复用，本接口不支持在多个子图之间共享同一条物理辅流。
	37	+
	38	+## 调用示例
	39	+
	40	+```c++
	41	+const auto aux_stream_id = ctx.RequestAttachedStream(ge::AscendString("aux"));
	42	+if (aux_stream_id == UINT32_MAX) {
	43	+ return ge::GRAPH_FAILED;
	44	+}
	45	+// 将aux_stream_id填入AnnotatedKernelLaunchInfo::stream_id，
	46	+// 再调用带predecessors参数的[`AddLaunch`](AddLaunch.md)重载在辅流上提交launch。
	47	+```

Mdocs/zh/api/graph_engine_api/cpp/gert/AnnotatedArgsContext/overview.md+2-1

	@@ -20,7 +20,8 @@ Context及其返回的Tensor指针只在当前`DeclareLaunchArgs`回调期间有
20	\| --- \| --- \|	20	\| --- \| --- \|
21	\| [`WorkspaceAddr MallocWorkSpace(size_t size)`](MallocWorkSpace.md) \| 申请逻辑workspace。 \|	21	\| [`WorkspaceAddr MallocWorkSpace(size_t size)`](MallocWorkSpace.md) \| 申请逻辑workspace。 \|
22	\| [`uint32_t GetStreamId() const`](GetStreamId.md) \| 获取节点的逻辑主stream ID。 \|	22	\| [`uint32_t GetStreamId() const`](GetStreamId.md) \| 获取节点的逻辑主stream ID。 \|
23	-\| [`ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args)`](AddLaunch.md) \| 添加一个kernel launch。 \|	23	+\| [`ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args)`](AddLaunch.md) \| 添加一个kernel launch；带`predecessors`参数的重载返回launch token并可声明前驱依赖。 \|
		24	+\| [`uint32_t RequestAttachedStream(const ge::AscendString &key)`](RequestAttachedStream.md) \| 按key申请当前子图内的声明式逻辑辅流。 \|
24	\| [`const Tensor *GetInputTensor(size_t index) const`](GetInputTensor.md) \| 按扁平实例索引获取输入Tensor。 \|	25	\| [`const Tensor *GetInputTensor(size_t index) const`](GetInputTensor.md) \| 按扁平实例索引获取输入Tensor。 \|
25	\| [`const Tensor *GetOutputTensor(size_t index) const`](GetOutputTensor.md) \| 按扁平实例索引获取输出Tensor。 \|	26	\| [`const Tensor *GetOutputTensor(size_t index) const`](GetOutputTensor.md) \| 按扁平实例索引获取输出Tensor。 \|
26	\| [`const Tensor *GetRequiredInputTensor(size_t ir_index) const`](GetRequiredInputTensor.md) \| 按IR原型索引获取必选输入Tensor。 \|	27	\| [`const Tensor *GetRequiredInputTensor(size_t ir_index) const`](GetRequiredInputTensor.md) \| 按IR原型索引获取必选输入Tensor。 \|

Mgraph_metadef/exe_graph/CMakeLists.txt+2-0

	@@ -147,6 +147,8 @@ target_link_libraries(stub_lowering
147	PRIVATE	147	PRIVATE
148	ge_intf_pub_Os	148	ge_intf_pub_Os
149	c_sec_headers	149	c_sec_headers
		150	+ runtime_headers
		151	+ slog
150	PUBLIC	152	PUBLIC
151	ge_metadef_headers	153	ge_metadef_headers
152	)	154	)

Mgraph_metadef/exe_graph/runtime/annotated_args_context.cc+64-0

 #include "exe_graph/runtime/gert_mem_allocator.h"
 #include "exe_graph/runtime/gert_mem_block.h"
 #include "graph/utils/args_format_desc_utils.h"
++#include "framework/common/debug/ge_log.h"
 namespace gert {
 namespace {
 constexpr uint32_t kInvalidStreamId = std::numeric_limits<uint32_t>::max();
++constexpr AnnotatedLaunchToken kInvalidToken = std::numeric_limits<AnnotatedLaunchToken>::max();
 }  // namespace
 WorkspaceAddr AnnotatedArgsContext::MallocWorkSpace(size_t size) {
   return static_cast<uint32_t>(stream_id);
 }
++uint32_t AnnotatedArgsContext::RequestAttachedStream(const ge::AscendString &key) {
++  const auto additional_output_start = GetAdditionalOutputStartIndex();
++  if (additional_output_start < 0) {
++    return kInvalidStreamId;
++  }
++  const int64_t handler_index = additional_output_start + static_cast<int64_t>(AdditionalOutputIndex::kArgsHandler);
++  auto *handler_chain = GetOutput(handler_index);
++  if (handler_chain == nullptr) {
++    return kInvalidStreamId;
++  }
++  auto *args_handler = handler_chain->GetValue<ArgsHandler *>();
++  auto *annotated_args_handler =
++      (args_handler == nullptr) ? nullptr : dynamic_cast<AnnotatedArgsHandler *>(args_handler);
++  return (annotated_args_handler == nullptr) ? kInvalidStreamId : annotated_args_handler->RequestAttachedStream(key);
++}
++
 ge::graphStatus AnnotatedArgsContext::AddLaunch(const AnnotatedKernelLaunchInfo &launch_info,
                                                 AnnotatedKernelArgs &&args) {
   GE_ASSERT_TRUE((launch_info.kernel_name != nullptr) && (launch_info.kernel_name[0] != '\0'),
   return ge::GRAPH_SUCCESS;
 }
++AnnotatedLaunchToken AnnotatedArgsContext::AddLaunch(const AnnotatedKernelLaunchInfo &launch_info,
++                                                     AnnotatedKernelArgs &&args,
++                                                     const std::vector<AnnotatedLaunchToken> &predecessors) {
++  if ((launch_info.kernel_name == nullptr) || (launch_info.kernel_name[0] == '\0') ||
++      (launch_info.kernel_bin == nullptr) || (launch_info.kernel_bin_size == 0U) || (launch_info.block_dim == 0U) ||
++      (launch_info.stream_id == kInvalidStreamId)) {
++    GELOGE(ge::PARAM_INVALID,
++           "[AnnotatedArgsContext] invalid launch info, kernel=%p, bin=%p, bin_size=%zu, block_dim=%u, stream_id=%u.",
++           launch_info.kernel_name, launch_info.kernel_bin, launch_info.kernel_bin_size, launch_info.block_dim,
++           launch_info.stream_id);
++    return kInvalidToken;
++  }
++  std::vector<uint8_t> args_data;
++  std::vector<ge::ArgDesc> arg_descs;
++  if (args.ExtractArgsData(args_data, arg_descs) != ge::GRAPH_SUCCESS) {
++    GELOGE(ge::GRAPH_FAILED, "[AnnotatedArgsContext] failed to extract annotated launch arguments.");
++    return kInvalidToken;
++  }
++  const auto additional_output_start = GetAdditionalOutputStartIndex();
++  if (additional_output_start < 0) {
++    GELOGE(ge::INTERNAL_ERROR, "[AnnotatedArgsContext] args handler output index is unavailable.");
++    return kInvalidToken;
++  }
++  const int64_t handler_index = additional_output_start + static_cast<int64_t>(AdditionalOutputIndex::kArgsHandler);
++  auto *handler_chain = GetOutput(handler_index);
++  if (handler_chain == nullptr) {
++    GELOGE(ge::INTERNAL_ERROR, "[AnnotatedArgsContext] args handler output chain is null, index=%ld.", handler_index);
++    return kInvalidToken;
++  }
++  auto *args_handler = handler_chain->GetValue<ArgsHandler *>();
++  auto *annotated_args_handler =
++      (args_handler == nullptr) ? nullptr : dynamic_cast<AnnotatedArgsHandler *>(args_handler);
++  if (annotated_args_handler == nullptr) {
++    GELOGE(ge::INTERNAL_ERROR, "[AnnotatedArgsContext] output handler is not AnnotatedArgsHandler.");
++    return kInvalidToken;
++  }
++  AnnotatedLaunchToken token = kInvalidToken;
++  if (annotated_args_handler->AddLaunchWithDependencies(
++          launch_info.kernel_name, launch_info.kernel_bin, launch_info.kernel_bin_size, launch_info.block_dim,
++          launch_info.stream_id, std::move(args_data), std::move(arg_descs), predecessors,
++          token) != ge::GRAPH_SUCCESS) {
++    return kInvalidToken;
++  }
++  return token;
++}
++
 }  // namespace gert

Mgraph_metadef/exe_graph/runtime/annotated_args_handler.cc+72-3

	@@ -11,20 +11,25 @@
11	#include "annotated_args_handler.h"	11	#include "annotated_args_handler.h"
12		12
13	#include <utility>	13	#include <utility>
		14	+#include <limits>
14		15
15	#include "graph/utils/args_format_desc_utils.h"	16	#include "graph/utils/args_format_desc_utils.h"
		17	+#include "framework/common/debug/ge_log.h"
16		18
17	namespace gert {	19	namespace gert {
18	AnnotatedArgsHandler::LaunchRecord::LaunchRecord(const char const kernel_name, const void const kernel_bin,	20	AnnotatedArgsHandler::LaunchRecord::LaunchRecord(const char const kernel_name, const void const kernel_bin,
19	const size_t kernel_bin_size, const uint32_t block_dim,	21	const size_t kernel_bin_size, const uint32_t block_dim,
20	const uint32_t stream_id, std::vector<uint8_t> &&args_data,	22	const uint32_t stream_id, std::vector<uint8_t> &&args_data,
21	- std::vector<ge::ArgDesc> &&arg_descs)	23	+ std::vector<ge::ArgDesc> &&arg_descs, const AnnotatedLaunchToken token,
		24	+ std::vector<AnnotatedLaunchToken> &&dependencies)
22	: kernel_name_(kernel_name),	25	: kernel_name_(kernel_name),
23	kernel_bin_(static_cast<const uint8_t >(kernel_bin), static_cast<const uint8_t >(kernel_bin) + kernel_bin_size),	26	kernel_bin_(static_cast<const uint8_t >(kernel_bin), static_cast<const uint8_t >(kernel_bin) + kernel_bin_size),
24	args_data_(std::move(args_data)),	27	args_data_(std::move(args_data)),
25	arg_descs_(std::move(arg_descs)),	28	arg_descs_(std::move(arg_descs)),
26	block_dim_(block_dim),	29	block_dim_(block_dim),
27	- stream_id_(stream_id) {}	30	+ stream_id_(stream_id),
		31	+ token_(token),
		32	+ dependencies_(std::move(dependencies)) {}
28		33
29	AnnotatedArgsHandler::LaunchRecord::LaunchRecord(LaunchRecord &&other) noexcept = default;	34	AnnotatedArgsHandler::LaunchRecord::LaunchRecord(LaunchRecord &&other) noexcept = default;
30		35
	@@ -77,6 +82,18 @@ uint32_t AnnotatedArgsHandler::LaunchRecord::GetStreamId() const {
77	return stream_id_;	82	return stream_id_;
78	}	83	}
79		84
		85	+AnnotatedLaunchToken AnnotatedArgsHandler::LaunchRecord::GetToken() const {
		86	+ return token_;
		87	+}
		88	+
		89	+size_t AnnotatedArgsHandler::LaunchRecord::GetDependencyCount() const {
		90	+ return dependencies_.size();
		91	+}
		92	+
		93	+const AnnotatedLaunchToken *AnnotatedArgsHandler::LaunchRecord::GetDependencies() const {
		94	+ return dependencies_.empty() ? nullptr : dependencies_.data();
		95	+}
		96	+
80	AnnotatedArgsHandler::~AnnotatedArgsHandler() = default;	97	AnnotatedArgsHandler::~AnnotatedArgsHandler() = default;
81		98
82	const KernelArgs AnnotatedArgsHandler::MallocReadOnlyDevArgs(void const host_args, const size_t args_size) {	99	const KernelArgs AnnotatedArgsHandler::MallocReadOnlyDevArgs(void const host_args, const size_t args_size) {
	@@ -102,11 +119,63 @@ const AnnotatedArgsHandler::LaunchRecord *AnnotatedArgsHandler::GetLaunch(const
102	return &launch_records_[index];	119	return &launch_records_[index];
103	}	120	}
104		121
		122	+void AnnotatedArgsHandler::SetAttachedStreamRequestFunc(AttachedStreamRequestFunc func) {
		123	+ attached_stream_request_func_ = std::move(func);
		124	+}
		125	+
		126	+uint32_t AnnotatedArgsHandler::RequestAttachedStream(const ge::AscendString &key) {
		127	+ if (!attached_stream_request_func_) {
		128	+ GELOGE(ge::INTERNAL_ERROR, "[AnnotatedArgsHandler] attached stream request callback is not configured.");
		129	+ return std::numeric_limits<uint32_t>::max();
		130	+ }
		131	+ const auto stream_id = attached_stream_request_func_(key);
		132	+ if (stream_id != std::numeric_limits<uint32_t>::max()) {
		133	+ attached_stream_ids_.emplace_back(stream_id);
		134	+ }
		135	+ return stream_id;
		136	+}
		137	+
		138	+const std::vector<uint32_t> &AnnotatedArgsHandler::GetAttachedStreamIds() const {
		139	+ return attached_stream_ids_;
		140	+}
		141	+
105	void AnnotatedArgsHandler::AddLaunch(const char const kernel_name, const void const kernel_bin,	142	void AnnotatedArgsHandler::AddLaunch(const char const kernel_name, const void const kernel_bin,
106	const size_t kernel_bin_size, const uint32_t block_dim, const uint32_t stream_id,	143	const size_t kernel_bin_size, const uint32_t block_dim, const uint32_t stream_id,
107	std::vector<uint8_t> &&args_data, std::vector<ge::ArgDesc> &&arg_descs) {	144	std::vector<uint8_t> &&args_data, std::vector<ge::ArgDesc> &&arg_descs) {
108	LaunchRecord launch_record(kernel_name, kernel_bin, kernel_bin_size, block_dim, stream_id, std::move(args_data),	145	LaunchRecord launch_record(kernel_name, kernel_bin, kernel_bin_size, block_dim, stream_id, std::move(args_data),
109	- std::move(arg_descs));	146	+ std::move(arg_descs), std::numeric_limits<AnnotatedLaunchToken>::max(), {});
110	(void)launch_records_.emplace_back(std::move(launch_record));	147	(void)launch_records_.emplace_back(std::move(launch_record));
111	}	148	}
		149	+
		150	+ge::graphStatus AnnotatedArgsHandler::AddLaunchWithDependencies(
		151	+ const char kernel_name, const void kernel_bin, const size_t kernel_bin_size, const uint32_t block_dim,
		152	+ const uint32_t stream_id, std::vector<uint8_t> &&args_data, std::vector<ge::ArgDesc> &&arg_descs,
		153	+ const std::vector<AnnotatedLaunchToken> &predecessors, AnnotatedLaunchToken &token) {
		154	+ const auto invalid_token = std::numeric_limits<AnnotatedLaunchToken>::max();
		155	+ if ((launch_records_.size() >= static_cast<size_t>(invalid_token)) \|\| (stream_id == invalid_token)) {
		156	+ GELOGE(ge::PARAM_INVALID, "[AnnotatedArgsHandler] invalid stream/token state, launch_count=%zu, stream_id=%u.",
		157	+ launch_records_.size(), stream_id);
		158	+ return ge::GRAPH_FAILED;
		159	+ }
		160	+ const auto new_token = static_cast<AnnotatedLaunchToken>(launch_records_.size());
		161	+ for (size_t i = 0U; i < predecessors.size(); ++i) {
		162	+ if (predecessors[i] >= new_token) {
		163	+ GELOGE(ge::PARAM_INVALID, "[AnnotatedArgsHandler] invalid dependency token=%u for new token=%u.", predecessors[i],
		164	+ new_token);
		165	+ return ge::GRAPH_FAILED;
		166	+ }
		167	+ for (size_t j = 0U; j < i; ++j) {
		168	+ if (predecessors[j] == predecessors[i]) {
		169	+ GELOGE(ge::PARAM_INVALID, "[AnnotatedArgsHandler] duplicate dependency token=%u for new token=%u.",
		170	+ predecessors[i], new_token);
		171	+ return ge::GRAPH_FAILED;
		172	+ }
		173	+ }
		174	+ }
		175	+ LaunchRecord launch_record(kernel_name, kernel_bin, kernel_bin_size, block_dim, stream_id, std::move(args_data),
		176	+ std::move(arg_descs), new_token, std::vector<AnnotatedLaunchToken>(predecessors));
		177	+ (void)launch_records_.emplace_back(std::move(launch_record));
		178	+ token = new_token;
		179	+ return ge::GRAPH_SUCCESS;
		180	+}
112	} // namespace gert	181	} // namespace gert

Mgraph_metadef/exe_graph/runtime/annotated_args_handler.h+23-2

	@@ -13,14 +13,17 @@
13		13
14	#include <cstddef>	14	#include <cstddef>
15	#include <cstdint>	15	#include <cstdint>
		16	+#include <functional>
16	#include <string>	17	#include <string>
17	#include <vector>	18	#include <vector>
18		19
		20	+#include "exe_graph/runtime/annotated_args_context.h"
19	#include "framework/runtime/args_handler.h"	21	#include "framework/runtime/args_handler.h"
20		22
21	namespace ge {	23	namespace ge {
22	struct ArgDesc;	24	struct ArgDesc;
23	-}	25	+class AscendString;
		26	+} // namespace ge
24		27
25	namespace gert {	28	namespace gert {
26	class AnnotatedArgsContext;	29	class AnnotatedArgsContext;
	@@ -45,11 +48,15 @@ class AnnotatedArgsHandler : public ArgsHandler {
45	size_t GetArgDescCount() const;	48	size_t GetArgDescCount() const;
46	uint32_t GetBlockDim() const;	49	uint32_t GetBlockDim() const;
47	uint32_t GetStreamId() const;	50	uint32_t GetStreamId() const;
		51	+ AnnotatedLaunchToken GetToken() const;
		52	+ size_t GetDependencyCount() const;
		53	+ const AnnotatedLaunchToken *GetDependencies() const;
48		54
49	private:	55	private:
50	LaunchRecord(const char const kernel_name, const void const kernel_bin, const size_t kernel_bin_size,	56	LaunchRecord(const char const kernel_name, const void const kernel_bin, const size_t kernel_bin_size,
51	const uint32_t block_dim, const uint32_t stream_id, std::vector<uint8_t> &&args_data,	57	const uint32_t block_dim, const uint32_t stream_id, std::vector<uint8_t> &&args_data,
52	- std::vector<ge::ArgDesc> &&arg_descs);	58	+ std::vector<ge::ArgDesc> &&arg_descs, AnnotatedLaunchToken token,
		59	+ std::vector<AnnotatedLaunchToken> &&dependencies);
53		60
54	std::string kernel_name_;	61	std::string kernel_name_;
55	std::vector<uint8_t> kernel_bin_;	62	std::vector<uint8_t> kernel_bin_;
	@@ -57,6 +64,8 @@ class AnnotatedArgsHandler : public ArgsHandler {
57	std::vector<ge::ArgDesc> arg_descs_;	64	std::vector<ge::ArgDesc> arg_descs_;
58	uint32_t block_dim_;	65	uint32_t block_dim_;
59	uint32_t stream_id_;	66	uint32_t stream_id_;
		67	+ AnnotatedLaunchToken token_;
		68	+ std::vector<AnnotatedLaunchToken> dependencies_;
60	friend class AnnotatedArgsHandler;	69	friend class AnnotatedArgsHandler;
61	};	70	};
62		71
	@@ -72,12 +81,24 @@ class AnnotatedArgsHandler : public ArgsHandler {
72	// Returned pointer is invalidated when another launch is appended or this handler is destroyed.	81	// Returned pointer is invalidated when another launch is appended or this handler is destroyed.
73	const LaunchRecord *GetLaunch(const size_t index) const;	82	const LaunchRecord *GetLaunch(const size_t index) const;
74		83
		84	+ using AttachedStreamRequestFunc = std::function<uint32_t(const ge::AscendString &)>;
		85	+ void SetAttachedStreamRequestFunc(AttachedStreamRequestFunc func);
		86	+ uint32_t RequestAttachedStream(const ge::AscendString &key);
		87	+ const std::vector<uint32_t> &GetAttachedStreamIds() const;
		88	+
75	private:	89	private:
76	void AddLaunch(const char const kernel_name, const void const kernel_bin, const size_t kernel_bin_size,	90	void AddLaunch(const char const kernel_name, const void const kernel_bin, const size_t kernel_bin_size,
77	const uint32_t block_dim, const uint32_t stream_id, std::vector<uint8_t> &&args_data,	91	const uint32_t block_dim, const uint32_t stream_id, std::vector<uint8_t> &&args_data,
78	std::vector<ge::ArgDesc> &&arg_descs);	92	std::vector<ge::ArgDesc> &&arg_descs);
		93	+ ge::graphStatus AddLaunchWithDependencies(const char kernel_name, const void kernel_bin, size_t kernel_bin_size,
		94	+ uint32_t block_dim, uint32_t stream_id, std::vector<uint8_t> &&args_data,
		95	+ std::vector<ge::ArgDesc> &&arg_descs,
		96	+ const std::vector<AnnotatedLaunchToken> &predecessors,
		97	+ AnnotatedLaunchToken &token);
79		98
80	std::vector<LaunchRecord> launch_records_;	99	std::vector<LaunchRecord> launch_records_;
		100	+ AttachedStreamRequestFunc attached_stream_request_func_;
		101	+ std::vector<uint32_t> attached_stream_ids_;
81	friend class AnnotatedArgsContext;	102	friend class AnnotatedArgsContext;
82	};	103	};
83		104

Mgraph_metadef/exe_graph/stub/stub_offline_launch_runtime.cc+14-0

	@@ -191,6 +191,11 @@ uint32_t AnnotatedArgsContext::GetStreamId() const {
191	return 0U;	191	return 0U;
192	}	192	}
193		193
		194	+uint32_t AnnotatedArgsContext::RequestAttachedStream(const ge::AscendString &key) {
		195	+ (void)key;
		196	+ return 0U;
		197	+}
		198	+
194	ge::graphStatus AnnotatedArgsContext::AddLaunch(const AnnotatedKernelLaunchInfo &launch_info,	199	ge::graphStatus AnnotatedArgsContext::AddLaunch(const AnnotatedKernelLaunchInfo &launch_info,
195	AnnotatedKernelArgs &&args) {	200	AnnotatedKernelArgs &&args) {
196	std::vector<uint8_t> args_data;	201	std::vector<uint8_t> args_data;
	@@ -227,6 +232,15 @@ ge::graphStatus ResourceUsageContext::ReportAttachedStream(const std::vector<ge:
227	return ge::GRAPH_FAILED;	232	return ge::GRAPH_FAILED;
228	}	233	}
229		234
		235	+AnnotatedLaunchToken AnnotatedArgsContext::AddLaunch(const AnnotatedKernelLaunchInfo &launch_info,
		236	+ AnnotatedKernelArgs &&args,
		237	+ const std::vector<AnnotatedLaunchToken> &predecessors) {
		238	+ (void)launch_info;
		239	+ (void)args;
		240	+ (void)predecessors;
		241	+ return 0U;
		242	+}
		243	+
230	Tensor *HostCpuOpExecutionContext::MallocOutputTensor(size_t index, const StorageShape &shape,	244	Tensor *HostCpuOpExecutionContext::MallocOutputTensor(size_t index, const StorageShape &shape,
231	const StorageFormat &format, ge::DataType dtype) {	245	const StorageFormat &format, ge::DataType dtype) {
232	(void)index;	246	(void)index;

Minc/graph_metadef/external/exe_graph/runtime/annotated_args_context.h+27-0

	@@ -20,12 +20,16 @@
20	#include "exe_graph/runtime/extended_kernel_context.h"	20	#include "exe_graph/runtime/extended_kernel_context.h"
21	#include "graph/error_codes.h"	21	#include "graph/error_codes.h"
22	#include "graph/tensor.h"	22	#include "graph/tensor.h"
		23	+#include "graph/ascend_string.h"
23		24
24	namespace ge {	25	namespace ge {
25	struct ArgDesc;	26	struct ArgDesc;
26	}	27	}
27		28
28	namespace gert {	29	namespace gert {
		30	+// kernel launch 任务的标识符，由带 predecessors 参数的 AddLaunch 重载返回。
		31	+// token 按提交顺序从 0 递增分配，用于后续 launch 声明前驱依赖关系。
		32	+using AnnotatedLaunchToken = uint32_t;
		H H hugo11127 天前加个注释，说明token的含义
29	/**	33	/**
30	* 逻辑输入地址描述符，用于 AnnotatedKernelArgs 追加输入地址参数。	34	* 逻辑输入地址描述符，用于 AnnotatedKernelArgs 追加输入地址参数。
31	* @since 9.2.0(2026-07)	35	* @since 9.2.0(2026-07)
	@@ -220,6 +224,29 @@ class AnnotatedArgsContext : public ExtendedKernelContext {
220	*/	224	*/
221	ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args);	225	ge::graphStatus AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args);
222		226
		227	+ /**
		228	+ * 添加一个 kernel launch 任务并返回其 token，可同时指定依赖的前驱 launch token。
		229	+ * 无依赖时 predecessors 传空 vector（{}）。
		230	+ * @param launch_info kernel launch 信息
		231	+ * @param args kernel launch 参数，通过 AnnotatedKernelArgs 构建
		232	+ * @param predecessors 依赖的前驱 kernel launch token 列表，可为空；
		233	+ * token 必须来自本上下文先前 AddLaunch 的返回值，不可重复、不可前向依赖
		234	+ * @return 新添加的 kernel launch token（从 0 递增），供后续 launch 声明依赖；
		235	+ * 失败时返回 UINT32_MAX
		236	+ * @since 9.3.0(2026-09)
		237	+ */
		238	+ AnnotatedLaunchToken AddLaunch(const AnnotatedKernelLaunchInfo &launch_info, AnnotatedKernelArgs &&args,
		239	+ const std::vector<AnnotatedLaunchToken> &predecessors);
		240	+
		241	+ /**
		242	+ * 由声明式算子在编译期调用申请一条辅流。
		243	+ * key 的作用域为算子所在的子图（对应一个编译子模型），同一子图内相同 key 复用同一辅流，不跨子图复用。
		H H hugo11127 天前机制上是否可以保证跨子图不会误复用的情况，比如条件算子的子图内算子，和子图外算子
		244	+ * @param key 辅流注册 key
		245	+ * @return 辅流逻辑 ID；异常时返回 UINT32_MAX。
		246	+ * @since 9.3.0(2026-09)
		247	+ */
		248	+ uint32_t RequestAttachedStream(const ge::AscendString &key);
		249	+
223	/**	250	/**
224	* 根据输入 index 获取输入 Tensor 指针。	251	* 根据输入 index 获取输入 Tensor 指针。
225	* @param index 输入索引	252	* @param index 输入索引

Mtests/ge/st/testcase/CMakeLists.txt+1-0

	@@ -108,6 +108,7 @@ add_executable(graph_engine_test
108	"${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_weighted_stream_merger_public_st_test.cc"	108	"${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_weighted_stream_merger_public_st_test.cc"
109	"${AIR_CODE_DIR}/tests/ge/st/testcase/graph/ir/named_io_node_builder_test.cc"	109	"${AIR_CODE_DIR}/tests/ge/st/testcase/graph/ir/named_io_node_builder_test.cc"
110	"${AIR_CODE_DIR}/tests/ge/st/testcase/test_multi_stream_tuning_step_recorder.cc"	110	"${AIR_CODE_DIR}/tests/ge/st/testcase/test_multi_stream_tuning_step_recorder.cc"
		111	+ "${AIR_CODE_DIR}/tests/ge/st/testcase/test_declarative_custom_op_attached_stream.cc"
111	)	112	)
112		113
113	add_dependencies(graph_engine_test	114	add_dependencies(graph_engine_test

Atests/ge/st/testcase/test_declarative_custom_op_attached_stream.cc+37-0

@@ -0,0 +1,37 @@
	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 <cstdint>
	12	+#include <gtest/gtest.h>
	13	+
	14	+#include "exe_graph/runtime/annotated_args_handler.h"
	15	+#include "ge/ge_api_types.h"
	16	+
	17	+namespace gert {
	18	+namespace {
	19	+
	20	+TEST(DeclarativeCustomOpAttachedStreamSt, RequestAttachedStreamKeepsPerKeyIds) {
	21	+ AnnotatedArgsHandler handler;
	22	+ handler.SetAttachedStreamRequestFunc(
	23	+ [](const ge::AscendString &key) { return key.GetLength() == 0U ? 0U : static_cast<uint32_t>(key.GetLength()); });
	24	+
	25	+ EXPECT_EQ(handler.RequestAttachedStream(ge::AscendString("aux")), 3U);
	26	+ EXPECT_EQ(handler.RequestAttachedStream(ge::AscendString("aux")), 3U);
	27	+ EXPECT_EQ(handler.GetAttachedStreamIds().size(), 2U);
	28	+}
	29	+
	30	+TEST(DeclarativeCustomOpAttachedStreamSt, MissingRequestCallbackReturnsInvalidStream) {
	31	+ AnnotatedArgsHandler handler;
	32	+ EXPECT_EQ(handler.RequestAttachedStream(ge::AscendString("aux")), UINT32_MAX);
	33	+ EXPECT_TRUE(handler.GetAttachedStreamIds().empty());
	34	+}
	35	+
	36	+} // namespace
	37	+} // namespace gert

Mtests/ge/ut/ge/common/custom_ops_kernel_info_store_unittest.cc+115-0

	@@ -39,6 +39,8 @@
39	#include "graph/utils/args_format_desc_utils.h"	39	#include "graph/utils/args_format_desc_utils.h"
40	#include "graph/args_format_desc.h"	40	#include "graph/args_format_desc.h"
41	#include "graph/utils/attr_utils.h"	41	#include "graph/utils/attr_utils.h"
		42	+#include "graph/build/stream/declarative_stream_registry.h"
		43	+#include "graph/build/task_generator_utils.h"
42	#include "graph/utils/tensor_utils.h"	44	#include "graph/utils/tensor_utils.h"
43	#include "runtime/custom_op/python_custom_op_adapter.h"	45	#include "runtime/custom_op/python_custom_op_adapter.h"
44	#include "securec.h"	46	#include "securec.h"
	@@ -603,6 +605,42 @@ class MockAnnotatedArgsWithMismatchStreamCustomOp : public AnnotatedArgsOp, publ
603	}	605	}
604	};	606	};
605		607
		608	+class MockAnnotatedArgsWithUnjoinedAttachedStreamCustomOp : public AnnotatedArgsOp, public MockPortableCustomOp {
		609	+ public:
		610	+ graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) override {
		611	+ static const uint8_t kBin[] = {0xA1U, 0xA2U};
		612	+ const auto *input = ctx.GetInputTensor(0U);
		613	+ GE_ASSERT_NOTNULL(input);
		614	+ gert::AnnotatedKernelArgs main_args(gert::InputAddr{0U, input->GetAddr()});
		615	+ GE_ASSERT_TRUE(ctx.AddLaunch(gert::AnnotatedKernelLaunchInfo{"unjoined_main_stream_kernel", kBin, sizeof(kBin), 1U,
		616	+ ctx.GetStreamId()},
		617	+ std::move(main_args), {}) != UINT32_MAX);
		618	+
		619	+ const auto attached_stream = ctx.RequestAttachedStream("unjoined");
		620	+ GE_ASSERT_TRUE(attached_stream != UINT32_MAX);
		621	+ gert::AnnotatedKernelLaunchInfo attached_info{"unjoined_attached_stream_kernel", kBin, sizeof(kBin), 1U,
		622	+ attached_stream};
		623	+ gert::AnnotatedKernelArgs attached_args(gert::InputAddr{0U, input->GetAddr()});
		624	+ const auto attached_token = ctx.AddLaunch(attached_info, std::move(attached_args), {});
		625	+ return (attached_token == UINT32_MAX) ? GRAPH_FAILED : GRAPH_SUCCESS;
		626	+ }
		627	+};
		628	+
		629	+class MockAnnotatedArgsWithOnlyAttachedStreamCustomOp : public AnnotatedArgsOp, public MockPortableCustomOp {
		630	+ public:
		631	+ graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) override {
		632	+ static const uint8_t kBin[] = {0xB1U, 0xB2U};
		633	+ const auto *input = ctx.GetInputTensor(0U);
		634	+ GE_ASSERT_NOTNULL(input);
		635	+ const auto attached_stream = ctx.RequestAttachedStream("only_attached");
		636	+ GE_ASSERT_TRUE(attached_stream != UINT32_MAX);
		637	+ gert::AnnotatedKernelLaunchInfo info{"only_attached_stream_kernel", kBin, sizeof(kBin), 1U, attached_stream};
		638	+ gert::AnnotatedKernelArgs args(gert::InputAddr{0U, input->GetAddr()});
		639	+ const auto token = ctx.AddLaunch(info, std::move(args), {});
		640	+ return (token == UINT32_MAX) ? GRAPH_FAILED : GRAPH_SUCCESS;
		641	+ }
		642	+};
		643	+
606	class MockAnnotatedArgsWithConstInputAndWorkspaceCustomOp : public AnnotatedArgsOp, public MockPortableCustomOp {	644	class MockAnnotatedArgsWithConstInputAndWorkspaceCustomOp : public AnnotatedArgsOp, public MockPortableCustomOp {
607	public:	645	public:
608	graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) override {	646	graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) override {
	@@ -1342,6 +1380,33 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskReusesImmutablePlanWithoutWork
1342	EXPECT_EQ(g_single_declare_no_workspace_count.load(), 1U);	1380	EXPECT_EQ(g_single_declare_no_workspace_count.load(), 1U);
1343	}	1381	}
1344		1382
		1383	+TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskRestoresSerializedAnnotatedArgsPlanWithoutCallback) {
		1384	+ const std::string kTestOpType = "TestSerializedPlanRestore_BuilderTest";
		1385	+ auto creator = []() -> std::unique_ptr<BaseCustomOp> {
		1386	+ return std::make_unique<MockSingleDeclareNoWorkspaceCustomOp>();
		1387	+ };
		1388	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(AscendString(kTestOpType.c_str()), creator), GRAPH_SUCCESS);
		1389	+
		1390	+ ComputeGraphPtr graph;
		1391	+ auto node = BuildStaticCustomNode(kTestOpType, graph);
		1392	+ ASSERT_NE(node, nullptr);
		1393	+ g_single_declare_no_workspace_count.store(0U);
		1394	+
		1395	+ std::vector<domi::TaskDef> first_tasks;
		1396	+ ASSERT_EQ(GenerateTaskForNode(node, first_tasks), SUCCESS);
		1397	+ ASSERT_EQ(first_tasks.size(), 1U);
		1398	+ EXPECT_EQ(g_single_declare_no_workspace_count.load(), 1U);
		1399	+
		1400	+ auto op_desc = node->GetOpDesc();
		1401	+ ge::Buffer encoded_plan;
		1402	+ ASSERT_TRUE(ge::AttrUtils::GetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan));
		1403	+ std::vector<domi::TaskDef> restored_tasks;
		1404	+ ASSERT_EQ(GenerateTaskForNode(node, restored_tasks), SUCCESS);
		1405	+ ASSERT_EQ(restored_tasks.size(), 1U);
		1406	+ EXPECT_EQ(g_single_declare_no_workspace_count.load(), 1U);
		1407	+ EXPECT_EQ(restored_tasks[0].kernel().context().args_format(), first_tasks[0].kernel().context().args_format());
		1408	+}
		1409	+
1345	TEST_F(UtestCustomOpsKernelInfoStore, CalcOpRunningParamStartsNewAnnotatedArgsPlanLifecycle) {	1410	TEST_F(UtestCustomOpsKernelInfoStore, CalcOpRunningParamStartsNewAnnotatedArgsPlanLifecycle) {
1346	const std::string kTestOpType = "TestSingleDeclareNewLifecycle_BuilderTest";	1411	const std::string kTestOpType = "TestSingleDeclareNewLifecycle_BuilderTest";
1347	auto creator = []() -> std::unique_ptr<BaseCustomOp> {	1412	auto creator = []() -> std::unique_ptr<BaseCustomOp> {
	@@ -1854,6 +1919,8 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskMultipleLaunchesWritesPrefixed
1854	ASSERT_NE(tbe_kernel, nullptr);	1919	ASSERT_NE(tbe_kernel, nullptr);
1855	EXPECT_EQ(tbe_kernel->GetName(), tasks[i].kernel().kernel_name());	1920	EXPECT_EQ(tbe_kernel->GetName(), tasks[i].kernel().kernel_name());
1856	}	1921	}
		1922	+ // 无依赖且无附着流的legacy计划物化时会刷新为节点当前主流，无需打标
		1923	+ EXPECT_FALSE(node->GetOpDesc()->HasAttr(kKeepGeneratedTasksAttr));
1857	}	1924	}
1858		1925
1859	TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskMultipleLaunchesProducesModelBuilderReadableKernelAttrs) {	1926	TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskMultipleLaunchesProducesModelBuilderReadableKernelAttrs) {
	@@ -1895,6 +1962,54 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskFailsWhenDeclareSetsMismatchSt
1895	ExpectGenerateTaskFailedWithStatus(kTestOpType, INTERNAL_ERROR);	1962	ExpectGenerateTaskFailedWithStatus(kTestOpType, INTERNAL_ERROR);
1896	}	1963	}
1897		1964
		1965	+TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskAllowsAttachedLaunchWithoutMainJoin) {
		1966	+ GetThreadLocalContext().SetGraphOption({{ge::SOC_VERSION, "Ascend910B"}});
		1967	+
		1968	+ const std::string kTestOpType = "TestUnjoinedAttachedStreamCustomOp_BuilderTest";
		1969	+ auto creator = []() -> std::unique_ptr<BaseCustomOp> {
		1970	+ return std::make_unique<MockAnnotatedArgsWithUnjoinedAttachedStreamCustomOp>();
		1971	+ };
		1972	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(AscendString(kTestOpType.c_str()), creator), GRAPH_SUCCESS);
		1973	+
		1974	+ ComputeGraphPtr graph;
		1975	+ auto node = BuildStaticCustomNode(kTestOpType, graph);
		1976	+ ASSERT_NE(node, nullptr);
		1977	+ auto registry = std::make_shared<DeclarativeStreamRegistry>(4U);
		1978	+ ASSERT_NE(registry->RequestAttachedStream(AscendString("unjoined")), UINT32_MAX);
		1979	+ ASSERT_TRUE(node->GetOwnerComputeGraph()->SetExtAttr(kDeclarativeStreamRegistryAttr, registry));
		1980	+
		1981	+ std::vector<domi::TaskDef> tasks;
		1982	+ EXPECT_EQ(GenerateTaskForNode(node, tasks), SUCCESS);
		1983	+ ASSERT_EQ(tasks.size(), 2U);
		1984	+ EXPECT_EQ(tasks[0].type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		1985	+ EXPECT_EQ(tasks[0].stream_id(), 3U);
		1986	+ EXPECT_EQ(tasks[1].type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		1987	+ EXPECT_EQ(tasks[1].stream_id(), 4U);
		1988	+ // 带附着流的计划需打标，避免拆流后二次生成任务时丢失事件任务与真实流号
		1989	+ bool keep_generated_tasks = false;
		1990	+ EXPECT_TRUE(AttrUtils::GetBool(node->GetOpDesc(), kKeepGeneratedTasksAttr, keep_generated_tasks));
		1991	+ EXPECT_TRUE(keep_generated_tasks);
		1992	+}
		1993	+
		1994	+TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskFailsWhenAnnotatedPlanHasNoMainLaunch) {
		1995	+ const std::string kTestOpType = "TestOnlyAttachedStreamCustomOp_BuilderTest";
		1996	+ auto creator = []() -> std::unique_ptr<BaseCustomOp> {
		1997	+ return std::make_unique<MockAnnotatedArgsWithOnlyAttachedStreamCustomOp>();
		1998	+ };
		1999	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(AscendString(kTestOpType.c_str()), creator), GRAPH_SUCCESS);
		2000	+
		2001	+ ComputeGraphPtr graph;
		2002	+ auto node = BuildStaticCustomNode(kTestOpType, graph);
		2003	+ ASSERT_NE(node, nullptr);
		2004	+ auto registry = std::make_shared<DeclarativeStreamRegistry>(4U);
		2005	+ ASSERT_NE(registry->RequestAttachedStream(AscendString("only_attached")), UINT32_MAX);
		2006	+ ASSERT_TRUE(node->GetOwnerComputeGraph()->SetExtAttr(kDeclarativeStreamRegistryAttr, registry));
		2007	+
		2008	+ std::vector<domi::TaskDef> tasks;
		2009	+ EXPECT_EQ(GenerateTaskForNode(node, tasks), INTERNAL_ERROR);
		2010	+ EXPECT_TRUE(tasks.empty());
		2011	+}
		2012	+
1898	TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskFinalDeclareUsesPlannedWorkspaceOffset) {	2013	TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskFinalDeclareUsesPlannedWorkspaceOffset) {
1899	const std::string kTestOpType = "TestAnnotatedArgsWithWorkspaceFinal_BuilderTest";	2014	const std::string kTestOpType = "TestAnnotatedArgsWithWorkspaceFinal_BuilderTest";
1900	auto creator = []() -> std::unique_ptr<BaseCustomOp> {	2015	auto creator = []() -> std::unique_ptr<BaseCustomOp> {

Mtests/ge/ut/ge/graph/build/stream_allocator_unittest.cc+201-0

	@@ -9,11 +9,16 @@
9	*/	9	*/
10		10
11	#include <string>	11	#include <string>
		12	+#include <limits>
		13	+#include <set>
12	#include <vector>	14	#include <vector>
13	#include <gtest/gtest.h>	15	#include <gtest/gtest.h>
14		16
15	#include "macro_utils/dt_public_scope.h"	17	#include "macro_utils/dt_public_scope.h"
16	#include "graph/build/stream/graph_stream_allocator.h"	18	#include "graph/build/stream/graph_stream_allocator.h"
		19	+#include "graph/build/stream/declarative_stream_registry.h"
		20	+#include "graph/build/stream/annotated_args_task_plan.h"
		21	+#include "graph/build/model_builder.h"
17	#include "graph/build/stream/stream_utils.h"	22	#include "graph/build/stream/stream_utils.h"
18	#include "graph/normal_graph/compute_graph_impl.h"	23	#include "graph/normal_graph/compute_graph_impl.h"
19	#include "macro_utils/dt_public_unscope.h"	24	#include "macro_utils/dt_public_unscope.h"
	@@ -29,11 +34,207 @@
29	#include "framework/ge_runtime_stub/include/stub/gert_runtime_stub.h"	34	#include "framework/ge_runtime_stub/include/stub/gert_runtime_stub.h"
30	#include "ge_running_env/op_reg.h"	35	#include "ge_running_env/op_reg.h"
31	#include "graph/build/stream/stream_info.h"	36	#include "graph/build/stream/stream_info.h"
		37	+#include "framework/omg/omg_inner_types.h"
32		38
33	extern std::string g_runtime_stub_mock;	39	extern std::string g_runtime_stub_mock;
34		40
35	namespace ge {	41	namespace ge {
36	namespace {	42	namespace {
		43	+TEST(DeclarativeStreamRegistryTest, DeduplicatesKeysAndRejectsEmptyKey) {
		44	+ DeclarativeStreamRegistry registry(3U);
		45	+ const auto first = registry.RequestAttachedStream(AscendString("compute"));
		46	+ EXPECT_EQ(first, 3U);
		47	+ EXPECT_EQ(registry.RequestAttachedStream(AscendString("compute")), 3U);
		48	+ EXPECT_EQ(registry.RequestAttachedStream(AscendString("copy")), 4U);
		49	+ EXPECT_EQ(registry.RequestAttachedStream(AscendString("")), std::numeric_limits<uint32_t>::max());
		50	+ EXPECT_EQ(registry.RequestAttachedStream(AscendString("CANN-FMK-internal")), 5U);
		51	+ EXPECT_EQ(registry.GetAllocatedStreamIds(), (std::vector<uint32_t>{3U, 4U, 5U}));
		52	+}
		53	+
		54	+TEST(DeclarativeStreamRegistryTest, RejectsOverflow) {
		55	+ DeclarativeStreamRegistry registry(std::numeric_limits<uint32_t>::max());
		56	+ EXPECT_EQ(registry.RequestAttachedStream(AscendString("compute")), std::numeric_limits<uint32_t>::max());
		57	+ EXPECT_TRUE(registry.GetAllocatedStreamIds().empty());
		58	+}
		59	+
		60	+TEST(AnnotatedArgsTaskPlanTest, RoundTripsLaunchDependenciesAndStreams) {
		61	+ AnnotatedArgsTaskPlan source;
		62	+ for (uint32_t i = 0U; i < 4U; ++i) {
		63	+ domi::TaskDef task;
		64	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		65	+ source.task_templates.emplace_back(std::move(task));
		66	+ }
		67	+ source.launch_stream_ids = {0U, 3U, 4U, 0U};
		68	+ source.dependencies = {{0U, 2U, 1U}, {1U, 2U, 2U}, {2U, 3U, 3U}};
		69	+ source.attached_stream_ids = {3U, 4U};
		70	+
		71	+ Buffer encoded;
		72	+ ASSERT_EQ(SerializeAnnotatedArgsTaskPlan(source, encoded), SUCCESS);
		73	+ AnnotatedArgsTaskPlan restored;
		74	+ ASSERT_EQ(DeserializeAnnotatedArgsTaskPlan(encoded, restored), SUCCESS);
		75	+ EXPECT_EQ(restored.launch_stream_ids, source.launch_stream_ids);
		76	+ EXPECT_EQ(restored.dependencies.size(), 3U);
		77	+ EXPECT_EQ(restored.attached_stream_ids, source.attached_stream_ids);
		78	+ EXPECT_EQ(restored.task_templates.size(), 4U);
		79	+}
		80	+
		81	+TEST(AnnotatedArgsTaskPlanTest, RejectsMissingLaunchStreamIds) {
		82	+ AnnotatedArgsTaskPlan plan;
		83	+ domi::TaskDef task;
		84	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		85	+ plan.task_templates.emplace_back(std::move(task));
		86	+
		87	+ Buffer encoded;
		88	+ EXPECT_NE(SerializeAnnotatedArgsTaskPlan(plan, encoded), SUCCESS);
		89	+}
		90	+
		91	+TEST(ModelBuilderTest, FinalizeDeclarativeAttachedStreamsRequiresRegistry) {
		92	+ auto graph = std::make_shared<ComputeGraph>("attached_stream_without_registry");
		93	+ auto op_desc = std::make_shared<OpDesc>("annotated", "AnnotatedArgs");
		94	+ auto node = graph->AddNode(op_desc);
		95	+ ASSERT_NE(node, nullptr);
		96	+
		97	+ AnnotatedArgsTaskPlan plan;
		98	+ domi::TaskDef task;
		99	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		100	+ plan.task_templates.emplace_back(std::move(task));
		101	+ plan.launch_stream_ids = {0U};
		102	+ plan.attached_stream_ids = {1U};
		103	+ Buffer encoded_plan;
		104	+ ASSERT_EQ(SerializeAnnotatedArgsTaskPlan(plan, encoded_plan), SUCCESS);
		105	+ ASSERT_TRUE(AttrUtils::SetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan));
		106	+
		107	+ ModelBuilder builder(0U, graph, Graph2SubGraphInfoList(), {}, false);
		108	+ EXPECT_NE(builder.FinalizeDeclarativeAttachedStreams(), SUCCESS);
		109	+}
		110	+
		111	+TEST(ModelBuilderTest, FinalizeDeclarativeAttachedStreamsRestoresAttachedStream) {
		112	+ auto graph = std::make_shared<ComputeGraph>("attached_stream_with_registry");
		113	+ auto op_desc = std::make_shared<OpDesc>("annotated", "AnnotatedArgs");
		114	+ auto node = graph->AddNode(op_desc);
		115	+ ASSERT_NE(node, nullptr);
		116	+
		117	+ auto registry = std::make_shared<DeclarativeStreamRegistry>(1U);
		118	+ ASSERT_EQ(registry->RequestAttachedStream(AscendString("attached")), 1U);
		119	+ ASSERT_TRUE(graph->SetExtAttr(kDeclarativeStreamRegistryAttr, registry));
		120	+
		121	+ AnnotatedArgsTaskPlan plan;
		122	+ domi::TaskDef task;
		123	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		124	+ plan.task_templates.emplace_back(std::move(task));
		125	+ plan.launch_stream_ids = {0U};
		126	+ plan.attached_stream_ids = {1U};
		127	+ Buffer encoded_plan;
		128	+ ASSERT_EQ(SerializeAnnotatedArgsTaskPlan(plan, encoded_plan), SUCCESS);
		129	+ ASSERT_TRUE(AttrUtils::SetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan));
		130	+
		131	+ ModelBuilder builder(0U, graph, Graph2SubGraphInfoList(), {}, false);
		132	+ ASSERT_EQ(builder.FinalizeDeclarativeAttachedStreams(), SUCCESS);
		133	+ EXPECT_EQ(op_desc->GetAttachedStreamIds(), (std::vector<int64_t>{1}));
		134	+}
		135	+
		136	+TEST(AnnotatedArgsTaskPlanTest, InsertsAndReconcilesCrossStreamEventTasks) {
		137	+ auto graph = std::make_shared<ComputeGraph>("annotated_event_graph");
		138	+ auto op_desc = std::make_shared<OpDesc>("annotated", "AnnotatedArgs");
		139	+ op_desc->SetStreamId(0);
		140	+ auto node = graph->AddNode(op_desc);
		141	+ ASSERT_NE(node, nullptr);
		142	+
		143	+ auto plan = std::make_shared<AnnotatedArgsTaskPlan>();
		144	+ for (size_t i = 0U; i < 3U; ++i) {
		145	+ domi::TaskDef task;
		146	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		147	+ plan->task_templates.emplace_back(std::move(task));
		148	+ }
		149	+ plan->launch_stream_ids = {0U, 1U, 2U};
		150	+ plan->dependencies = {{0U, 2U, 0U}};
		151	+ Buffer encoded_plan;
		152	+ ASSERT_EQ(SerializeAnnotatedArgsTaskPlan(*plan, encoded_plan), SUCCESS);
		153	+ ASSERT_TRUE(AttrUtils::SetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan));
		154	+
		155	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> tasks;
		156	+ for (const auto stream : plan->launch_stream_ids) {
		157	+ domi::TaskDef task;
		158	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		159	+ task.set_stream_id(stream);
		160	+ tasks[op_desc->GetId()].emplace_back(std::move(task));
		161	+ }
		162	+ StreamAllocator allocator(graph, Graph2SubGraphInfoList());
		163	+ int64_t event_num = 4;
		164	+ ASSERT_EQ(allocator.InsertAnnotatedArgsEventTasks(tasks, event_num), SUCCESS);
		165	+ EXPECT_EQ(event_num, 5);
		166	+ ASSERT_EQ(allocator.ReconcileAnnotatedArgsEventTasks(tasks, event_num), SUCCESS);
		167	+ size_t records = 0U;
		168	+ size_t waits = 0U;
		169	+ size_t kernels = 0U;
		170	+ for (const auto &task : tasks[op_desc->GetId()]) {
		171	+ kernels += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		172	+ records += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD));
		173	+ waits += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT));
		174	+ if (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD) \|\|
		175	+ task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT)) {
		176	+ EXPECT_EQ(task.event_id(), 4U);
		177	+ }
		178	+ }
		179	+ EXPECT_EQ(records, 1U);
		180	+ EXPECT_EQ(waits, 1U);
		181	+ EXPECT_EQ(kernels, 3U);
		182	+ EXPECT_EQ(tasks[op_desc->GetId()].size(), 5U);
		183	+ EXPECT_EQ(event_num, 5);
		184	+}
		185	+
		186	+TEST(AnnotatedArgsTaskPlanTest, PrunesSameStreamAndReconcilesMultipleDependencies) {
		187	+ auto graph = std::make_shared<ComputeGraph>("annotated_event_multi_graph");
		188	+ auto op_desc = std::make_shared<OpDesc>("annotated_multi", "AnnotatedArgs");
		189	+ auto node = graph->AddNode(op_desc);
		190	+ ASSERT_NE(node, nullptr);
		191	+
		192	+ auto plan = std::make_shared<AnnotatedArgsTaskPlan>();
		193	+ for (size_t i = 0U; i < 4U; ++i) {
		194	+ domi::TaskDef task;
		195	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		196	+ plan->task_templates.emplace_back(std::move(task));
		197	+ }
		198	+ plan->launch_stream_ids = {0U, 1U, 1U, 0U};
		199	+ plan->dependencies = {{0U, 1U, 0U}, {1U, 2U, 1U}, {2U, 3U, 2U}};
		200	+ Buffer encoded_plan;
		201	+ ASSERT_EQ(SerializeAnnotatedArgsTaskPlan(*plan, encoded_plan), SUCCESS);
		202	+ ASSERT_TRUE(AttrUtils::SetBytes(op_desc, "_custom_annotated_args_task_plan_bytes", encoded_plan));
		203	+
		204	+ std::unordered_map<int64_t, std::vector<domi::TaskDef>> tasks;
		205	+ for (const auto stream : plan->launch_stream_ids) {
		206	+ domi::TaskDef task;
		207	+ task.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		208	+ task.set_stream_id(stream);
		209	+ tasks[op_desc->GetId()].emplace_back(std::move(task));
		210	+ }
		211	+ StreamAllocator allocator(graph, Graph2SubGraphInfoList());
		212	+ int64_t event_num = 2;
		213	+ ASSERT_EQ(allocator.InsertAnnotatedArgsEventTasks(tasks, event_num), SUCCESS);
		214	+ ASSERT_EQ(allocator.ReconcileAnnotatedArgsEventTasks(tasks, event_num), SUCCESS);
		215	+
		216	+ size_t records = 0U;
		217	+ size_t waits = 0U;
		218	+ size_t kernels = 0U;
		219	+ std::set<uint32_t> event_ids;
		220	+ for (const auto &task : tasks[op_desc->GetId()]) {
		221	+ kernels += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		222	+ if (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD)) {
		223	+ ++records;
		224	+ event_ids.insert(task.event_id());
		225	+ } else if (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT)) {
		226	+ ++waits;
		227	+ event_ids.insert(task.event_id());
		228	+ }
		229	+ }
		230	+ EXPECT_EQ(records, 2U);
		231	+ EXPECT_EQ(waits, 2U);
		232	+ EXPECT_EQ(kernels, 4U);
		233	+ EXPECT_EQ(tasks[op_desc->GetId()].size(), 8U);
		234	+ EXPECT_EQ(event_ids, (std::set<uint32_t>{2U, 3U}));
		235	+ EXPECT_EQ(event_num, 4);
		236	+}
		237	+
37	void AddTaskDefWithStreamId(std::vector<domi::TaskDef> &task_defs, uint32_t stream_id) {	238	void AddTaskDefWithStreamId(std::vector<domi::TaskDef> &task_defs, uint32_t stream_id) {
38	auto task = domi::TaskDef();	239	auto task = domi::TaskDef();
39	task.set_stream_id(stream_id);	240	task.set_stream_id(stream_id);

Mtests/ge/ut/ge/graph/build/task_generator_unittest.cc+98-0

	@@ -26,6 +26,7 @@
26	#include "api/gelib/gelib.h"	26	#include "api/gelib/gelib.h"
27	#include "engines/manager/opskernel_manager/ops_kernel_builder_manager.h"	27	#include "engines/manager/opskernel_manager/ops_kernel_builder_manager.h"
28	#include "graph/build/task_generator.h"	28	#include "graph/build/task_generator.h"
		29	+#include "graph/build/task_generator_utils.h"
29	#include "graph/manager/mem_manager.h"	30	#include "graph/manager/mem_manager.h"
30	#include "graph/manager/graph_var_manager.h"	31	#include "graph/manager/graph_var_manager.h"
31	#include "ge/ut/ge/ffts_plus_proto_tools.h"	32	#include "ge/ut/ge/ffts_plus_proto_tools.h"
	@@ -35,6 +36,7 @@ using namespace std;
35	using namespace testing;	36	using namespace testing;
36	using namespace ge;	37	using namespace ge;
37	namespace {	38	namespace {
		39	+bool g_emit_attached_tasks = false;
38	const char *const kIsInputVar = "INPUT_IS_VAR";	40	const char *const kIsInputVar = "INPUT_IS_VAR";
39	const char *const kIsOutputVar = "OUTPUT_IS_VAR";	41	const char *const kIsOutputVar = "OUTPUT_IS_VAR";
40	const char *const kKernelInfoNameHccl = "ops_kernel_info_hccl";	42	const char *const kKernelInfoNameHccl = "ops_kernel_info_hccl";
	@@ -62,6 +64,27 @@ class UtestTaskGeneratorTest : public testing::Test {
62	};	64	};
63	Status GenerateTask(const Node &node, RunContext &context, std::vector<domi::TaskDef> &tasks) override {	65	Status GenerateTask(const Node &node, RunContext &context, std::vector<domi::TaskDef> &tasks) override {
64	domi::TaskDef task_def;	66	domi::TaskDef task_def;
		67	+ if (g_emit_attached_tasks) {
		68	+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		69	+ task_def.set_stream_id(node.GetOpDesc()->GetStreamId());
		70	+ task_def.set_sqe_num(1U);
		71	+ task_def.mutable_kernel()->mutable_context()->set_op_index(node.GetOpDesc()->GetId());
		72	+ tasks.push_back(task_def);
		73	+ task_def.Clear();
		74	+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD));
		75	+ task_def.set_stream_id(node.GetOpDesc()->GetStreamId());
		76	+ task_def.set_event_id(0U);
		77	+ task_def.mutable_event_ex()->set_op_index(node.GetOpDesc()->GetId());
		78	+ tasks.push_back(task_def);
		79	+ task_def.Clear();
		80	+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT));
		81	+ task_def.set_stream_id(node.GetOpDesc()->GetStreamId());
		82	+ task_def.set_event_id(0U);
		83	+ task_def.set_sqe_num(2U);
		84	+ task_def.mutable_event_ex()->set_op_index(node.GetOpDesc()->GetId());
		85	+ tasks.push_back(task_def);
		86	+ return SUCCESS;
		87	+ }
65	tasks.push_back(task_def);	88	tasks.push_back(task_def);
66	return SUCCESS;	89	return SUCCESS;
67	};	90	};
	@@ -598,6 +621,47 @@ TEST_F(UtestTaskGeneratorTest, ReGetTaskInfo_success) {
598	EXPECT_EQ(ret, SUCCESS);	621	EXPECT_EQ(ret, SUCCESS);
599	}	622	}
600		623
		624	+TEST_F(UtestTaskGeneratorTest, ReGetTaskInfo_keep_generated_tasks_of_marked_node) {
		625	+ OpsKernelManager::GetInstance().ops_kernel_store_[kKernelInfoNameHccl] = MakeShared<FakeOpsKernelInfoStore>();
		626	+ OpsKernelBuilderRegistry::GetInstance().kernel_builders_[kKernelInfoNameHccl] = MakeShared<FakeOpsKernelBuilder>();
		627	+
		628	+ auto graph = BuildHcclGraph();
		629	+ RunContext run_context;
		630	+ TaskGenerator task_generator(nullptr, 0, &run_context);
		631	+ Model model;
		632	+ // 首次生成kernel+事件任务，模拟声明式自定义算子带依赖/附着流的场景
		633	+ g_emit_attached_tasks = true;
		634	+ ASSERT_EQ(task_generator.GenerateTask(graph, model), SUCCESS);
		635	+ g_emit_attached_tasks = false;
		636	+
		637	+ auto hccl_node = graph->FindNode("hccl_phony_node");
		638	+ ASSERT_NE(hccl_node, nullptr);
		639	+ const auto node_id = hccl_node->GetOpDesc()->GetId();
		640	+ ASSERT_EQ(task_generator.node_id_2_node_tasks_[node_id].size(), 3U);
		641	+ ASSERT_TRUE(AttrUtils::SetBool(hccl_node->GetOpDesc(), kKeepGeneratedTasksAttr, true));
		642	+
		643	+ // 新增节点触发二次生成，首尾节点标记会把hccl节点一并纳入重生成集合
		644	+ auto add_op_desc = std::make_shared<OpDesc>("Add_new", ADD);
		645	+ auto add = graph->AddNode(add_op_desc);
		646	+ add->GetOpDesc()->SetOpKernelLibName(kKernelLibName);
		647	+ add->GetOpDesc()->SetOpEngineName(kKernelLibName);
		648	+ EXPECT_EQ(task_generator.ReGetTaskInfo(graph), SUCCESS);
		649	+
		650	+ // 打标节点保留首次生成的任务，事件任务不被清空
		651	+ const auto &tasks = task_generator.node_id_2_node_tasks_[node_id];
		652	+ ASSERT_EQ(tasks.size(), 3U);
		653	+ size_t record_num = 0U;
		654	+ size_t wait_num = 0U;
		655	+ for (const auto &task : tasks) {
		656	+ record_num += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD));
		657	+ wait_num += (task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT));
		658	+ }
		659	+ EXPECT_EQ(record_num, 1U);
		660	+ EXPECT_EQ(wait_num, 1U);
		661	+ OpsKernelBuilderRegistry::GetInstance().kernel_builders_.erase(kKernelInfoNameHccl);
		662	+ OpsKernelManager::GetInstance().ops_kernel_store_.erase(kKernelInfoNameHccl);
		663	+}
		664	+
601	TEST_F(UtestTaskGeneratorTest, AutoFindFpOpIndex) {	665	TEST_F(UtestTaskGeneratorTest, AutoFindFpOpIndex) {
602	auto graph = BuildGraphFpProfiling();	666	auto graph = BuildGraphFpProfiling();
603	SetGraphNodeKernel(graph);	667	SetGraphNodeKernel(graph);
	@@ -768,6 +832,40 @@ TEST_F(UtestTaskGeneratorTest, GenerateTask) {
768	OpsKernelBuilderRegistry::GetInstance().kernel_builders_.erase(kKernelInfoNameHccl);	832	OpsKernelBuilderRegistry::GetInstance().kernel_builders_.erase(kKernelInfoNameHccl);
769	}	833	}
770		834
		835	+TEST_F(UtestTaskGeneratorTest, GenerateTask_AnnotatedKernelRecordWait) {
		836	+ map<string, string> options;
		837	+ ASSERT_EQ(ge::GELib::Initialize(options), SUCCESS);
		838	+ auto info_store = MakeShared<FakeOpsKernelInfoStore>();
		839	+ OpsKernelManager::GetInstance().ops_kernel_store_[kKernelInfoNameHccl] = info_store;
		840	+ auto builder = MakeShared<FakeOpsKernelBuilder>();
		841	+ OpsKernelBuilderRegistry::GetInstance().kernel_builders_[kKernelInfoNameHccl] = builder;
		842	+
		843	+ g_emit_attached_tasks = true;
		844	+ auto graph = BuildHcclGraph();
		845	+ RunContext run_context;
		846	+ TaskGenerator task_generator(nullptr, 0, &run_context);
		847	+ Model model;
		848	+ const auto generate_ret = task_generator.GenerateTask(graph, model);
		849	+ g_emit_attached_tasks = false;
		850	+ ASSERT_EQ(generate_ret, SUCCESS);
		851	+
		852	+ size_t task_num = 0U;
		853	+ size_t wait_num = 0U;
		854	+ size_t record_num = 0U;
		855	+ for (const auto &entry : task_generator.node_id_2_node_tasks_) {
		856	+ task_num += entry.second.size();
		857	+ for (const auto &task : entry.second) {
		858	+ wait_num += task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT);
		859	+ record_num += task.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD);
		860	+ }
		861	+ }
		862	+ EXPECT_EQ(task_num, 3U);
		863	+ EXPECT_EQ(record_num, 1U);
		864	+ EXPECT_EQ(wait_num, 1U);
		865	+ OpsKernelBuilderRegistry::GetInstance().kernel_builders_.erase(kKernelInfoNameHccl);
		866	+ OpsKernelManager::GetInstance().ops_kernel_store_.erase(kKernelInfoNameHccl);
		867	+}
		868	+
771	TEST_F(UtestTaskGeneratorTest, SetFpBpByOptions) {	869	TEST_F(UtestTaskGeneratorTest, SetFpBpByOptions) {
772	std::string fp_name("data");	870	std::string fp_name("data");
773	std::string bp_name("addn1");	871	std::string bp_name("addn1");

Mtests/ge/ut/ge/graph/load/custom_task_info_unittest.cc+92-0

	@@ -15,6 +15,8 @@
15	#include "common/model/ge_model.h"	15	#include "common/model/ge_model.h"
16	#include "graph/debug/ge_attr_define.h"	16	#include "graph/debug/ge_attr_define.h"
17	#include "graph/load/model_manager/task_info/ge/custom_task_info.h"	17	#include "graph/load/model_manager/task_info/ge/custom_task_info.h"
		18	+#include "graph/load/model_manager/task_info/rts/event_record_task_info.h"
		19	+#include "graph/load/model_manager/task_info/rts/event_wait_task_info.h"
18	#include "graph/load/model_manager/davinci_model.h"	20	#include "graph/load/model_manager/davinci_model.h"
19	#include "graph/load/model_manager/reusable_stream_allocator.h"	21	#include "graph/load/model_manager/reusable_stream_allocator.h"
20	#include "graph/load/model_manager/model_args_manager.h"	22	#include "graph/load/model_manager/model_args_manager.h"
	@@ -61,6 +63,89 @@ class AclMockAnnotatedLaunch : public AclRuntimeStub {
61		63
62	class UtestCustomTaskInfo : public testing::Test {	64	class UtestCustomTaskInfo : public testing::Test {
63	protected:	65	protected:
		66	+ std::vector<domi::TaskDef> MakeAnnotatedKernelRecordWaitTaskDefs() const {
		67	+ std::vector<domi::TaskDef> task_defs;
		68	+ domi::TaskDef kernel;
		69	+ kernel.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
		70	+ kernel.set_stream_id(0U);
		71	+ kernel.set_sqe_num(1U);
		72	+ kernel.mutable_kernel()->mutable_context()->set_op_index(0U);
		73	+ task_defs.emplace_back(kernel);
		74	+
		75	+ domi::TaskDef record;
		76	+ record.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD));
		77	+ record.set_stream_id(0U);
		78	+ record.set_event_id(0U);
		79	+ record.mutable_event_ex()->set_op_index(0U);
		80	+ task_defs.emplace_back(record);
		81	+
		82	+ domi::TaskDef wait;
		83	+ wait.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT));
		84	+ wait.set_stream_id(0U);
		85	+ wait.set_event_id(0U);
		86	+ wait.set_sqe_num(2U);
		87	+ wait.mutable_event_ex()->set_op_index(0U);
		88	+ task_defs.emplace_back(wait);
		89	+ return task_defs;
		90	+ }
		91	+
		92	+ Status LoadTaskDefs(const std::vector<domi::TaskDef> &task_defs) const {
		93	+ DavinciModel model(0, nullptr);
		94	+ model.reusable_stream_allocator_ = ReusableStreamAllocator::Create();
		95	+ rtStream_t stream = nullptr;
		96	+ if (model.reusable_stream_allocator_->GetOrCreateRtStream(stream, 0, 0, 0) != SUCCESS) {
		97	+ return FAILED;
		98	+ }
		99	+ model.stream_list_ = {stream};
		100	+ rtEvent_t event = nullptr;
		101	+ if (aclrtCreateEvent(&event) != ACL_SUCCESS) {
		102	+ return FAILED;
		103	+ }
		104	+ model.event_list_ = {event};
		105	+ model.op_list_[0] = CreateOpDesc("annotated_custom", "CustomOp");
		106	+ for (const auto &task_def : task_defs) {
		107	+ if (task_def.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD)) {
		108	+ EventRecordTaskInfo info;
		109	+ if (info.Init(task_def, &model) != SUCCESS \|\| info.Distribute() != SUCCESS) {
		110	+ return FAILED;
		111	+ }
		112	+ } else if (task_def.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT)) {
		113	+ EventWaitTaskInfo info;
		114	+ if (info.Init(task_def, &model) != SUCCESS \|\| info.Distribute() != SUCCESS) {
		115	+ return FAILED;
		116	+ }
		117	+ }
		118	+ }
		119	+ return SUCCESS;
		120	+ }
		121	+
		122	+ const domi::TaskDef &FindWaitTask(const std::vector<domi::TaskDef> &task_defs) const {
		123	+ for (const auto &task_def : task_defs) {
		124	+ if (task_def.type() == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT)) {
		125	+ return task_def;
		126	+ }
		127	+ }
		128	+ return task_defs.front();
		129	+ }
		130	+
		131	+ uint32_t GetTaskNumOfTaskDef(const domi::TaskDef &task_def) const {
		132	+ return task_def.sqe_num();
		133	+ }
		134	+
		135	+ size_t CountUnexpectedEventTaskTypes(const std::vector<domi::TaskDef> &task_defs) const {
		136	+ size_t count = 0U;
		137	+ for (const auto &task_def : task_defs) {
		138	+ const auto type = task_def.type();
		139	+ if (type == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_RECORD) \|\|
		140	+ type == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_EVENT_WAIT) \|\|
		141	+ type == static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL)) {
		142	+ continue;
		143	+ }
		144	+ ++count;
		145	+ }
		146	+ return count;
		147	+ }
		148	+
64	void SetUp() override {	149	void SetUp() override {
65	RTS_STUB_SETUP();	150	RTS_STUB_SETUP();
66	auto acl_mock_memcpy = [](void dst, size_t dest_max, const void src, size_t count,	151	auto acl_mock_memcpy = [](void dst, size_t dest_max, const void src, size_t count,
	@@ -166,6 +251,13 @@ TEST_F(UtestCustomTaskInfo, Release_AclrtGetCurrentContextFailStillSuccess) {
166	AclRuntimeStub::GetInstance()->SetErrorResultApiName("");	251	AclRuntimeStub::GetInstance()->SetErrorResultApiName("");
167	}	252	}
168		253
		254	+TEST_F(UtestCustomTaskInfo, AnnotatedEventWaitUsesExistingWaitAndResetSemantics) {
		255	+ const auto task_defs = MakeAnnotatedKernelRecordWaitTaskDefs();
		256	+ ASSERT_EQ(LoadTaskDefs(task_defs), SUCCESS);
		257	+ EXPECT_EQ(CountUnexpectedEventTaskTypes(task_defs), 0U);
		258	+ EXPECT_EQ(GetTaskNumOfTaskDef(FindWaitTask(task_defs)), 2U);
		259	+}
		260	+
169	TEST_F(UtestCustomTaskInfo, InsertDumpOp_ReturnSuccessWhenOpNeedDumpIsFalse) {	261	TEST_F(UtestCustomTaskInfo, InsertDumpOp_ReturnSuccessWhenOpNeedDumpIsFalse) {
170	DavinciModel model(0, nullptr);	262	DavinciModel model(0, nullptr);
171	CustomTaskInfo task_info;	263	CustomTaskInfo task_info;

Mtests/graph_metadef/ut/exe_graph/annotated_args_context_unittest.cc+40-0

	@@ -137,6 +137,46 @@ TEST(AnnotatedArgsContextUT, RejectsInvalidLaunchParameters) {
137	EXPECT_EQ(args_handler.GetLaunchCount(), 0U);	137	EXPECT_EQ(args_handler.GetLaunchCount(), 0U);
138	}	138	}
139		139
		140	+TEST(AnnotatedArgsContextUT, AddLaunchWithTokenValidatesAndCopiesDependencies) {
		141	+ AllocatorFaker allocator;
		142	+ constexpr uint32_t kStreamId = 3U;
		143	+ allocator.SetStreamId(kStreamId);
		144	+ Tensor input_tensor;
		145	+ Tensor output_tensor;
		146	+ std::vector<GertMemBlock *> workspace_mems;
		147	+ AnnotatedArgsHandler args_handler;
		148	+ auto holder = KernelRunContextBuilder()
		149	+ .Inputs({&input_tensor, &allocator})
		150	+ .Outputs({&output_tensor, &workspace_mems, static_cast<ArgsHandler *>(&args_handler)})
		151	+ .Build(MakeOpDesc());
		152	+ auto context = reinterpret_cast<AnnotatedArgsContext >(holder.GetKernelContext());
		153	+ ASSERT_NE(context, nullptr);
		154	+ const uint8_t bin[] = {0x11U};
		155	+ auto make_args = []() { return AnnotatedKernelArgs(InputAddr{0U, reinterpret_cast<void *>(0x1000U)}); };
		156	+
		157	+ const auto token_a =
		158	+ context->AddLaunch(AnnotatedKernelLaunchInfo{"a", bin, sizeof(bin), 1U, kStreamId}, make_args(), {});
		159	+ ASSERT_NE(token_a, std::numeric_limits<AnnotatedLaunchToken>::max());
		160	+ std::vector<AnnotatedLaunchToken> predecessors{token_a};
		161	+ const auto token_b =
		162	+ context->AddLaunch(AnnotatedKernelLaunchInfo{"b", bin, sizeof(bin), 1U, kStreamId}, make_args(), predecessors);
		163	+ ASSERT_NE(token_b, std::numeric_limits<AnnotatedLaunchToken>::max());
		164	+ predecessors[0] = std::numeric_limits<AnnotatedLaunchToken>::max();
		165	+ ASSERT_EQ(args_handler.GetLaunchCount(), 2U);
		166	+ EXPECT_EQ(args_handler.GetLaunch(1U)->GetToken(), token_b);
		167	+ ASSERT_EQ(args_handler.GetLaunch(1U)->GetDependencyCount(), 1U);
		168	+ EXPECT_EQ(args_handler.GetLaunch(1U)->GetDependencies()[0], token_a);
		169	+
		170	+ const AnnotatedKernelLaunchInfo invalid_info{"c", bin, sizeof(bin), 1U, kStreamId};
		171	+ EXPECT_EQ(context->AddLaunch(invalid_info, make_args(), {2U}), std::numeric_limits<AnnotatedLaunchToken>::max());
		172	+ EXPECT_EQ(context->AddLaunch(invalid_info, make_args(), {3U}), std::numeric_limits<AnnotatedLaunchToken>::max());
		173	+ EXPECT_EQ(context->AddLaunch(invalid_info, make_args(), {token_a, token_a}),
		174	+ std::numeric_limits<AnnotatedLaunchToken>::max());
		175	+ const AnnotatedKernelLaunchInfo bad_name_info{nullptr, bin, sizeof(bin), 1U, kStreamId};
		176	+ EXPECT_EQ(context->AddLaunch(bad_name_info, make_args(), {}), std::numeric_limits<AnnotatedLaunchToken>::max());
		177	+ EXPECT_EQ(args_handler.GetLaunchCount(), 2U);
		178	+}
		179	+
140	TEST(AnnotatedArgsContextUT, RecordsMultipleLaunchesInOrder) {	180	TEST(AnnotatedArgsContextUT, RecordsMultipleLaunchesInOrder) {
141	AllocatorFaker allocator;	181	AllocatorFaker allocator;
142	constexpr uint32_t kStreamId = 3U;	182	constexpr uint32_t kStreamId = 3U;

	@@ -107,6 +107,7 @@ set(COMPILER_SRC_LIST
107	"graph/build/graph_builder.cc"	107	"graph/build/graph_builder.cc"
108	"graph/build/label_allocator.cc"	108	"graph/build/label_allocator.cc"
109	"graph/build/stream/assign_attached_stream_pass.cc"	109	"graph/build/stream/assign_attached_stream_pass.cc"
		110	+ "graph/build/stream/annotated_args_task_plan.cc"
		T T tang-haojie9月23日 pr缺描述
110	"graph/build/stream/assign_attached_notify_pass.cc"	111	"graph/build/stream/assign_attached_notify_pass.cc"
111	"graph/build/stream/assign_attached_event_pass.cc"	112	"graph/build/stream/assign_attached_event_pass.cc"
112	"graph/build/stream/attach_resource_assign_helper.cc"	113	"graph/build/stream/attach_resource_assign_helper.cc"
	@@ -154,6 +155,7 @@ set(COMPILER_SRC_LIST
154	"graph/build/model_builder.cc"	155	"graph/build/model_builder.cc"
155	"graph/build/input_h2d_overlap_plan.cc"	156	"graph/build/input_h2d_overlap_plan.cc"
156	"graph/build/input_h2d_overlap_planner.cc"	157	"graph/build/input_h2d_overlap_planner.cc"
		158	+ "graph/build/stream/declarative_stream_registry.cc"
157	"graph/build/run_context_util.cc"	159	"graph/build/run_context_util.cc"
158	"graph/build/stream_graph_optimizer.cc"	160	"graph/build/stream_graph_optimizer.cc"
159	"graph/build/task_generator.cc"	161	"graph/build/task_generator.cc"