Add Instruction and Function class. (PaddlePaddle#17)

* Add Instruction data structure.

* Add more methods for Instruction.

* Update the cmake dependency of note_ir.

* Add Function class for Note IR.

* Print attribute values of instruction.

* Add the assert statement in FunctionToString UT.

* Add the unit test for Function and Instruction.

* Add the default action for VisitParameter.

* Fix errors of layout_test.cc and shape_test.cc.

* Add comments for Instruction and Function.

* Add non-const Accept for Instruction class.

* Add comments of constructors in Instruction and Function.

Author: wzzju

paddle/fluid/compiler/piano/backends/note_visitor_base.h

@@ -11,14 +11,15 @@
 // 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
 
 namespace paddle {
 namespace piano {
 
 namespace note {
 class Instruction;
-}
+}  // namespace note
 
 namespace backends {
 
@@ -28,6 +29,8 @@ class NoteVisitorBase {
 
   // Scalar op
   virtual void VisitConstant(const note::Instruction&) = 0;
+  // TODO(sunli): use the pure virtual function instead
+  virtual void VisitParameter(const note::Instruction&) {}
 
   // ops can be replaced by library
   virtual void VisitBatchNormGrad(const note::Instruction&) = 0;

paddle/fluid/compiler/piano/layout.cc

@@ -41,7 +41,7 @@ std::string Layout::ToString() const {
                  std::back_inserter(dim_names),
                  [](const auto& dim) { return std::to_string(dim); });
   return paddle::string::format_string(
-      "{%s}", paddle::string::join_strings(dim_names, ',').c_str());
+      "{%s}", paddle::string::join_strings(dim_names, ", ").c_str());
 }
 
 bool Layout::Valid() const { return !minor_to_major().empty(); }

paddle/fluid/compiler/piano/layout_test.cc

@@ -36,7 +36,7 @@ TEST_F(LayoutTest, LayoutTransWithProto) {
 }
 
 TEST_F(LayoutTest, LayoutToString) {
-  ASSERT_EQ("{3,2,1,0}", layout_.ToString());
+  ASSERT_EQ("{3, 2, 1, 0}", layout_.ToString());
 }
 
 }  // namespace piano

paddle/fluid/compiler/piano/note/CMakeLists.txt

@@ -3,3 +3,6 @@ cc_test(note_opcode_test SRCS opcode_test.cc DEPS note_opcode)
 
 proto_library(note_proto SRCS note.proto)
 target_compile_options(note_proto PUBLIC "-Wno-extra")
+
+cc_library(note_ir SRCS instruction.cc function.cc DEPS note_opcode note_proto piano_data_description)
+cc_test(note_ir_test SRCS note_ir_test.cc DEPS note_ir)

paddle/fluid/compiler/piano/note/function.cc

@@ -0,0 +1,103 @@
+/* Copyright (c) 2021 PaddlePaddle 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 "paddle/fluid/compiler/piano/note/function.h"
+#include <memory>
+#include <sstream>
+#include <unordered_map>
+#include <vector>
+#include "paddle/fluid/compiler/piano/note/instruction.h"
+#include "paddle/fluid/compiler/piano/note/opcode.h"
+#include "paddle/fluid/string/string_helper.h"
+
+namespace paddle {
+namespace piano {
+namespace note {
+
+Function::Function(
+    const FunctionProto& proto,
+    const std::unordered_map<std::int64_t, Function*>& func_index)
+    : name_(proto.name()),
+      signature_(proto.signature()),
+      global_id_(proto.id()) {
+  // the map used to record `id -> Instruction*`
+  std::unordered_map<std::int64_t, Instruction*> instr_index;
+
+  // the map used to record `Instruction* -> id`, which is opposite
+  // to the instr_index map
+  std::unordered_map<Instruction*, std::int64_t> inverted_index;
+  for (const auto& instr_proto : proto.instructions()) {
+    auto instr =
+        std::make_unique<Instruction>(instr_proto, instr_index, func_index);
+    instr->set_parent(this);
+    // set parameter(input) instructions field
+    if (instr->opcode() == OpCode::kParameter) {
+      param_instrs_.push_back(instr.get());
+    }
+    instr_index[instr_proto.id()] = instr.get();
+    inverted_index[instr.get()] = instr_proto.id();
+    instructions_.emplace_back(std::move(instr));
+  }
+  PADDLE_ENFORCE_EQ(
+      proto.return_id() >= 0 && instr_index.count(proto.return_id()), true,
+      platform::errors::PreconditionNotMet(
+          "The return instruction id is %ld, and it is not "
+          "included in this function.",
+          proto.return_id()));
+
+  // set the returned instruction field
+  return_instr_ = instr_index[proto.return_id()];
+  std::sort(instructions_.begin(), instructions_.end(),
+            [&inverted_index](const std::unique_ptr<Instruction>& l,
+                              const std::unique_ptr<Instruction>& r) {
+              return inverted_index[l.get()] < inverted_index[r.get()];
+            });
+}
+
+FunctionProto Function::ToProto() const {
+  FunctionProto proto;
+  proto.set_name(name_);
+  *proto.mutable_signature() = signature_.ToProto();
+  proto.set_id(global_id_);
+  proto.set_return_id(return_instr_->global_id());
+  // serialize instruction protos
+  for (const auto& instr : instructions_) {
+    *proto.add_instructions() = instr->ToProto();
+  }
+  return proto;
+}
+
+std::string Function::ToString() const {
+  std::ostringstream out_str;
+  // get the function name and signature
+  out_str << "def %" << name_ << signature_.ToString() << " {\n";
+  const std::string tab(2, ' ');
+
+  // get the string value of each instruction
+  std::size_t num = instructions_.size();
+  for (decltype(instructions_.size()) i = 0; i < num; i++) {
+    if (num - 1 == i) {
+      out_str << tab << "return ";
+    } else {
+      out_str << tab;
+    }
+    out_str << instructions_[i]->ToString() << "\n";
+  }
+  out_str << "}";
+  return out_str.str();
+}
+
+}  // namespace note
+}  // namespace piano
+}  // namespace paddle

paddle/fluid/compiler/piano/note/function.h

@@ -0,0 +1,120 @@
+/* Copyright (c) 2021 PaddlePaddle 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 <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <vector>
+#include "paddle/fluid/compiler/piano/note/note.pb.h"
+#include "paddle/fluid/compiler/piano/shape.h"
+
+namespace paddle {
+namespace piano {
+namespace note {
+
+class Instruction;
+// class Module;
+
+class Function {
+ public:
+  // Construct a Function object with a given FunctionProto value.
+  // 'func_index' is used to transform function id into Function pointer,
+  // which is used to construct instructions in this function.
+  Function(const FunctionProto &proto,
+           const std::unordered_map<std::int64_t, Function *> &func_index);
+
+  FunctionProto ToProto() const;
+
+  std::string ToString() const;
+
+  // return the name of this function
+  const std::string &name() const { return name_; }
+
+  // return instructions owned by this function
+  std::vector<Instruction *> instructions() const {
+    std::vector<Instruction *> instrs;
+    instrs.reserve(instructions_.size());
+    std::transform(
+        instructions_.cbegin(), instructions_.cend(),
+        std::back_inserter(instrs),
+        [](const std::unique_ptr<Instruction> &instr) { return instr.get(); });
+    return instrs;
+  }
+
+  const Instruction *instruction(std::int64_t idx) const {
+    return instructions_.at(idx).get();
+  }
+
+  Instruction *mutable_instruction(std::int64_t idx) {
+    return instructions_.at(idx).get();
+  }
+
+  // return the function signature
+  const Signature &signature() const { return signature_; }
+
+  Signature *mutable_signature() { return &signature_; }
+
+  std::int64_t global_id() const { return global_id_; }
+
+  // return the returned instruction of this function
+  const Instruction *return_instr() const { return return_instr_; }
+
+  // const Module *parent() const { return parent_; }
+
+  // Module *mutable_parent() { return parent_; }
+
+  // void set_parent(Module *module) { parent_ = module; }
+
+  const std::vector<Instruction *> &param_instrs() const {
+    return param_instrs_;
+  }
+
+  // return parameter instructions of this function
+  const Instruction *param_instr(std::int64_t idx) const {
+    return param_instrs_.at(idx);
+  }
+
+  // return the parameter(input) number of this function
+  std::size_t params_num() const { return param_instrs_.size(); }
+
+ private:
+  // the name of this function
+  std::string name_;
+  // instructions owned by this function
+  std::vector<std::unique_ptr<Instruction>> instructions_;
+  // the function signature, including parameter and return types
+  Signature signature_;
+  // the global id of this function in a module
+  std::int64_t global_id_;
+  // the returned instruction of this function
+  Instruction *return_instr_;
+
+  // TODO(wzzju): Add Module class.
+  // the module where this function is contained
+  // Module *parent_{nullptr};
+
+  // parameter instructions of this function,
+  // which denote input parameters
+  std::vector<Instruction *> param_instrs_;
+
+  DISABLE_COPY_AND_ASSIGN(Function);
+};
+
+}  // namespace note
+}  // namespace piano
+}  // namespace paddle