forked from OSSInnovation/mindspore
add cust aicpu path
This commit is contained in:
parent
6c26629404
commit
83b156b287
|
@ -32,6 +32,7 @@ using AicpuTaskInfoPtr = std::shared_ptr<ge::model_runner::AicpuTaskInfo>;
|
|||
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<AddressPtr> &inputs
|
|||
const std::vector<AddressPtr> &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<void *> 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";
|
||||
|
||||
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;
|
||||
(void)std::transform(std::begin(inputs), std::end(inputs), std::back_inserter(input_data_addrs),
|
||||
[](const AddressPtr &input) -> void * { return input->addr; });
|
||||
|
|
|
@ -19,6 +19,7 @@
|
|||
#include <cstdint>
|
||||
#include <vector>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#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<std::string> kCustAiCpuKernelOps{kTopK, kEditDistance};
|
||||
|
||||
struct AicpuParamHead {
|
||||
uint32_t length; // Total length: include cunstom message
|
||||
|
|
Loading…
Reference in New Issue