feat: 导出分配内存和运行推理的接口

Signed-off-by: YdrMaster <ydrml@hotmail.com>
This commit is contained in:
YdrMaster 2023-02-23 11:08:00 +08:00
parent d7e52054e6
commit 6871fff02b
2 changed files with 9 additions and 1 deletions

View File

@ -76,6 +76,12 @@ class GraphHandlerObj {
const optional<vector<int>> &steps);
Tensor pad(Tensor input, Tensor output, const vector<int> &pads,
const optional<vector<int>> &axes);
//------ runtime
inline void data_malloc() { g->dataMalloc(); }
inline void run() { g->getRuntime()->run(g); }
};
} // namespace infini

View File

@ -102,7 +102,9 @@ void init_graph_builder(py::module &m) {
.def("pad",
py::overload_cast<Tensor, Tensor, const vector<int> &,
const optional<vector<int>> &>(&Handler::pad),
policy::move);
policy::move)
.def("data_malloc", &Handler::data_malloc, policy::automatic)
.def("run", &Handler::run, policy::automatic);
}
} // namespace infini