execution_engine.h

Go to the documentation of this file.

#pragma once

#include <map>
#include <memory>
#include <string>
#include <vector>

#include <dlpack/dlpack.h>

#include "tc/core/mapping_options.h"
#include "tc/core/tc_executor.h"

namespace tc {

class ExecutionEngine {
 public:
  struct ExecutorInfo {
    ExecutorInfo(
        std::string id,
        std::vector<const DLTensor*> inputsInfo,
        std::unique_ptr<MappingOptions> options,
        lang::TreeRef tc,
        size_t handle)
        : identifier(id),
          inputsInfo(dlutils::makeDLTensorVector(inputsInfo)),
          options(std::move(options)),
          exec(tc, inputsInfo),
          objectLocalHandle(handle) {}

    void clear() {
      exec.clearRTC();
    }

    std::string identifier;
    std::vector<dlutils::DLTensorUPtr> inputsInfo;
    std::unique_ptr<MappingOptions> options;
    TcExecutor exec;
    size_t objectLocalHandle;
  };

  ExecutionEngine() = default;

  void define(const std::string& language);

  void define(const std::vector<lang::TreeRef>& treeRefs);

  void addTC(const std::string& tc);

  std::vector<const DLTensor*> inferOutputTensorInfo(
      const std::string& name,
      const std::vector<const DLTensor*>& inTensorPtrs);

  lang::TreeRef treeForFunction(const std::string& name) {
    return tcNameMap_.at(name);
  }

  // TODO: Pass autotuning info (none by default, otherwise some struct with
  //       maxtime and other things)

  size_t compile(
      const std::string& name,
      const std::vector<const DLTensor*>& inputs,
      const MappingOptions& options);

  // TODO: sanity check on name and input / output sizes.
  Duration run(
      size_t handle,
      const std::vector<const DLTensor*>& inputs,
      const std::vector<DLTensor*>& outputs,
      bool profile = false,
      std::function<bool(const ExecutorInfo*)> pruningFunction =
          [](const ExecutorInfo*) { return false; });

  void uncheckedRun(
      size_t handle,
      const std::vector<const void*>& inputs,
      const std::vector<void*>& outputs);

  void clear(size_t handle);

 private:
  size_t getHandle(
      const std::string& name,
      const std::vector<const DLTensor*>& inputsInfo,
      const MappingOptions& options);
  std::unique_ptr<ExecutorInfo> makeExecutorInfo(
      const std::string& name,
      const std::vector<const DLTensor*>& inputsInfo,
      const MappingOptions& options);
  size_t emplaceExecutor(std::unique_ptr<ExecutorInfo> p);

  std::mutex executorInfoMutex;

  // XXX:if ExecutorInfo is moved/copied (even when the vector's underlying
  // storage is extended) something inside isl segfaults,  unique_ptr is used as
  // a workaround
  std::vector<std::unique_ptr<ExecutorInfo>> executors_;
  std::map<std::string, lang::TreeRef> tcNameMap_;

  size_t uidCounter = 0;
};

} // namespace tc