From 6871fff02bd642c52674e620fe5364ee46afebe5 Mon Sep 17 00:00:00 2001
From: YdrMaster <ydrml@hotmail.com>
Date: Thu, 23 Feb 2023 11:08:00 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AF=BC=E5=87=BA=E5=88=86=E9=85=8D?=
 =?UTF-8?q?=E5=86=85=E5=AD=98=E5=92=8C=E8=BF=90=E8=A1=8C=E6=8E=A8=E7=90=86?=
 =?UTF-8?q?=E7=9A=84=E6=8E=A5=E5=8F=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: YdrMaster <ydrml@hotmail.com>
---
 include/core/graph_handler.h | 6 ++++++
 src/ffi/ffi_infinitensor.cc  | 4 +++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/include/core/graph_handler.h b/include/core/graph_handler.h
index 22897c4c..b74f8e13 100644
--- a/include/core/graph_handler.h
+++ b/include/core/graph_handler.h
@@ -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
diff --git a/src/ffi/ffi_infinitensor.cc b/src/ffi/ffi_infinitensor.cc
index b48a9e1c..9300cdb8 100644
--- a/src/ffi/ffi_infinitensor.cc
+++ b/src/ffi/ffi_infinitensor.cc
@@ -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