feat(python_ffi): 支持从 executor 存取 Device::Blob

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

src/03runtime/include/runtime/stream.h

@@ -42,9 +42,11 @@ namespace refactor::runtime {
                decltype(_device));
 
         decltype(_graph) const &graph() const noexcept { return _graph; }
-        void setData(count_t, void const *, size_t);
+        auto setData(count_t, size_t) -> Arc<hardware::Device::Blob>;
         void setData(count_t, Arc<hardware::Device::Blob>);
-        bool getData(count_t, void *, size_t) const;
+        auto getData(count_t) -> Arc<hardware::Device::Blob> const;
+        void setData(count_t, void const *, size_t);
+        bool copyData(count_t, void *, size_t) const;
         void run();
         auto bench(void (*sync)()) -> std::vector<std::chrono::nanoseconds>;
         void trace(std::function<void(count_t, void const *const *, void const *const *)>);

src/03runtime/src/stream.cc

@@ -18,15 +18,21 @@ namespace refactor::runtime {
               std::move(edges),
           } {}
 
+    auto Stream::setData(count_t i, size_t size) -> Arc<hardware::Device::Blob> {
+        return _graph.edges[i].blob = _device->malloc(size);
+    }
+    void Stream::setData(count_t i, Arc<hardware::Device::Blob> blob) {
+        _graph.edges[i].blob = std::move(blob);
+    }
     void Stream::setData(count_t i, void const *data, size_t size) {
         auto blob = _device->malloc(size);
         blob->copyFromHost(data, size);
         _graph.edges[i].blob = std::move(blob);
     }
-    void Stream::setData(count_t i, Arc<hardware::Device::Blob> blob) {
-        _graph.edges[i].blob = std::move(blob);
+    auto Stream::getData(count_t i) -> Arc<hardware::Device::Blob> const {
+        return _graph.edges[i].blob;
     }
-    bool Stream::getData(count_t i, void *data, size_t size) const {
+    bool Stream::copyData(count_t i, void *data, size_t size) const {
         if (!_graph.edges[i].blob) { return false; }
         _graph.edges[i].blob->copyToHost(data, size);
         return true;

src/09python_ffi/src/executor.cc

@@ -26,7 +26,7 @@ namespace refactor::python_ffi {
         for (auto i : graph.topology.globalInputs()) {
             auto size = graph.edges[i].tensor->bytesSize();
             buffer.resize(size);
-            if (stream.getData(i, buffer.data(), size)) {
+            if (stream.copyData(i, buffer.data(), size)) {
                 _stream.setData(i, buffer.data(), size);
             }
         }
@@ -35,24 +35,36 @@ namespace refactor::python_ffi {
     void Executor::setInput(count_t i, pybind11::array data) {
         i = _stream.graph().topology.globalInputs().at(i);
 
-        auto const &name = _stream.graph().edges[i].name;
-        auto const &edges = _graph.internal().contiguous().edges;
-        auto const &tensor = *std::find_if(edges.begin(), edges.end(), [&](auto const &e) { return e.name == name; })->tensor;
+        auto const &tensor = *_graph.internal().contiguous().edges[i].tensor;
         ASSERT(tensor.bytesSize() == static_cast<size_t>(data.nbytes()), "input size mismatch");
         _stream.setData(i, data.data(), data.nbytes());
     }
 
     auto Executor::getOutput(count_t i) -> pybind11::array {
         i = _stream.graph().topology.globalOutputs().at(i);
 
-        auto const &name = _stream.graph().edges[i].name;
-        auto const &edges = _graph.internal().contiguous().edges;
-        auto const &tensor = *std::find_if(edges.begin(), edges.end(), [&](auto const &e) { return e.name == name; })->tensor;
+        auto const &tensor = *_graph.internal().contiguous().edges[i].tensor;
         auto ans = pybind11::array(buildNumpyDType(tensor.dataType), std::move(tensor.shape));
-        _stream.getData(i, ans.mutable_data(), ans.nbytes());
+        _stream.copyData(i, ans.mutable_data(), ans.nbytes());
         return ans;
     }
 
+    auto Executor::pin(count_t i) -> Arc<hardware::Device::Blob> {
+        i = _stream.graph().topology.globalInputs().at(i);
+
+        if (auto pinned = _stream.getData(i); pinned) {
+            return pinned;
+        } else {
+            auto const &tensor = *_graph.internal().contiguous().edges[i].tensor;
+            return _stream.setData(i, tensor.bytesSize());
+        }
+    }
+    void Executor::setPinned(count_t i, Arc<hardware::Device::Blob> pinned) {
+        i = _stream.graph().topology.globalInputs().at(i);
+
+        _stream.setData(i, std::move(pinned));
+    }
+
     void Executor::run() {
         _stream.run();
     }

src/09python_ffi/src/executor.h

@@ -16,6 +16,8 @@ namespace refactor::python_ffi {
         void dispatch(Arc<hardware::Device>, std::string allocator);
         void setInput(count_t, pybind11::array);
         auto getOutput(count_t) -> pybind11::array;
+        auto pin(count_t) -> Arc<hardware::Device::Blob>;
+        void setPinned(count_t, Arc<hardware::Device::Blob>);
         void run();
         void bench(bool sync);
         void trace(std::string path, std::string format);

src/09python_ffi/src/main.cpp

@@ -21,6 +21,7 @@ namespace refactor::python_ffi {
         py::class_<Tensor      , Arc<Tensor>      >(m, "Tensor"      );
         py::class_<OpBox       , Arc<OpBox>       >(m, "Operator"    );
         py::class_<Device      , Arc<Device>      >(m, "Device"      );
+        py::class_<Device::Blob, Arc<Device::Blob>>(m, "Pinned"      );
 
         m   .def("config_log"      , &configLog                  , return_::automatic )
             .def("find_device"     , &findDevice                 , return_::move      )
@@ -44,6 +45,8 @@ namespace refactor::python_ffi {
             .def("dispatch"        , &Executor::dispatch         , return_::automatic )
             .def("set_input"       , &Executor::setInput         , return_::automatic )
             .def("get_output"      , &Executor::getOutput        , return_::move      )
+            .def("pin"             , &Executor::pin              , return_::move      )
+            .def("set_pinned"      , &Executor::setPinned        , return_::automatic )
             .def("run"             , &Executor::run              , return_::automatic )
             .def("bench"           , &Executor::bench            , return_::automatic )
             .def("trace"           , &Executor::trace            , return_::automatic )