add decomposer registry (PaddlePaddle#458)

* add decomposer registry

Co-authored-by: Yan Chunwei <yanchunwei@outlook.com>

Author: Zhang Ting

cinn/common/target.cc

@@ -2,6 +2,8 @@
 
 #include <glog/logging.h>
 
+#include <sstream>
+
 #include "cinn/runtime/cinn_runtime.h"
 
 namespace cinn {
@@ -49,6 +51,12 @@ int Target::get_target_bits() const {
   return -1;
 }
 
+std::string Target::arch_str() const {
+  std::ostringstream oss;
+  oss << arch;
+  return oss.str();
+}
+
 std::ostream &operator<<(std::ostream &os, const Target &target) {
   os << "Target<";
   switch (target.os) {

cinn/common/target.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ostream>
+#include <string>
 #include <vector>
 
 namespace cinn {
@@ -69,6 +70,8 @@ struct Target {
 
   std::vector<Lib> get_target_libs() const;
 
+  std::string arch_str() const;
+
   bool operator==(const Target& other) const;
   bool operator!=(const Target& other) const { return !(*this == other); }
   friend std::ostream& operator<<(std::ostream& os, const Target& target);

cinn/frontend/CMakeLists.txt

@@ -26,5 +26,8 @@ else()
 endif()
 
 cc_test(test_net_builder SRCS net_builder_test.cc DEPS cinncore)
+cc_test(test_decomposer_registry
+        SRCS decomposer_registry_test.cc DEPS cinncore)
 
 add_subdirectory(paddle)
+add_subdirectory(decomposer)

cinn/frontend/decomposer/CMakeLists.txt

@@ -0,0 +1,5 @@
+core_gather_headers()
+
+core_gather_srcs(SRCS
+    activation.cc
+    )

cinn/frontend/decomposer/activation.cc

@@ -0,0 +1,17 @@
+#include "cinn/frontend/decomposer_registry.h"
+
+namespace cinn {
+namespace frontend {
+namespace decomposer {
+
+void relu(const Instruction& instr, const DecomposerContext& context) { LOG(FATAL) << "not implemented"; }
+
+}  // namespace decomposer
+}  // namespace frontend
+}  // namespace cinn
+
+CINN_REGISTER_HELPER(activation) {
+  CINN_DECOMPOSER_REGISTER(relu, ::cinn::common::DefaultHostTarget()).set_body(cinn::frontend::decomposer::relu);
+
+  return true;
+}

cinn/frontend/decomposer/use_decomposer.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include "cinn/common/macros.h"
+
+CINN_USE_REGISTER(activation)

cinn/frontend/decomposer_registry.h

@@ -0,0 +1,70 @@
+#pragma once
+
+#include <functional>
+#include <string>
+#include <unordered_map>
+
+#include "cinn/common/target.h"
+#include "cinn/frontend/syntax.h"
+
+namespace cinn {
+namespace frontend {
+
+class Decomposer;
+
+class DecomposerContext {
+ public:
+  explicit DecomposerContext(Program* prog) : program(prog) {}
+
+  Program* program{nullptr};
+};
+
+class InstrDecomposerRegistry : public Registry<Decomposer> {
+ public:
+  static InstrDecomposerRegistry* Global() {
+    static InstrDecomposerRegistry x;
+    return &x;
+  }
+
+  inline const Decomposer* Get(const std::string& op_name, const common::Target& target) {
+    const Decomposer* decomposer = Find(op_name, target);
+    CHECK(decomposer) << "Decomposer for [" << op_name << ", " << target << "] is not registered";
+    return decomposer;
+  }
+
+  inline const Decomposer* Find(const std::string& name, const common::Target& target) {
+    return Registry<Decomposer>::Find(name + "_" + target.arch_str());
+  }
+
+  inline Decomposer& __REGISTER__(const std::string& name, const common::Target& target) {
+    return Registry<Decomposer>::__REGISTER__(name + "_" + target.arch_str());
+  }
+
+ private:
+  InstrDecomposerRegistry() = default;
+  CINN_DISALLOW_COPY_AND_ASSIGN(InstrDecomposerRegistry);
+};
+
+class Decomposer {
+ public:
+  using DecomposerKernel = std::function<void(const Instruction& instr, const DecomposerContext&)>;
+
+  Decomposer& set_body(const DecomposerKernel& kernel) {
+    kernel_ = kernel;
+    return *this;
+  }
+
+  void Run(const Instruction& instr, const DecomposerContext& context) { kernel_(instr, context); }
+
+  std::string name;
+
+ private:
+  DecomposerKernel kernel_;
+};
+
+#define CINN_DECOMPOSER_REGISTER(name, target)                                                \
+  static ::cinn::frontend::Decomposer& CINN_STR_CONCAT(__make_Decomposer_name, __COUNTER__) = \
+      ::cinn::frontend::InstrDecomposerRegistry::Global()->__REGISTER__(#name, target)
+
+}  // namespace frontend
+}  // namespace cinn

cinn/frontend/decomposer_registry_test.cc

@@ -0,0 +1,15 @@
+#include "cinn/frontend/decomposer_registry.h"
+
+#include <gtest/gtest.h>
+
+#include "cinn/frontend/decomposer/use_decomposer.h"
+
+namespace cinn::frontend {
+
+TEST(InstrDecomposerRegistry, basic) {
+  common::Target target = common::DefaultHostTarget();
+  ASSERT_EQ(InstrDecomposerRegistry::Global()->Find("conv", target), nullptr);
+  ASSERT_NE(InstrDecomposerRegistry::Global()->Find("relu", target), nullptr);
+}
+
+}  // namespace cinn::frontend