Program Listing for File kernel_api.h

↰ Return to documentation for file (include/runtime/include/api/kernel_api.h)

#ifndef MINDSPORE_INCLUDE_API_KERNEL_API_H
#define MINDSPORE_INCLUDE_API_KERNEL_API_H
#include <vector>
#include <string>
#include <utility>
#include <map>
#include "include/api/types.h"
namespace mindspore {
class Context;
namespace kernel {
template <typename Primitive>
class MS_API IKernel {
 public:
  IKernel() = default;
  IKernel(const std::vector<mindspore::MSTensor> &inputs, const std::vector<mindspore::MSTensor> &outputs,
          const Primitive *primitive, const mindspore::Context *ctx)
      : context_(ctx), inputs_(std::move(inputs)), outputs_(std::move(outputs)), primitive_(primitive) {}
  virtual ~IKernel() = 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_; }
  const 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 &section) 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_;
  const Primitive *primitive_ = nullptr;
  std::map<std::string, std::string> attrs_;
  const std::map<std::string, std::map<std::string, std::string>> *config_ = nullptr;
};
}  // namespace kernel
}  // namespace mindspore
#endif  // MINDSPORE_INCLUDE_API_KERNEL_API_H