Pass api

cinn/CMakeLists.txt

@@ -2,6 +2,7 @@ if (WITH_TESTING)
   cc_library(cinn_gtest_main SRCS gtest_main.cc DEPS gtest)
 endif()
 
+add_subdirectory(api)
 add_subdirectory(auto_schedule)
 add_subdirectory(common)
 add_subdirectory(utils)

cinn/api/CMakeLists.txt

@@ -0,0 +1,8 @@
+core_gather_headers()
+
+gather_srcs(cinnapi_src SRCS
+    op_node.cc
+    tensor_node.cc
+    )
+
+message(STATUS "srcs: ${cinnapi_src}")

cinn/api/README.md

@@ -0,0 +1,44 @@
+The classes in this directory are the interface of group fusion pass, you can use these apis to build the stragey for group fusion.
+
+The Class and APIs are following:
+
+`OpGroup` : A set of op nodes, which will pass to cinn backend for generating kernel code. Two groups can fuse togather according to the rule of merging written in the passes.
+
+`OpNode` : Map the op in the program.
+
+`TensorNode` : Map the tensor in the program.
+
+`Shape` : The shape infomation of tensor
+
+`FusePassCtx` : The context is the parameter for the pass, it hold the data all you need in the pass.
+
+`FuseHelper` : We provide some util methods such as `DetectCycleIfFuse` in fuse_helper to simplify development of pass.
+
+| Class      | method | description |
+| :--:       | :--: | :--: |
+| `OpGroup`  | kind()| Get the Kind of group |
+|            | producers()| Get producer groups of current group |
+|            | consumers() | Get consumer groups of current group |
+|            | WalkOpNodes(const std::function<void(const OpNode&)>& VisitOpNode) | Visit the op_nodes in the group and execute the VisitOpNode function for each OpNode |
+|  |  |  |
+| `OpNode`   | kind() | Get the Kind of op_node |
+|            | inputs() | Get input tensors of op_node |
+|            | outputs() | Get output tensors of op_node |
+|            | GetAttr(const std::string& attr_name) | Get attribute of op_node by attr name |
+|  |  |  |
+| `TensorNode` | shape() | Get shape of tensor |
+|            | producer() | Get the producer op_node of tensor |
+|            | consumers() | Get the consumer op_nodes of tensor |
+|  |  |  |
+| `Shape`    | numel() | Get total number of elements in the shape |
+|            | other methods are same with std::vector<int64_t> | |
+|  |  |  |
+| `LightwareFusePassCtx` | PickOpGroup() | Get the current group in the pass context |
+|                      | void EnableFuse(const OpGroup& first, const OpGroup& second) | Mark the two groups which can fuse togather |
+|  | fuse_helper()     | Get the fuse_helper provided by pass context  |
+|  |  |  |
+| `InputFusePassCtx`   | PickConsumersWithSameInputs() | Get all consumer groups for input tensors of graph |
+|                      | void EnableFuse(const OpGroup& first, const OpGroup& second) | Mark the two groups which can fuse togather |
+|  | fuse_helper()     | Get the fuse_helper provided by pass context  |
+|  |  |  |
+| `FuseHelper` | DetectCycleIfFuse(const OpGroup& first, const OpGroup& second) | Whether there is cycle in graph after fusing two groups |

cinn/api/op_group.h

