feat: 语言无关Custom Task Info 适配 OM2-ge-AtomGit

Mbase/common/helper/om2/om2_package_contants.h+1-0

	@@ -28,6 +28,7 @@ namespace ge {
28	DO(OM2_RUNTIME_DIR_FORMAT, "data/model_%s/runtime/"); \	28	DO(OM2_RUNTIME_DIR_FORMAT, "data/model_%s/runtime/"); \
29	DO(OM2_DEBUG_DIR_FORMAT, "data/model_%s/debug/"); \	29	DO(OM2_DEBUG_DIR_FORMAT, "data/model_%s/debug/"); \
30	DO(OM2_OP_ATTR_PATH_FORMAT, "data/model_%s/debug/op_attr.json"); \	30	DO(OM2_OP_ATTR_PATH_FORMAT, "data/model_%s/debug/op_attr.json"); \
		31	+ DO(OM2_CUSTOM_KERNELS_DIR_FORMAT, "data/custom_ops/%s/"); \
31	DO(OM2_KERNELS_DIR_FORMAT, "data/kernels_%s/"); \	32	DO(OM2_KERNELS_DIR_FORMAT, "data/kernels_%s/"); \
32	DO(OM2_CONSTANTS_DIR, "data/constants/"); \	33	DO(OM2_CONSTANTS_DIR, "data/constants/"); \
33	DO(OM2_CONSTANTS_FILE_PREFIX, "constant_"); \	34	DO(OM2_CONSTANTS_FILE_PREFIX, "constant_"); \

Mbase/common/helper/om2/om2_zip_saver.cc+25-2

	@@ -208,6 +208,26 @@ Status SerializeKernelBinaries(const gert::Om2ModelData &model_data,
208	return SUCCESS;	208	return SUCCESS;
209	}	209	}
210		210
		211	+Status SerializeCustomKernelBinaries(const gert::Om2ModelData &model_data,
		212	+ const std::shared_ptr<ZipArchiveWriter> &zip_writer) {
		213	+ const auto kernel_bin_dir = FormatOm2Path(OM2_CUSTOM_KERNELS_DIR_FORMAT, "binaries_npu_arch");
		214	+ for (const auto &kb : model_data.custom_kernel_binaries) {
		215	+ const auto entry_path = kernel_bin_dir + kb.name;
		216	+ GE_ASSERT_TRUE(zip_writer->WriteBytes(entry_path, kb.data.get(), kb.data_size, false));
		217	+ }
		218	+ return SUCCESS;
		219	+}
		220	+
		221	+Status SerializeCustomKernelSharedLibs(const gert::Om2ModelData &model_data,
		222	+ const std::shared_ptr<ZipArchiveWriter> &zip_writer) {
		223	+ const auto kernel_bin_dir = FormatOm2Path(OM2_CUSTOM_KERNELS_DIR_FORMAT, "shared_libs");
		224	+ for (const auto &kb : model_data.custom_shared_libs) {
		225	+ const auto entry_path = kernel_bin_dir + kb.name;
		226	+ GE_ASSERT_TRUE(zip_writer->WriteBytes(entry_path, kb.data.get(), kb.data_size, false));
		227	+ }
		228	+ return SUCCESS;
		229	+}
		230	+
