Merge pull request PaddlePaddle#48 from Superjomn/fea/make-lowered_fu…

…nc-ir

fea/make lowered func ir

.gitignore

@@ -5,3 +5,4 @@
 cmake-build*
 build*
 .idea*
+*.html

cinn/backends/codegen_c.cc

@@ -1,18 +1,20 @@
 #include "cinn/backends/codegen_c.h"
 
+#include "cinn/ir/lowered_func.h"
+
 namespace cinn {
 namespace backends {
 
 CodeGenC::CodeGenC(std::ostream &os, Target target) : ir::IrPrinter(os), target_(target) {}
 
 void CodeGenC::Compile(const lang::Module &module) {}
-void CodeGenC::Compile(const lang::LoweredFunc &function) {
-  os() << "void " << function.name;
+void CodeGenC::Compile(const ir::LoweredFunc &function) {
+  os() << "void " << function->name;
 
   // output arguments
   os() << "(";
 
-  auto print_arg = [&](const lang::Argument &arg) {
+  auto print_arg = [&](const ir::Argument &arg) {
     if (arg.is_buffer()) {
       os() << "struct cinn_buffer_t *";
     } else if (arg.is_scalar()) {
@@ -22,20 +24,20 @@ void CodeGenC::Compile(const lang::LoweredFunc &function) {
     os() << arg.name;
   };
 
-  for (int i = 0; i < function.args.size() - 1; i++) {
-    print_arg(function.args[i]);
+  for (int i = 0; i < function->args.size() - 1; i++) {
+    print_arg(function->args[i]);
     os() << ", ";
   }
-  if (function.args.size() >= 1) {
-    print_arg(function.args.back());
+  if (function->args.size() >= 1) {
+    print_arg(function->args.back());
   }
 
   os() << ")";
 
   DoIndent();
   os() << "{\n";
 
-  Print(function.body);
+  Print(function->body);
 
   DoIndent();
   os() << "}";

cinn/backends/codegen_c.h

@@ -7,6 +7,7 @@
 #include "cinn/ir/function.h"
 #include "cinn/ir/ir.h"
 #include "cinn/ir/ir_printer.h"
+#include "cinn/ir/lowered_func.h"
 #include "cinn/lang/module.h"
 
 namespace cinn {
@@ -24,7 +25,7 @@ class CodeGenC : public ir::IrPrinter {
   void Compile(const lang::Module& module);
 
  protected:
-  void Compile(const lang::LoweredFunc& function);
+  void Compile(const ir::LoweredFunc& function);
   void Compile(const ir::Buffer& buffer);
 
   std::string PrintType(Type type);

cinn/ir/CMakeLists.txt

@@ -7,6 +7,7 @@ set(srcs
     ir_mutator.cc
     function.cc
     function_definition.cc
+    lowered_func.cc
     ir_operators.cc
     buffer.cc
     function_base.cc

cinn/ir/ir_mutator.cc

@@ -114,5 +114,10 @@ void IRMutator::Visit(const _Tensor_ *expr, Expr *op) {
   }
 }
 
+void IRMutator::Visit(const _LoweredFunc_ *expr, Expr *op) {
+  auto *node = op->As<_LoweredFunc_>();
+  IRVisitorBase::Visit(&node->body, &node->body);
+}
+
 }  // namespace ir
 }  // namespace cinn

cinn/ir/ir_printer.cc

@@ -2,7 +2,10 @@
 
 #include <vector>
 
+#include "cinn/ir/lowered_func.h"
+#include "cinn/lang/module.h"
 #include "cinn/lang/tensor.h"
+#include "cinn/utils/string.h"
 
 namespace cinn {
 namespace ir {
@@ -173,6 +176,27 @@ void IrPrinter::Visit(const _Tensor_ *x) {
   }
   os_ << ")";
 }
+void IrPrinter::Visit(const _LoweredFunc_ *f) {
+  os_ << "function " << f->name << " ";
+
+  std::vector<std::string> arg_names;
+  for (auto &arg : f->args) {
+    arg_names.push_back(arg.name);
+  }
+  os_ << "(" << utils::Join(arg_names, ", ");
+
+  DoIndent();
+  os_ << "{";
+
+  IncIndent();
+
+  Print(f->body);
+
+  DecIndent();
+
+  DoIndent();
+  os_ << "}";
+}
 std::ostream &operator<<(std::ostream &os, Expr a) {
   std::stringstream ss;
   IrPrinter printer(ss);
@@ -181,5 +205,9 @@ std::ostream &operator<<(std::ostream &os, Expr a) {
   return os;
 }
 
+std::ostream &operator<<(std::ostream &os, const ir::LoweredFunc &f) {}
+
+std::ostream &operator<<(std::ostream &os, const lang::Module &m);
+
 }  // namespace ir
 }  // namespace cinn

cinn/ir/ir_printer.h

@@ -7,6 +7,12 @@
 #include "cinn/ir/ir_visitor.h"
 
 namespace cinn {
+
+namespace lang {
+class Module;
+class LoweredFunc;
+}  // namespace lang
+
 namespace ir {
 
 struct IrPrinter : public IRVisitor {
@@ -72,6 +78,7 @@ struct IrPrinter : public IRVisitor {
   void Visit(const _IterVar_ *x) override {}
   void Visit(const _Buffer_ *x) override;
   void Visit(const _Tensor_ *x) override;
+  void Visit(const _LoweredFunc_ *x) override;
 
  private:
   std::ostream &os_;
@@ -80,6 +87,7 @@ struct IrPrinter : public IRVisitor {
 };
 
 std::ostream &operator<<(std::ostream &os, Expr a);
+std::ostream &operator<<(std::ostream &os, const lang::Module &m);
 
 }  // namespace ir
 }  // namespace cinn

cinn/ir/ir_visitor.h

@@ -4,6 +4,7 @@
 
 #include "cinn/ir/buffer.h"
 #include "cinn/ir/ir.h"
+#include "cinn/ir/lowered_func.h"
 #include "cinn/lang/tensor.h"
 
 namespace cinn {

cinn/ir/lowered_func.cc

@@ -0,0 +1,30 @@
+#include "cinn/ir/lowered_func.h"
+
+#include "cinn/common/common.h"
+
+namespace cinn {
+namespace ir {
+
+const _LoweredFunc_* LoweredFunc::operator->() const { return As<_LoweredFunc_>(); }
+_LoweredFunc_* LoweredFunc::operator->() { return As<_LoweredFunc_>(); }
+
+LoweredFunc _LoweredFunc_::Make(const std::string& name, const std::vector<Argument>& args, const Expr& body) {
+  auto* n = make_shared<_LoweredFunc_>();
+  n->name = name;
+  n->args = args;
+  n->body = body;
+  return LoweredFunc(n);
+}
+
+LoweredFunc _LoweredFunc_::Make(const std::string& name,
+                                const std::vector<Argument>& args,
+                                const std::vector<Expr>& body) {
+  CHECK_EQ(body.size(), 1);
+  return Make(name, args, body.front());
+}
+
+std::vector<Expr*> _LoweredFunc_::expr_fields() { return {&body}; }
+std::vector<const Expr*> _LoweredFunc_::expr_fields() const { return {&body}; }
+
+}  // namespace ir
+}  // namespace cinn

cinn/ir/lowered_func.h

@@ -0,0 +1,70 @@
+#pragma once
+#include "cinn/ir/buffer.h"
+#include "cinn/ir/node.h"
+
+namespace cinn {
+namespace ir {
+
+class _LoweredFunc_;
+
+/**
+ * A struct representing an argument to a lowered function. Used for specifying the function signature of generated
+ * code.
+ */
+struct Argument {
+  //! The name of the argument.
+  std::string name;
+
+  enum class Kind { kScalar = 0, kBuffer } kind{Kind::kScalar};
+
+  //! Number of the dimensions of buffer.
+  uint32_t ndims{0};
+
+  //! The type of the buffer or scalar.
+  Type type;
+
+  bool is_buffer() const { return kind == Kind::kBuffer; }
+  bool is_scalar() const { return kind == Kind::kScalar; }
+
+  Argument() {}
+  Argument(const std::string& name, Kind kind, const Type& type, int ndims)
+      : name(name), kind(kind), type(type), ndims(ndims) {}
+
+  explicit Argument(const ir::Buffer& buffer) : name(buffer->name), type(buffer->type()), ndims(buffer->shape.size()) {}
+};
+
+//! Wrapper for _LoweredFunc_
+class LoweredFunc : public IrNodeRef {
+ public:
+  LoweredFunc() = default;
+  explicit LoweredFunc(IrNode* n) : IrNodeRef(n) {}
+
+  const _LoweredFunc_* operator->() const;
+  _LoweredFunc_* operator->();
+};
+
+/**
+ * Definition of a lowered function. Note that, it should be functional.
+ */
+struct _LoweredFunc_ : ExprNode<_LoweredFunc_> {
+  //! The name of this function.
+  std::string name;
+
+  //! The Arguments used in the body of the function.
+  std::vector<Argument> args;
+
+  //! Body of this function.
+  Expr body;
+
+  static LoweredFunc Make(const std::string& name, const std::vector<Argument>& args, const Expr& body);
+
+  static LoweredFunc Make(const std::string& name, const std::vector<Argument>& args, const std::vector<Expr>& body);
+
+  std::vector<Expr*> expr_fields() override;
+  std::vector<const Expr*> expr_fields() const override;
+
+  static const IrNodeTy _node_type_ = IrNodeTy::_LoweredFunc_;
+};
+
+}  // namespace ir
+}  // namespace cinn

cinn/ir/node.h

@@ -66,6 +66,7 @@ class IRVisitor;
   macro__(_IterVar_)                        \
   macro__(_Buffer_)                         \
   macro__(_Tensor_)                         \
+  macro__(_LoweredFunc_)                    \
 
 #define NODETY_FORALL(__m)              \
   NODETY_PRIMITIVE_TYPE_FOR_EACH(__m)   \

cinn/lang/lower.cc

@@ -45,7 +45,7 @@ Expr LowerGroup(const poly::detail::Group& group, const std::map<std::string, Ex
   return e;
 }
 
-std::vector<LoweredFunc> Lower(const std::string& name, const std::vector<Tensor>& args) {
+std::vector<ir::LoweredFunc> Lower(const std::string& name, const std::vector<Tensor>& args) {
   // make sure the graph's start-points in the args.
 
   auto stages = poly::GatherStagesInTensors(args);
@@ -106,12 +106,12 @@ std::vector<LoweredFunc> Lower(const std::string& name, const std::vector<Tensor
   optim::RemoveNestedBlock(&block);
 
   // prepare arguments
-  std::vector<Argument> arguments;
+  std::vector<ir::Argument> arguments;
   for (auto& arg : args) {
-    arguments.emplace_back(arg->name, Argument::Kind::kBuffer, arg->type(), arg->shape.size());
+    arguments.emplace_back(arg->name, ir::Argument::Kind::kBuffer, arg->type(), arg->shape.size());
   }
 
-  return {LoweredFunc(name, arguments, block)};
+  return {ir::_LoweredFunc_::Make(name, arguments, block)};
 }
 
 }  // namespace lang

cinn/lang/lower.h

@@ -16,7 +16,7 @@ namespace cinn {
 namespace lang {
 using ir::Tensor;
 
-std::vector<LoweredFunc> Lower(const std::string& name, const std::vector<Tensor>& args);
+std::vector<ir::LoweredFunc> Lower(const std::string& name, const std::vector<Tensor>& args);
 
 }  // namespace lang
 }  // namespace cinn

cinn/lang/lower_test.cc

@@ -42,7 +42,7 @@ TEST(lower, basic) {
   }
 }
 )ROC";
-  TEST_SOUTPUT(lower_funcs.front().body, out);
+  TEST_SOUTPUT(lower_funcs.front()->body, out);
 }
 
 TEST(lower, more_complex) {

cinn/lang/module.cc

@@ -10,7 +10,7 @@ struct _Module_ : Object {
   std::string name;
   Target target;
   std::vector<ir::Buffer> buffers;
-  std::vector<ir::PackedFunc> functions;
+  std::vector<ir::LoweredFunc> functions;
   std::vector<Module> submodules;
 
   const char *type_info() const override { return "_Module_"; }
@@ -30,23 +30,17 @@ const std::string &Module::name() const { return self()->name; }
 
 const std::vector<ir::Buffer> &Module::buffers() const { return self()->buffers; }
 
-const std::vector<ir::PackedFunc> &Module::functions() const { return self()->functions; }
+const std::vector<ir::LoweredFunc> &Module::functions() const { return self()->functions; }
 
 const std::vector<Module> &Module::submodules() const { return self()->submodules; }
 
-void Module::Append(const ir::Buffer &buffer) { self()->buffers.push_back(buffer); }
+void Module::Append(const Buffer &buffer) { self()->buffers.push_back(buffer.buffer()); }
 
-void Module::Append(const ir::PackedFunc &function) { self()->functions.push_back(function); }
+void Module::Append(const ir::LoweredFunc &function) { self()->functions.push_back(function); }
 
 void Module::Append(const Module &module) { self()->submodules.push_back(module); }
 
 void Module::Compile(const backends::Outputs &outputs) const {}
 
-LoweredFunc::LoweredFunc(const std::string &name, const std::vector<Argument> &args, const std::vector<Expr> &body) {
-  this->name = name;
-  this->args = args;
-  this->body = ir::Block::Make(body);
-}
-
 }  // namespace lang
 }  // namespace cinn