@@ -0,0 +1,155 @@
+// Copyright (c) 2023 CINN Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+
+#include "cinn/api/op_node.h"
+#include "cinn/hlir/framework/graph.h"
+#include "cinn/hlir/pass/fusion_helper_base.h"
+
+namespace cinn {
+namespace api {
+
+using Comparator = hlir::framework::Graph::Group::SharedGroupComparator;
+using Hasher     = hlir::framework::Graph::Group::SharedGroupHasher;
+
+class OpGroup {
+ public:
+  OpGroup(const std::shared_ptr<hlir::framework::Graph::Group>& group) : group_(group) {}
+
+  OpGroup(const OpGroup& other) = default;
+
+  class OpGroupListIterator {
+   public:
+    OpGroupListIterator(
+        std::unordered_set<std::shared_ptr<hlir::framework::Graph::Group>, Hasher, Comparator>::const_iterator it)
+        : iter_(it) {}
+
+    OpGroupListIterator& operator++() {
+      ++iter_;
+      return *this;
+    }
+
+    OpGroupListIterator operator++(int) {
+      OpGroupListIterator tmp = *this;
+      ++iter_;
+      return tmp;
+    }
+
+    bool operator==(const OpGroupListIterator& other) const { return iter_ == other.iter_; }
+
+    bool operator!=(const OpGroupListIterator& other) const { return !(*this == other); }
+
+    OpGroup operator*() const { return OpGroup(*iter_); }
+
+   private:
+    std::unordered_set<std::shared_ptr<hlir::framework::Graph::Group>, Hasher, Comparator>::const_iterator iter_;
+  };
+
+  class ProducerOpGroupListView {
+   public:
+    ProducerOpGroupListView(const std::weak_ptr<hlir::framework::Graph::Group>& group) : group_(group) {}
+
+    ProducerOpGroupListView(const ProducerOpGroupListView& other) = delete;
+    ProducerOpGroupListView(ProducerOpGroupListView&& other)      = delete;
+
+    ProducerOpGroupListView& operator=(const ProducerOpGroupListView& other) = delete;
+
+    using const_iterator = OpGroupListIterator;
+
+    size_t size() const { return group_.lock()->producer_groups().size(); }
+
+    const_iterator begin() const { return const_iterator(group_.lock()->producer_groups().begin()); }
+
+    const_iterator end() const { return const_iterator(group_.lock()->producer_groups().end()); }
+
+   private:
+    const std::weak_ptr<hlir::framework::Graph::Group> group_;
+  };
+
+  class ConsumerOpGroupListView {
+   public:
+    ConsumerOpGroupListView(const std::weak_ptr<hlir::framework::Graph::Group>& group) : group_(group) {}
+
+    ConsumerOpGroupListView(const ConsumerOpGroupListView& other) = delete;
+    ConsumerOpGroupListView(ConsumerOpGroupListView&& other)      = delete;
+
+    ConsumerOpGroupListView& operator=(const ConsumerOpGroupListView& other) = delete;
+
+    using const_iterator = OpGroupListIterator;
+
+    size_t size() const { return group_.lock()->consumer_groups().size(); }
+
+    const_iterator begin() const { return const_iterator(group_.lock()->consumer_groups().begin()); }
+
+    const_iterator end() const { return const_iterator(group_.lock()->consumer_groups().end()); }
+
+   private:
+    const std::weak_ptr<hlir::framework::Graph::Group> group_;
+  };
+
+  const std::string& group_id() const { return group_.lock()->group_id; }
+
+  hlir::framework::OpPatternKind kind() const { return group_.lock()->kind(); }
+
+  // The WalkOpNodes function is used to traverse the op_nodes in the group and execute
+  // the VisitOpNode function for each OpNode. This function is equivalent to for loop
+  // for op_nodes in graph.
+  //
+  // In order to avoid unnecessary memory copies, we use WalkOpNodes function instead of
+  // providing a function to get all op_nodes directly.
+  //
+  // Example: Get the all Reduction op_nodes in the group.
+  //   OpGroup group = ...;
+  //   std::set<api::OpNode> reduce_op_set;
+  //   group.WalkOpNodes([&reduce_op_set](const api::OpNode& op){
+  //     // The lambda funtion of VisitOpNode to get reduction op_nodes.
+  //     if (op.kind() == OpPatternKind::kReduction) {
+  //       reduce_op_set.insert(op);
+  //     }
+  //   });
+  void WalkOpNodes(const std::function<void(const OpNode&)>& VisitOpNode) const {
+    group_.lock()->WalkNodes(
+        [&](const hlir::framework::Node* node) { VisitOpNode(OpNode(node, group_.lock()->graph_)); });
+  }
+
+  ProducerOpGroupListView producers() const { return ProducerOpGroupListView(group_); }
+
+  ConsumerOpGroupListView consumers() const { return ConsumerOpGroupListView(group_); }
+
+  std::shared_ptr<hlir::framework::Graph::Group> GetGroup() const { return group_.lock(); }
+
+  bool operator==(const OpGroup& other) const { return group_.lock().get() == other.group_.lock().get(); }
+
+  bool operator<(const OpGroup& other) const { return group_.lock().get() < other.group_.lock().get(); }
+
+ private:
+  const std::weak_ptr<hlir::framework::Graph::Group> group_;
+};
+
+}  // namespace api
+}  // namespace cinn
+
+namespace std {
+
+template <>
+struct hash<cinn::api::OpGroup> {
+  size_t operator()(const cinn::api::OpGroup& obj) const {
+    return std::hash<int64_t>()(reinterpret_cast<uint64_t>(obj.GetGroup().get()));
+  }
+};
+
+}  // namespace std

cinn/api/op_node.cc

