diff --git a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_kernel_mod.cc b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_kernel_mod.cc index c7d7a3f1a2..d6fafbcf43 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_kernel_mod.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_kernel_mod.cc @@ -32,6 +32,7 @@ using AicpuTaskInfoPtr = std::shared_ptr; namespace mindspore { namespace kernel { constexpr auto AICPU_OPS_SO_NAME = "libaicpu_kernels.so"; +constexpr auto CUST_AICPU_OPS_SO_NAME = "libcpu_kernels.so"; AicpuOpKernelMod::AicpuOpKernelMod() : anf_node_(nullptr) {} @@ -66,8 +67,12 @@ void AicpuOpKernelMod::CreateCpuKernelInfo(const std::vector &inputs const std::vector &outputs) { MS_LOG(INFO) << "CreateCpuKernelInfoOffline start"; - node_so_ = AICPU_OPS_SO_NAME; - + if (kCustAiCpuKernelOps.find(node_name_) != kCustAiCpuKernelOps.end()) { + node_so_ = CUST_AICPU_OPS_SO_NAME; + node_name_ = kCustRunApi; + } else { + node_so_ = AICPU_OPS_SO_NAME; + } // InputOutputAddr vector io_addrs; (void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(io_addrs), @@ -148,7 +153,12 @@ std::vector AicpuOpKernelMod::GenTask(const std::vector MS_LOG(INFO) << "AicpuOpKernelMod GenTask start"; stream_id_ = stream_id; - node_so_ = AICPU_OPS_SO_NAME; + if (kCustAiCpuKernelOps.find(node_name_) != kCustAiCpuKernelOps.end()) { + node_so_ = CUST_AICPU_OPS_SO_NAME; + node_name_ = kCustRunApi; + } else { + node_so_ = AICPU_OPS_SO_NAME; + } std::vector input_data_addrs; (void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(input_data_addrs), [](const AddressPtr &input) -> void * { return input->addr; }); diff --git a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h index 01a8f57718..ecb9644cdf 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h +++ b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include "backend/kernel_compiler/kernel.h" namespace mindspore { @@ -41,6 +42,9 @@ constexpr auto kSeed1 = "Seed1"; constexpr auto kSeed2 = "seed2"; constexpr auto kTopK = "TopK"; constexpr auto kTopKV2 = "TopKV2"; +constexpr auto kEditDistance = "EditDistance"; +constexpr auto kCustRunApi = "RunCpuKernel"; +const std::set kCustAiCpuKernelOps{kTopK, kEditDistance}; struct AicpuParamHead { uint32_t length; // Total length: include cunstom message