refactor(computation): 现在计算图下降到 kernel 时会直接从图结构删除无意义的节点

Signed-off-by: YdrMaster <ydrml@hotmail.com>
InfiniTensor · Dec 6, 2023 · e2ba442 · e2ba442
1 parent 1600518
commit e2ba442
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 18 deletions.
diff --git a/src/00common/include/common/range.h b/src/00common/include/common/range.h
@@ -16,7 +16,7 @@ namespace refactor {
 
         bool empty() const noexcept { return end_ == begin_; }
         size_t size() const noexcept { return end_ - begin_; }
-        t at(size_t i) const noexcept {
+        t at(size_t i) const {
             ASSERT(i < size(), "Index out of range");
             return operator[](i);
         }

diff --git a/src/01graph_topo/include/graph_topo/linked_graph.hpp b/src/01graph_topo/include/graph_topo/linked_graph.hpp
@@ -37,7 +37,7 @@ namespace refactor::graph_topo {
         Rc<Node> pushNode(TN, std::vector<Rc<Edge>>);
         void eraseNode(count_t);
         void eraseNode(Rc<Node>);
-        size_t cleanup(bool useless(TE const &) = nullptr);
+        size_t cleanup(bool useful(TE const &) = nullptr);
         bool sort();
     };
 
@@ -86,6 +86,7 @@ namespace refactor::graph_topo {
             node.disconnect(i);
         }
         for (auto const &out : node._outputs) {
+            ASSERT(out->_targets.empty(), "Output edge should not have targets");
             out->_source = nullptr;
         }
     }
@@ -180,23 +181,22 @@ namespace refactor::graph_topo {
         _nodes.erase(it);
     }
 
-    LINKED_GRAPH_FN cleanup(bool useless(TE const &))->size_t {
+    LINKED_GRAPH_FN cleanup(bool useful(TE const &))->size_t {
         std::unordered_set<Edge *> outputs;
         outputs.reserve(_outputs.size());
         std::transform(_outputs.begin(), _outputs.end(), std::inserter(outputs, outputs.end()), [](auto const &e) { return e.get(); });
-        auto useful = [&](Rc<Edge> const &e) {
-            return !e->_targets.empty() ||     // 还有节点连接到这个边
-                   outputs.contains(e.get()) ||// 这个边是全图输出
-                   !useless ||                 // 不需要其他判断
-                   !useless(e->_info);         // 这个边其他原因有用
+        auto useful_ = [&](Rc<Edge> const &e) {
+            return !e->_targets.empty() ||      // 还有节点连接到这个边
+                   outputs.contains(e.get()) || // 这个边是全图输出
+                   (useful && useful(e->_info));// 这个边其他原因有用
         };
 
         auto before = _nodes.size();
         while (true) {
             auto endit = std::remove_if(
                 _nodes.begin(), _nodes.end(),
                 [&, this](auto &n) {
-                    auto useless_ = std::none_of(n->_outputs.begin(), n->_outputs.end(), useful);
+                    auto useless_ = std::none_of(n->_outputs.begin(), n->_outputs.end(), useful_);
                     if (useless_) { _cleanupNode(*n); }
                     return useless_;
                 });

diff --git a/src/03runtime/include/runtime/stream.h b/src/03runtime/include/runtime/stream.h
@@ -24,6 +24,7 @@ namespace refactor::runtime {
     struct Edge {
         Arc<hardware::Device::Blob> blob;
         size_t stackOffset;
+        std::string name;
     };
 
     class Stream {
@@ -39,6 +40,8 @@ namespace refactor::runtime {
                std::vector<Node>,
                std::vector<Edge>,
                decltype(_device));
+
+        decltype(_graph) const &graph() const noexcept { return _graph; }
         void setData(count_t, void const *, size_t);
         void setData(count_t, Arc<hardware::Device::Blob>);
         bool getData(count_t, void *, size_t) const;

diff --git a/src/04kernel/src/graph.cc b/src/04kernel/src/graph.cc
@@ -33,6 +33,7 @@ namespace refactor::kernel {
 
         for (auto i : range0_(edges_.size())) {
             auto const &edge = _internal.edges[i];
+            edges_[i].name = edge.name;
             if (edge.data) {
                 auto blob = device->malloc(edge.size);
                 blob->copyFromHost(edge.data->get<void>());

diff --git a/src/05computation/src/graph.cc b/src/05computation/src/graph.cc
@@ -58,7 +58,16 @@ namespace refactor::computation {
 
         auto modifier = graph_topo::InplaceModifier(graph.topology);
         modifier.reconnect(identities);
-        return kernel::Graph(modifier.take(), std::move(nodes), std::move(edges));
+
+        auto temp = graph_topo::LinkedGraph(graph_topo::Graph{
+            modifier.take(),
+            std::move(nodes),
+            std::move(edges),
+        });
+        temp.cleanup();
+        auto [topo__, nodes__, edges__] = temp.intoGraph();
+
+        return kernel::Graph(std::move(topo__), std::move(nodes__), std::move(edges__));
     }
 
     auto Graph::internal() const -> decltype(_internal) const & { return _internal; }

diff --git a/src/09python_ffi/src/executor.cc b/src/09python_ffi/src/executor.cc
@@ -33,21 +33,21 @@ namespace refactor::python_ffi {
     }
 
     void Executor::setInput(count_t i, pybind11::array data) {
-        auto globalInputs = _graph.internal().contiguous().topology.globalInputs();
-        ASSERT(i < globalInputs.size(), "input index out of range");
-        i = globalInputs[i];
+        i = _stream.graph().topology.globalInputs().at(i);
 
-        auto const &tensor = *_graph.internal().contiguous().edges[i].tensor;
+        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;
         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 {
-        auto globalOutputs = _graph.internal().contiguous().topology.globalOutputs();
-        ASSERT(i < globalOutputs.size(), "output index out of range");
-        i = globalOutputs[i];
+        i = _stream.graph().topology.globalOutputs().at(i);
 
-        auto const &tensor = *_graph.internal().contiguous().edges[i].tensor;
+        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 ans = pybind11::array(buildNumpyDType(tensor.dataType), std::move(tensor.shape));
         _stream.getData(i, ans.mutable_data(), ans.nbytes());
         return ans;