@@ -0,0 +1,29 @@
+// Copyright (c) 2023 CINN Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "cinn/api/op_node.h"
+
+namespace cinn {
+namespace api {
+
+TensorNode OpNode::InputTensorListView::operator[](size_t index) const {
+  return TensorNode(edges_[index]->source()->safe_as<hlir::framework::NodeData>(), graph_);
+}
+
+TensorNode OpNode::OutputTensorListView::operator[](size_t index) const {
+  return TensorNode(edges_[index]->sink()->safe_as<hlir::framework::NodeData>(), graph_);
+}
+
+}  // namespace api
+}  // namespace cinn

cinn/api/op_node.h

@@ -0,0 +1,133 @@
+// Copyright (c) 2023 CINN Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include "cinn/hlir/framework/graph.h"
+#include "cinn/hlir/pass/fusion_helper_base.h"
+#include "cinn/api/tensor_node.h"
+#include "cinn/hlir/framework/op.h"
+
+namespace cinn {
+namespace api {
+
+using OpPatternKind = cinn::hlir::framework::OpPatternKind;
+using Attribute = cinn::utils::Attribute;
+
+class OpNode {
+ public:
+  OpNode(const hlir::framework::Node* node, const hlir::framework::Graph* graph) : node_(node), graph_(graph), input_tensors_(node->inlinks_in_order(), graph_), output_tensors_(node->outlinks_in_order(), graph_) {
+    VLOG(1) << "[OpNode] node: " << node->id();
+  }
+
+  OpNode(const OpNode& other) : node_(other.node_), graph_(other.graph_), input_tensors_(node_->inlinks_in_order(), graph_), output_tensors_(node_->outlinks_in_order(), graph_) {}
+
+  OpPatternKind kind () const {
+    thread_local const static hlir::framework::OpValueType<OpPatternKind>& op_pattern_dict = hlir::framework::Operator::GetAttrs<OpPatternKind>("OpPattern");
+    auto kind = op_pattern_dict[node_->op()];
+
+    if (kind == hlir::framework::kBroadcast) {
+      // As binary op was defined as broadcast, actually it should be element-wise.
+      if (node_->op()->name != "broadcast_to") {
+        return hlir::framework::kElementWise;
+      }
+    }
+    return kind;
+  }
+
+  class InputTensorListView {
+   public:
+    InputTensorListView(const std::vector<common::Shared<common::GraphEdge>>& edges, const hlir::framework::Graph* graph) : edges_(edges), graph_(graph) {}
+
+    InputTensorListView(const InputTensorListView& other) = delete;
+    InputTensorListView(InputTensorListView&& other) = delete;
+
+    InputTensorListView& operator=(const InputTensorListView& other) = delete;
+
+    size_t size() const { return edges_.size(); }
+
+    TensorNode operator[](size_t index) const;
+
+   private:
+    std::vector<common::Shared<common::GraphEdge>> edges_;
+    const hlir::framework::Graph* graph_;
+  };
+
+  class OutputTensorListView {
+   public:
+    OutputTensorListView(const std::vector<common::Shared<common::GraphEdge>>& edges, const hlir::framework::Graph* graph) : edges_(edges), graph_(graph) {}
+
+    OutputTensorListView(const OutputTensorListView& other) = delete;
+    OutputTensorListView(OutputTensorListView&& other) = delete;
+
+    OutputTensorListView& operator=(const OutputTensorListView& other) = delete;
+
+    size_t size() const { return edges_.size(); }
+
+    TensorNode operator[](size_t index) const;
+
+   private:
+    std::vector<common::Shared<common::GraphEdge>> edges_;
+    const hlir::framework::Graph* graph_;
+  };
+
+  bool operator == (const OpNode& other) const {
+    return node_ == other.node_;
+  }
+
+  bool operator < (const OpNode& other) const {
+    return node_ < other.node_;
+  }
+
+  const InputTensorListView& inputs() const {
+    return input_tensors_;
+  }
+
+  const OutputTensorListView& outputs() const {
+    return output_tensors_;
+  }
+
+  template <typename T>
+  const T& GetAttr(const std::string& attr_name) const {
+    return absl::get<T>(GetAttr(attr_name));
+  }
+
+ private:
+  const Attribute& GetAttr(const std::string& attr_name) const {
+    return node_->attrs.attr_store.at(attr_name);
+  }
+
+  friend struct std::hash<OpNode>;
+
+  const hlir::framework::Node* node_;
+  const hlir::framework::Graph* graph_;
+
+  const InputTensorListView input_tensors_;
+  const OutputTensorListView output_tensors_;
+};
+
+}  // namespace api
+}  // namespace cinn
+
+namespace std {
+
+template <>
+struct hash<cinn::api::OpNode> {
+  size_t operator()(const cinn::api::OpNode& obj) const {
+    return std::hash<int64_t>()(reinterpret_cast<uint64_t>(obj.node_));
+  }
+};
+
+} // namespace std