fix: 完善 HostCPU 自定义算子的引擎选择-ge-AtomGit

已合并

fix: 完善 HostCPU 自定义算子的引擎选择 #4630

duhua创建于 5 天前从du-hua1024/ge： new 合入到 cann/ge： develop

fix: 完善 HostCPU 自定义算子的引擎选择 #4630

已合并

duhua创建于 5 天前

从du-hua1024/ge new合入到cann/ge develop

共 5 个文件变更+83-2

Mcompiler/engines/manager/engine_manager/dnnengine_manager.cc+7-0

	@@ -466,6 +466,13 @@ std::string DNNEngineManager::GetDNNEngineName(const ge::NodePtr &node_ptr,
466	return "";	466	return "";
467	}	467	}
468		468
		469	+void DNNEngineManager::TrySelectHostCpuCustomOp(const OpDescPtr &op_desc, OpInfo &matched_op_info) const {
		470	+ if ((op_desc == nullptr) \|\| (op_desc->GetOpEngineName() != kHostCpuEngineName) \|\| !IsHostCpuCustomOp(op_desc)) {
		471	+ return;
		472	+ }
		473	+ SetHostCpuCustomOp(op_desc, matched_op_info);
		474	+}
		475	+
469	std::string DNNEngineManager::GetCompositeEngineName(const ge::NodePtr &node_ptr, uint32_t recursive_depth) {	476	std::string DNNEngineManager::GetCompositeEngineName(const ge::NodePtr &node_ptr, uint32_t recursive_depth) {
470	// op_desc of node should not be null	477	// op_desc of node should not be null
471	const auto &op_desc = node_ptr->GetOpDesc();	478	const auto &op_desc = node_ptr->GetOpDesc();

Mcompiler/engines/manager/engine_manager/dnnengine_manager.h+1-0

	@@ -68,6 +68,7 @@ class DNNEngineManager {
68	std::string GetDNNEngineName(const ge::NodePtr &node_ptr);	68	std::string GetDNNEngineName(const ge::NodePtr &node_ptr);
69	std::string GetDNNEngineName(const ge::NodePtr &node_ptr, const std::set<std::string> &exclude_engines,	69	std::string GetDNNEngineName(const ge::NodePtr &node_ptr, const std::set<std::string> &exclude_engines,
70	OpInfo &matched_op_info);	70	OpInfo &matched_op_info);
		71	+ void TrySelectHostCpuCustomOp(const OpDescPtr &op_desc, OpInfo &matched_op_info) const;
71	std::string GetCompositeEngineName(const ge::NodePtr &node_ptr, uint32_t recursive_depth = 1);	72	std::string GetCompositeEngineName(const ge::NodePtr &node_ptr, uint32_t recursive_depth = 1);
72	std::string GetCompositeEngineName(const std::string &atomic_engine_name);	73	std::string GetCompositeEngineName(const std::string &atomic_engine_name);
73	std::string GetCompositeEngineKernelLibName(const std::string &composite_engine_name) const;	74	std::string GetCompositeEngineKernelLibName(const std::string &composite_engine_name) const;

Mcompiler/graph/partition/engine_place.cc+3-0

	@@ -75,6 +75,9 @@ Status EnginePlacer::SelectEngine(const NodePtr &node, const std::set<std::strin
75	matched_op_info.flagAsync = false;	75	matched_op_info.flagAsync = false;
76	}	76	}
77		77
		78	+ // Host CPU custom operators keep the custom engine identity and use Host CPU lowering.
		79	+ DNNEngineManager::GetInstance().TrySelectHostCpuCustomOp(op_desc, matched_op_info);
		80	+
78	// Record the node assigned atomic_engine name	81	// Record the node assigned atomic_engine name
79	GELOGD("Assigning DNNEngine %s to node %s, op type %s", op_desc->GetOpEngineName().c_str(), node->GetName().c_str(),	82	GELOGD("Assigning DNNEngine %s to node %s, op type %s", op_desc->GetOpEngineName().c_str(), node->GetName().c_str(),
80	node->GetType().c_str());	83	node->GetType().c_str());

Mexamples/fusion_pass/pattern_base_pass/1_fuse_matmul_add_pass/cpp/src/fuse_matmul_add_pass.cpp+2-2

	@@ -50,8 +50,8 @@ class FuseMatMulAndAddPass : public PatternFusionPass {
50	std::cout << "Define replacement for FuseMatMulAndAddPass" << std::endl;	50	std::cout << "Define replacement for FuseMatMulAndAddPass" << std::endl;
51	auto replace_graph_builder = es::EsGraphBuilder("replacement");	51	auto replace_graph_builder = es::EsGraphBuilder("replacement");
52	auto [r_a, r_b, r_c] = replace_graph_builder.CreateInputs<3>();	52	auto [r_a, r_b, r_c] = replace_graph_builder.CreateInputs<3>();
53	- auto alpha_const = replace_graph_builder.CreateScalar(1);	53	+ auto alpha_const = replace_graph_builder.CreateScalar(1.0f);
54	- auto beta_const = replace_graph_builder.CreateScalar(1);	54	+ auto beta_const = replace_graph_builder.CreateScalar(1.0f);
55	auto gemm = es::GEMM(r_a, r_b, r_c, alpha_const, beta_const);	55	auto gemm = es::GEMM(r_a, r_b, r_c, alpha_const, beta_const);
56	return replace_graph_builder.BuildAndReset({gemm});	56	return replace_graph_builder.BuildAndReset({gemm});
57	}	57	}

Mtests/ge/ut/ge/graph/partition/engine_place_unittest.cc+70-0

	@@ -13,6 +13,7 @@
13	#include "compute_graph.h"	13	#include "compute_graph.h"
14	#include "graph/partition/engine_place.h"	14	#include "graph/partition/engine_place.h"
15	#include "engines/manager/engine_manager/dnnengine_manager.h"	15	#include "engines/manager/engine_manager/dnnengine_manager.h"
		16	+#include "graph/custom_op_factory.h"
16	#include "graph/debug/ge_attr_define.h"	17	#include "graph/debug/ge_attr_define.h"
17		18
18	namespace ge {	19	namespace ge {
	@@ -115,6 +116,75 @@ TEST_F(UtestEnginePlace, select_engine_when_opdesc_confilct_with_attr) {
115	ASSERT_EQ(op_desc->GetOpKernelLibName(), op_kernel_name);	116	ASSERT_EQ(op_desc->GetOpKernelLibName(), op_kernel_name);
116	}	117	}
117		118
		119	+TEST_F(UtestEnginePlace, select_engine_overrides_legacy_host_cpu_with_host_custom) {
		120	+ const AscendString op_type("EnginePlaceHostCpuCustom");
		121	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(op_type, OpBackend::kHostCPU,
		122	+ []() -> std::unique_ptr<BaseCustomOp> { return nullptr; }),
		123	+ GRAPH_SUCCESS);
		124	+ ComputeGraphPtr graph = std::make_shared<ComputeGraph>("default");
		125	+ auto op_desc = std::make_shared<OpDesc>("mock_op_name", op_type.GetString());
		126	+ op_desc->SetOpEngineName("DNN_VM_HOST_CPU");
		127	+ op_desc->SetOpKernelLibName("DNN_VM_HOST_CPU_OP_STORE");
		128	+ auto node_ptr = graph->AddNode(op_desc);
		129	+
		130	+ EnginePlacer engine_place(graph);
		131	+ bool is_check_support_success = true;
		132	+ OpInfo op_info;
		133	+ EXPECT_EQ(engine_place.SelectEngine(node_ptr, {}, is_check_support_success, op_info), SUCCESS);
		134	+ EXPECT_EQ(op_desc->GetOpEngineName(), kEngineNameCustom);
		135	+ EXPECT_EQ(op_desc->GetOpKernelLibName(), kCustomOpKernelLibName);
		136	+ EXPECT_EQ(op_info.engine, kEngineNameCustom);
		137	+ EXPECT_EQ(op_info.opKernelLib, kCustomOpKernelLibName);
		138	+ std::string lowering_func;
		139	+ ASSERT_TRUE(AttrUtils::GetStr(op_desc, kAttrLowingFunc, lowering_func));
		140	+ EXPECT_EQ(lowering_func, kHostCpuCustomOpLowerFunc);
		141	+ CustomOpFactory::RemoveCustomOps({op_type});
		142	+}
		143	+
		144	+TEST_F(UtestEnginePlace, select_engine_overrides_host_cpu_attr_with_host_custom) {
		145	+ const AscendString op_type("EnginePlaceHostCpuCustomAttr");
		146	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(op_type, OpBackend::kHostCPU,
		147	+ []() -> std::unique_ptr<BaseCustomOp> { return nullptr; }),
		148	+ GRAPH_SUCCESS);
		149	+ ComputeGraphPtr graph = std::make_shared<ComputeGraph>("default");
		150	+ auto op_desc = std::make_shared<OpDesc>("mock_op_name", op_type.GetString());
		151	+ AttrUtils::SetStr(op_desc, ATTR_NAME_ENGINE_NAME_FOR_LX, "DNN_VM_HOST_CPU");
		152	+ AttrUtils::SetStr(op_desc, ATTR_NAME_KKERNEL_LIB_NAME_FOR_LX, "DNN_VM_HOST_CPU_OP_STORE");
		153	+ auto node_ptr = graph->AddNode(op_desc);
		154	+
		155	+ EnginePlacer engine_place(graph);
		156	+ bool is_check_support_success = true;
		157	+ OpInfo op_info;
		158	+ EXPECT_EQ(engine_place.SelectEngine(node_ptr, {}, is_check_support_success, op_info), SUCCESS);
		159	+ EXPECT_EQ(op_desc->GetOpEngineName(), kEngineNameCustom);
		160	+ EXPECT_EQ(op_info.engine, kEngineNameCustom);
		161	+ EXPECT_EQ(op_info.opKernelLib, kCustomOpKernelLibName);
		162	+ std::string lowering_func;
		163	+ ASSERT_TRUE(AttrUtils::GetStr(op_desc, kAttrLowingFunc, lowering_func));
		164	+ EXPECT_EQ(lowering_func, kHostCpuCustomOpLowerFunc);
		165	+ CustomOpFactory::RemoveCustomOps({op_type});
		166	+}
		167	+
		168	+TEST_F(UtestEnginePlace, select_engine_keeps_device_engine_with_host_custom) {
		169	+ const AscendString op_type("EnginePlaceHostCpuCustomDeviceEngine");
		170	+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(op_type, OpBackend::kHostCPU,
		171	+ []() -> std::unique_ptr<BaseCustomOp> { return nullptr; }),
		172	+ GRAPH_SUCCESS);
		173	+ ComputeGraphPtr graph = std::make_shared<ComputeGraph>("default");
		174	+ auto op_desc = std::make_shared<OpDesc>("mock_op_name", op_type.GetString());
		175	+ op_desc->SetOpEngineName("AIcoreEngine");
		176	+ op_desc->SetOpKernelLibName("AiCoreLib");
		177	+ auto node_ptr = graph->AddNode(op_desc);
		178	+
		179	+ EnginePlacer engine_place(graph);
		180	+ bool is_check_support_success = true;
		181	+ OpInfo op_info;
		182	+ EXPECT_EQ(engine_place.SelectEngine(node_ptr, {}, is_check_support_success, op_info), SUCCESS);
		183	+ EXPECT_EQ(op_desc->GetOpEngineName(), "AIcoreEngine");
		184	+ EXPECT_EQ(op_desc->GetOpKernelLibName(), "AiCoreLib");
		185	+ CustomOpFactory::RemoveCustomOps({op_type});
		186	+}
		187	+
118	TEST_F(UtestEnginePlace, check_when_graph_is_null) {	188	TEST_F(UtestEnginePlace, check_when_graph_is_null) {
119	EnginePlacer engine_place(nullptr);	189	EnginePlacer engine_place(nullptr);
120	EXPECT_EQ(engine_place.Check(), FAILED);	190	EXPECT_EQ(engine_place.Check(), FAILED);