aten_compiler.h

Go to the documentation of this file.

#pragma once

#include <chrono>
#include <string>
#include <vector>

#include <ATen/ATen.h>
#include <ATen/DLConvertor.h>

#include "tc/core/execution_engine.h"
#include "tc/core/utils/cuda_info.h"
#include "tc/lang/parser.h"

namespace tc {


class ATenCompilationUnit {
 public:
  explicit ATenCompilationUnit();

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

  // TODO: Pass struct to allow autotuning
  size_t compile(
      const std::string& name,
      const std::vector<at::Tensor>& inputs,
      const MappingOptions& options);

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

  Duration run(
      const std::string& name,
      const std::vector<at::Tensor>& inputs,
      std::vector<at::Tensor>& outputs,
      size_t handle,
      bool profile = false);

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

 private:
  std::unique_ptr<ExecutionEngine> executionEngine_;
};

std::pair<std::vector<DLTensor*>, std::vector<DLManagedTensor*>>
toDlpackTensors(const std::vector<at::Tensor>& tensors);

std::pair<std::vector<const DLTensor*>, std::vector<DLManagedTensor*>>
toConstDlpackTensors(const std::vector<at::Tensor>& tensors);

void deleteDlmTensors(std::vector<DLManagedTensor*>& tensors);

} // namespace tc