[CINN] Refactor pass api of group fusion in CINN

cmake/cinn/core.cmake

@@ -435,6 +435,28 @@ function(download_and_uncompress INSTALL_DIR URL FILENAME)
     INSTALL_COMMAND "")
 endfunction()
 
+set(fusion_pass_file
+    ${CMAKE_CURRENT_BINARY_DIR}/paddle/cinn/hlir/pass/use_general_pass.h
+    CACHE INTERNAL "use_general_pass.h file")
+file(
+  WRITE ${fusion_pass_file}
+  "#include \"paddle/cinn/common/macros.h\" // Generated by the paddle/cinn/hlir/pass/CMakeLists.txt.  DO NOT EDIT!\n\n"
+)
+
+function(find_fusion_pass_register FILENAME ADD_PATH PATTERN)
+  # set op_name to OUTPUT
+  file(READ ${FILENAME} CONTENT)
+  string(REGEX MATCHALL "${PATTERN}\\([a-zA-Z0-9_]*," fusion_pass_patterns
+               "${CONTENT}")
+  if(NOT fusion_pass_patterns STREQUAL "")
+    foreach(pass_pattern ${fusion_pass_patterns})
+      string(REPLACE "${PATTERN}(" "" pass_pattern "${pass_pattern}")
+      string(REPLACE "," "" pass_pattern "${pass_pattern}")
+      file(APPEND ${ADD_PATH} "USE_FUSION_PASS(${pass_pattern});\n")
+    endforeach()
+  endif()
+endfunction()
+
 function(gather_srcs SRC_GROUP)
   set(options)
   set(oneValueArgs)
@@ -444,6 +466,8 @@ function(gather_srcs SRC_GROUP)
     set(${SRC_GROUP}
         "${${SRC_GROUP}};${CMAKE_CURRENT_SOURCE_DIR}/${cpp}"
         CACHE INTERNAL "")
+    find_fusion_pass_register("${CMAKE_CURRENT_SOURCE_DIR}/${cpp}"
+                              ${fusion_pass_file} "CINN_REGISTER_FUSION_PASS")
   endforeach()
 endfunction()
 

paddle/cinn/CMakeLists.txt

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

paddle/cinn/api/CMakeLists.txt

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

paddle/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 |

paddle/cinn/api/fuse_pass_context.h

@@ -0,0 +1,131 @@
+// 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 "paddle/cinn/api/op_group.h"
+
+namespace cinn {
+namespace api {
+
+class FusePassContext {
+ public:
+  virtual ~FusePassCtx() {}
+
+  virtual void EnableFuse(const OpGroup& first, const OpGroup& second) = 0;
+
+ protected:
+  FusePassCtx() = default;
+};
+
+class LightwareFusePassCtx : public FusePassContext {
+ public:
+  virtual ~LightwareFusePassCtx() {}
+
+  virtual const OpGroup& PickOpGroup() const = 0;
+
+  virtual const FuseHelper& fuse_helper() const = 0;
+
+  virtual void EnableFuse(const OpGroup& first, const OpGroup& second) = 0;
+
+ protected:
+  LightwareFusePassCtx() = default;
+};
+
+class GraphGroupLightwareFusePassCtx final : public LightwareFusePassCtx {
+ public:
+  GraphGroupLightwareFusePassCtx(
+      const FusionHelperBase* graph_group_fusion_helper,
+      const OpGroup& group,
+      const std::function<void(const OpGroup& first, const OpGroup& second)>&
+          EnableFuse)
+      : graph_group_fusion_helper_(graph_group_fusion_helper),
+        group_(group),
+        EnableFuse_(EnableFuse),
+        fuse_helper_(
+            new GraphGroupFuseHelper<GraphGroupLightwareFusePassCtx>(this)) {}
+
+  const OpGroup& PickOpGroup() const override { return group_; }
+
+  const FuseHelper& fuse_helper() const override { return *fuse_helper_; }
+
+  void EnableFuse(const OpGroup& first, const OpGroup& second) override {
+    EnableFuse_(first, second);
+  }
+
+  const FusionHelperBase& graph_group_fusion_helper() const {
+    return *graph_group_fusion_helper_;
+  }
+
+ private:
+  const FusionHelperBase* graph_group_fusion_helper_;
+  const OpGroup& group_;
+  const std::function<void(const OpGroup& first, const OpGroup& second)>
+      EnableFuse_;
+  const std::unique_ptr<const FuseHelper> fuse_helper_;
+};
+
+class InputFusePassCtx : public FusePassCtx {
+ public:
+  virtual ~InputFusePassCtx() {}
+
+  virtual const OpGroupList& PickConsumersWithSameInputs() const = 0;
+
+  virtual const FuseHelper& fuse_helper() const = 0;
+
+  virtual void EnableFuse(const OpGroup& first, const OpGroup& second) = 0;
+
+ protected:
+  InputFusePassCtx() = default;
+};
+
+class GraphGroupInputFusePassCtx final : public InputFusePassCtx {
+ public:
+  GraphGroupInputFusePassCtx(
+      const FusionHelperBase* graph_group_fusion_helper,
+      const OpGroupList& groups,
+      const std::function<void(const OpGroup& first, const OpGroup& second)>&
+          EnableFuse)
+      : graph_group_fusion_helper_(graph_group_fusion_helper),
+        groups_(groups),
+        EnableFuse_(EnableFuse),
+        fuse_helper_(
+            new GraphGroupFuseHelper<GraphGroupInputFusePassCtx>(this)) {}
+
+  const OpGroupList& PickConsumersWithSameInputs() const override {
+    return groups_;
+  }
+
+  const FuseHelper& fuse_helper() const override { return *fuse_helper_; }
+
+  void EnableFuse(const OpGroup& first, const OpGroup& second) override {
+    EnableFuse_(first, second);
+  }
+
+  const FusionHelperBase& graph_group_fusion_helper() const {
+    return *graph_group_fusion_helper_;
+  }
+
+ private:
+  const FusionHelperBase* graph_group_fusion_helper_;
+  const OpGroupList& groups_;
+  const std::function<void(const OpGroup& first, const OpGroup& second)>
+      EnableFuse_;
+  const std::unique_ptr<const FuseHelper> fuse_helper_;
+};
+
+}  // namespace api
+}  // namespace cinn