Program Listing for File kernel.h
↰ Return to documentation for file (include/include/api/kernel.h
)
#ifndef MINDSPORE_INCLUDE_API_KERNEL_H
#define MINDSPORE_INCLUDE_API_KERNEL_H
#include <vector>
#include <string>
#include <utility>
#include <map>
#include "schema/model_generated.h"
#include "include/api/types.h"
#include "include/api/context.h"
namespace mindspore::kernel {
class Kernel {
public:
Kernel() = default;
Kernel(const std::vector<mindspore::MSTensor> &inputs, const std::vector<mindspore::MSTensor> &outputs,
const schema::Primitive *primitive, const mindspore::Context *ctx)
: context_(ctx), inputs_(std::move(inputs)), outputs_(std::move(outputs)), primitive_(primitive) {
Initialize();
}
virtual ~Kernel() = default;
virtual int Prepare() = 0;
virtual int Execute() = 0;
virtual int ReSize() = 0;
virtual void set_inputs(const std::vector<mindspore::MSTensor> &in_tensors) { this->inputs_ = in_tensors; }
virtual void set_input(mindspore::MSTensor in_tensor, int index) { this->inputs_[index] = in_tensor; }
virtual void set_outputs(const std::vector<mindspore::MSTensor> &out_tensors) { this->outputs_ = out_tensors; }
virtual void set_output(mindspore::MSTensor out_tensor, int index) { this->outputs_[index] = out_tensor; }
virtual const std::vector<mindspore::MSTensor> &inputs() { return this->inputs_; }
virtual const std::vector<mindspore::MSTensor> &outputs() { return this->outputs_; }
std::string name() const { return this->name_; }
void set_name(const std::string &name) { this->name_ = name; }
const mindspore::Context *context() const { return this->context_; }
virtual schema::PrimitiveType type() const { return type_; }
virtual schema::QuantType quant_type() const { return quant_type_; }
const schema::Primitive *primitive() const { return this->primitive_; }
std::string GetAttr(const std::string &key) const {
auto iter = attrs_.find(key);
if (iter != attrs_.end()) {
return iter->second;
}
return "";
}
void SetConfig(const std::map<std::string, std::map<std::string, std::string>> *config) {
config_ = config;
}
std::map<std::string, std::string> GetConfig(const std::string §ion) const {
if (config_ == nullptr) {
return std::map<std::string, std::string>();
}
auto iter = config_->find(section);
if (iter != config_->end()) {
return iter->second;
}
return std::map<std::string, std::string>();
}
protected:
void SetAttr(const std::string &key, const std::string &value) { attrs_[key] = value; }
std::string name_;
const mindspore::Context *context_ = nullptr;
std::vector<mindspore::MSTensor> inputs_;
std::vector<mindspore::MSTensor> outputs_;
schema::PrimitiveType type_ = schema::PrimitiveType_NONE;
const schema::Primitive *primitive_ = nullptr;
std::map<std::string, std::string> attrs_;
const std::map<std::string, std::map<std::string, std::string>> *config_;
schema::QuantType quant_type_ = schema::QuantType_QUANT_NONE;
private:
void Initialize();
};
} // namespace mindspore::kernel
#endif // MINDSPORE_INCLUDE_API_KERNEL_H