[BYOC] Enhance partitioning and external codegen

src/relay/backend/contrib/codegen_c/codegen.cc

@@ -40,35 +40,39 @@ using namespace backend;
  * purpose. Only several binary options are covered. Users
  * may need to extend them to cover more operators.
  */
-class CodegenC : public ExprVisitor, public CodegenCBase {
+class CodegenC : public ExprFunctor<std::vector<Output>(const Expr&)>,
+                 public CodegenCBase {
  public:
   explicit CodegenC(const std::string& id) { this->ext_func_id_ = id; }
 
-  void VisitExpr_(const VarNode* node) final {
+  std::vector<Output> VisitExpr(const Expr& expr) final {
+    if (visited_.count(expr)) return visited_.at(expr);
+    std::vector<Output> output = ExprFunctor::VisitExpr(expr);
+    visited_[expr] = output;
+    return output;
+  }
+
+  std::vector<Output> VisitExprDefault_(const Object* op) final {
+    LOG(FATAL) << "C codegen doesn't support: " << op->GetTypeKey();
+    return {};
+  }
+
+  std::vector<Output> VisitExpr_(const VarNode* node) final {
     ext_func_args_.push_back(GetRef<Var>(node));
-    out_.clear();
     Output output;
     output.name = node->name_hint();
-    out_.push_back(output);
+    return {output};
   }
 
-  void VisitExpr_(const ConstantNode* cn) final {
-    Constant constant = GetRef<Constant>(cn);
-    if (visited_.count(constant)) {
-      // Note this is for demostration purpose. ConstantNode doesn't necessarily
-      // belong to calls. We need to revisit this when tuples come into play.
-      out_.push_back(visited_[constant]);
-      return;
-    }
+  std::vector<Output> VisitExpr_(const ConstantNode* cn) final {
+    // Note this is for demonstration purpose. ConstantNode doesn't necessarily
+    // belong to calls. We need to revisit this when tuples come into play.
 
     std::ostringstream decl_stream;
     std::ostringstream buf_stream;
 
-    out_.clear();
     Output output;
     output.name = "const_" + std::to_string(const_idx_++);
-    out_.push_back(output);
-    visited_[constant] = output;
 
     runtime::NDArray array = cn->data;
     const auto& shape = array.Shape();
@@ -99,9 +103,11 @@ class CodegenC : public ExprVisitor, public CodegenCBase {
     }
     buf_stream << "};";
     ext_func_body.insert(ext_func_body.begin(), buf_stream.str());
+
+    return {output};
   }
 
-  void VisitExpr_(const CallNode* call) final {
+  std::vector<Output> VisitExpr_(const CallNode* call) final {
     std::ostringstream macro_stream;
     std::ostringstream decl_stream;
     std::ostringstream buf_stream;
@@ -138,8 +144,8 @@ class CodegenC : public ExprVisitor, public CodegenCBase {
     bool first = true;
     decl_stream << func_name << "(";
     for (size_t i = 0; i < call->args.size(); ++i) {
-      VisitExpr(call->args[i]);
-      for (auto out : out_) {
+      auto res = VisitExpr(call->args[i]);
+      for (auto out : res) {
         if (!first) {
           decl_stream << ", ";
         }
@@ -162,26 +168,27 @@ class CodegenC : public ExprVisitor, public CodegenCBase {
     ext_func_body.push_back(decl_stream.str());
 
     // Update output buffer
-    out_.clear();
+    // Note C codegen only handles TensorType. Therefore, we don't flatten
+    // tuples and only return a single vaule.
     Output output;
     output.name = out;
     output.dtype = dtype;
     output.need_copy = true;
     output.size = out_size;
-    out_.push_back(output);
+    return {output};
   }
 
   /*!
    * \brief Emit the source code that invokes C compiler compatible wrappers.
    *
    * \return The emitted code.
    */
-  std::string JIT() {
+  std::string JIT(const std::vector<Output>& out) {
     // Write function macros
     for (auto decl : func_decl_) {
       code_stream_ << decl << "\n";
     }
-    return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out_);
+    return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out);
   }
 
  private:
@@ -202,9 +209,7 @@ class CodegenC : public ExprVisitor, public CodegenCBase {
   /*! \brief The declaration statements of buffers. */
   std::vector<std::string> buf_decl_;
   /*! \brief The name and index pairs for output. */
-  std::vector<Output> out_;
-  /*! \brief The cached expressions. */
-  std::unordered_map<Expr, Output, ObjectHash, ObjectEqual> visited_;
+  std::unordered_map<Expr, std::vector<Output>, ObjectHash, ObjectEqual> visited_;
 };
 
 class CSourceCodegen : public CSourceModuleCodegenBase {
@@ -216,8 +221,8 @@ class CSourceCodegen : public CSourceModuleCodegenBase {
     auto sid = GetExtSymbol(func);
 
     CodegenC builder(sid);
-    builder.VisitExpr(func->body);
-    code_stream_ << builder.JIT();
+    auto out = builder.VisitExpr(func->body);
+    code_stream_ << builder.JIT(out);
   }
 
   runtime::Module CreateCSourceModule(const ObjectRef& ref) override {

src/relay/backend/contrib/codegen_c/codegen_c.h

@@ -165,9 +165,11 @@ class CodegenCBase {
   /*!
    * \brief Emit the code for external runtime.
    *
+   * \param out The outputs.
+   *
    * \return The code string.
    */
-  virtual std::string JIT() = 0;
+  virtual std::string JIT(const std::vector<Output>& out) = 0;
 
   /*!
    * \brief A common interface that is used by various external runtime to

src/relay/backend/contrib/dnnl/codegen.cc

@@ -128,42 +128,43 @@ std::vector<std::string> Add(const CallNode* call) {
 
 // TODO(@zhiics, @comaniac): This is a basic implementation. We should implement
 // all utilities and make a base class for users to implement.
-class CodegenDNNL : public ExprVisitor, public CodegenCBase {
+class CodegenDNNL : public ExprFunctor<std::vector<Output>(const Expr&)>,
+                    public CodegenCBase {
  public:
   explicit CodegenDNNL(const std::string& id) { this->ext_func_id_ = id; }
 
-  void VisitExpr_(const VarNode* node) final {
+  std::vector<Output> VisitExpr(const Expr& expr) final {
+    if (visited_.count(expr)) return visited_.at(expr);
+    std::vector<Output> output = ExprFunctor::VisitExpr(expr);
+    visited_[expr] = output;
+    return output;
+  }
+
+  std::vector<Output> VisitExprDefault_(const Object* op) final {
+    LOG(FATAL) << "DNNL codegen doesn't support: " << op->GetTypeKey();
+    return {};
+  }
+
+  std::vector<Output> VisitExpr_(const VarNode* node) final {
     ext_func_args_.push_back(GetRef<Var>(node));
-    out_.clear();
     Output output;
     output.name = node->name_hint();
-    out_.push_back(output);
+    return {output};
   }
 
-  void VisitExpr_(const TupleGetItemNode* op) final {
-    VisitExpr(op->tuple);
-    CHECK(out_.size() > static_cast<size_t>(op->index));
+  std::vector<Output> VisitExpr_(const TupleGetItemNode* op) final {
+    auto res = VisitExpr(op->tuple);
+    CHECK_GT(res.size(), static_cast<size_t>(op->index));
 
     // Only keep the item we want for the child node.
     // FIXME(@comaniac): The other items should still be requried for the primary outputs.
-    auto item = out_[op->index];
-    out_.clear();
-    out_.push_back(item);
+    return {res[op->index]};
   }
 
-  void VisitExpr_(const ConstantNode* cn) final {
-    Constant constant = GetRef<Constant>(cn);
-    if (visited_.count(constant)) {
-      out_.push_back(visited_[constant]);
-      return;
-    }
-
-    out_.clear();
+  std::vector<Output> VisitExpr_(const ConstantNode* cn) final {
     Output output;
     output.name = "const_" + std::to_string(const_idx_++);
     output.dtype = "float";
-    out_.push_back(output);
-    visited_[constant] = output;
 
     runtime::NDArray array = cn->data;
 
@@ -176,33 +177,37 @@ class CodegenDNNL : public ExprVisitor, public CodegenCBase {
     CHECK_EQ(GetDtypeString(type_node), "float") << "Only float is supported for now.";
 
     std::ostringstream buf_stream;
-    buf_stream << "float* " << output.name << " = (float*)std::malloc(4 * " << num_elems << ");\n";
     const float* ptr = static_cast<float*>(array.ToDLPack()->dl_tensor.data);
-    for (int64_t i = 0; i < num_elems; i++) {
-      buf_stream << "  " << output.name << "[" << i << "] = " << ptr[i] << ";\n";
+
+    // Allocate large arrays on the static section to avoid stakc overflow.
+    // Note that this would probably increase compilation time as the source
+    // file could be really large.
+    buf_stream << "static float " << output.name << "[" << num_elems <<"] = {";
+    for (int64_t i = 0; i < num_elems - 1; i++) {
+      buf_stream << ptr[i] << ",";
     }
+    if (num_elems > 0) buf_stream << ptr[num_elems - 1];
+    buf_stream << "};\n";
 
     ext_func_body.insert(ext_func_body.begin(), buf_stream.str());
+    return {output};
   }
 
-  void VisitExpr_(const CallNode* call) final {
+  std::vector<Output> VisitExpr_(const CallNode* call) final {
     GenerateBodyOutput ret;
     if (const auto* func = call->op.as<FunctionNode>()) {
       ret = GenerateCompositeFunctionCall(func, call);
     } else {
       ret = GenerateOpCall(call);
     }
 
-    out_.clear();
-    for (size_t i = 0; i < ret.outputs.size(); ++i) {
-      buf_decl_.push_back(ret.buffers[i]);
-      out_.push_back(ret.outputs[i]);
-    }
+    buf_decl_.insert(buf_decl_.end(), ret.buffers.begin(), ret.buffers.end());
     ext_func_body.push_back(ret.decl);
+    return ret.outputs;
   }
 
-  std::string JIT(void) {
-    return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out_);
+  std::string JIT(const std::vector<Output>& out) {
+    return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out);
   }
 
  private:
@@ -215,8 +220,8 @@ class CodegenDNNL : public ExprVisitor, public CodegenCBase {
   std::vector<std::string> GetArgumentNames(const CallNode* call) {
     std::vector<std::string> arg_names;
     for (size_t i = 0; i < call->args.size(); ++i) {
-      VisitExpr(call->args[i]);
-      for (auto out : out_) {
+      auto res = VisitExpr(call->args[i]);
+      for (const auto& out : res) {
         arg_names.push_back(out.name);
       }
     }
@@ -331,17 +336,15 @@ class CodegenDNNL : public ExprVisitor, public CodegenCBase {
    */
   int buf_idx_{0};
   /*! \brief The index of global constants. */
-  int const_idx_ = 0;
+  int const_idx_{0};
   /*! \brief The arguments used by a wrapped function that calls DNNL kernels. */
   Array<Var> ext_func_args_;
   /*! \brief statement of the function that will be compiled using DNNL kernels. */
   std::vector<std::string> ext_func_body;
   /*! \brief The declaration of intermeidate buffers. */
   std::vector<std::string> buf_decl_;
-  /*! \brief The name of the the outputs. */
-  std::vector<Output> out_;
   /*! \brief The cached expressions. */
-  std::unordered_map<Expr, Output, ObjectHash, ObjectEqual> visited_;
+  std::unordered_map<Expr, std::vector<Output>, ObjectHash, ObjectEqual> visited_;
 };
 
 /*!
@@ -361,8 +364,8 @@ class DNNLModuleCodegen : public CSourceModuleCodegenBase {
     auto sid = GetExtSymbol(func);
 
     CodegenDNNL builder(sid);
-    builder.VisitExpr(func->body);
-    code_stream_ << builder.JIT();
+    auto out = builder.VisitExpr(func->body);
+    code_stream_ << builder.JIT(out);
   }
 
   /*!

src/relay/transforms/partition_graph.cc

@@ -148,25 +148,42 @@ class Partitioner : public ExprMutator {
       CHECK_EQ(call->args.size(), 1U);
 
       // Traverse the rest graph.
-      auto input_expr = VisitExpr(call->args[0]);
+      Expr parent = call->args[0];
+      auto input_expr = VisitExpr(parent);
+
+      // Backtrace the parent to find the first ancestor node that is not a begin or end op
+      while (const auto* parent_call = parent.as<CallNode>()) {
+        if (parent_call->op == compiler_begin_op ||
+            parent_call->op == compiler_end_op) {
+          parent = parent_call->args[0];
+        } else {
+          break;
+        }
+      }
 
       AnnotatedRegion sg = GetRegion(GetRef<Call>(call));
       int index = GetArgIdx(sg, GetRef<Call>(call));
       CHECK_NE(index, -1);
-      // The type of the created variable is the same as the compiler_begin
-      // node.
-      std::string target = call->attrs.as<CompilerAttrs>()->compiler;
-      std::string varname =
-          target + "_" + std::to_string(sg->GetID()) + "_i" + std::to_string(index);
-      auto var = Var(varname, GetRef<Call>(call)->checked_type_);
-
-      auto cand = std::make_pair(var, input_expr);
-      if (std::find(region_args[sg].begin(), region_args[sg].end(), cand) ==
-          region_args[sg].end()) {
-        region_args[sg].push_back(cand);
-      }
 
-      return std::move(var);
+      if (shared_output_.count(parent) && shared_output_[parent].count(sg)) {
+        return shared_output_[parent][sg];
+      } else {
+        // The type of the created variable is the same as the compiler_begin
+        // node.
+        std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+        std::string varname =
+            target + "_" + std::to_string(sg->GetID()) + "_i" + std::to_string(index);
+        auto var = Var(varname, GetRef<Call>(call)->checked_type_);
+
+        std::pair<Var, Expr> cand = std::make_pair(var, input_expr);
+
+        if (std::find(region_args[sg].begin(), region_args[sg].end(), cand) ==
+            region_args[sg].end()) {
+          region_args[sg].push_back(cand);
+        }
+        shared_output_[parent][sg] = var;
+        return std::move(var);
+      }
     } else {
       CHECK_EQ(call->op, compiler_end_op);
       // The annotation node is inserted on edge so it must have only one
@@ -474,6 +491,12 @@ class Partitioner : public ExprMutator {
    * belongs to
    */
   std::unordered_map<AnnotatedRegionSet, BaseFunc, ObjectHash, ObjectEqual> regions_sets_;
+
+  /*!\brief Cache the output that is shared by different nodes. */
+  using RegionOutputMap = std::unordered_map<AnnotatedRegion, Var, ObjectHash, ObjectEqual>;
+  std::unordered_map<Expr, RegionOutputMap, ObjectHash, ObjectEqual> shared_output_;
+
+  /*!\brief The IRModule used for partitioning. */
   IRModule module_;
 };