211	JsonFile SerializeAippDimsToJson(const std::vector<ge::InputOutputDims> &dims_list, const std::string &fmt_str,	231	JsonFile SerializeAippDimsToJson(const std::vector<ge::InputOutputDims> &dims_list, const std::string &fmt_str,
212	const std::string &dt_str) {	232	const std::string &dt_str) {
213	JsonFile::json arr = JsonFile::json::array();	233	JsonFile::json arr = JsonFile::json::array();
	@@ -396,10 +416,11 @@ Status Om2ZipSaver::Save(const gert::Om2ModelData &model_data, ModelBufferData &
396	const std::string &writer_path) {	416	const std::string &writer_path) {
397	GELOGI(	417	GELOGI(
398	"[OM2] Begin to serialize Om2ModelData to ZIP, model_name:%s, root_graph:%s, "	418	"[OM2] Begin to serialize Om2ModelData to ZIP, model_name:%s, root_graph:%s, "
399	- "inputs:%zu, outputs:%zu, kernels:%zu, weight_size:%zu",	419	+ "inputs:%zu, outputs:%zu, kernels:%zu, custom kernels: %zu, weight_size:%zu",
400	model_data.model_meta.model_name.c_str(), model_data.model_meta.root_graph_name.c_str(),	420	model_data.model_meta.model_name.c_str(), model_data.model_meta.root_graph_name.c_str(),
401	model_data.model_meta.input_desc.size(), model_data.model_meta.output_desc.size(),	421	model_data.model_meta.input_desc.size(), model_data.model_meta.output_desc.size(),
402	- model_data.kernel_binaries.size(), model_data.constants_data.internal_weight_size);	422	+ model_data.kernel_binaries.size(), model_data.custom_kernel_binaries.size(),
		423	+ model_data.constants_data.internal_weight_size);
403	const std::string path = writer_path.empty() ? "om2_model" : writer_path;	424	const std::string path = writer_path.empty() ? "om2_model" : writer_path;
404	auto zip_writer = std::make_shared<ZipArchiveWriter>(path);	425	auto zip_writer = std::make_shared<ZipArchiveWriter>(path);
405	GE_ASSERT_NOTNULL(zip_writer);	426	GE_ASSERT_NOTNULL(zip_writer);
	@@ -411,6 +432,8 @@ Status Om2ZipSaver::Save(const gert::Om2ModelData &model_data, ModelBufferData &
411	GE_ASSERT_SUCCESS(SerializeVarResource(model_data, zip_writer));	432	GE_ASSERT_SUCCESS(SerializeVarResource(model_data, zip_writer));
412	GE_ASSERT_SUCCESS(SerializeVarMetas(model_data, zip_writer, "0"));	433	GE_ASSERT_SUCCESS(SerializeVarMetas(model_data, zip_writer, "0"));
413	GE_ASSERT_SUCCESS(SerializeKernelBinaries(model_data, zip_writer));	434	GE_ASSERT_SUCCESS(SerializeKernelBinaries(model_data, zip_writer));
		435	+ GE_ASSERT_SUCCESS(SerializeCustomKernelBinaries(model_data, zip_writer));
		436	+ GE_ASSERT_SUCCESS(SerializeCustomKernelSharedLibs(model_data, zip_writer));
414	GE_ASSERT_SUCCESS(SerializeModelMeta(model_data, zip_writer));	437	GE_ASSERT_SUCCESS(SerializeModelMeta(model_data, zip_writer));
415	GE_ASSERT_SUCCESS(SerializeDebugInfo(model_data, zip_writer));	438	GE_ASSERT_SUCCESS(SerializeDebugInfo(model_data, zip_writer));
416	GE_ASSERT_SUCCESS(SerializeManifest(model_data, zip_writer));	439	GE_ASSERT_SUCCESS(SerializeManifest(model_data, zip_writer));

Mbase/common/helper/om2_package_helper.cc+169-26

	@@ -26,11 +26,12 @@
26	#include "graph/utils/type_utils.h"	26	#include "graph/utils/type_utils.h"
27	#include "graph/utils/tensor_utils.h"	27	#include "graph/utils/tensor_utils.h"
28	#include "graph_metadef/graph/utils/file_utils.h"	28	#include "graph_metadef/graph/utils/file_utils.h"
29	-#include "graph/model.h"	29	+#include "graph/custom_op_factory.h"
30	#include "common/helper/visual_json_converter.h"	30	#include "common/helper/visual_json_converter.h"
31	#include "graph/ge_context.h"	31	#include "graph/ge_context.h"
32	#include "graph/manager/graph_var_manager.h"	32	#include "graph/manager/graph_var_manager.h"
33	#include "common/helper/om2/rt_var_resource_builder.h"	33	#include "common/helper/om2/rt_var_resource_builder.h"
		34	+#include "common/op_so_store/op_so_store_utils.h"
34		35
35	namespace ge {	36	namespace ge {
36	namespace {	37	namespace {
	@@ -578,11 +579,11 @@ Status Om2PackageHelper::ExtractVisualJson(const void *model_data, size_t model_
578	return SUCCESS;	579	return SUCCESS;
579	}	580	}
580		581
581	-Status Om2PackageHelper::BuildProgramBody(const GeModelPtr &ge_model, gert::Om2ProgramBody &body,	582	+Status Om2PackageHelper::BuildProgramBody(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) {
582	- std::vector<Om2ConstMeta> &const_metas, std::vector<Om2VarMeta> &var_metas) {
583	GELOGI("[OM2] Begin to build program body");	583	GELOGI("[OM2] Begin to build program body");
		584	+ auto &body = model_data.program_body;
584	Om2Codegen codegen;	585	Om2Codegen codegen;
585	- GE_ASSERT_SUCCESS(codegen.Om2CodegenAndCompile(ge_model, body.source_artifacts, const_metas, var_metas));	586	+ GE_ASSERT_SUCCESS(codegen.Om2CodegenAndCompile(ge_model, model_data));
586	GE_ASSERT_TRUE(!body.source_artifacts.empty());	587	GE_ASSERT_TRUE(!body.source_artifacts.empty());
587		588
588	for (const auto &artifact : body.source_artifacts) {	589	for (const auto &artifact : body.source_artifacts) {
	@@ -592,16 +593,159 @@ Status Om2PackageHelper::BuildProgramBody(const GeModelPtr &ge_model, gert::Om2P
592	}	593	}
593	}	594	}
594		595
		596	+ auto &const_metas = model_data.constants_data.consts;
		597	+ auto &var_metas = model_data.var_metas;
595	GELOGI("[OM2] Successfully built program body, artifacts count=%zu, const_metas count=%zu, var_metas count=%zu",	598	GELOGI("[OM2] Successfully built program body, artifacts count=%zu, const_metas count=%zu, var_metas count=%zu",
596	body.source_artifacts.size(), const_metas.size(), var_metas.size());	599	body.source_artifacts.size(), const_metas.size(), var_metas.size());
597	return SUCCESS;	600	return SUCCESS;
598	}	601	}
599		602
600	-Status Om2PackageHelper::BuildKernelBinaries(const GeModelPtr &ge_model,	603	+Status Om2PackageHelper::BuildCustomSharedLibs(const GeRootModelPtr &ge_root_model, gert::Om2ModelData &model_data) {
601	- std::vector<gert::Om2KernelBinary> &kernel_binaries) {	604	+ GELOGI("[OM2] Begin to build custom kernel shared libraries");
		605	+ GE_ASSERT_NOTNULL(ge_root_model);
		606	+ if (!OpSoStoreUtils::IsSoBinType(ge_root_model->GetSoInOmFlag(), SoBinType::kCustomOp)) {
		607	+ return SUCCESS;
		608	+ }
		609	+ GE_ASSERT_SUCCESS(ReadCustomOpSoToBuffer(ge_root_model->GetCustomOpSoSet(), model_data.custom_shared_libs));
		610	+ GELOGI("[OM2] Save %zu custom op so to OpSoStore success.", ge_root_model->GetCustomOpSoSet().size());
		611	+ return SUCCESS;
		612	+}
		613	+
		614	+Status Om2PackageHelper::ReadCustomOpSoToBuffer(const std::unordered_set<std::string> &ops_so_set,
		615	+ std::vector<gert::Om2KernelBinary> &shared_lib_binaries) {
		616	+ for (const auto &op_so : ops_so_set) {
		617	+ uint32_t bin_len = 0U;
		618	+ auto op_so_bin = GetBinDataFromFile(op_so, bin_len);
		619	+ GE_ASSERT_NOTNULL(op_so_bin, "open so fail, path=%s", op_so.c_str());
		620	+ const auto &pos = op_so.find_last_of("/");
		621	+ GE_ASSERT_TRUE(pos != std::string::npos);
		622	+ const auto &so_name = op_so.substr(pos + 1UL);
		623	+ const auto &vendor_name = op_so.substr(0, pos);
		624	+ const size_t hash_id = std::hash<std::string>{}(std::string(op_so_bin.get(), op_so_bin.get() + bin_len));
		625	+ gert::Om2KernelBinary kb;
		626	+ kb.name = std::to_string(bin_len) + "_" + std::to_string(hash_id) + "_" + so_name;
		627	+ kb.data = ge::ReadonlyByteBuffer(reinterpret_cast<uint8_t *>(op_so_bin.release()), ge::ConditionalDeleter{true});
		628	+ kb.data_size = bin_len;
		629	+ shared_lib_binaries.emplace_back(std::move(kb));
		630	+
		631	+ GELOGD("[OM2] Serialized custom op so '%s', bin size:%zu", so_name.c_str(), bin_len);
		632	+ }
		633	+ return SUCCESS;
		634	+}
		635	+
		636	+Status Om2PackageHelper::CollectUsedCustomOpTypes(const GeRootModelPtr &ge_root_model,
		637	+ std::set<std::string> &used_custom_op_types) {
		638	+ if (ge_root_model->GetRootGraph() != nullptr) {
		639	+ const auto &root_graph = ge_root_model->GetRootGraph();
		640	+ for (const auto &node : root_graph->GetAllNodes()) {
		641	+ const std::string op_type = node->GetType();
		642	+ if (CustomOpFactory::IsExistOp(AscendString(op_type.c_str()))) {
		643	+ used_custom_op_types.insert(op_type);
		644	+ }
		645	+ }
		646	+ }
		647	+
		648	+ // subgraph_instance_name_to_model_ 中的 GeModel 可能持有独立的 ComputeGraph 对象，
		649	+ // 这些子图未必通过 AddSubgraph 挂入 root_graph 的子图树，因此无法被上方 GetAllNodes() 遍历到。
		650	+ // 典型场景：编译分区后各分区独立持有自己的 ComputeGraph，或反序列化时每个 GeModel 单独还原图对象。
		651	+ // 此处需额外遍历，以确保这类游离子图中的自定义算子也被收集到。
		652	+ const auto &subgraph_map = ge_root_model->GetSubgraphInstanceNameToModel();
		653	+ for (const auto &subgraph_pair : subgraph_map) {
		654	+ const auto &ge_model = subgraph_pair.second;
		655	+ if (ge_model == nullptr \|\| ge_model->GetGraph() == nullptr) {
		656	+ continue;
		657	+ }
		658	+ const auto &graph = ge_model->GetGraph();
		659	+ if (graph == ge_root_model->GetRootGraph()) {
		660	+ continue;
		661	+ }
		662	+ for (const auto &node : graph->GetAllNodes()) {
		663	+ const std::string op_type = node->GetType();
		664	+ if (CustomOpFactory::IsExistOp(AscendString(op_type.c_str()))) {
		665	+ used_custom_op_types.insert(op_type);
		666	+ }
		667	+ }
		668	+ }
		669	+ return SUCCESS;
		670	+}
		671	+
		672	+Status Om2PackageHelper::BuildCustomKernelBinaries(const GeRootModelPtr &ge_root_model,
		673	+ gert::Om2ModelData &model_data) {
		674	+ GELOGI("[OM2] Begin to build custom kernel binaries");
		675	+ auto &kernel_binaries = model_data.custom_kernel_binaries;
		676	+
		677	+ std::set<std::string> used_custom_op_types;
		678	+ GE_ASSERT_SUCCESS(CollectUsedCustomOpTypes(ge_root_model, used_custom_op_types));
		679	+
		680	+ if (used_custom_op_types.empty()) {
		681	+ GELOGI("[OM2] No custom ops used in graph, skip building custom kernels.");
		682	+ return SUCCESS;
		683	+ }
		684	+
		685	+ bool has_serializable_custom_op = false;
		686	+ bool has_non_serializable_custom_op = false;
		687	+ std::vector<std::pair<std::string, PortableOp *>> serializable_ops;
		688	+ serializable_ops.reserve(used_custom_op_types.size());
		689	+ for (const auto &op_type_str : used_custom_op_types) {
		690	+ auto op = CustomOpFactory::CreateOrGetCustomOp(AscendString(op_type_str.c_str()));
		691	+ if (op == nullptr) {
		692	+ GELOGE(FAILED, "[OM2] create custom op failed, op_type:%s", op_type_str.c_str());
		693	+ return FAILED;
		694	+ }
		695	+ auto serializable_op = dynamic_cast<PortableOp >(op);
		696	+ if (serializable_op == nullptr) {
		697	+ has_non_serializable_custom_op = true;
		698	+ } else {
		699	+ has_serializable_custom_op = true;
		700	+ serializable_ops.emplace_back(op_type_str, serializable_op);
		701	+ }
		702	+ if (has_serializable_custom_op && has_non_serializable_custom_op) {
		703	+ GELOGE(FAILED, "[OM2] graph contains both serializable and non-serializable custom ops.");
		704	+ return FAILED;
		705	+ }
		706	+ }
		707	+
		708	+ for (const auto &[op_type, serializable_op] : serializable_ops) {
		709	+ if (serializable_op == nullptr) {
		710	+ GELOGE(FAILED, "[OM2] serializable custom op is null, op_type:%s", op_type.c_str());
		711	+ return FAILED;
		712	+ }
		713	+
		714	+ std::vector<uint8_t> buffer;
		715	+ const auto ret = serializable_op->Serialize(buffer);
		716	+ if (ret != GRAPH_SUCCESS) {
		717	+ GELOGE(ret, "[OM2] serialize failed, op_type:%s", op_type.c_str());
		718	+ return ret;
		719	+ }
		720	+ if (buffer.empty()) {
		721	+ GELOGW("[OM2] serialized buffer is empty, skip, op_type:%s", op_type.c_str());
		722	+ continue;
		723	+ }
		724	+
		725	+ auto bin_data = new (std::nothrow) uint8_t[buffer.size()];
		726	+ if (bin_data == nullptr) {
		727	+ GELOGE(FAILED, "[Allocate][Mem]Allocate mem failed");
		728	+ return FAILED;
		729	+ }
		730	+ GE_ASSERT_EOK(memcpy_s(bin_data, buffer.size(), buffer.data(), buffer.size()));
		731	+ gert::Om2KernelBinary kb;
		732	+ kb.data = ge::ReadonlyByteBuffer(bin_data, ge::ConditionalDeleter{true});
		733	+ kb.data_size = buffer.size();
		734	+ const size_t hash_id = std::hash<std::string>{}(std::string(kb.data.get(), kb.data.get() + kb.data_size));
		735	+ const auto entry_path = op_type + "_" + std::to_string(hash_id) + "_CustomKernel.bin";
		736	+ kb.name = op_type + "_" + std::to_string(hash_id) + "_CustomKernel.bin";
		737	+ kernel_binaries.push_back(std::move(kb));
		738	+ GELOGD("[OM2] Serialized custom op '%s', bin size:%zu", op_type.c_str(), kb.data_size);
		739	+ }
		740	+ GELOGI("[OM2] Successfully built custom kernel binaries, count=%zu", kernel_binaries.size());
		741	+ return SUCCESS;
		742	+}
		743	+
		744	+Status Om2PackageHelper::BuildKernelBinaries(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) {
602	GELOGI("[OM2] Begin to build kernel binaries");	745	GELOGI("[OM2] Begin to build kernel binaries");
603	const auto &graph = ge_model->GetGraph();	746	const auto &graph = ge_model->GetGraph();
604	GE_ASSERT_NOTNULL(graph);	747	GE_ASSERT_NOTNULL(graph);
		748	+ std::vector<gert::Om2KernelBinary> &kernel_binaries = model_data.kernel_binaries;
605		749
606	// Collect TBE kernels	750	// Collect TBE kernels
607	const auto &tbe_kernel_store = ge_model->GetTBEKernelStore();	751	const auto &tbe_kernel_store = ge_model->GetTBEKernelStore();
	@@ -667,10 +811,11 @@ Status Om2PackageHelper::BuildKernelBinaries(const GeModelPtr &ge_model,
667	return SUCCESS;	811	return SUCCESS;
668	}	812	}
669		813
670	-Status Om2PackageHelper::BuildModelMeta(const GeModelPtr &ge_model, gert::Om2ModelMeta &model_meta) {	814	+Status Om2PackageHelper::BuildModelMeta(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) {
671	GELOGI("[OM2] Begin to build model meta");	815	GELOGI("[OM2] Begin to build model meta");
672	const auto &graph = ge_model->GetGraph();	816	const auto &graph = ge_model->GetGraph();
673	GE_ASSERT_NOTNULL(graph);	817	GE_ASSERT_NOTNULL(graph);
		818	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
674		819
675	ModelIoNodes io_nodes;	820	ModelIoNodes io_nodes;
676	GE_ASSERT_SUCCESS(CollectModelIoNodes(graph, io_nodes));	821	GE_ASSERT_SUCCESS(CollectModelIoNodes(graph, io_nodes));
	@@ -800,9 +945,11 @@ Status Om2PackageHelper::BuildModelMeta(const GeModelPtr &ge_model, gert::Om2Mod
800	return SUCCESS;	945	return SUCCESS;
801	}	946	}
802		947
803	-Status Om2PackageHelper::BuildConstantsData(const GeModelPtr &ge_model, const std::vector<Om2ConstMeta> &const_metas,	948	+Status Om2PackageHelper::BuildConstantsData(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) {
804	- gert::Om2ConstantsData &data) {
805	GELOGI("[OM2] Begin to build constants data");	949	GELOGI("[OM2] Begin to build constants data");
		950	+ gert::Om2ConstantsData &data = model_data.constants_data;
		951	+ auto &const_metas = data.consts;
		952	+
806	bool has_internal_const = false;	953	bool has_internal_const = false;
807	for (const auto &const_meta : const_metas) {	954	for (const auto &const_meta : const_metas) {
808	if (const_meta.type == "INTERNAL") {	955	if (const_meta.type == "INTERNAL") {
	@@ -812,11 +959,6 @@ Status Om2PackageHelper::BuildConstantsData(const GeModelPtr &ge_model, const st
812	}	959	}
813		960
814	data.internal_weight_size = has_internal_const ? ge_model->GetWeightSize() : 0U;	961	data.internal_weight_size = has_internal_const ? ge_model->GetWeightSize() : 0U;
815	-
816	- for (const auto &const_meta : const_metas) {
817	- data.consts.push_back(const_meta);
818	- }
819	-
820	if (has_internal_const) {	962	if (has_internal_const) {
821	const uint8_t *weight_ptr = ge_model->GetWeightData();	963	const uint8_t *weight_ptr = ge_model->GetWeightData();
822	GE_ASSERT_NOTNULL(weight_ptr, "[OM2] Weight data pointer is null");	964	GE_ASSERT_NOTNULL(weight_ptr, "[OM2] Weight data pointer is null");
	@@ -828,10 +970,11 @@ Status Om2PackageHelper::BuildConstantsData(const GeModelPtr &ge_model, const st
828	return SUCCESS;	970	return SUCCESS;
829	}	971	}
830		972
831	-Status Om2PackageHelper::BuildDebugInfo(const GeModelPtr &ge_model, gert::Om2DebugInfo &debug_info) {	973	+Status Om2PackageHelper::BuildDebugInfo(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) {
832	GELOGI("[OM2] Begin to build debug info");	974	GELOGI("[OM2] Begin to build debug info");
833	const auto &graph = ge_model->GetGraph();	975	const auto &graph = ge_model->GetGraph();
834	GE_ASSERT_NOTNULL(graph);	976	GE_ASSERT_NOTNULL(graph);
		977	+ gert::Om2DebugInfo &debug_info = model_data.debug_info;
835		978
836	// Build op_attr_map	979	// Build op_attr_map
837	for (const auto &node : graph->GetNodes(graph->GetGraphUnknownFlag())) {	980	for (const auto &node : graph->GetNodes(graph->GetGraphUnknownFlag())) {
	@@ -859,9 +1002,9 @@ Status Om2PackageHelper::BuildDebugInfo(const GeModelPtr &ge_model, gert::Om2Deb
859	return SUCCESS;	1002	return SUCCESS;
860	}	1003	}
861		1004
862	-Status Om2PackageHelper::BuildManifest(const GeRootModelPtr &ge_root_model,	1005	+Status Om2PackageHelper::BuildManifest(const GeRootModelPtr &ge_root_model, gert::Om2ModelData &model_data) {
863	- std::map<std::string, std::string> &manifest) {
864	GELOGI("[OM2] Begin to build manifest");	1006	GELOGI("[OM2] Begin to build manifest");
		1007	+ std::map<std::string, std::string> &manifest = model_data.manifest;
865	manifest[OM2_ARCHIVE_VERSION] = OM2_ARCHIVE_VERSION_VALUE;	1008	manifest[OM2_ARCHIVE_VERSION] = OM2_ARCHIVE_VERSION_VALUE;
866	if (ge_root_model != nullptr) {	1009	if (ge_root_model != nullptr) {
867	manifest[OM2_MODEL_NUM] = std::to_string(ge_root_model->GetSubgraphInstanceNameToModel().size());	1010	manifest[OM2_MODEL_NUM] = std::to_string(ge_root_model->GetSubgraphInstanceNameToModel().size());
	@@ -890,13 +1033,13 @@ Status Om2PackageHelper::BuildOm2ModelData(const GeModelPtr &ge_model, gert::Om2
890	GELOGW("[OM2] Ge model set opp version unsuccessful!");	1033	GELOGW("[OM2] Ge model set opp version unsuccessful!");
891	}	1034	}
892		1035
893	- std::vector<Om2ConstMeta> const_metas;	1036	+ GE_ASSERT_SUCCESS(BuildCustomKernelBinaries(ge_root_model, model_data));
894	- std::vector<Om2VarMeta> var_metas;	1037	+ GE_ASSERT_SUCCESS(BuildCustomSharedLibs(ge_root_model, model_data));
895	- GE_ASSERT_SUCCESS(BuildProgramBody(ge_model, model_data.program_body, const_metas, var_metas));	1038	+ GE_ASSERT_SUCCESS(BuildProgramBody(ge_model, model_data));
896	- GE_ASSERT_SUCCESS(BuildKernelBinaries(ge_model, model_data.kernel_binaries));	1039	+ GE_ASSERT_SUCCESS(BuildKernelBinaries(ge_model, model_data));
897	- GE_ASSERT_SUCCESS(BuildModelMeta(ge_model, model_data.model_meta));	1040	+ GE_ASSERT_SUCCESS(BuildModelMeta(ge_model, model_data));
898	- GE_ASSERT_SUCCESS(BuildConstantsData(ge_model, const_metas, model_data.constants_data));	1041	+ GE_ASSERT_SUCCESS(BuildConstantsData(ge_model, model_data));
899	- model_data.var_metas = std::move(var_metas);	1042	+
900	const auto compute_graph = ge_model->GetGraph();	1043	const auto compute_graph = ge_model->GetGraph();
901	model_data.graph_id = (compute_graph != nullptr) ? compute_graph->GetGraphID() : 0U;	1044	model_data.graph_id = (compute_graph != nullptr) ? compute_graph->GetGraphID() : 0U;
902	const auto session_id = GetContext().SessionId();	1045	const auto session_id = GetContext().SessionId();
	@@ -905,8 +1048,8 @@ Status Om2PackageHelper::BuildOm2ModelData(const GeModelPtr &ge_model, gert::Om2
905	GE_ASSERT_SUCCESS(	1048	GE_ASSERT_SUCCESS(
906	gert::BuildRTVarResource(*var_manager, ge_model->GetGraph(), model_data.var_metas, model_data.rt_var_resource));	1049	gert::BuildRTVarResource(*var_manager, ge_model->GetGraph(), model_data.var_metas, model_data.rt_var_resource));
907	}	1050	}
908	- GE_ASSERT_SUCCESS(BuildDebugInfo(ge_model, model_data.debug_info));	1051	+ GE_ASSERT_SUCCESS(BuildDebugInfo(ge_model, model_data));
909	- GE_ASSERT_SUCCESS(BuildManifest(ge_root_model, model_data.manifest));	1052	+ GE_ASSERT_SUCCESS(BuildManifest(ge_root_model, model_data));
910		1053
911	GELOGI("[OM2] Successfully built Om2ModelData");	1054	GELOGI("[OM2] Successfully built Om2ModelData");
912	return SUCCESS;	1055	return SUCCESS;

Mbase/common/om2/codegen/ast/ast_nodes.h+2-0

	@@ -103,6 +103,8 @@ enum class StablePartId : uint8_t {
103	kCreateLabelListForLabelGotoEx,	103	kCreateLabelListForLabelGotoEx,
104	kOpDefStructs,	104	kOpDefStructs,
105	kOm2LogMacros,	105	kOm2LogMacros,
		106	+ kCreateClassCustomOpFactory,
		107	+ kCustomTaskHelpers,
106	};	108	};
107		109
108	enum class StablePartPlacement : uint8_t {	110	enum class StablePartPlacement : uint8_t {

Mbase/common/om2/codegen/emitter/stable_parts/stable_part_provider.cc+225-0

	@@ -763,6 +763,15 @@ Status ResolveStablePart(StablePartId id, std::string &output) {
763	" uint32_t task_type;\n"	763	" uint32_t task_type;\n"
764	"};\n"	764	"};\n"
765	"\n"	765	"\n"
		766	+ "struct CustomDispatchInfo {\n"
		767	+ " const OpArgInfo *args_info; // IO 地址解析数组\n"
		768	+ " uint32_t args_info_num; // args_info 数组长度\n"
		769	+ " const char *op_type; // 算子类型名，用于 Report 上报\n"
		770	+ " uint32_t args_idx; // 参数表索引，用于 GetArgsInfo 查找\n"
		771	+ " uint32_t stream_id; // 执行流索引\n"
		772	+ " uint32_t task_type;\n"
		773	+ "};\n"
		774	+ "\n"
766	"struct DsaDispatchInfo {\n"	775	"struct DsaDispatchInfo {\n"
767	" const OpArgInfo *args_info; // IO 地址解析数组\n"	776	" const OpArgInfo *args_info; // IO 地址解析数组\n"
768	" const char *op_type; // 算子类型名，用于 Report 上报\n"	777	" const char *op_type; // 算子类型名，用于 Report 上报\n"
	@@ -858,6 +867,7 @@ Status ResolveStablePart(StablePartId id, std::string &output) {
858	" union {\n"	867	" union {\n"
859	" AicoreDispatchInfo aicore;\n"	868	" AicoreDispatchInfo aicore;\n"
860	" AicpuDispatchInfo aicpu;\n"	869	" AicpuDispatchInfo aicpu;\n"
		870	+ " CustomDispatchInfo custom;\n"
861	" CmoDispatchInfo cmo;\n"	871	" CmoDispatchInfo cmo;\n"
862	" MemcpyAsyncDispatchInfo memcpy_async;\n"	872	" MemcpyAsyncDispatchInfo memcpy_async;\n"
863	" MemcpyAddrDispatchInfo memcpy_addr;\n"	873	" MemcpyAddrDispatchInfo memcpy_addr;\n"
	@@ -959,6 +969,221 @@ Status ResolveStablePart(StablePartId id, std::string &output) {
959	" } \\\n"	969	" } \\\n"
960	" } while (false)\n";	970	" } while (false)\n";
961	return SUCCESS;	971	return SUCCESS;
		972	+ case StablePartId::kCreateClassCustomOpFactory:
		973	+ output =
		974	+ "class CustomOpFactory {\n"
		975	+ " public:\n"
		976	+ " static graphStatus RegisterCustomOpCreator(const AscendString &op_type, const BaseOpCreator "
		977	+ "&op_creator);\n"
		978	+ " static BaseCustomOp *CreateOrGetCustomOp(const AscendString &op_type);\n"
		979	+ " static graphStatus GetAllRegisteredOps(std::vector<AscendString> &all_registered_ops);\n"
		980	+ " static bool IsExistOp(const AscendString &op_type);\n"
		981	+ " static graphStatus LoadCustomOpsPartition(const uint8_t *data, size_t len);\n"
		982	+ "};\n";
		983	+ return SUCCESS;
		984	+ case StablePartId::kCustomTaskHelpers:
		985	+ output =
		986	+ "using GeTensorPtr = std::shared_ptr<gert::Tensor>;\n"
		987	+ "\n"
		988	+ "GeTensorPtr BuildGeTensor(const Om2Tensor &om2_tensor) {\n"
		989	+ " gert::StorageShape storage_shape;\n"
		990	+ " for (auto i = 0U; i < om2_tensor.shape_dims_num; ++i) {\n"
		991	+ " auto dim = om2_tensor.shape_dims[i];\n"
		992	+ " (void)storage_shape.MutableStorageShape().AppendDim(dim);\n"
		993	+ " (void)storage_shape.MutableOriginShape().AppendDim(dim);\n"
		994	+ " }\n"
		995	+ " gert::StorageFormat storage_format(static_cast<ge::Format>(om2_tensor.format),\n"
		996	+ " static_cast<ge::Format>(om2_tensor.format), {});\n"
		997	+ " return std::make_shared<gert::Tensor>(storage_shape, storage_format, gert::kOnDeviceHbm,\n"
		998	+ " static_cast<ge::DataType>(om2_tensor.data_type),\n"
		999	+ " reinterpret_cast<gert::TensorAddress>(om2_tensor.device_address));\n"
		1000	+ "}\n"
		1001	+ "\n"
		1002	+ "class CustKernelContextHolder {\n"
		1003	+ " public:\n"
		1004	+ " CustKernelContextHolder() = default;\n"
		1005	+ " CustKernelContextHolder(CustKernelContextHolder &&holder) {\n"
		1006	+ " context_holder_ = std::move(holder.context_holder_);\n"
		1007	+ " value_holder_ = std::move(holder.value_holder_);\n"
		1008	+ " compute_node_extend_holder_ = std::move(holder.compute_node_extend_holder_);\n"
		1009	+ " context_ = holder.context_;\n"
		1010	+ " holder.context_ = nullptr;\n"
		1011	+ " }\n"
		1012	+ "\n"
		1013	+ " CustKernelContextHolder &operator=(CustKernelContextHolder &&holder) {\n"
		1014	+ " context_holder_ = std::move(holder.context_holder_);\n"
		1015	+ " value_holder_ = std::move(holder.value_holder_);\n"
		1016	+ " compute_node_extend_holder_ = std::move(holder.compute_node_extend_holder_);\n"
		1017	+ " context_ = holder.context_;\n"
		1018	+ " holder.context_ = nullptr;\n"
		1019	+ " return *this;\n"
		1020	+ " }\n"
		1021	+ "\n"
		1022	+ " ~CustKernelContextHolder() {\n"
		1023	+ " for (auto &value : value_holder_) {\n"
		1024	+ " value.Set(nullptr, nullptr);\n"
		1025	+ " }\n"
		1026	+ " }\n"
		1027	+ "\n"
		1028	+ " std::unique_ptr<uint8_t[]> context_holder_;\n"
		1029	+ " std::vector<gert::Chain> value_holder_;\n"
		1030	+ " std::unique_ptr<uint8_t[]> compute_node_extend_holder_;\n"
		1031	+ " gert::KernelContext *context_;\n"
		1032	+ "};\n"
		1033	+ "\n"
		1034	+ "CustKernelContextHolder BuildKernelContextHolder(const char_t op_name, const char_t op_type,\n"
		1035	+ " const std::vector<GeTensorPtr> &inputs,\n"
		1036	+ " const std::vector<GeTensorPtr> &outputs, void *allocator,\n"
		1037	+ " void *stream) {\n"
		1038	+ " std::vector<std::pair<void *, gert::Chain::Deleter>> holder_inputs;\n"
		1039	+ " std::vector<std::pair<void *, gert::Chain::Deleter>> holder_outputs;\n"
		1040	+ " std::vector<void *> ws_vec;\n"
		1041	+ " for (auto &elem : inputs) {\n"
		1042	+ " holder_inputs.emplace_back(elem.get(), nullptr);\n"
		1043	+ " }\n"
		1044	+ " holder_inputs.emplace_back(allocator, nullptr);\n"
		1045	+ " holder_inputs.emplace_back(stream, nullptr);\n"
		1046	+ " for (auto &elem : outputs) {\n"
		1047	+ " holder_outputs.emplace_back(elem.get(), nullptr);\n"
		1048	+ " }\n"
		1049	+ " holder_outputs.emplace_back(&ws_vec, nullptr);\n"
		atomgit-bot7 天前 🟡 Medium Priority
		1050	+ "\n"
		1051	+ " CustKernelContextHolder holder;\n"
		1052	+ " size_t context_size =\n"
		1053	+ " sizeof(KernelRunContext) + sizeof(gert::Chain ) (holder_inputs.size() + holder_outputs.size());\n"
		1054	+ " holder.context_holder_ = std::make_unique<uint8_t[]>(context_size);\n"
		1055	+ " if (holder.context_holder_ == nullptr) {\n"
		1056	+ " printf(\"Create context holder failed.\");\n"
		1057	+ " return holder;\n"
		1058	+ " }\n"
		1059	+ " holder.context_ = PtrToPtr<uint8_t, gert::KernelContext>(holder.context_holder_.get());\n"
		1060	+ "\n"
		1061	+ " holder.compute_node_extend_holder_ = std::make_unique<uint8_t[]>(sizeof(gert::ComputeNodeInfo));\n"
		1062	+ " if (holder.compute_node_extend_holder_ == nullptr) {\n"
		1063	+ " printf(\"Create compute node holder failed.\");\n"
		1064	+ " return holder;\n"
		1065	+ " }\n"
		1066	+ " auto compute_node_info = PtrToPtr<uint8_t, "
		1067	+ "gert::ComputeNodeInfo>(holder.compute_node_extend_holder_.get());\n"
		1068	+ " compute_node_info->Init(inputs.size(), inputs.size(), outputs.size(), op_name, op_type);\n"
		1069	+ "\n"
		1070	+ " auto kernel_run_context = holder.context_->GetContext();\n"
		1071	+ " kernel_run_context->input_size = holder_inputs.size();\n"
		1072	+ " kernel_run_context->output_size = holder_outputs.size();\n"
		1073	+ " kernel_run_context->compute_node_info = compute_node_info;\n"
		1074	+ " kernel_run_context->output_start = &(kernel_run_context->values[kernel_run_context->input_size]);\n"
		1075	+ " holder.value_holder_.resize(holder_inputs.size() + holder_outputs.size());\n"
		1076	+ " for (size_t i = 0UL; i < holder.value_holder_.size(); ++i) {\n"
		1077	+ " kernel_run_context->values[i] = PtrToPtr<gert::Chain, AsyncAnyValue>(&holder.value_holder_[i]);\n"
		1078	+ " }\n"
		1079	+ " for (size_t i = 0UL; i < holder_inputs.size(); ++i) {\n"
		1080	+ " holder.value_holder_[i].Set(holder_inputs[i].first, holder_inputs[i].second);\n"
		1081	+ " }\n"
		1082	+ " for (size_t i = 0UL; i < holder_outputs.size(); ++i) {\n"
		1083	+ " holder.value_holder_[holder_inputs.size() + i].Set(holder_outputs[i].first, holder_outputs[i].second);\n"
		1084	+ " }\n"
		1085	+ "\n"
		1086	+ " return holder;\n"
		1087	+ "}\n"
		1088	+ "\n"
		1089	+ "class AllocatorFaker : public gert::GertAllocator {\n"
		1090	+ " gert::GertMemBlock Malloc(size_t / size */) override {\n"
		1091	+ " return nullptr;\n"
		1092	+ " }\n"
		1093	+ " void Free(gert::GertMemBlock * /* block */) override {}\n"
		1094	+ " gert::GertTensorData MallocTensorData(size_t /* size */) override {\n"
		1095	+ " return gert::GertTensorData();\n"
		1096	+ " }\n"
		1097	+ " gert::TensorData MallocTensorDataFromL1(size_t /* size */) override {\n"
		1098	+ " return gert::TensorData();\n"
		1099	+ " }\n"
		1100	+ " ge::graphStatus ShareFromTensorData(const gert::TensorData & /* td /, gert::GertTensorData & / gtd */) "
		1101	+ "override {\n"
		1102	+ " return 0;\n"
		1103	+ " }\n"
		1104	+ " ge::graphStatus SetL1Allocator(ge::Allocator * /* allocator */) override {\n"
		1105	+ " return ge::GRAPH_SUCCESS;\n"
		1106	+ " }\n"
		1107	+ " ge::graphStatus FreeAt(int64_t /* stream_id /, gert::GertMemBlock /* block */) override {\n"
		1108	+ " return ge::GRAPH_SUCCESS;\n"
		1109	+ " }\n"
		1110	+ " int64_t GetStreamNum() override {\n"
		1111	+ " return 0;\n"
		1112	+ " }\n"
		1113	+ "};\n"
		1114	+ "\n"
		1115	+ "aclError KernelCustTaskDistribute(const char_t op_name, const char_t op_type, const "
		1116	+ "std::vector<Om2Tensor> &inputs,\n"
		1117	+ " const std::vector<Om2Tensor> &outputs, aclrtStream stream) {\n"
		1118	+ " std::vector<GeTensorPtr> op_inputs;\n"
		1119	+ " std::vector<GeTensorPtr> op_outputs;\n"
		1120	+ " for (auto &elem : inputs) {\n"
		1121	+ " op_inputs.emplace_back(BuildGeTensor(elem));\n"
		1122	+ " }\n"
		1123	+ " for (auto &elem : outputs) {\n"
		1124	+ " op_outputs.emplace_back(BuildGeTensor(elem));\n"
		1125	+ " }\n"
		1126	+ " auto allocator = std::make_shared<AllocatorFaker>();\n"
		1127	+ " auto eager_context_holder =\n"
		1128	+ " BuildKernelContextHolder(op_name, op_type, op_inputs, op_outputs, allocator.get(), stream);\n"
		1129	+ " auto eager_context = reinterpret_cast<gert::EagerOpExecutionContext *>(eager_context_holder.context_);\n"
		1130	+ " auto custom_op_ptr = ge::CustomOpFactory::CreateOrGetCustomOp(op_type);\n"
		1131	+ " OM2_CHK_NOTNULL(custom_op_ptr);\n"
		1132	+ " auto eager_execute_op_ptr = dynamic_cast<ge::EagerExecuteOp >(custom_op_ptr);\n"
		1133	+ " if (eager_execute_op_ptr == nullptr) {\n"
		1134	+ " OM2_LOGE(\"%s is custom op but did not implement EagerExecuteOp\", eager_context->GetNodeType());\n"
		1135	+ " return ACL_ERROR_FAILURE;\n"
		1136	+ " }\n"
		1137	+ " OM2_CHK_STATUS(eager_execute_op_ptr->Execute(eager_context));\n"
		1138	+ "\n"
		1139	+ " return ACL_SUCCESS;\n"
		1140	+ "}\n"
		1141	+ "\n"
		1142	+ "aclError DeserializeCustKernelBinaries(std::unordered_map<std::string, BinDataInfo> &bin_info_map) {\n"
		1143	+ " for (auto &[name, bin_info] : bin_info_map) {\n"
		1144	+ " auto post_fix_pos = name.rfind(\"_CustomKernel.bin\");\n"
		1145	+ " if (post_fix_pos == std::string::npos \|\| post_fix_pos == 0) {\n"
		1146	+ " continue;\n"
		1147	+ " }\n"
		1148	+ "\n"
		1149	+ " OM2_LOGI(\"[OM2] Begin to deserialize kernel binary '%s'\", name.c_str());\n"
		1150	+ " auto kernel_name_end_pos = name.rfind('_', post_fix_pos - 1);\n"
		1151	+ " if (kernel_name_end_pos == std::string::npos \|\| kernel_name_end_pos == 0) {\n"
		1152	+ " OM2_LOGW(\"[OM2] Name is invalid, ignore kernel binary '%s'\", name.c_str());\n"
		1153	+ " continue;\n"
		1154	+ " }\n"
		1155	+ "\n"
		1156	+ " std::string kernel_name = name.substr(0, kernel_name_end_pos);\n"
		1157	+ " OM2_LOGI(\"[OM2] Extract kernel name `%s` from '%s'\", kernel_name.c_str(), name.c_str());\n"
		1158	+ "\n"
		1159	+ " // Create operator by kernel name\n"
		1160	+ " auto op = ge::CustomOpFactory::CreateOrGetCustomOp(ge::AscendString(kernel_name.c_str()));\n"
		1161	+ " if (op == nullptr) {\n"
		1162	+ " OM2_LOGE(\"[OM2] Custom op '%s' not found in registry\", kernel_name.c_str());\n"
		1163	+ " return ACL_ERROR_FAILURE;\n"
		1164	+ " }\n"
		1165	+ "\n"
		1166	+ " // Cast to PortableOp\n"
		1167	+ " auto serializable_op = dynamic_cast<ge::PortableOp >(op);\n"
		1168	+ " if (serializable_op == nullptr) {\n"
		1169	+ " OM2_LOGE(\"[OM2] Custom op '%s' is not PortableOp\", kernel_name.c_str());\n"
		1170	+ " return ACL_ERROR_FAILURE;\n"
		1171	+ " }\n"
		1172	+ "\n"
		1173	+ " // Deserialize kernel binary\n"
		1174	+ " uint8_t data_addr = reinterpret_cast<uint8_t >(const_cast<void *>(bin_info.data));\n"
		1175	+ " const std::vector<uint8_t> kernel_bin_buffer(data_addr, data_addr + bin_info.size);\n"
		1176	+ " const auto ret = serializable_op->Deserialize(kernel_bin_buffer);\n"
		1177	+ " if (ret != ge::GRAPH_SUCCESS) {\n"
		1178	+ " OM2_LOGE(\"[OM2] Failed to deserialize custom op '%s'\", kernel_name.c_str());\n"
		1179	+ " return ACL_ERROR_FAILURE;\n"
		1180	+ " }\n"
		1181	+ "\n"
		1182	+ " OM2_LOGI(\"[OM2] Successfully deserialized custom op '%s'\", kernel_name.c_str());\n"
		1183	+ " }\n"
		1184	+ " return ACL_SUCCESS;\n"
		1185	+ "}\n";
		1186	+ return SUCCESS;
962	default:	1187	default:
963	output.clear();	1188	output.clear();
964	return FAILED;	1189	return FAILED;

Mbase/common/om2/codegen/file_code_generator/load_and_run_file_code_generator.cc+4-0

	@@ -49,6 +49,7 @@ const std::map<std::string, std::vector<OpDispatchType::Value>> &GetDispatchFunc
49	{"DispatchCmoAddr", {OpDispatchType::DISPATCH_CMO_ADDR}},	49	{"DispatchCmoAddr", {OpDispatchType::DISPATCH_CMO_ADDR}},
50	{"DispatchDsa", {OpDispatchType::DISPATCH_DSA}},	50	{"DispatchDsa", {OpDispatchType::DISPATCH_DSA}},
51	{"DispatchKernelEx", {OpDispatchType::DISPATCH_KERNEL_EX}},	51	{"DispatchKernelEx", {OpDispatchType::DISPATCH_KERNEL_EX}},
		52	+ {"DispatchCustomKernel", {OpDispatchType::DISPATCH_CUSTOM_KERNEL}},
52	};	53	};
53	return kMap;	54	return kMap;
54	}	55	}
	@@ -92,6 +93,9 @@ MethodDef *LoadAndRunFileCodeGenerator::BuildGetRtModelHandleMethod() const {
92	Status LoadAndRunFileCodeGenerator::BuildLoadBody(std::vector<BodyItem> &body, const Om2CodegenModel &codegen_model,	93	Status LoadAndRunFileCodeGenerator::BuildLoadBody(std::vector<BodyItem> &body, const Om2CodegenModel &codegen_model,
93	const std::vector<TaskCodeBuilderPtr> &task_code_builders) {	94	const std::vector<TaskCodeBuilderPtr> &task_code_builders) {
94	body.push_back(ast_.Call("OM2_LOGI", {ast_.Str("Load begin")}));	95	body.push_back(ast_.Call("OM2_LOGI", {ast_.Str("Load begin")}));
		96	+ if (has_custom_kernel_) {
		97	+ body.push_back(ChkStatus(ast_.Call("DeserializeCustKernelBinaries", {ast_.Var("", "bin_info_map_")})));
		98	+ }
95	body.push_back(dev_ext_info_mem_ptrs_.Resize(codegen_model.aicpu_task_count));	99	body.push_back(dev_ext_info_mem_ptrs_.Resize(codegen_model.aicpu_task_count));
96		100
97	// 公共 DispatchOpContext 初始化列表（for 循环版和展开版共享）	101	// 公共 DispatchOpContext 初始化列表（for 循环版和展开版共享）

Mbase/common/om2/codegen/file_code_generator/load_and_run_file_code_generator.h+4-0

	@@ -33,6 +33,9 @@ class LoadAndRunFileCodeGenerator : public Om2ModelClassGeneratorBase {
33	FunctionDef *BuildCommitProfUnit() const;	33	FunctionDef *BuildCommitProfUnit() const;
34	DeclNode *BuildOpDefTable(const Om2CodegenModel &codegen_model,	34	DeclNode *BuildOpDefTable(const Om2CodegenModel &codegen_model,
35	const std::vector<TaskCodeBuilderPtr> &task_code_builders) const;	35	const std::vector<TaskCodeBuilderPtr> &task_code_builders) const;
		36	+ void SetHasCustomKernel(bool value) {
		37	+ has_custom_kernel_ = value;
		38	+ }
36		39
37	private:	40	private:
38	Status BuildLoadBody(std::vector<BodyItem> &body, const Om2CodegenModel &codegen_model,	41	Status BuildLoadBody(std::vector<BodyItem> &body, const Om2CodegenModel &codegen_model,
	@@ -54,6 +57,7 @@ class LoadAndRunFileCodeGenerator : public Om2ModelClassGeneratorBase {
54	Status BuildAclrtMallocFunction(std::vector<DeclNode *> &items) const;	57	Status BuildAclrtMallocFunction(std::vector<DeclNode *> &items) const;
55	Status BuildDispatchOp(std::vector<DeclNode *> &items,	58	Status BuildDispatchOp(std::vector<DeclNode *> &items,
56	const std::map<uint32_t, std::string> &type_to_func_name) const;	59	const std::map<uint32_t, std::string> &type_to_func_name) const;
		60	+ bool has_custom_kernel_ = false;
57	};	61	};
58	} // namespace ge	62	} // namespace ge
59		63

Mbase/common/om2/codegen/om2_codegen.cc+6-10

	@@ -16,7 +16,6 @@
16	#include "common/om2/codegen/ast/ast_build_context.h"	16	#include "common/om2/codegen/ast/ast_build_context.h"
17	#include "common/om2/codegen/ast/ast_context.h"	17	#include "common/om2/codegen/ast/ast_context.h"
18	#include "common/om2/codegen/om2_codegen_model_builder.h"	18	#include "common/om2/codegen/om2_codegen_model_builder.h"
19	-#include "common/om2/codegen/om2_codegen_utils.h"
20	#include "program_generator.h"	19	#include "program_generator.h"
21	#include "om2_code_printer.h"	20	#include "om2_code_printer.h"
22		21
	@@ -55,8 +54,11 @@ void DumpGeneratedFiles(const Om2CodegenArtifacts &artifacts) {
55	}	54	}
56	} // namespace	55	} // namespace
57		56
58	-Status Om2Codegen::Om2CodegenAndCompile(const ge::GeModelPtr &ge_model, Om2CodegenArtifacts &artifacts,	57	+Status Om2Codegen::Om2CodegenAndCompile(const ge::GeModelPtr &ge_model, gert::Om2ModelData &model_data) const {
59	- Om2ConstMetas &const_metas, std::vector<Om2VarMeta> &var_metas) const {	58	+ Om2CodegenArtifacts &artifacts = model_data.program_body.source_artifacts;
		59	+ auto &const_metas = model_data.constants_data.consts;
		60	+ std::vector<Om2VarMeta> &var_metas = model_data.var_metas;
		61	+ bool has_custom_kernel = !model_data.custom_kernel_binaries.empty();
60	artifacts.clear();	62	artifacts.clear();
61	const_metas.clear();	63	const_metas.clear();
62	var_metas.clear();	64	var_metas.clear();
	@@ -68,7 +70,7 @@ Status Om2Codegen::Om2CodegenAndCompile(const ge::GeModelPtr &ge_model, Om2Codeg
68	Om2CodegenModelBuilder builder;	70	Om2CodegenModelBuilder builder;
69	GE_ASSERT_SUCCESS(builder.Build(ge_model, task_code_builders, codegen_model, const_metas));	71	GE_ASSERT_SUCCESS(builder.Build(ge_model, task_code_builders, codegen_model, const_metas));
70	var_metas = codegen_model.var_metas;	72	var_metas = codegen_model.var_metas;
71	- ProgramGenerator generator(ast, task_code_builders, codegen_model);	73	+ ProgramGenerator generator(ast, task_code_builders, codegen_model, has_custom_kernel);
72		74
73	Om2CodePrinter code_printer(ge_model->GetName());	75	Om2CodePrinter code_printer(ge_model->GetName());
74	GE_ASSERT_SUCCESS(generator.GenerateProgram(code_printer));	76	GE_ASSERT_SUCCESS(generator.GenerateProgram(code_printer));
	@@ -89,10 +91,4 @@ Status Om2Codegen::Om2CodegenAndCompile(const ge::GeModelPtr &ge_model, Om2Codeg
89	artifacts.push_back(std::move(so_artifact));	91	artifacts.push_back(std::move(so_artifact));
90	return SUCCESS;	92	return SUCCESS;
91	}	93	}
92	-
93	-Status Om2Codegen::Om2CodegenAndCompile(const ge::GeModelPtr &ge_model, Om2CodegenArtifacts &artifacts,
94	- Om2ConstMetas &const_metas) const {
95	- std::vector<Om2VarMeta> var_metas;
96	- return Om2CodegenAndCompile(ge_model, artifacts, const_metas, var_metas);
97	-}
98	} // namespace ge	94	} // namespace ge

Mbase/common/om2/codegen/om2_codegen.h+2-6

	@@ -7,21 +7,17 @@
7	* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.	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.	8	* See LICENSE in the root of the software repository for the full text of the License.
9	*/	9	*/
10	-
11	#ifndef AIR_CXX_BASE_COMMON_OM2_CODEGEN_OM2_CODEGEN_H_	10	#ifndef AIR_CXX_BASE_COMMON_OM2_CODEGEN_OM2_CODEGEN_H_
12	#define AIR_CXX_BASE_COMMON_OM2_CODEGEN_OM2_CODEGEN_H_	11	#define AIR_CXX_BASE_COMMON_OM2_CODEGEN_OM2_CODEGEN_H_
13		12
14	#include "common/model/ge_model.h"	13	#include "common/model/ge_model.h"
15	#include "ge_common/ge_common_api_types.h"	14	#include "ge_common/ge_common_api_types.h"
16	-#include "common/om2/codegen/om2_codegen_types.h"	15	+#include "common/om2/om2_model_data.h"
17		16
18	namespace ge {	17	namespace ge {
19	class Om2Codegen {	18	class Om2Codegen {
20	public:	19	public:
21	- Status Om2CodegenAndCompile(const GeModelPtr &ge_model, Om2CodegenArtifacts &artifacts,	20	+ Status Om2CodegenAndCompile(const GeModelPtr &ge_model, gert::Om2ModelData &model_data) const;
22	- Om2ConstMetas &const_metas) const;
23	- Status Om2CodegenAndCompile(const GeModelPtr &ge_model, Om2CodegenArtifacts &artifacts, Om2ConstMetas &const_metas,
24	- std::vector<Om2VarMeta> &var_metas) const;
25	};	21	};
26	} // namespace ge	22	} // namespace ge
27		23

Mbase/common/om2/codegen/om2_codegen_types.h+3-2

	@@ -290,6 +290,7 @@ struct OpDispatchType {
290	DISPATCH_CMO_ADDR = 18,	290	DISPATCH_CMO_ADDR = 18,
291	DISPATCH_DSA = 19,	291	DISPATCH_DSA = 19,
292	DISPATCH_KERNEL_EX = 20,	292	DISPATCH_KERNEL_EX = 20,
		293	+ DISPATCH_CUSTOM_KERNEL = 21,
293	};	294	};
294		295
295	static std::string ToString() {	296	static std::string ToString() {
	@@ -303,9 +304,9 @@ struct OpDispatchType {
303	"DISPATCH_BARRIER", "DISPATCH_CMO",	304	"DISPATCH_BARRIER", "DISPATCH_CMO",
304	"DISPATCH_MEMCPY_ASYNC", "DISPATCH_MEMCPY_ADDR_ASYNC",	305	"DISPATCH_MEMCPY_ASYNC", "DISPATCH_MEMCPY_ADDR_ASYNC",
305	"DISPATCH_CMO_ADDR", "DISPATCH_DSA",	306	"DISPATCH_CMO_ADDR", "DISPATCH_DSA",
306	- "DISPATCH_KERNEL_EX"};	307	+ "DISPATCH_KERNEL_EX", "DISPATCH_CUSTOM_KERNEL"};
307	std::string code = "enum OpDispatchType : uint32_t {\n";	308	std::string code = "enum OpDispatchType : uint32_t {\n";
308	- for (uint32_t i = 0U; i <= static_cast<uint32_t>(DISPATCH_KERNEL_EX); i++) {	309	+ for (uint32_t i = 0U; i < sizeof(kNames) / sizeof(kNames[0]); i++) {
309	code += " " + std::string(kNames[i]) + " = " + std::to_string(i) + ",\n";	310	code += " " + std::string(kNames[i]) + " = " + std::to_string(i) + ",\n";
310	}	311	}
311	code += " DISPATCH_TYPE_COUNT\n};\n";	312	code += " DISPATCH_TYPE_COUNT\n};\n";

Mbase/common/om2/codegen/om2_codegen_utils.cc+1-0

	@@ -44,6 +44,7 @@ const std::unordered_set<ModelTaskType> kSupportedTaskTypes = {
44	ModelTaskType::MODEL_TASK_CMO,	44	ModelTaskType::MODEL_TASK_CMO,
45	ModelTaskType::MODEL_TASK_CMO_ADDR,	45	ModelTaskType::MODEL_TASK_CMO_ADDR,
46	ModelTaskType::MODEL_TASK_BARRIER,	46	ModelTaskType::MODEL_TASK_BARRIER,
		47	+ ModelTaskType::MODEL_TASK_CUSTOM_KERNEL,
47	};	48	};
48	} // namespace	49	} // namespace
49	std::string Om2CodegenUtils::GetKernelNameWithExtension(const std::string &kernel_name) {	50	std::string Om2CodegenUtils::GetKernelNameWithExtension(const std::string &kernel_name) {

Mbase/common/om2/codegen/program_generator.cc+26-13

	@@ -188,23 +188,36 @@ Status ProgramGenerator::GenerateKernelRegSource(Om2CodePrinter &code_printer) {
188		188
189	Status ProgramGenerator::GenerateLoadAndRunSource(Om2CodePrinter &code_printer) {	189	Status ProgramGenerator::GenerateLoadAndRunSource(Om2CodePrinter &code_printer) {
190	LoadAndRunFileCodeGenerator load_and_run_handler(ast_);	190	LoadAndRunFileCodeGenerator load_and_run_handler(ast_);
		191	+ load_and_run_handler.SetHasCustomKernel(has_custom_kernel_);
191	auto anonymous_items = load_and_run_handler.BuildAnonymousNamespaceItems(codegen_model_, task_code_builder_list_);	192	auto anonymous_items = load_and_run_handler.BuildAnonymousNamespaceItems(codegen_model_, task_code_builder_list_);
		193	+ if (has_custom_kernel_) {
		194	+ anonymous_items.insert(anonymous_items.begin(),
		195	+ ast_.StablePart(StablePartId::kCustomTaskHelpers, StablePartPlacement::kNamespace));
		196	+ }
192	(void)anonymous_items.insert(anonymous_items.begin(),	197	(void)anonymous_items.insert(anonymous_items.begin(),
193	ast_.StablePart(StablePartId::kLoadAndRunDumpHelpers, StablePartPlacement::kNamespace));	198	ast_.StablePart(StablePartId::kLoadAndRunDumpHelpers, StablePartPlacement::kNamespace));
194	anonymous_items.push_back(load_and_run_handler.BuildOpDefTable(codegen_model_, task_code_builder_list_));	199	anonymous_items.push_back(load_and_run_handler.BuildOpDefTable(codegen_model_, task_code_builder_list_));
195	- auto *translation_unit = ast_.File({	200	+ std::vector<DeclNode *> body_items = {ast_.Include(codegen_model_.model_name + "_interface.h")};
196	- ast_.Include(codegen_model_.model_name + "_interface.h"),	201	+ if (has_custom_kernel_) {
197	- ast_.Space(),	202	+ body_items.emplace_back(ast_.Include("graph/custom_op.h"));
198	- ast_.Namespace("om2",	203	+ body_items.emplace_back(ast_.Include("exe_graph/runtime/gert_mem_allocator.h"));
199	- {	204	+ }
200	- ast_.Namespace("", anonymous_items),	205	+ body_items.emplace_back(ast_.Space());
201	- load_and_run_handler.BuildGetRtModelHandleMethod(),	206	+ if (has_custom_kernel_) {
202	- load_and_run_handler.BuildLoadMethod(codegen_model_, task_code_builder_list_),	207	+ body_items.emplace_back(ast_.Namespace(
203	- load_and_run_handler.BuildRunAsyncMethod(codegen_model_),	208	+ "ge", {ast_.StablePart(StablePartId::kCreateClassCustomOpFactory, StablePartPlacement::kNamespace)}));
204	- load_and_run_handler.BuildRunMethod(codegen_model_),	209	+ body_items.emplace_back(ast_.Space());
205	- }),	210	+ }
206	- ast_.StablePart(StablePartId::kLoadAndRunExternalApis),	211	+ body_items.emplace_back(
207	- });	212	+ ast_.Namespace("om2", {
		213	+ ast_.Namespace("", anonymous_items),
		214	+ load_and_run_handler.BuildGetRtModelHandleMethod(),
		215	+ load_and_run_handler.BuildLoadMethod(codegen_model_, task_code_builder_list_),
		216	+ load_and_run_handler.BuildRunAsyncMethod(codegen_model_),
		217	+ load_and_run_handler.BuildRunMethod(codegen_model_),
		218	+ }));
		219	+ body_items.emplace_back(ast_.StablePart(StablePartId::kLoadAndRunExternalApis));
		220	+ auto *translation_unit = ast_.File(body_items);
208	GE_ASSERT_SUCCESS(EmitFile(GeneratedFileIndex::kLoadingAndRunningFile, translation_unit, code_printer));	221	GE_ASSERT_SUCCESS(EmitFile(GeneratedFileIndex::kLoadingAndRunningFile, translation_unit, code_printer));
209	GELOGD("[OM2] Load and run source file code is generated.");	222	GELOGD("[OM2] Load and run source file code is generated.");
210	return SUCCESS;	223	return SUCCESS;

Mbase/common/om2/codegen/program_generator.h+6-2

	@@ -28,8 +28,11 @@ namespace ge {
28	class ProgramGenerator {	28	class ProgramGenerator {
29	public:	29	public:
30	ProgramGenerator(AstBuildContext &ast, const std::vector<TaskCodeBuilderPtr> &task_code_builders,	30	ProgramGenerator(AstBuildContext &ast, const std::vector<TaskCodeBuilderPtr> &task_code_builders,
31	- const Om2CodegenModel &codegen_model)	31	+ const Om2CodegenModel &codegen_model, bool has_custom_kernel = false)
32	- : ast_(ast), task_code_builder_list_(task_code_builders), codegen_model_(codegen_model) {}	32	+ : ast_(ast),
		33	+ task_code_builder_list_(task_code_builders),
		34	+ codegen_model_(codegen_model),
		35	+ has_custom_kernel_(has_custom_kernel) {}
33	Status GenerateProgram(Om2CodePrinter &code_printer);	36	Status GenerateProgram(Om2CodePrinter &code_printer);
34		37
35	private:	38	private:
	@@ -45,6 +48,7 @@ class ProgramGenerator {
45	std::vector<TaskCodeBuilderPtr> task_code_builder_list_;	48	std::vector<TaskCodeBuilderPtr> task_code_builder_list_;
46	uint64_t args_table_index_ = 0U;	49	uint64_t args_table_index_ = 0U;
47	Om2CodegenModel codegen_model_;	50	Om2CodegenModel codegen_model_;
		51	+ bool has_custom_kernel_ = false;
48	};	52	};
49	} // namespace ge	53	} // namespace ge
50		54

Abase/common/om2/codegen/task_code_builder/fe/custom_task_code_builder.cc+252-0

@@ -0,0 +1,252 @@
	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 "custom_task_code_builder.h"
	12	+#include "common/om2/codegen/om2_model_utils.h"
	13	+#include "common/om2/codegen/task_code_builder_factory.h"
	14	+#include "common/om2/codegen/task_code_builder/task_code_builder_util.h"
	15	+#include "opskernel/ops_kernel_info_types.h"
	16	+#include "graph/utils/args_format_desc_utils.h"
	17	+#include "graph/args_format_desc.h"
	18	+#include "common/checker.h"
	19	+
	20	+namespace ge {
	21	+namespace {
	22	+constexpr uint32_t kAddressLen = static_cast<uint32_t>(sizeof(uint64_t));
	23	+} // namespace
	24	+
	25	+int64_t CustomTaskCodeBuilder::ParseOpIndex(const domi::TaskDef &task_def) {
	26	+ const domi::KernelDef &kernel_def = task_def.kernel();
	27	+ domi::KernelContext context = kernel_def.context();
	28	+ return static_cast<int64_t>(context.op_index());
	29	+}
	30	+
	31	+Status CustomTaskCodeBuilder::Contribute(TaskSemanticContributeContext &context) {
	32	+ GE_ASSERT_SUCCESS(TaskCodeBuilder::Contribute(context));
	33	+ GE_ASSERT_NOTNULL(context.next_args_table_index);
	34	+ GE_ASSERT_NOTNULL(context.next_host_args_offset);
	35	+ GE_ASSERT_NOTNULL(context.op_desc);
	36	+
	37	+ build_data_.semantic.task_type = context.task_type;
	38	+ build_data_.semantic.kernel_type = static_cast<ccKernelType>(context.task_def.kernel().context().kernel_type());
	39	+
	40	+ GE_ASSERT_SUCCESS(Om2ModelUtils::ResolveWorkspaceAddrs(context, build_data_.semantic.workspace_addrs));
	41	+ GE_ASSERT_SUCCESS(Om2ModelUtils::ResolveInputAddrs(context, build_data_.semantic.input_addrs));
	42	+ GE_ASSERT_SUCCESS(Om2ModelUtils::ResolveOutputAddrs(context, true, build_data_.semantic.output_addrs));
	43	+ AssignTaskLocalIoNames();
	44	+
	45	+ dispatch_type_ = OpDispatchType::DISPATCH_CUSTOM_KERNEL;
	46	+
	47	+ // parse args format
	48	+ std::vector<ArgDesc> arg_descs;
	49	+ domi::KernelContext kernel_context = context.task_def.kernel().context();
	50	+ GE_ASSERT_SUCCESS(ArgsFormatDesc::Parse(context.op_desc, kernel_context.args_format(), arg_descs),
	51	+ "[OM2] Formatted args [%s] parsed failed.", kernel_context.args_format().c_str());
	52	+
	53	+ // calc args size
	54	+ size_t args_size = 0U;
	55	+ for (const auto &arg_desc : arg_descs) {
	56	+ (void)ArgsFormatDesc::GetArgSize(context.op_desc, arg_desc, args_size);
	57	+ }
	58	+ InitArgsTableEntry(context, args_size);
	59	+
	60	+ // update values of context.next_args_table_index and context.next_host_args_offset
	61	+ if (build_data_.semantic.args_table_entry.has_value()) {
	62	+ ++(*context.next_args_table_index);
	63	+ *context.next_host_args_offset +=
	64	+ Om2ModelUtils::ArgsSizeAlign8(static_cast<size_t>(build_data_.semantic.args_table_entry->args_size));
	65	+ }
	66	+
	67	+ // construct build_data_.ordered_args
	68	+ uint64_t current_args_offset = 0U;
	69	+ auto append_args = [this, &current_args_offset](const std::vector<AddrSemantic> &addrs) {
	70	+ for (const auto &addr : addrs) {
	71	+ OpArgDesc arg = TaskCodeBuilderUtil::ConvertAddrDesc(addr);
	72	+ arg.args_offset = current_args_offset;
	73	+ current_args_offset += kAddressLen;
	74	+ build_data_.ordered_args.push_back(std::move(arg));
	75	+ }
	76	+ };
	77	+ append_args(build_data_.semantic.input_addrs);
	78	+ append_args(build_data_.semantic.output_addrs);
	79	+ append_args(build_data_.semantic.workspace_addrs);
	80	+
	81	+ return SUCCESS;
	82	+}
	83	+
	84	+Status CustomTaskCodeBuilder::RenderDistHelper(std::vector<DeclNode *> &items) {
	85	+ auto op = ast_.Var("const TaskDispatchInfo *", "op");
	86	+ auto ctx = ast_.Var("const DispatchOpContext &", "ctx");
	87	+ GE_ASSERT_SUCCESS(RenderDispatchCustomKernel(op, ctx, items));
	88	+ return SUCCESS;
	89	+}
	90	+
	91	+Status CustomTaskCodeBuilder::RenderOpDefTableFields(std::vector<std::pair<std::string, Arg>> &fields) {
	92	+ GELOGI("[OM2] BuildOpDefTable: op=%s, func_idx=%u", header_.op_name.c_str(),
	93	+ build_data_.semantic.launch.func_handle_index);
	94	+ fields.push_back({"dispatch_type", ast_.StaticCast("OpDispatchType", static_cast<int64_t>(dispatch_type_))});
	95	+ fields.push_back({"op_name", Arg::StringLiteral(header_.op_name)});
	96	+
	97	+ auto custom_fields = std::vector<std::pair<std::string, Arg>>{
	98	+ {"args_info", TaskCodeBuilderUtil::RenderOpArgDesc(ast_, build_data_.ordered_args)},
	99	+ {"args_info_num", static_cast<int64_t>(build_data_.ordered_args.size())},
	100	+ {"op_type", Arg::StringLiteral(header_.op_type)},
	101	+ {"args_idx", static_cast<int64_t>(build_data_.semantic.args_table_entry->table_index)},
	102	+ {"stream_id", static_cast<uint32_t>(header_.stream_id)},
	103	+ {"task_type", static_cast<int64_t>(build_data_.semantic.task_type)},
	104	+ };
	105	+ auto custom_dispatch = ast_.DesignatedInit({{"custom", ast_.DesignatedInit(custom_fields)}});
	106	+ fields.emplace_back("dispatch_info", custom_dispatch);
	107	+
	108	+ return SUCCESS;
	109	+}
	110	+
	111	+std::string CustomTaskCodeBuilder::GetFuncName() const {
	112	+ return kDispatchFuncName;
	113	+}
	114	+
	115	+void CustomTaskCodeBuilder::AssignTaskLocalIoNames() {
	116	+ const std::string task_prefix = "op" + std::to_string(header_.op_index);
	117	+ for (size_t i = 0U; i < build_data_.semantic.input_addrs.size(); ++i) {
	118	+ if (build_data_.semantic.input_addrs[i].tensor_info.has_value()) {
	119	+ build_data_.semantic.input_addrs[i].symbol_hint = task_prefix + "_input" + std::to_string(i);
	120	+ }
	121	+ }
	122	+ for (size_t i = 0U; i < build_data_.semantic.output_addrs.size(); ++i) {
	123	+ if (build_data_.semantic.output_addrs[i].tensor_info.has_value()) {
	124	+ build_data_.semantic.output_addrs[i].symbol_hint = task_prefix + "_output" + std::to_string(i);
	125	+ }
	126	+ }
	127	+}
	128	+
	129	+void CustomTaskCodeBuilder::InitArgsTableEntry(const TaskSemanticContributeContext &context, const uint32_t args_size) {
	130	+ (void)build_data_.semantic.args_table_entry.emplace();
	131	+ build_data_.semantic.args_table_entry->table_index = *context.next_args_table_index;
	132	+ build_data_.semantic.args_table_entry->args_size = args_size;
	133	+ build_data_.semantic.args_table_entry->host_offset = *context.next_host_args_offset;
	134	+ args_table_entry_ = &(*build_data_.semantic.args_table_entry);
	135	+}
	136	+
	137	+Status CustomTaskCodeBuilder::RenderDispatchCustomKernel(const VarRef &op, const VarRef &ctx,
	138	+ std::vector<DeclNode *> &items) {
	139	+ std::vector<BodyItem> body;
	140	+ auto setup = RenderDispatchSetup(op, ctx);
	141	+ body.insert(body.end(), setup.begin(), setup.end());
	142	+ body.push_back(RenderDispatchLoop(op, ctx));
	143	+ auto distribution = RenderDistribution(op, ctx);
	144	+ body.insert(body.end(), distribution.begin(), distribution.end());
	145	+ return TaskCodeBuilderUtil::RenderDispatchFunc(ast_, "DispatchCustomKernel", body, items);
	146	+}
	147	+
	148	+std::vector<BodyItem> CustomTaskCodeBuilder::RenderDispatchSetup(const VarRef &op, const VarRef &ctx) {
	149	+ return {
	150	+ ast_.VarDecl(
	151	+ ast_.Var("ArgsInfo *", "args_info"),
	152	+ ctx.Attr("args_table").Attr("GetArgsInfo")(op.Arrow("dispatch_info").Attr("custom").Attr("args_idx"))),
	153	+ ChkNotNull(ast_.Var("", "args_info")),
	154	+ // -- 声明 ordered_io_addrs 和 Report IO 向量 --
	155	+ ast_.VarDecl(ast_.Var("std::vector<Om2Tensor>", "io_tensors")),
	156	+ ast_.VarDecl(ast_.Var("std::vector<Om2Tensor>", "input_tensors")),
	157	+ ast_.VarDecl(ast_.Var("std::vector<Om2Tensor>", "output_tensors")),
	158	+ ast_.Call(
	159	+ "",
	160	+ {ast_.Var("", "io_tensors").Attr("reserve")(op.Arrow("dispatch_info").Attr("custom").Attr("args_info_num"))}),
	161	+ ast_.VarDecl(ast_.Var("std::vector<Om2TaskIoEntry>", "report_inputs")),
	162	+ ast_.VarDecl(ast_.Var("std::vector<Om2TaskIoEntry>", "report_outputs")),
	163	+ ast_.VarDecl(ast_.Var("std::vector<uint64_t>", "report_workspace_addrs")),
	164	+ ast_.VarDecl(ast_.Var("std::vector<uint64_t>", "report_workspace_sizes")),
	165	+ };
	166	+}
	167	+
	168	+BodyItem CustomTaskCodeBuilder::RenderDispatchLoop(const VarRef &op, const VarRef &ctx) {
	169	+ auto a = ast_.Var("const auto &", "a");
	170	+ return ast_.For(ast_.VarDecl("uint32_t", "j", ast_.UInt(0)),
	171	+ ast_.Var("", "j") < op.Arrow("dispatch_info").Attr("custom").Attr("args_info_num"),
	172	+ ast_.PostInc(ast_.Var("", "j")),
	173	+ std::initializer_list<BodyItem>{
	174	+ ast_.VarDecl(a, op.Arrow("dispatch_info").Attr("custom").Attr("args_info")[ast_.Var("", "j")]),
	175	+ ast_.VarDecl(ast_.Var("uint64_t", "_addr"), ast_.UInt(0)),
	176	+ ast_.Switch(ast_.Var("", "a").Attr("type"),
	177	+ std::vector<BodyItem>{
	178	+ // INPUT / OUTPUT / CONST_TENSOR → 共享 handler（内部根据 a.type 区分）
	179	+ ast_.Case(ast_.Var("", "OP_ARG_INPUT")),
	180	+ ast_.Case(ast_.Var("", "OP_ARG_OUTPUT")),
	181	+ ast_.Case(ast_.Var("", "OP_ARG_CONST_TENSOR")),
	182	+ ast_.Block(HandleInputOutputArg(a, ctx)),
	183	+ ast_.Case(Arg(nullptr)),
	184	+ ast_.Block({
	185	+ ast_.Break(),
	186	+ }),
	187	+ }),
	188	+ });
	189	+}
	190	+
	191	+std::vector<BodyItem> CustomTaskCodeBuilder::RenderDistribution(const VarRef &op, const VarRef &ctx) {
	192	+ auto custom = op.Arrow("dispatch_info").Attr("custom");
	193	+ auto dispatch_type = ast_.StaticCast("uint32_t", op.Arrow("dispatch_type"));
	194	+ auto stream = ctx.Attr("stream_list")[custom.Attr("stream_id")];
	195	+
	196	+ return {
	197	+ ChkStatus(ast_.Call(
	198	+ "ReportOm2TaskPreprocess",
	199	+ {op.Arrow("op_name"), custom.Attr("op_type"),
	200	+ ast_.UInt(0), // op_desc_id
	201	+ ast_.ReinterpretCast("uintptr_t", ast_.Var("", "args_info").Arrow("dev_addr")),
	202	+ ast_.Var("", "args_info").Arrow("size"), ast_.Var("", "report_inputs"), ast_.Var("", "report_outputs"),
	203	+ ast_.Var("", "report_workspace_addrs"), ast_.Var("", "report_workspace_sizes"), dispatch_type,
	204	+ ast_.Var("", "0"), stream, ast_.Var("", "nullptr"), ctx.Attr("model_id"), ctx.Attr("instance_handle")})),
	205	+ ast_.VarDecl(ast_.Var("uint64_t", "_launch_begin"), ast_.Call("MsprofSysCycleTime", {})),
	206	+ ChkStatus(ast_.Call("KernelCustTaskDistribute",
	207	+ {ast_.Var("", "op->op_name"), ast_.Var("", "op->dispatch_info.custom.op_type"),
	208	+ ast_.Var("", "input_tensors"), ast_.Var("", "output_tensors"), stream})),
	209	+ ChkStatus(ast_.Call(
	210	+ "ReportLaunchedOm2Task",
	211	+ {op.Arrow("op_name"), custom.Attr("op_type"),
	212	+ ast_.UInt(0), // op_desc_id
	213	+ ast_.ReinterpretCast("uintptr_t", ast_.Var("", "args_info").Arrow("dev_addr")),
	214	+ ast_.Var("", "args_info").Arrow("size"), ast_.Var("", "report_inputs").Data(),
	215	+ ast_.StaticCast("uint64_t", ast_.Var("", "report_inputs").Size()), ast_.Var("", "report_outputs").Data(),
	216	+ ast_.StaticCast("uint32_t", ast_.Var("", "report_outputs").Size()),
	217	+ ast_.Var("", "report_workspace_addrs").Data(), ast_.Var("", "report_workspace_sizes").Data(),
	218	+ ast_.StaticCast("uint32_t", ast_.Var("", "report_workspace_sizes").Size()), dispatch_type, ast_.Var("", "0"),
	atomgit-bot7 天前 🟡 Medium Priority 建议：将 203 行与 218 行的 `dispatch_type` 实参替换为 `custom.Attr("task_type")`，并删除 193 行的 `dispatch_type` 局部变量；`task_type` 形参应传 ModelTaskType，而非 OpDispatchType。
	219	+ stream, ctx.Attr("model_id"), ctx.Attr("instance_handle"), ast_.UInt(0U),
	220	+ ast_.Var("uint64_t", "_launch_begin")})),
	221	+ };
	222	+}
	223	+
	224	+std::vector<BodyItem> CustomTaskCodeBuilder::HandleInputOutputArg(const VarRef &a, const VarRef &ctx) {
	225	+ return {
	226	+ ast_.Assign(
	227	+ ast_.Var("", "_addr"),
	228	+ ast_.ReinterpretCast("uint64_t",
	229	+ ast_.Call("ResolveOpAddr", {a.Attr("addr").Attr("mem_src"), a.Attr("addr").Attr("index"),
	230	+ a.Attr("addr").Attr("offset"), ctx.Attr("total_dev_mem_ptr"),
	231	+ ctx.Attr("session_scope_mem_ptr"), ctx.Attr("constants"),
	232	+ ctx.Attr("var_addrs")}))),
	233	+ ast_.Var("", "io_tensors")
	234	+ .PushBack(ast_.Call(
	235	+ "BuildOm2Tensor",
	236	+ {ast_.ReinterpretCast("void *", ast_.Var("", "_addr")), a.Attr("data").Attr("tensor").Attr("size"),
	237	+ a.Attr("data").Attr("tensor").Attr("data_type"), a.Attr("data").Attr("tensor").Attr("format"),
	238	+ a.Attr("data").Attr("tensor").Attr("shape"), a.Attr("data").Attr("tensor").Attr("shape_dims")})),
	239	+ ast_.VarDecl(ast_.Var("Om2TaskIoEntry", "_entry"),
	240	+ ast_.InitList({ast_.Var("", "io_tensors").Attr("back")().Addr(),
	241	+ a.Attr("data").Attr("tensor").Attr("args_offset")})),
	242	+ ast_.If(a.Attr("type") == ast_.Var("", "OP_ARG_INPUT") \|\| a.Attr("type") == ast_.Var("", "OP_ARG_CONST_TENSOR"),
	243	+ {ast_.Var("", "report_inputs").PushBack(ast_.Var("", "_entry")),
	244	+ ast_.Var("", "input_tensors").PushBack(ast_.Var("", "io_tensors.back()"))},
	245	+ {ast_.Var("", "report_outputs").PushBack(ast_.Var("", "_entry")),
	246	+ ast_.Var("", "output_tensors").PushBack(ast_.Var("", "io_tensors.back()"))}),
	247	+ ast_.Break(),
	248	+ };
	249	+}
	250	+
	251	+REGISTER_TASK_CODE_BUILDER(MODEL_TASK_CUSTOM_KERNEL, CustomTaskCodeBuilder);
	252	+} // namespace ge

Abase/common/om2/codegen/task_code_builder/fe/custom_task_code_builder.h+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	+#ifndef AIR_CXX_BASE_COMMON_OM2_CODEGEN_TASK_CODE_GENERATOR_FE_CUSTOM_TASK_CODE_GENERATOR_H_
	12	+#define AIR_CXX_BASE_COMMON_OM2_CODEGEN_TASK_CODE_GENERATOR_FE_CUSTOM_TASK_CODE_GENERATOR_H_
	13	+
	14	+#include "common/om2/codegen/task_code_builder/task_code_builder.h"
	15	+
	16	+namespace ge {
	17	+struct CustomBuildData {
	18	+ std::vector<OpArgDesc> ordered_args;
	19	+ KernelTaskSemantic semantic{};
	20	+};
	21	+
	22	+class CustomTaskCodeBuilder : public TaskCodeBuilder {
	23	+ static constexpr const char *kDispatchFuncName = "DispatchCustomKernel";
	24	+
	25	+ public:
	26	+ explicit CustomTaskCodeBuilder(AstBuildContext &ast) : TaskCodeBuilder(ast) {}
	27	+
	28	+ // ── Public overrides & accessors ──
	29	+ int64_t ParseOpIndex(const domi::TaskDef &task_def) override;
	30	+ Status Contribute(TaskSemanticContributeContext &context) override;
	31	+ Status RenderDistHelper(std::vector<DeclNode *> &items) override;
	32	+ Status RenderOpDefTableFields(std::vector<std::pair<std::string, Arg>> &fields) override;
	33	+ std::string GetFuncName() const override;
	34	+
	35	+ private:
	36	+ // ── Build data assembly ──
	37	+ Status RenderDispatchCustomKernel(const VarRef &op, const VarRef &ctx, std::vector<DeclNode *> &items);
	38	+ std::vector<BodyItem> RenderDispatchSetup(const VarRef &op, const VarRef &ctx);
	39	+ BodyItem RenderDispatchLoop(const VarRef &op, const VarRef &ctx);
	40	+ std::vector<BodyItem> RenderDistribution(const VarRef &op, const VarRef &ctx);
	41	+ std::vector<BodyItem> HandleInputOutputArg(const VarRef &a, const VarRef &ctx);
	42	+ void AssignTaskLocalIoNames();
	43	+ void InitArgsTableEntry(const TaskSemanticContributeContext &context, const uint32_t args_size);
	44	+
	45	+ // ── Member variables ──
	46	+ CustomBuildData build_data_;
	47	+ OpDispatchType::Value dispatch_type_{OpDispatchType::DISPATCH_CUSTOM_KERNEL};
	48	+};
	49	+} // namespace ge
	50	+
	51	+#endif // AIR_CXX_BASE_COMMON_OM2_CODEGEN_TASK_CODE_GENERATOR_FE_CUSTOM_TASK_CODE_GENERATOR_H_

Mbase/common/om2/om2_model_data.h+2-0

	@@ -85,6 +85,8 @@ struct Om2ModelData {
85	Om2ModelMeta model_meta;	85	Om2ModelMeta model_meta;
86	Om2ConstantsData constants_data;	86	Om2ConstantsData constants_data;
87	std::vector<Om2KernelBinary> kernel_binaries;	87	std::vector<Om2KernelBinary> kernel_binaries;
		88	+ std::vector<Om2KernelBinary> custom_kernel_binaries;
		89	+ std::vector<Om2KernelBinary> custom_shared_libs;
88	Om2DebugInfo debug_info;	90	Om2DebugInfo debug_info;
89	std::map<std::string, std::string> manifest;	91	std::map<std::string, std::string> manifest;
90	std::unique_ptr<RTVarResource> rt_var_resource;	92	std::unique_ptr<RTVarResource> rt_var_resource;

Minc/framework/common/helper/om2_package_helper.h+13-12

	@@ -12,9 +12,6 @@
12	#define INC_FRAMEWORK_COMMON_HELPER_OM2_PACKAGE_HELPER_H	12	#define INC_FRAMEWORK_COMMON_HELPER_OM2_PACKAGE_HELPER_H
13		13
14	#include "framework/common/helper/model_save_helper.h"	14	#include "framework/common/helper/model_save_helper.h"
15	-#include "common/om2/codegen/om2_codegen_types.h"
16	-#include <map>
17	-#include <memory>
18	#include <string>	15	#include <string>
19		16
20	namespace gert {	17	namespace gert {
	@@ -48,6 +45,8 @@ class GE_FUNC_VISIBILITY Om2PackageHelper : public ModelSaveHelper {
48		45
49	static Status RelocateExternalWeights(const std::string &output_file_name, const ModelBufferData &model,	46	static Status RelocateExternalWeights(const std::string &output_file_name, const ModelBufferData &model,
50	ModelBufferData &relocated_model, bool &relocated);	47	ModelBufferData &relocated_model, bool &relocated);
		48	+ static Status ReadCustomOpSoToBuffer(const std::unordered_set<std::string> &ops_so_set,
		49	+ std::vector<gert::Om2KernelBinary> &shared_lib_binaries);
51		50
52	/// @brief 从 OM2 ZIP 模型内提取 visual JSON 内容。	51	/// @brief 从 OM2 ZIP 模型内提取 visual JSON 内容。
53	/// @param model_data OM2 ZIP 数据内存地址。	52	/// @param model_data OM2 ZIP 数据内存地址。
	@@ -56,16 +55,18 @@ class GE_FUNC_VISIBILITY Om2PackageHelper : public ModelSaveHelper {
56	static Status ExtractVisualJson(const void *model_data, size_t model_len, std::string &json_out);	55	static Status ExtractVisualJson(const void *model_data, size_t model_len, std::string &json_out);
57		56
58	private:	57	private:
59	- static Status BuildProgramBody(const GeModelPtr &ge_model, gert::Om2ProgramBody &body,	58	+ static Status BuildProgramBody(const GeModelPtr &ge_model, gert::Om2ModelData &model_data);
60	- std::vector<Om2ConstMeta> &const_metas, std::vector<Om2VarMeta> &var_metas);	59	+ static Status BuildKernelBinaries(const GeModelPtr &ge_model, gert::Om2ModelData &model_data);
61	- static Status BuildKernelBinaries(const GeModelPtr &ge_model, std::vector<gert::Om2KernelBinary> &kernel_binaries);	60	+ static Status BuildModelMeta(const GeModelPtr &ge_model, gert::Om2ModelData &model_data);
62	- static Status BuildModelMeta(const GeModelPtr &ge_model, gert::Om2ModelMeta &model_meta);	61	+ static Status BuildConstantsData(const GeModelPtr &ge_model, gert::Om2ModelData &model_data);
63	- static Status BuildConstantsData(const GeModelPtr &ge_model, const std::vector<Om2ConstMeta> &const_metas,	62	+ static Status BuildDebugInfo(const GeModelPtr &ge_model, gert::Om2ModelData &model_data);
64	- gert::Om2ConstantsData &data);	63	+ static Status BuildManifest(const GeRootModelPtr &ge_root_model, gert::Om2ModelData &model_data);
65	- static Status BuildDebugInfo(const GeModelPtr &ge_model, gert::Om2DebugInfo &debug_info);	64	+
66	- static Status BuildManifest(const GeRootModelPtr &ge_root_model, std::map<std::string, std::string> &manifest);	65	+ static Status CollectUsedCustomOpTypes(const GeRootModelPtr &ge_root_model,
		66	+ std::set<std::string> &used_custom_op_types);
		67	+ static Status BuildCustomKernelBinaries(const GeRootModelPtr &ge_root_model, gert::Om2ModelData &model_data);
		68	+ static Status BuildCustomSharedLibs(const GeRootModelPtr &ge_root_model, gert::Om2ModelData &model_data);
67		69
68	- private:
69	bool is_offline_{true};	70	bool is_offline_{true};
70	};	71	};
71	} // namespace ge	72	} // namespace ge

Mruntime/om2/om2_model_executor.cc+47-4

	@@ -25,6 +25,7 @@
25	#include "mmpa/mmpa_api.h"	25	#include "mmpa/mmpa_api.h"
26	#include "../../inc/framework/runtime/om2_context.h"	26	#include "../../inc/framework/runtime/om2_context.h"
27	#include "graph/utils/type_utils_inner.h"	27	#include "graph/utils/type_utils_inner.h"
		28	+#include "graph/custom_op_factory.h"
28	#include "graph_metadef/common/ge_common/util.h"	29	#include "graph_metadef/common/ge_common/util.h"
29	#include "rt_external_mem.h"	30	#include "rt_external_mem.h"
30	#include "common/helper/om2/json_file.h"	31	#include "common/helper/om2/json_file.h"
	@@ -39,6 +40,8 @@
39	#include "zip_archive_reader.h"	40	#include "zip_archive_reader.h"
40	#include "common/om2/om2_model_data.h"	41	#include "common/om2/om2_model_data.h"
41	#include "om2_aipp_utils.h"	42	#include "om2_aipp_utils.h"
		43	+#include <fstream>
		44	+#include <vector>
42		45
43	namespace gert {	46	namespace gert {
44	namespace {	47	namespace {
	@@ -54,6 +57,12 @@ using DestroyFunc = ge::graphStatus ()(Om2ModelHandle );
54	using RunFunc = ge::graphStatus ()(Om2ModelHandle , int, void , int, void , int32_t, Om2ProfInfos *);	57	using RunFunc = ge::graphStatus ()(Om2ModelHandle , int, void , int, void , int32_t, Om2ProfInfos *);
55	using RunAsyncFunc = ge::graphStatus ()(Om2ModelHandle , rtStream_t, int, void , int, void , Om2ProfInfos *);	58	using RunAsyncFunc = ge::graphStatus ()(Om2ModelHandle , rtStream_t, int, void , int, void , Om2ProfInfos *);
56		59
		60	+struct CustSharedLibInfo {
		61	+ std::string so_file;
		62	+ int32_t so_fd = -1;
		63	+ void *so_handle = nullptr;
		64	+};
		65	+
57	struct RunModelInfo {	66	struct RunModelInfo {
58	std::string so_file;	67	std::string so_file;
59	int32_t so_fd = -1;	68	int32_t so_fd = -1;
	@@ -61,6 +70,7 @@ struct RunModelInfo {
61	void *so_handle = nullptr;	70	void *so_handle = nullptr;
62	std::string model_name;	71	std::string model_name;
63	std::string root_graph_name;	72	std::string root_graph_name;
		73	+ std::vector<CustSharedLibInfo> cust_shared_libs;
64	Om2ModelHandle model_handle = nullptr;	74	Om2ModelHandle model_handle = nullptr;
65	rtModel_t rt_model_handle = nullptr;	75	rtModel_t rt_model_handle = nullptr;
66	CreateFunc create_func = nullptr;	76	CreateFunc create_func = nullptr;
	@@ -389,8 +399,8 @@ ge::Status DeserializeVariablesConfigEntry(const ge::RAIIZipArchive &archive, co
389	return ge::SUCCESS;	399	return ge::SUCCESS;
390	}	400	}
391		401
392	-ge::Status DeserializeKernelEntry(const ge::RAIIZipArchive &archive, const std::string &entry,	402	+ge::Status DeserializeBinaryEntry(const ge::RAIIZipArchive &archive, const std::string &entry,
393	- gert::Om2ModelData &model_data) {	403	+ std::vector<Om2KernelBinary> &binaries) {
394	gert::Om2KernelBinary kernel_binary;	404	gert::Om2KernelBinary kernel_binary;
395	kernel_binary.name = ExtractParentDirAndFileName(entry).second;	405	kernel_binary.name = ExtractParentDirAndFileName(entry).second;
396	size_t buffer_size{0U};	406	size_t buffer_size{0U};
	@@ -399,10 +409,15 @@ ge::Status DeserializeKernelEntry(const ge::RAIIZipArchive &archive, const std::
399	GE_ASSERT_TRUE(buffer_size > 0U, "[OM2] Empty archive entry %s", entry.c_str());	409	GE_ASSERT_TRUE(buffer_size > 0U, "[OM2] Empty archive entry %s", entry.c_str());
400	kernel_binary.data = std::move(buffer);	410	kernel_binary.data = std::move(buffer);
401	kernel_binary.data_size = buffer_size;	411	kernel_binary.data_size = buffer_size;
402	- (void)model_data.kernel_binaries.emplace_back(std::move(kernel_binary));	412	+ (void)binaries.emplace_back(std::move(kernel_binary));
403	return ge::SUCCESS;	413	return ge::SUCCESS;
404	}	414	}
405		415
		416	+ge::Status DeserializeKernelEntry(const ge::RAIIZipArchive &archive, const std::string &entry,
		417	+ gert::Om2ModelData &model_data) {
		418	+ return DeserializeBinaryEntry(archive, entry, model_data.kernel_binaries);
		419	+}
		420	+
406	ge::Status DeserializeModelMetaEntry(const ge::RAIIZipArchive &archive, const std::string &entry,	421	ge::Status DeserializeModelMetaEntry(const ge::RAIIZipArchive &archive, const std::string &entry,
407	gert::Om2ModelData &model_data) {	422	gert::Om2ModelData &model_data) {
408	size_t buff_size = 0U;	423	size_t buff_size = 0U;
	@@ -512,6 +527,15 @@ ge::Status HandleArchiveEntry(const ge::RAIIZipArchive &archive, const std::stri
512	}	527	}
513	if (entry.find("data/kernels_") != std::string::npos && IsFileNameEndsWith(entry, ".o")) {	528	if (entry.find("data/kernels_") != std::string::npos && IsFileNameEndsWith(entry, ".o")) {
514	GE_ASSERT_SUCCESS(DeserializeKernelEntry(archive, entry, model_data));	529	GE_ASSERT_SUCCESS(DeserializeKernelEntry(archive, entry, model_data));
		530	+ return ge::SUCCESS;
		531	+ }
		532	+ if (entry.find("data/custom_ops/shared_libs") != std::string::npos && IsFileNameEndsWith(entry, ".so")) {
		533	+ GE_ASSERT_SUCCESS(DeserializeBinaryEntry(archive, entry, model_data.custom_shared_libs));
		534	+ return ge::SUCCESS;
		535	+ }
		536	+ if (entry.find("data/custom_ops/binaries_") != std::string::npos && IsFileNameEndsWith(entry, ".bin")) {
		537	+ GE_ASSERT_SUCCESS(DeserializeBinaryEntry(archive, entry, model_data.custom_kernel_binaries));
		538	+ return ge::SUCCESS;
515	}	539	}
516	return ge::SUCCESS;	540	return ge::SUCCESS;
517	}	541	}
	@@ -743,6 +767,26 @@ class Om2ModelExecutor::Impl {
743	has_model_ = true;	767	has_model_ = true;
744		768
745	GE_ASSERT_SUCCESS(BuildKernelBinInfoFromStruct(om2_data.kernel_binaries, kernel_bin_info));	769	GE_ASSERT_SUCCESS(BuildKernelBinInfoFromStruct(om2_data.kernel_binaries, kernel_bin_info));
		770	+
		771	+ // Add custom kernel binaries
		772	+ GE_ASSERT_SUCCESS(BuildKernelBinInfoFromStruct(om2_data.custom_kernel_binaries, kernel_bin_info));
		773	+
		774	+ // dlopen custom kernel shared libraries
		775	+ for (auto &kb : om2_data.custom_shared_libs) {
		776	+ CustSharedLibInfo so_info;
		777	+ GE_ASSERT_SUCCESS(CreateSoMemFd(kb.name, kb.data.get(), kb.data_size, so_info.so_file, so_info.so_fd));
		778	+ so_info.so_handle = mmDlopen(so_info.so_file.c_str(), MMPA_RTLD_NOW);
		779	+ if (so_info.so_handle == nullptr) {
		780	+ CloseMemFd(so_info.so_fd);
		781	+ const char_t *error = mmDlerror();
		782	+ error = (error == nullptr) ? "" : error;
		783	+ GELOGE(ge::FAILED, "[OM2][Invoke][DlOpen] Failed to load so, path = [%s], error = [%s]",
		784	+ so_info.so_file.c_str(), error);
		785	+ return ge::FAILED;
		786	+ }
		atomgit-bot7 天前 🟡 Medium Priority 变更行 `runtime/om2/om2_model_executor.cc:777-785`：`CreateSoMemFd(...)` 成功返回后（此时 `so_info.so_fd` 已是一个打开的 memfd，且 `CreateSoMemFd` 内部的 `GE_DISMISS_GUARD` 已被解除），紧接着 `mmDlopen` 若返回 `nullptr`，代码直接 `return ge::FAILED`。失败模式：`so_info.so_fd` 指向的 memfd 未调用 `CloseMemFd` 即被丢弃，加载失败一次泄漏一个 fd。多次加载失败同样会累积 fd。 F fary86 7 天前评论： 7 天前评论：
		787	+ run_model_info_.cust_shared_libs.emplace_back(so_info);
		788	+ }
		atomgit-bot7 天前 🟡 Medium Priority 变更行（`runtime/om2/om2_model_executor.cc:73` 新增 `RunModelInfo::cust_shared_libs`，并在 `:775-787` 对每个自定义 so `mmDlopen` 后 `emplace_back` 到该向量）→ 这些 `CustSharedLibInfo` 中的 `so_handle`（`mmDlopen` 结果）和 `so_fd`（`CreateSoMemFd` 创建的 memfd）在整个生命周期内没有任何释放路径。证据链：失败模式：每次模型卸载或重复 `Load` 都会泄漏 N 个 dlopen 句柄和 N 个文件描述符（N = 自定义 so 数量）。长期反复加载/卸载会累积，最终导致 fd 耗尽（进程无法再创建 memfd / 打开文件）或共享库引用计数永不下降。这是确定的资源泄漏，无任何清理分支。建议：在 Cleanup() 中释放 cust_shared_libs 的 so_handle 与 so_fd，并在 LoadFromOm2ModelData 重建 run_model_info_ 前先释放旧值。例如在 Cleanup() 末尾追加：for (auto &lib : run_model_info_.cust_shared_libs) { if (lib.so_handle) { (void)mmDlclose(lib.so_handle); } CloseMemFd(lib.so_fd); } run_model_info_.cust_shared_libs.clear(); F fary86 7 天前评论： 7 天前评论：
		789	+
746	GE_ASSERT_SUCCESS(LoadSoFromBuffer(om2_data.program_body.so_artifact));	790	GE_ASSERT_SUCCESS(LoadSoFromBuffer(om2_data.program_body.so_artifact));
747	GE_ASSERT_TRUE(!run_model_info_.so_file.empty(), "[OM2] Om2 compiled so not found in Om2ModelData.");	791	GE_ASSERT_TRUE(!run_model_info_.so_file.empty(), "[OM2] Om2 compiled so not found in Om2ModelData.");
748		792
	@@ -785,7 +829,6 @@ class Om2ModelExecutor::Impl {
785		829
786	ge::Status BuildKernelBinInfoFromStruct(const std::vector<gert::Om2KernelBinary> &kernels,	830	ge::Status BuildKernelBinInfoFromStruct(const std::vector<gert::Om2KernelBinary> &kernels,
787	std::vector<KernelBinInfo> &info) {	831	std::vector<KernelBinInfo> &info) {
788	- info.clear();
789	for (const auto &k : kernels) {	832	for (const auto &k : kernels) {
790	KernelBinInfo bin_info;	833	KernelBinInfo bin_info;
791	bin_info.file = k.name;	834	bin_info.file = k.name;

Atests/framework/ge_runtime_stub/include/faker/custom_taskdef_faker.h+36-0

@@ -0,0 +1,36 @@
	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 INC_CUSTOM_TASKDEF_FAKER_H
	12	+#define INC_CUSTOM_TASKDEF_FAKER_H
	13	+
	14	+#include "task_def_faker.h"
	15	+namespace gert {
	16	+
	17	+struct CustomTaskDefFaker : public TaskDefFaker {
	18	+ CustomTaskDefFaker(std::string stub_name = "");
	19	+ CustomTaskDefFaker &BinData(uint64_t data);
	20	+ CustomTaskDefFaker &ArgsFormat(const std::string &args_format);
	21	+
	22	+ private:
	23	+ vector<domi::TaskDef> CreateTaskDef(uint64_t op_index = 0) override;
	24	+ std::unique_ptr<TaskDefFaker> Clone() const override;
	25	+
	26	+ private:
	27	+ void Init();
	28	+ bool inited_;
	29	+ uint64_t bin_data = 0;
	30	+ std::string stub_name_;
	31	+ std::string args_format_;
	32	+};
	33	+
	34	+} // namespace gert
	35	+
	36	+#endif

Mtests/framework/ge_runtime_stub/include/faker/task_def_faker.h+1-0

	@@ -28,6 +28,7 @@ class TaskDefFaker {
28	kRts, // ModelTaskType::MODEL_TASK_MEMCPY_ASYNC	28	kRts, // ModelTaskType::MODEL_TASK_MEMCPY_ASYNC
29	kLabelSwitch, // ModelTaskType::MODEL_TASK_STREAM_LABEL_SWITCH_BY_INDEX	29	kLabelSwitch, // ModelTaskType::MODEL_TASK_STREAM_LABEL_SWITCH_BY_INDEX
30	kEvent, // ModelTaskType::MODEL_TASK_EVENT_RECORD	30	kEvent, // ModelTaskType::MODEL_TASK_EVENT_RECORD
		31	+ kCustom, // ModelTaskType::MODEL_TASK_CUSTOM_KERNEL
31	kTaskTypeEnd	32	kTaskTypeEnd
32	};	33	};
33		34

Atests/framework/ge_runtime_stub/src/custom_taskdef_faker.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 "faker/custom_taskdef_faker.h"
	12	+#include "faker/task_def_faker.h"
	13	+#include "framework/common/debug/ge_log.h"
	14	+namespace gert {
	15	+
	16	+CustomTaskDefFaker::CustomTaskDefFaker(std::string stub_name) : inited_(false), stub_name_(stub_name) {}
	17	+
	18	+vector<domi::TaskDef> CustomTaskDefFaker::CreateTaskDef(uint64_t op_index) {
	19	+ Init();
	20	+ auto task_def = TaskDefFaker::CreateTaskDef(op_index);
	21	+ task_def[0].mutable_kernel()->set_stub_func(stub_name_);
	22	+ if (!args_format_.empty()) {
	23	+ task_def[0].mutable_kernel()->mutable_context()->set_args_format(args_format_);
	24	+ }
	25	+ GELOGD("CreateTaskDef size:%zu.", task_def.size());
	26	+ return task_def;
	27	+}
	28	+
	29	+std::unique_ptr<TaskDefFaker> CustomTaskDefFaker::Clone() const {
	30	+ return std::unique_ptr<CustomTaskDefFaker>(new CustomTaskDefFaker(*this));
	31	+}
	32	+
	33	+CustomTaskDefFaker &CustomTaskDefFaker::BinData(uint64_t data) {
	34	+ bin_data = data;
	35	+ return *this;
	36	+}
	37	+
	38	+CustomTaskDefFaker &CustomTaskDefFaker::ArgsFormat(const std::string &args_format) {
	39	+ args_format_ = args_format;
	40	+ return *this;
	41	+}
	42	+
	43	+void CustomTaskDefFaker::Init() {
	44	+ if (inited_) {
	45	+ return;
	46	+ }
	47	+ AddTask({kCustom, kTE_AiCore, bin_data});
	48	+ inited_ = true;
	49	+}
	50	+
	51	+}; // namespace gert

Mtests/framework/ge_runtime_stub/src/task_def_faker.cc+2-1

	@@ -30,6 +30,7 @@ std::array<ge::ModelTaskType, TaskDefFaker::kTaskTypeEnd> task_types_map = {
30	ge::ModelTaskType::MODEL_TASK_MEMCPY_ASYNC, // kRts	30	ge::ModelTaskType::MODEL_TASK_MEMCPY_ASYNC, // kRts
31	ge::ModelTaskType::MODEL_TASK_STREAM_LABEL_SWITCH_BY_INDEX, // kLabelSwitch	31	ge::ModelTaskType::MODEL_TASK_STREAM_LABEL_SWITCH_BY_INDEX, // kLabelSwitch
32	ge::ModelTaskType::MODEL_TASK_EVENT_RECORD, // kEvent	32	ge::ModelTaskType::MODEL_TASK_EVENT_RECORD, // kEvent
		33	+ ge::ModelTaskType::MODEL_TASK_CUSTOM_KERNEL, // kCustom
33	};	34	};
34	} // namespace	35	} // namespace
35		36
	@@ -53,7 +54,7 @@ std::vector<domi::TaskDef> TaskDefFaker::CreateTaskDef(uint64_t op_index) {
53	task_defs[i].mutable_kernel_with_handle()->mutable_context()->set_args_offset(&offset, sizeof(offset));	54	task_defs[i].mutable_kernel_with_handle()->mutable_context()->set_args_offset(&offset, sizeof(offset));
54	task_defs[i].mutable_kernel_with_handle()->set_args(addrs, sizeof(addrs));	55	task_defs[i].mutable_kernel_with_handle()->set_args(addrs, sizeof(addrs));
55	task_defs[i].mutable_kernel_with_handle()->set_block_dim(8);	56	task_defs[i].mutable_kernel_with_handle()->set_block_dim(8);
56	- } else if (task_type == kWithoutHandle \|\| task_type == kCCAicpu) {	57	+ } else if (task_type == kWithoutHandle \|\| task_type == kCCAicpu \|\| task_type == kCustom) {
57	task_defs[i].mutable_kernel()->mutable_context()->set_op_index(op_index);	58	task_defs[i].mutable_kernel()->mutable_context()->set_op_index(op_index);
58	task_defs[i].mutable_kernel()->mutable_context()->set_kernel_type(	59	task_defs[i].mutable_kernel()->mutable_context()->set_kernel_type(
59	static_cast<uint32_t>(kernel_types_map[kernel_type]));	60	static_cast<uint32_t>(kernel_types_map[kernel_type]));

Mtests/ge/st/testcase/test_om2.cc+148-0

	@@ -52,7 +52,10 @@
52	#include "ge_runtime_stub/include/faker/ge_model_builder.h"	52	#include "ge_runtime_stub/include/faker/ge_model_builder.h"
53	#include "ge_runtime_stub/include/faker/aicore_taskdef_faker.h"	53	#include "ge_runtime_stub/include/faker/aicore_taskdef_faker.h"
54	#include "ge_runtime_stub/include/faker/aicpu_taskdef_faker.h"	54	#include "ge_runtime_stub/include/faker/aicpu_taskdef_faker.h"
		55	+#include "ge_runtime_stub/include/faker/custom_taskdef_faker.h"
55		56
		57	+#include "graph/custom_op_factory.h"
		58	+#include "graph/custom_op.h"
56	#include <cinttypes>	59	#include <cinttypes>
57	#include <securec.h>	60	#include <securec.h>
58	#include <sys/stat.h>	61	#include <sys/stat.h>
	@@ -1412,6 +1415,99 @@ GeRootModelPtr CreateGeRootModelWithCmoAddrTask(bool with_explicit_format = fals
1412	return ge_root_model;	1415	return ge_root_model;
1413	}	1416	}
1414		1417
		1418	+class TestPortableCustomOp : public PortableOp, public EagerExecuteOp {
		1419	+ public:
		1420	+ graphStatus Execute(gert::EagerOpExecutionContext *ctx) override {
		1421	+ return SUCCESS;
		1422	+ }
		1423	+
		1424	+ graphStatus Serialize(std::vector<uint8_t> &buffer) override {
		1425	+ const std::string payload = "test_portable_custom_op_kernel_bin";
		1426	+ buffer.assign(payload.begin(), payload.end());
		1427	+ return GRAPH_SUCCESS;
		1428	+ }
		1429	+
		1430	+ graphStatus Deserialize(const std::vector<uint8_t> &buffer) override {
		1431	+ return GRAPH_SUCCESS;
		1432	+ }
		1433	+};
		1434	+
		1435	+static ComputeGraphPtr BuildCustomOpGraph() {
		1436	+ auto graph = std::make_shared<ComputeGraph>("custom_op_om2_graph");
		1437	+ GeTensorDesc tensor_desc(GeShape({2, 2, 2}), FORMAT_ND, DT_FLOAT);
		1438	+
		1439	+ auto data0_desc = std::make_shared<OpDesc>("data0", DATA);
		1440	+ (void)data0_desc->AddInputDesc(tensor_desc);
		1441	+ (void)data0_desc->AddOutputDesc(tensor_desc);
		1442	+ AttrUtils::SetInt(data0_desc, ATTR_NAME_INDEX, 0);
		1443	+ auto data0 = graph->AddNode(data0_desc);
		1444	+
		1445	+ auto data1_desc = std::make_shared<OpDesc>("data1", DATA);
		1446	+ (void)data1_desc->AddInputDesc(tensor_desc);
		1447	+ (void)data1_desc->AddOutputDesc(tensor_desc);
		1448	+ AttrUtils::SetInt(data1_desc, ATTR_NAME_INDEX, 1);
		1449	+ auto data1 = graph->AddNode(data1_desc);
		1450	+
		1451	+ auto custom_op_desc = std::make_shared<OpDesc>("custom_op", "TestPortableOp");
		1452	+ (void)custom_op_desc->AddInputDesc("x0", tensor_desc);
		1453	+ (void)custom_op_desc->AddInputDesc("x1", tensor_desc);
		1454	+ (void)custom_op_desc->AddOutputDesc("y", tensor_desc);
		1455	+ custom_op_desc->AppendIrInput("x0", kIrInputRequired);
		1456	+ custom_op_desc->AppendIrInput("x1", kIrInputRequired);
		1457	+ custom_op_desc->AppendIrOutput("y", kIrOutputRequired);
		1458	+ auto custom_op_node = graph->AddNode(custom_op_desc);
		1459	+
		1460	+ auto netoutput_desc = std::make_shared<OpDesc>("netoutput", NETOUTPUT);
		1461	+ (void)netoutput_desc->AddInputDesc(tensor_desc);
		1462	+ auto netoutput = graph->AddNode(netoutput_desc);
		1463	+
		1464	+ GraphUtils::AddEdge(data0->GetOutDataAnchor(0), custom_op_node->GetInDataAnchor(0));
		1465	+ GraphUtils::AddEdge(data1->GetOutDataAnchor(0), custom_op_node->GetInDataAnchor(1));
		1466	+ GraphUtils::AddEdge(custom_op_node->GetOutDataAnchor(0), netoutput->GetInDataAnchor(0));
		1467	+ netoutput_desc->SetSrcName({"custom_op"});
		1468	+ netoutput_desc->SetSrcIndex({0});
		1469	+ graph->TopologicalSorting();
		1470	+ return graph;
		1471	+}
		1472	+
		1473	+GeRootModelPtr CreateGeRootModelWithCustomOp() {
		1474	+ auto graph = BuildCustomOpGraph();
		1475	+
		1476	+ gert::GeModelBuilder builder(graph);
		1477	+ auto ge_root_model =
		1478	+ builder
		1479	+ .AddTaskDef(
		1480	+ "custom_op",
		1481	+ gert::CustomTaskDefFaker("custom_op_stub").ArgsFormat("{i_instance0}{i_instance1}{o_instance0*}"))
		1482	+ .FakeTbeBin({"custom_op"})
		1483	+ .BuildGeRootModel();
		1484	+ auto &compute_graph = ge_root_model->GetRootGraph();
		1485	+ compute_graph->SetGraphUnknownFlag(false);
		1486	+ for (const auto &node : compute_graph->GetDirectNode()) {
		1487	+ auto op_desc = node->GetOpDesc();
		1488	+ if (op_desc == nullptr) {
		1489	+ return nullptr;
		1490	+ }
		1491	+ if (op_desc->GetType() == DATA) {
		1492	+ op_desc->SetOutputOffset({1024});
		1493	+ } else if (op_desc->GetType() == NETOUTPUT) {
		1494	+ op_desc->SetInputOffset({3072});
		1495	+ } else {
		1496	+ op_desc->SetInputOffset(std::vector<int64_t>(op_desc->GetInputsSize(), 1024));
		1497	+ op_desc->SetOutputOffset(std::vector<int64_t>(op_desc->GetOutputsSize(), 1024));
		1498	+ }
		1499	+ }
		1500	+
		1501	+ const auto ge_model = ge_root_model->GetSubgraphInstanceNameToModel().begin()->second;
		1502	+ std::vector<uint8_t> weights_value(512, 1U);
		1503	+ ge_model->SetWeight(Buffer::CopyFrom(weights_value.data(), weights_value.size()));
		1504	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_MEMORY_SIZE, 2048);
		1505	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_WEIGHT_SIZE, weights_value.size());
		1506	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_STREAM_NUM, 1);
		1507	+
		1508	+ return ge_root_model;
		1509	+}
		1510	+
1415	} // namespace	1511	} // namespace
1416		1512
1417	class Om2St : public testing::Test {	1513	class Om2St : public testing::Test {
	@@ -3473,4 +3569,56 @@ TEST_F(Om2VarSt, GenOm2WithoutVarNodes_NoVarResourceFiles) {
3473	}	3569	}
3474	}	3570	}
3475		3571
		3572	+/**
		3573	+ * 用例描述：包含自定义算子的计算图打包为 OM2 模型，验证自定义算子 kernel bin 被序列化写入 OM2 包
		3574	+ *
		3575	+ * 预置条件：
		3576	+ * 1. 注册一个实现 PortableOp 接口的自定义算子 TestPortableOp
		3577	+ * 2. 构造计算图：data0 -> custom_op <- data1, custom_op -> netoutput
		3578	+ *
		3579	+ * 测试步骤：
		3580	+ * 1. 通过 CustomOpFactory 注册 TestPortableOp
		3581	+ * 2. 构造包含 TestPortableOp 节点的 GeRootModel
		3582	+ * 3. 调用 Om2PackageHelper::SaveToOmRootModel 打包为 OM2
		3583	+ * 4. 验证 OM2 包中存在自定义算子 kernel bin 文件
		3584	+ *
		3585	+ * 预期结果：
		3586	+ * 1. OM2 打包成功
		3587	+ * 2. OM2 包中包含 custom_op 的序列化 kernel bin 文件
		3588	+ */
		3589	+TEST_F(Om2St, ConvertOm2Model_Ok_GenOm2WithCustomOp) {
		3590	+ const AscendString kOpType("TestPortableOp");
		3591	+ CustomOpFactory::RegisterCustomOpCreator(
		3592	+ kOpType, []() -> std::unique_ptr<BaseCustomOp> { return std::make_unique<TestPortableCustomOp>(); });
		3593	+
		3594	+ Om2PackageHelper om2_packager;
		3595	+ const auto ge_root_model = CreateGeRootModelWithCustomOp();
		3596	+ ASSERT_NE(ge_root_model, nullptr);
		3597	+ ModelBufferData model_data;
		3598	+ const std::string output_file = PathUtils::Join({test_work_dir, kZipFileBaseName + "_custom_op.om2"});
		3599	+ SyncKernelNameForAllModels(ge_root_model);
		3600	+ ASSERT_EQ(om2_packager.SaveToOmRootModel(ge_root_model, output_file, model_data, false), SUCCESS);
		3601	+ ASSERT_EQ(mmAccess2(output_file.c_str(), M_F_OK), EOK);
		3602	+
		3603	+ uint32_t model_buf_size = 0;
		3604	+ const auto model_buf = GetBinDataFromFile(output_file, model_buf_size);
		3605	+ ASSERT_NE(model_buf, nullptr);
		3606	+ ASSERT_GT(model_buf_size, 0U);
		3607	+
		3608	+ RAIIZipArchive archive(reinterpret_cast<const uint8_t *>(model_buf.get()), model_buf_size);
		3609	+ ASSERT_TRUE(archive.IsGood());
		3610	+ const auto file_names = archive.ListFiles();
		3611	+
		3612	+ bool has_custom_kernel_bin = false;
		3613	+ for (const auto &f : file_names) {
		3614	+ if (f.find("CustomKernel.bin") != std::string::npos) {
		3615	+ has_custom_kernel_bin = true;
		3616	+ break;
		3617	+ }
		3618	+ }
		3619	+ EXPECT_TRUE(has_custom_kernel_bin) << "OM2 archive should contain custom op kernel binary";
		3620	+
		3621	+ CustomOpFactory::RemoveCustomOps({kOpType});
		3622	+}
		3623	+
3476	} // namespace ge	3624	} // namespace ge

Mtests/ge/ut/ge/common/control_task_code_builder_unittest.cc+11-0

	@@ -775,6 +775,7 @@ enum OpDispatchType : uint32_t {
775	DISPATCH_CMO_ADDR = 18,	775	DISPATCH_CMO_ADDR = 18,
776	DISPATCH_DSA = 19,	776	DISPATCH_DSA = 19,
777	DISPATCH_KERNEL_EX = 20,	777	DISPATCH_KERNEL_EX = 20,
		778	+ DISPATCH_CUSTOM_KERNEL = 21,
778	DISPATCH_TYPE_COUNT	779	DISPATCH_TYPE_COUNT
779	};	780	};
780		781
	@@ -886,6 +887,15 @@ struct AicpuDispatchInfo {
886	uint32_t task_type;	887	uint32_t task_type;
887	};	888	};
888		889
		890	+struct CustomDispatchInfo {
		891	+ const OpArgInfo *args_info; // IO 地址解析数组
		892	+ uint32_t args_info_num; // args_info 数组长度
		893	+ const char *op_type; // 算子类型名，用于 Report 上报
		894	+ uint32_t args_idx; // 参数表索引，用于 GetArgsInfo 查找
		895	+ uint32_t stream_id; // 执行流索引
		896	+ uint32_t task_type;
		897	+};
		898	+
889	struct DsaDispatchInfo {	899	struct DsaDispatchInfo {
890	const OpArgInfo *args_info; // IO 地址解析数组	900	const OpArgInfo *args_info; // IO 地址解析数组
891	const char *op_type; // 算子类型名，用于 Report 上报	901	const char *op_type; // 算子类型名，用于 Report 上报
	@@ -981,6 +991,7 @@ struct TaskDispatchInfo {
981	union {	991	union {
982	AicoreDispatchInfo aicore;	992	AicoreDispatchInfo aicore;
983	AicpuDispatchInfo aicpu;	993	AicpuDispatchInfo aicpu;
		994	+ CustomDispatchInfo custom;
984	CmoDispatchInfo cmo;	995	CmoDispatchInfo cmo;
985	MemcpyAsyncDispatchInfo memcpy_async;	996	MemcpyAsyncDispatchInfo memcpy_async;
986	MemcpyAddrDispatchInfo memcpy_addr;	997	MemcpyAddrDispatchInfo memcpy_addr;

Mtests/ge/ut/ge/common/program_generator_unittest.cc+68-5

	@@ -1394,13 +1394,13 @@ GeRootModelPtr CreateGeRootModelWithNoTaskConcatOutputReuseDimOne() {
1394	return ge_root_model;	1394	return ge_root_model;
1395	}	1395	}
1396		1396
1397	-ProgramGenerator CreateProgramGenerator(GeRootModelPtr &ge_root_model) {	1397	+ProgramGenerator CreateProgramGenerator(GeRootModelPtr &ge_root_model, bool has_custom_kernel = false) {
1398	static AstContext ast_ctx;	1398	static AstContext ast_ctx;
1399	static AstBuildContext ast(ast_ctx);	1399	static AstBuildContext ast(ast_ctx);
1400	const auto &name_to_ge_model = ge_root_model->GetSubgraphInstanceNameToModel();	1400	const auto &name_to_ge_model = ge_root_model->GetSubgraphInstanceNameToModel();
1401	if (name_to_ge_model.empty()) {	1401	if (name_to_ge_model.empty()) {
1402	ADD_FAILURE() << "[OM2] No subgraphs found in ge_root_model";	1402	ADD_FAILURE() << "[OM2] No subgraphs found in ge_root_model";
1403	- return ProgramGenerator(ast, {}, Om2CodegenModel());	1403	+ return ProgramGenerator(ast, {}, Om2CodegenModel(), has_custom_kernel);
1404	}	1404	}
1405	const auto &ge_model = name_to_ge_model.begin()->second;	1405	const auto &ge_model = name_to_ge_model.begin()->second;
1406	SyncKernelNameFromOpDesc(ge_model);	1406	SyncKernelNameFromOpDesc(ge_model);
	@@ -1408,15 +1408,15 @@ ProgramGenerator CreateProgramGenerator(GeRootModelPtr &ge_root_model) {
1408	Om2CodegenModel codegen_model;	1408	Om2CodegenModel codegen_model;
1409	if (Om2CodegenModelBuilder::CreateTaskCodeBuilders(ge_model, ast, task_code_builders, codegen_model) != SUCCESS) {	1409	if (Om2CodegenModelBuilder::CreateTaskCodeBuilders(ge_model, ast, task_code_builders, codegen_model) != SUCCESS) {
1410	ADD_FAILURE() << "[OM2] Failed to create task code handlers";	1410	ADD_FAILURE() << "[OM2] Failed to create task code handlers";
1411	- return ProgramGenerator(ast, {}, Om2CodegenModel());	1411	+ return ProgramGenerator(ast, {}, Om2CodegenModel(), has_custom_kernel);
1412	}	1412	}
1413	Om2CodegenModelBuilder builder;	1413	Om2CodegenModelBuilder builder;
1414	Om2ConstMetas const_metas;	1414	Om2ConstMetas const_metas;
1415	if (builder.Build(ge_model, task_code_builders, codegen_model, const_metas) != SUCCESS) {	1415	if (builder.Build(ge_model, task_code_builders, codegen_model, const_metas) != SUCCESS) {
1416	ADD_FAILURE() << "[OM2] Failed to build om2 codegen model";	1416	ADD_FAILURE() << "[OM2] Failed to build om2 codegen model";
1417	- return ProgramGenerator(ast, {}, Om2CodegenModel());	1417	+ return ProgramGenerator(ast, {}, Om2CodegenModel(), has_custom_kernel);
1418	}	1418	}
1419	- return ProgramGenerator(ast, task_code_builders, codegen_model);	1419	+ return ProgramGenerator(ast, task_code_builders, codegen_model, has_custom_kernel);
1420	}	1420	}
1421		1421
1422	const std::map<GeneratedFileIndex, std::string> kGeneratedFileNames = {	1422	const std::map<GeneratedFileIndex, std::string> kGeneratedFileNames = {
	@@ -2110,6 +2110,7 @@ enum OpDispatchType : uint32_t {
2110	DISPATCH_CMO_ADDR = 18,	2110	DISPATCH_CMO_ADDR = 18,
2111	DISPATCH_DSA = 19,	2111	DISPATCH_DSA = 19,
2112	DISPATCH_KERNEL_EX = 20,	2112	DISPATCH_KERNEL_EX = 20,
		2113	+ DISPATCH_CUSTOM_KERNEL = 21,
2113	DISPATCH_TYPE_COUNT	2114	DISPATCH_TYPE_COUNT
2114	};	2115	};
2115		2116
	@@ -2221,6 +2222,15 @@ struct AicpuDispatchInfo {
2221	uint32_t task_type;	2222	uint32_t task_type;
2222	};	2223	};
2223		2224
		2225	+struct CustomDispatchInfo {
		2226	+ const OpArgInfo *args_info; // IO 地址解析数组
		2227	+ uint32_t args_info_num; // args_info 数组长度
		2228	+ const char *op_type; // 算子类型名，用于 Report 上报
		2229	+ uint32_t args_idx; // 参数表索引，用于 GetArgsInfo 查找
		2230	+ uint32_t stream_id; // 执行流索引
		2231	+ uint32_t task_type;
		2232	+};
		2233	+
2224	struct DsaDispatchInfo {	2234	struct DsaDispatchInfo {
2225	const OpArgInfo *args_info; // IO 地址解析数组	2235	const OpArgInfo *args_info; // IO 地址解析数组
2226	const char *op_type; // 算子类型名，用于 Report 上报	2236	const char *op_type; // 算子类型名，用于 Report 上报
	@@ -2316,6 +2326,7 @@ struct TaskDispatchInfo {
2316	union {	2326	union {
2317	AicoreDispatchInfo aicore;	2327	AicoreDispatchInfo aicore;
2318	AicpuDispatchInfo aicpu;	2328	AicpuDispatchInfo aicpu;
		2329	+ CustomDispatchInfo custom;
2319	CmoDispatchInfo cmo;	2330	CmoDispatchInfo cmo;
2320	MemcpyAsyncDispatchInfo memcpy_async;	2331	MemcpyAsyncDispatchInfo memcpy_async;
2321	MemcpyAddrDispatchInfo memcpy_addr;	2332	MemcpyAddrDispatchInfo memcpy_addr;
	@@ -7610,4 +7621,56 @@ TEST_F(ProgramGeneratorUt, CopyTilingDataIfNeeded_EmptyTilingData_SkipsTiling) {
7610	std::map<GeneratedFileIndex, std::string> outputs;	7621	std::map<GeneratedFileIndex, std::string> outputs;
7611	ASSERT_EQ(GenerateProgramFiles(generator, outputs), SUCCESS);	7622	ASSERT_EQ(GenerateProgramFiles(generator, outputs), SUCCESS);
7612	}	7623	}
		7624	+
		7625	+// 校验 has_custom_kernel = true 时，load_and_run 源文件生成自定义算子内核所需代码
		7626	+TEST_F(ProgramGeneratorUt, GenerateLoadAndRunSource_WithCustomKernel_Ok) {
		7627	+ GeRootModelPtr ge_root_model = CreateGeRootModelWithAicoreOp();
		7628	+ auto generator = CreateProgramGenerator(ge_root_model, true);
		7629	+ std::map<GeneratedFileIndex, std::string> outputs;
		7630	+ ASSERT_EQ(GenerateProgramFiles(generator, outputs), SUCCESS);
		7631	+
		7632	+ const auto &load_run = outputs[GeneratedFileIndex::kLoadingAndRunningFile];
		7633	+
		7634	+ // has_custom_kernel 为 true 时应额外包含自定义算子相关头文件
		7635	+ EXPECT_NE(load_run.find("#include \"graph/custom_op.h\""), std::string::npos);
		7636	+ EXPECT_NE(load_run.find("#include \"exe_graph/runtime/gert_mem_allocator.h\""), std::string::npos);
		7637	+
		7638	+ // ge 命名空间内应包含 CustomOpFactory 类定义（kCreateClassCustomOpFactory）
		7639	+ EXPECT_NE(load_run.find("namespace ge {"), std::string::npos);
		7640	+ EXPECT_NE(load_run.find("class CustomOpFactory"), std::string::npos);
		7641	+ EXPECT_NE(load_run.find("CreateOrGetCustomOp"), std::string::npos);
		7642	+ EXPECT_NE(load_run.find("} // namespace ge"), std::string::npos);
		7643	+
		7644	+ // om2 匿名命名空间内应包含 kCustomTaskHelpers 提供的自定义内核辅助代码
		7645	+ EXPECT_NE(load_run.find("BuildGeTensor(const Om2Tensor &om2_tensor)"), std::string::npos);
		7646	+ EXPECT_NE(load_run.find("class CustKernelContextHolder"), std::string::npos);
		7647	+ EXPECT_NE(load_run.find("CustKernelContextHolder BuildKernelContextHolder"), std::string::npos);
		7648	+ EXPECT_NE(load_run.find("class AllocatorFaker : public gert::GertAllocator"), std::string::npos);
		7649	+ EXPECT_NE(load_run.find("aclError KernelCustTaskDistribute("), std::string::npos);
		7650	+ EXPECT_NE(load_run.find("aclError DeserializeCustKernelBinaries("), std::string::npos);
		7651	+
		7652	+ // Load 方法中应调用 DeserializeCustKernelBinaries 反序列化自定义内核二进制
		7653	+ EXPECT_NE(load_run.find("DeserializeCustKernelBinaries(bin_info_map_)"), std::string::npos);
		7654	+
		7655	+ // ge 命名空间应在 om2 命名空间之前生成
		7656	+ EXPECT_LT(load_run.find("namespace ge {"), load_run.find("namespace om2 {"));
		7657	+}
		7658	+
		7659	+// 校验 has_custom_kernel 默认为 false 时，不生成自定义算子内核相关代码
		7660	+TEST_F(ProgramGeneratorUt, GenerateLoadAndRunSource_WithoutCustomKernel_Omitted) {
		7661	+ GeRootModelPtr ge_root_model = CreateGeRootModelWithAicoreOp();
		7662	+ auto generator = CreateProgramGenerator(ge_root_model);
		7663	+ std::map<GeneratedFileIndex, std::string> outputs;
		7664	+ ASSERT_EQ(GenerateProgramFiles(generator, outputs), SUCCESS);
		7665	+
		7666	+ const auto &load_run = outputs[GeneratedFileIndex::kLoadingAndRunningFile];
		7667	+
		7668	+ // has_custom_kernel 为 false 时不应包含自定义算子相关头文件与辅助代码
		7669	+ EXPECT_EQ(load_run.find("#include \"graph/custom_op.h\""), std::string::npos);
		7670	+ EXPECT_EQ(load_run.find("#include \"exe_graph/runtime/gert_mem_allocator.h\""), std::string::npos);
		7671	+ EXPECT_EQ(load_run.find("class CustomOpFactory"), std::string::npos);
		7672	+ EXPECT_EQ(load_run.find("BuildGeTensor(const Om2Tensor &om2_tensor)"), std::string::npos);
		7673	+ EXPECT_EQ(load_run.find("KernelCustTaskDistribute"), std::string::npos);
		7674	+ EXPECT_EQ(load_run.find("DeserializeCustKernelBinaries"), std::string::npos);
		7675	+}
7613	} // namespace ge	7676	} // namespace ge

Mtests/ge/ut/ge/graph/load/om2_package_helper_unittest.cc+197-31

	@@ -8,6 +8,7 @@
8	* See LICENSE in the root of the software repository for the full text of the License.	8	* See LICENSE in the root of the software repository for the full text of the License.
9	*/	9	*/
10		10
		11	+#include <string>
11	#include "framework/common/helper/model_save_helper.h"	12	#include "framework/common/helper/model_save_helper.h"
12	#include "common/helper/om2/zip_archive_writer.h"	13	#include "common/helper/om2/zip_archive_writer.h"
13	#include "runtime/om2/zip_archive_reader.h"	14	#include "runtime/om2/zip_archive_reader.h"
	@@ -29,6 +30,7 @@
29	#include <gtest/gtest.h>	30	#include <gtest/gtest.h>
30	#include <filesystem>	31	#include <filesystem>
31	#include <fstream>	32	#include <fstream>
		33	+#include <filesystem> // for std::filesystem::remove
32	#include "common/env_path.h"	34	#include "common/env_path.h"
33	#include "mmpa/mmpa_api.h"	35	#include "mmpa/mmpa_api.h"
34	#include "graph/debug/ge_attr_define.h"	36	#include "graph/debug/ge_attr_define.h"
	@@ -36,6 +38,7 @@
36	#include "graph/utils/tensor_utils.h"	38	#include "graph/utils/tensor_utils.h"
37	#include "graph/utils/file_utils.h"	39	#include "graph/utils/file_utils.h"
38	#include "graph/utils/graph_utils.h"	40	#include "graph/utils/graph_utils.h"
		41	+#include "graph/custom_op_factory.h"
39	#include <cstdio>	42	#include <cstdio>
40	#include <sstream>	43	#include <sstream>
41	#include <system_error>	44	#include <system_error>
	@@ -43,6 +46,7 @@
43	#include "ge_runtime_stub/include/common/share_graph.h"	46	#include "ge_runtime_stub/include/common/share_graph.h"
44	#include "ge_runtime_stub/include/faker/ge_model_builder.h"	47	#include "ge_runtime_stub/include/faker/ge_model_builder.h"
45	#include "ge_runtime_stub/include/faker/aicore_taskdef_faker.h"	48	#include "ge_runtime_stub/include/faker/aicore_taskdef_faker.h"
		49	+#include "ge_runtime_stub/include/faker/custom_taskdef_faker.h"
46	#include "common/tbe_handle_store/tbe_kernel_store.h"	50	#include "common/tbe_handle_store/tbe_kernel_store.h"
47		51
48	namespace ge {	52	namespace ge {
	@@ -263,6 +267,109 @@ void ExpectVisualJsonCanLoad(const Archive &archive, const std::string &expected
263	EXPECT_FALSE(pb_json["graph"][0]["op"].empty());	267	EXPECT_FALSE(pb_json["graph"][0]["op"].empty());
264	}	268	}
265		269
		270	+class TestPortableCustomOp : public PortableOp, public EagerExecuteOp {
		271	+ public:
		272	+ graphStatus Execute(gert::EagerOpExecutionContext *ctx) override {
		273	+ return SUCCESS;
		274	+ }
		275	+
		276	+ graphStatus Serialize(std::vector<uint8_t> &buffer) override {
		277	+ const std::string payload = "test_portable_custom_op_kernel_bin";
		278	+ buffer.assign(payload.begin(), payload.end());
		279	+ return GRAPH_SUCCESS;
		280	+ }
		281	+
		282	+ graphStatus Deserialize(const std::vector<uint8_t> &buffer) override {
		283	+ return GRAPH_SUCCESS;
		284	+ }
		285	+};
		286	+
		287	+static ComputeGraphPtr BuildCustomOpGraph() {
		288	+ auto graph = std::make_shared<ComputeGraph>("custom_op_om2_graph");
		289	+ GeTensorDesc tensor_desc(GeShape({2, 2, 2}), FORMAT_ND, DT_FLOAT);
		290	+
		291	+ auto data0_desc = std::make_shared<OpDesc>("data0", DATA);
		292	+ (void)data0_desc->AddInputDesc(tensor_desc);
		293	+ (void)data0_desc->AddOutputDesc(tensor_desc);
		294	+ AttrUtils::SetInt(data0_desc, ATTR_NAME_INDEX, 0);
		295	+ auto data0 = graph->AddNode(data0_desc);
		296	+
		297	+ auto data1_desc = std::make_shared<OpDesc>("data1", DATA);
		298	+ (void)data1_desc->AddInputDesc(tensor_desc);
		299	+ (void)data1_desc->AddOutputDesc(tensor_desc);
		300	+ AttrUtils::SetInt(data1_desc, ATTR_NAME_INDEX, 1);
		301	+ auto data1 = graph->AddNode(data1_desc);
		302	+
		303	+ auto custom_op_desc = std::make_shared<OpDesc>("custom_op", "TestPortableOp");
		304	+ (void)custom_op_desc->AddInputDesc("x0", tensor_desc);
		305	+ (void)custom_op_desc->AddInputDesc("x1", tensor_desc);
		306	+ (void)custom_op_desc->AddOutputDesc("y", tensor_desc);
		307	+ custom_op_desc->AppendIrInput("x0", kIrInputRequired);
		308	+ custom_op_desc->AppendIrInput("x1", kIrInputRequired);
		309	+ custom_op_desc->AppendIrOutput("y", kIrOutputRequired);
		310	+ auto custom_op_node = graph->AddNode(custom_op_desc);
		311	+
		312	+ auto netoutput_desc = std::make_shared<OpDesc>("netoutput", NETOUTPUT);
		313	+ (void)netoutput_desc->AddInputDesc(tensor_desc);
		314	+ auto netoutput = graph->AddNode(netoutput_desc);
		315	+
		316	+ GraphUtils::AddEdge(data0->GetOutDataAnchor(0), custom_op_node->GetInDataAnchor(0));
		317	+ GraphUtils::AddEdge(data1->GetOutDataAnchor(0), custom_op_node->GetInDataAnchor(1));
		318	+ GraphUtils::AddEdge(custom_op_node->GetOutDataAnchor(0), netoutput->GetInDataAnchor(0));
		319	+ netoutput_desc->SetSrcName({"custom_op"});
		320	+ netoutput_desc->SetSrcIndex({0});
		321	+ graph->TopologicalSorting();
		322	+ return graph;
		323	+}
		324	+
		325	+GeRootModelPtr CreateGeRootModelWithCustomOp() {
		326	+ auto graph = BuildCustomOpGraph();
		327	+
		328	+ gert::GeModelBuilder builder(graph);
		329	+ auto ge_root_model =
		330	+ builder
		331	+ .AddTaskDef(
		332	+ "custom_op",
		333	+ gert::CustomTaskDefFaker("custom_op_stub").ArgsFormat("{i_instance0}{i_instance1}{o_instance0*}"))
		334	+ .FakeTbeBin({"custom_op"})
		335	+ .BuildGeRootModel();
		336	+ auto &compute_graph = ge_root_model->GetRootGraph();
		337	+ compute_graph->SetGraphUnknownFlag(false);
		338	+ for (const auto &node : compute_graph->GetDirectNode()) {
		339	+ auto op_desc = node->GetOpDesc();
		340	+ if (op_desc == nullptr) {
		341	+ return nullptr;
		342	+ }
		343	+ if (op_desc->GetType() == DATA) {
		344	+ op_desc->SetOutputOffset({1024});
		345	+ } else if (op_desc->GetType() == NETOUTPUT) {
		346	+ op_desc->SetInputOffset({3072});
		347	+ } else {
		348	+ op_desc->SetInputOffset(std::vector<int64_t>(op_desc->GetInputsSize(), 1024));
		349	+ op_desc->SetOutputOffset(std::vector<int64_t>(op_desc->GetOutputsSize(), 1024));
		350	+ }
		351	+ }
		352	+
		353	+ const auto ge_model = ge_root_model->GetSubgraphInstanceNameToModel().begin()->second;
		354	+ std::vector<uint8_t> weights_value(512, 1U);
		355	+ ge_model->SetWeight(Buffer::CopyFrom(weights_value.data(), weights_value.size()));
		356	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_MEMORY_SIZE, 2048);
		357	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_WEIGHT_SIZE, weights_value.size());
		358	+ (void)AttrUtils::SetInt(ge_model, ATTR_MODEL_STREAM_NUM, 1);
		359	+
		360	+ return ge_root_model;
		361	+}
		362	+
		363	+int WriteBinFile(const char *file_name, const std::string &text) {
		364	+ std::ofstream outFile(file_name, std::ios::out \| std::ios::binary);
		365	+ if (!outFile.is_open()) {
		366	+ return 1;
		367	+ }
		368	+ outFile.write(text.c_str(), text.size());
		369	+ outFile.close();
		370	+ return 0;
		371	+}
		372	+
266	} // namespace	373	} // namespace
267		374
268	class Om2PackageHelperUt : public testing::Test {	375	class Om2PackageHelperUt : public testing::Test {
	@@ -644,10 +751,12 @@ TEST_F(Om2PackageHelperUt, Om2CodegenAndCompile_Fail_DumpGeneratedFiles) {
644	}	751	}
645	}	752	}
646		753
647	- Om2CodegenArtifacts artifacts;	754	+ gert::Om2ModelData model_data;
648	- Om2ConstMetas const_metas;	755	+ Om2CodegenArtifacts &artifacts = model_data.program_body.source_artifacts;
		756	+ Om2ConstMetas &const_metas = model_data.constants_data.consts;
		757	+ ;
649	Om2Codegen codegen;	758	Om2Codegen codegen;
650	- ASSERT_NE(codegen.Om2CodegenAndCompile(ge_model, artifacts, const_metas), SUCCESS);	759	+ ASSERT_NE(codegen.Om2CodegenAndCompile(ge_model, model_data), SUCCESS);
651		760
652	if (!IsTmpDirExists()) {	761	if (!IsTmpDirExists()) {
653	return;	762	return;
	@@ -1033,8 +1142,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_SpecialInputSize) {
1033	}	1142	}
1034		1143
1035	Om2PackageHelper om2_packager;	1144	Om2PackageHelper om2_packager;
1036	- gert::Om2ModelMeta model_meta;	1145	+ gert::Om2ModelData model_data;
1037	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1146	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1147	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1038	ASSERT_FALSE(model_meta.input_desc.empty());	1148	ASSERT_FALSE(model_meta.input_desc.empty());
1039	EXPECT_EQ(model_meta.input_desc[0].GetByteSize(), 2048U);	1149	EXPECT_EQ(model_meta.input_desc[0].GetByteSize(), 2048U);
1040	}	1150	}
	@@ -1056,8 +1166,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_InputDimsAttr) {
1056	}	1166	}
1057		1167
1058	Om2PackageHelper om2_packager;	1168	Om2PackageHelper om2_packager;
1059	- gert::Om2ModelMeta model_meta;	1169	+ gert::Om2ModelData model_data;
1060	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1170	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1171	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1061	ASSERT_FALSE(model_meta.input_desc_v2.empty());	1172	ASSERT_FALSE(model_meta.input_desc_v2.empty());
1062	EXPECT_EQ(model_meta.input_desc_v2[0].GetShape(), std::vector<int64_t>({2, 8}));	1173	EXPECT_EQ(model_meta.input_desc_v2[0].GetShape(), std::vector<int64_t>({2, 8}));
1063	}	1174	}
	@@ -1081,8 +1192,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_SpecialOutputSize) {
1081	}	1192	}
1082		1193
1083	Om2PackageHelper om2_packager;	1194	Om2PackageHelper om2_packager;
1084	- gert::Om2ModelMeta model_meta;	1195	+ gert::Om2ModelData model_data;
1085	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1196	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1197	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1086	ASSERT_FALSE(model_meta.output_desc.empty());	1198	ASSERT_FALSE(model_meta.output_desc.empty());
1087	EXPECT_EQ(model_meta.output_desc[0].GetByteSize(), 4096U);	1199	EXPECT_EQ(model_meta.output_desc[0].GetByteSize(), 4096U);
1088	}	1200	}
	@@ -1104,8 +1216,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_DynamicOutputDims) {
1104	}	1216	}
1105		1217
1106	Om2PackageHelper om2_packager;	1218	Om2PackageHelper om2_packager;
1107	- gert::Om2ModelMeta model_meta;	1219	+ gert::Om2ModelData model_data;
1108	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1220	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1221	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1109	EXPECT_EQ(model_meta.dynamic_output_shape.size(), 2U);	1222	EXPECT_EQ(model_meta.dynamic_output_shape.size(), 2U);
1110	EXPECT_EQ(model_meta.dynamic_output_shape[0], "1,4");	1223	EXPECT_EQ(model_meta.dynamic_output_shape[0], "1,4");
1111	EXPECT_EQ(model_meta.dynamic_output_shape[1], "2,8");	1224	EXPECT_EQ(model_meta.dynamic_output_shape[1], "2,8");
	@@ -1128,8 +1241,9 @@ TEST_F(Om2PackageHelperUt, BuildDebugInfo_DumpOriginOpNames) {
1128	}	1241	}
1129		1242
1130	Om2PackageHelper om2_packager;	1243	Om2PackageHelper om2_packager;
1131	- gert::Om2DebugInfo debug_info;	1244	+ gert::Om2ModelData model_data;
1132	- ASSERT_EQ(om2_packager.BuildDebugInfo(ge_model, debug_info), SUCCESS);	1245	+ gert::Om2DebugInfo &debug_info = model_data.debug_info;
		1246	+ ASSERT_EQ(om2_packager.BuildDebugInfo(ge_model, model_data), SUCCESS);
1133	ASSERT_FALSE(debug_info.op_attr_map.empty());	1247	ASSERT_FALSE(debug_info.op_attr_map.empty());
1134		1248
1135	bool found = false;	1249	bool found = false;
	@@ -1147,8 +1261,9 @@ TEST_F(Om2PackageHelperUt, BuildDebugInfo_DumpOriginOpNames) {
1147		1261
1148	TEST_F(Om2PackageHelperUt, BuildManifest_NullRootModel) {	1262	TEST_F(Om2PackageHelperUt, BuildManifest_NullRootModel) {
1149	Om2PackageHelper om2_packager;	1263	Om2PackageHelper om2_packager;
1150	- std::map<std::string, std::string> manifest;	1264	+ gert::Om2ModelData model_data;
1151	- ASSERT_EQ(om2_packager.BuildManifest(nullptr, manifest), SUCCESS);	1265	+ std::map<std::string, std::string> &manifest = model_data.manifest;
		1266	+ ASSERT_EQ(om2_packager.BuildManifest(nullptr, model_data), SUCCESS);
1152		1267
1153	ASSERT_EQ(manifest.count(OM2_MODEL_NUM), 1U);	1268	ASSERT_EQ(manifest.count(OM2_MODEL_NUM), 1U);
1154	EXPECT_EQ(manifest[OM2_MODEL_NUM], "1");	1269	EXPECT_EQ(manifest[OM2_MODEL_NUM], "1");
	@@ -1210,10 +1325,11 @@ TEST_F(Om2PackageHelperUt, Om2CodegenAndCompile_Success) {
1210	const auto &ge_model = ge_root_model->GetSubgraphInstanceNameToModel().begin()->second;	1325	const auto &ge_model = ge_root_model->GetSubgraphInstanceNameToModel().begin()->second;
1211	ASSERT_NE(ge_model, nullptr);	1326	ASSERT_NE(ge_model, nullptr);
1212		1327
1213	- Om2CodegenArtifacts artifacts;	1328	+ gert::Om2ModelData model_data;
1214	- Om2ConstMetas const_metas;	1329	+ Om2CodegenArtifacts &artifacts = model_data.program_body.source_artifacts;
		1330	+ Om2ConstMetas &const_metas = model_data.constants_data.consts;
1215	Om2Codegen codegen;	1331	Om2Codegen codegen;
1216	- ASSERT_EQ(codegen.Om2CodegenAndCompile(ge_model, artifacts, const_metas), SUCCESS);	1332	+ ASSERT_EQ(codegen.Om2CodegenAndCompile(ge_model, model_data), SUCCESS);
1217	EXPECT_FALSE(artifacts.empty());	1333	EXPECT_FALSE(artifacts.empty());
1218	bool found_so = false;	1334	bool found_so = false;
1219	for (const auto &artifact : artifacts) {	1335	for (const auto &artifact : artifacts) {
	@@ -1233,10 +1349,11 @@ TEST_F(Om2PackageHelperUt, Om2CodegenAndCompile_InvalidModel_Fail) {
1233	const auto &ge_model = name_to_ge_model.begin()->second;	1349	const auto &ge_model = name_to_ge_model.begin()->second;
1234	ASSERT_NE(ge_model, nullptr);	1350	ASSERT_NE(ge_model, nullptr);
1235		1351
1236	- Om2CodegenArtifacts artifacts;	1352	+ gert::Om2ModelData model_data;
1237	- Om2ConstMetas const_metas;	1353	+ Om2CodegenArtifacts &artifacts = model_data.program_body.source_artifacts;
		1354	+ Om2ConstMetas &const_metas = model_data.constants_data.consts;
1238	Om2Codegen codegen;	1355	Om2Codegen codegen;
1239	- EXPECT_NE(codegen.Om2CodegenAndCompile(ge_model, artifacts, const_metas), SUCCESS);	1356	+ EXPECT_NE(codegen.Om2CodegenAndCompile(ge_model, model_data), SUCCESS);
1240	}	1357	}
1241		1358
1242	// ============================================================================	1359	// ============================================================================
	@@ -1599,9 +1716,10 @@ TEST_F(Om2PackageHelperUt, BuildKernelBinaries_WithAtomicKernel_Ok) {
1599	graph->SetGraphUnknownFlag(false);	1716	graph->SetGraphUnknownFlag(false);
1600	ge_model->SetGraph(graph);	1717	ge_model->SetGraph(graph);
1601		1718
1602	- std::vector<gert::Om2KernelBinary> kernel_binaries;	1719	+ gert::Om2ModelData model_data;
1603	- ASSERT_EQ(Om2PackageHelper::BuildKernelBinaries(ge_model, kernel_binaries), SUCCESS);	1720	+ ASSERT_EQ(Om2PackageHelper::BuildKernelBinaries(ge_model, model_data), SUCCESS);
1604		1721
		1722	+ auto &kernel_binaries = model_data.kernel_binaries;
1605	ASSERT_EQ(kernel_binaries.size(), 2U);	1723	ASSERT_EQ(kernel_binaries.size(), 2U);
1606	EXPECT_EQ(kernel_binaries[0].name, "normal_kernel.o");	1724	EXPECT_EQ(kernel_binaries[0].name, "normal_kernel.o");
1607	EXPECT_NE(kernel_binaries[0].data, nullptr);	1725	EXPECT_NE(kernel_binaries[0].data, nullptr);
	@@ -1636,8 +1754,9 @@ TEST_F(Om2PackageHelperUt, BuildKernelBinaries_WithAtomicKernel_Success) {
1636	ge_model->SetGraph(graph);	1754	ge_model->SetGraph(graph);
1637		1755
1638	Om2PackageHelper om2_packager;	1756	Om2PackageHelper om2_packager;
1639	- std::vector<gert::Om2KernelBinary> kernel_binaries;	1757	+ gert::Om2ModelData model_data;
1640	- ASSERT_EQ(om2_packager.BuildKernelBinaries(ge_model, kernel_binaries), SUCCESS);	1758	+ std::vector<gert::Om2KernelBinary> &kernel_binaries = model_data.kernel_binaries;
		1759	+ ASSERT_EQ(om2_packager.BuildKernelBinaries(ge_model, model_data), SUCCESS);
1641	EXPECT_FALSE(kernel_binaries.empty());	1760	EXPECT_FALSE(kernel_binaries.empty());
1642	}	1761	}
1643		1762
	@@ -1664,8 +1783,9 @@ TEST_F(Om2PackageHelperUt, BuildKernelBinaries_WithCustAicpuKernel_Success) {
1664	ge_model->SetGraph(graph);	1783	ge_model->SetGraph(graph);
1665		1784
1666	Om2PackageHelper om2_packager;	1785	Om2PackageHelper om2_packager;
1667	- std::vector<gert::Om2KernelBinary> kernel_binaries;	1786	+ gert::Om2ModelData model_data;
1668	- ASSERT_EQ(om2_packager.BuildKernelBinaries(ge_model, kernel_binaries), SUCCESS);	1787	+ std::vector<gert::Om2KernelBinary> &kernel_binaries = model_data.kernel_binaries;
		1788	+ ASSERT_EQ(om2_packager.BuildKernelBinaries(ge_model, model_data), SUCCESS);
1669	bool found_cust = false;	1789	bool found_cust = false;
1670	for (const auto &kb : kernel_binaries) {	1790	for (const auto &kb : kernel_binaries) {
1671	if (kb.name.find("_CustAicpuKernel.o") != std::string::npos) {	1791	if (kb.name.find("_CustAicpuKernel.o") != std::string::npos) {
	@@ -1687,8 +1807,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_WithOutputNameContainingColon_Success)
1687	AttrUtils::SetListStr(ge_model, ATTR_MODEL_OUT_NODES_NAME, out_node_names);	1807	AttrUtils::SetListStr(ge_model, ATTR_MODEL_OUT_NODES_NAME, out_node_names);
1688		1808
1689	Om2PackageHelper om2_packager;	1809	Om2PackageHelper om2_packager;
1690	- gert::Om2ModelMeta model_meta;	1810	+ gert::Om2ModelData model_data;
1691	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1811	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1812	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1692	ASSERT_FALSE(model_meta.output_desc.empty());	1813	ASSERT_FALSE(model_meta.output_desc.empty());
1693	EXPECT_EQ(model_meta.output_desc[0].GetName(), "add1:0");	1814	EXPECT_EQ(model_meta.output_desc[0].GetName(), "add1:0");
1694	}	1815	}
	@@ -1704,8 +1825,9 @@ TEST_F(Om2PackageHelperUt, BuildModelMeta_WithOutputNameWithoutColon_Success) {
1704	AttrUtils::SetListStr(ge_model, ATTR_MODEL_OUT_NODES_NAME, out_node_names);	1825	AttrUtils::SetListStr(ge_model, ATTR_MODEL_OUT_NODES_NAME, out_node_names);
1705		1826
1706	Om2PackageHelper om2_packager;	1827	Om2PackageHelper om2_packager;
1707	- gert::Om2ModelMeta model_meta;	1828	+ gert::Om2ModelData model_data;
1708	- ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_meta), SUCCESS);	1829	+ gert::Om2ModelMeta &model_meta = model_data.model_meta;
		1830	+ ASSERT_EQ(om2_packager.BuildModelMeta(ge_model, model_data), SUCCESS);
1709	ASSERT_FALSE(model_meta.output_desc.empty());	1831	ASSERT_FALSE(model_meta.output_desc.empty());
1710	EXPECT_NE(model_meta.output_desc[0].GetName().find(":"), std::string::npos);	1832	EXPECT_NE(model_meta.output_desc[0].GetName().find(":"), std::string::npos);
1711	}	1833	}
	@@ -1717,4 +1839,48 @@ TEST_F(Om2PackageHelperUt, SetSaveMode_False) {
1717	helper.SetSaveMode(true);	1839	helper.SetSaveMode(true);
1718	EXPECT_TRUE(helper.is_offline_);	1840	EXPECT_TRUE(helper.is_offline_);
1719	}	1841	}
		1842	+
		1843	+TEST_F(Om2PackageHelperUt, SaveToOmModel_WithCustomKernel) {
		1844	+ const AscendString kOpType("TestPortableOp");
		1845	+ CustomOpFactory::RegisterCustomOpCreator(
		1846	+ kOpType, []() -> std::unique_ptr<BaseCustomOp> { return std::make_unique<TestPortableCustomOp>(); });
		1847	+ Om2PackageHelper om2_packager;
		1848	+ om2_packager.SetSaveMode(false);
		1849	+ const auto ge_root_model = CreateGeRootModelWithCustomOp();
		1850	+ ASSERT_NE(ge_root_model, nullptr);
		1851	+ ModelBufferData model_data;
		1852	+ const std::string output_file = PathUtils::Join({test_work_dir, kZipFileBaseName + "_buffer.om2"});
		1853	+ SyncKernelNameForAllModels(ge_root_model);
		1854	+ ASSERT_EQ(om2_packager.SaveToOmRootModel(ge_root_model, output_file, model_data, false), SUCCESS);
		1855	+ EXPECT_NE(mmAccess2(output_file.c_str(), M_F_OK), EOK);
		1856	+ ASSERT_NE(model_data.data, nullptr);
		1857	+ ASSERT_GT(model_data.length, 0U);
		1858	+
		1859	+ SimpleZipArchiveReader archive(model_data.data.get(), model_data.length);
		1860	+ ASSERT_TRUE(archive.IsGood());
		1861	+ const auto file_names = archive.ListFiles();
		1862	+ bool has_custom_op_binary = false;
		1863	+ for (const auto &file_name : file_names) {
		1864	+ if ((file_name.find("/custom_ops/binaries_npu_arch/TestPortableOp_") != std::string::npos) &&
		1865	+ (file_name.find("_CustomKernel.bin") != std::string::npos)) {
		1866	+ has_custom_op_binary = true;
		1867	+ break;
		1868	+ }
		1869	+ }
		1870	+ EXPECT_TRUE(has_custom_op_binary);
		1871	+ CustomOpFactory::RemoveCustomOps({kOpType});
		1872	+}
		1873	+
		1874	+TEST_F(Om2PackageHelperUt, ReadCustomOpSoFiles) {
		1875	+ std::string so_file = "/tmp/libcusom_op_" + std::to_string(getpid()) + ".so";
		1876	+ std::string text = "fake custom op so content";
		1877	+ EXPECT_EQ(WriteBinFile(so_file.c_str(), text), 0);
		1878	+ std::unordered_set<std::string> ops_so_set = {so_file};
		1879	+ std::vector<gert::Om2KernelBinary> shared_lib_binaries;
		1880	+ EXPECT_EQ(Om2PackageHelper::ReadCustomOpSoToBuffer(ops_so_set, shared_lib_binaries), 0);
		1881	+ EXPECT_EQ(ops_so_set.size(), shared_lib_binaries.size());
		1882	+ EXPECT_EQ(text.size(), shared_lib_binaries[0].data_size);
		1883	+ EXPECT_EQ(memcmp(text.data(), shared_lib_binaries[0].data.get(), text.size()), 0);
		1884	+ std::filesystem::remove(so_file);
		1885	+}
1720	} // namespace ge	1886	} // namespace ge