add cust aicpu path

pull/6634/head
wuxuejian 4 years ago
parent 6c26629404
commit 83b156b287

@ -32,6 +32,7 @@ using AicpuTaskInfoPtr = std::shared_ptr<ge::model_runner::AicpuTaskInfo>;
namespace mindspore { namespace mindspore {
namespace kernel { namespace kernel {
constexpr auto AICPU_OPS_SO_NAME = "libaicpu_kernels.so"; constexpr auto AICPU_OPS_SO_NAME = "libaicpu_kernels.so";
constexpr auto CUST_AICPU_OPS_SO_NAME = "libcpu_kernels.so";
AicpuOpKernelMod::AicpuOpKernelMod() : anf_node_(nullptr) {} AicpuOpKernelMod::AicpuOpKernelMod() : anf_node_(nullptr) {}
@ -66,8 +67,12 @@ void AicpuOpKernelMod::CreateCpuKernelInfo(const std::vector<AddressPtr> &inputs
const std::vector<AddressPtr> &outputs) { const std::vector<AddressPtr> &outputs) {
MS_LOG(INFO) << "CreateCpuKernelInfoOffline start"; 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 // InputOutputAddr
vector<void *> io_addrs; vector<void *> io_addrs;
(void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(io_addrs), (void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(io_addrs),
@ -148,7 +153,12 @@ std::vector<TaskInfoPtr> AicpuOpKernelMod::GenTask(const std::vector<AddressPtr>
MS_LOG(INFO) << "AicpuOpKernelMod GenTask start"; MS_LOG(INFO) << "AicpuOpKernelMod GenTask start";
stream_id_ = stream_id; 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<void *> input_data_addrs; std::vector<void *> input_data_addrs;
(void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(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; }); [](const AddressPtr &input) -> void * { return input->addr; });

@ -19,6 +19,7 @@
#include <cstdint> #include <cstdint>
#include <vector> #include <vector>
#include <map> #include <map>
#include <set>
#include <string> #include <string>
#include "backend/kernel_compiler/kernel.h" #include "backend/kernel_compiler/kernel.h"
namespace mindspore { namespace mindspore {
@ -41,6 +42,9 @@ constexpr auto kSeed1 = "Seed1";
constexpr auto kSeed2 = "seed2"; constexpr auto kSeed2 = "seed2";
constexpr auto kTopK = "TopK"; constexpr auto kTopK = "TopK";
constexpr auto kTopKV2 = "TopKV2"; constexpr auto kTopKV2 = "TopKV2";
constexpr auto kEditDistance = "EditDistance";
constexpr auto kCustRunApi = "RunCpuKernel";
const std::set<std::string> kCustAiCpuKernelOps{kTopK, kEditDistance};
struct AicpuParamHead { struct AicpuParamHead {
uint32_t length; // Total length: include cunstom message uint32_t length; // Total length: include cunstom message

Loading…
Cancel
Save