[IR] add argument/constant for BasicBlock class

Author: youbeny

Diff for: include/halo/lib/ir/argument.h

@@ -39,8 +39,8 @@ class Argument final : public IRObject {
   explicit Argument(GlobalContext& context, const std::string& name,
                     const Type& type);
 
-  /// Get the parent function of this argument.
-  Function* GetParent() const noexcept { return parent_function_; }
+  /// Returns the parent object that could be a Function or a BasicBlock.
+  IRObject* GetParent() const noexcept { return parent_; }
 
   /// Set the type of this argument.
   void SetType(const Type& type);
@@ -55,11 +55,11 @@ class Argument final : public IRObject {
   void Print(std::ostream& os) const override;
 
  private:
-  Function* parent_function_ = nullptr;
+  IRObject* parent_ = nullptr;
 
   friend class ArgumentBuilder;
 };
 
 } // namespace halo
 
-#endif // HALO_LIB_IR_ARGUMENT_H_
+#endif // HALO_LIB_IR_ARGUMENT_H_

Diff for: include/halo/lib/ir/basic_block.h

@@ -20,10 +20,14 @@
 
 #include <list>
 
+#include "halo/lib/ir/argument.h"
+#include "halo/lib/ir/basic_block.h"
+#include "halo/lib/ir/constant.h"
 #include "halo/lib/ir/instruction.h"
 
 namespace halo {
 
+class Argument;
 class Function;
 
 /// This class defines a basic block in IR.
@@ -77,8 +81,36 @@ class BasicBlock final : public IRObject {
   InstructionList& Instructions() noexcept { return instructions_; }
   const InstructionList& Instructions() const noexcept { return instructions_; }
 
+  /// Arguments
+  using ArgumentList = std::list<std::unique_ptr<Argument>>;
+  using arg_iterator = ArgumentList::iterator;
+  using const_arg_iterator = ArgumentList::const_iterator;
+  arg_iterator arg_begin() noexcept { return args_.begin(); }
+  const_arg_iterator arg_begin() const noexcept { return args_.begin(); }
+  arg_iterator arg_end() noexcept { return args_.end(); }
+  const_arg_iterator arg_end() const noexcept { return args_.end(); }
+  ArgumentList& Args() noexcept { return args_; }
+  const ArgumentList& Args() const noexcept { return args_; }
+  Argument* arg_front() const { return args_.front().get(); }
+  Argument* arg_back() const { return args_.back().get(); }
+
+  /// Constants
+  using ConstantList = std::list<std::unique_ptr<Constant>>;
+  using constant_iterator = ConstantList::iterator;
+  using const_constant_iterator = ConstantList::const_iterator;
+  constant_iterator constant_begin() noexcept { return constants_.begin(); }
+  const_constant_iterator constant_begin() const noexcept {
+    return constants_.begin();
+  }
+  constant_iterator constant_end() noexcept { return constants_.end(); }
+  const_constant_iterator constant_end() const noexcept {
+    return constants_.end();
+  }
+  ConstantList& Constants() noexcept { return constants_; }
+  const ConstantList& Constants() const noexcept { return constants_; }
+
   /// Return the parent function to which this basic block belongs.
-  Function* GetParent() noexcept { return parent_function_; }
+  Function* GetParent() noexcept { return parent_; }
 
   Kind GetKind() const noexcept override { return Kind::BasicBlock; }
 
@@ -91,12 +123,14 @@ class BasicBlock final : public IRObject {
   void Print(std::ostream& os) const override;
 
  private:
-  Function* parent_function_ = nullptr;
+  Function* parent_ = nullptr;
   InstructionList instructions_;
+  ArgumentList args_;
+  ConstantList constants_;
 
   friend class BasicBlockBuilder;
 };
 
 } // namespace halo
 
-#endif // HALO_LIB_IR_BASIC_BLOCK_H_
+#endif // HALO_LIB_IR_BASIC_BLOCK_H_

Diff for: include/halo/lib/ir/ir_builder.h

@@ -154,11 +154,13 @@ class FunctionBuilder final
 
 /// This class provides the APIs to create arguments for the current function.
 class ArgumentBuilder final
-    : public IRObjectBuilder<Function, Function::ArgumentList> {
+    : public IRObjectBuilder<IRObject, Function::ArgumentList> {
  public:
   explicit ArgumentBuilder(Function* function)
       : IRObjectBuilder(function->GetGlobalContext(), function,
                         function->Args()) {}
+  explicit ArgumentBuilder(BasicBlock* bb)
+      : IRObjectBuilder(bb->GetGlobalContext(), bb, bb->Args()) {}
 
   /// Create a new argument and insert it to the end of current argument list.
   Argument* CreateArgument(const std::string& name);
@@ -189,6 +191,9 @@ class ConstantBuilder final
   explicit ConstantBuilder(Module* parent)
       : IRObjectBuilder(parent->GetGlobalContext(), parent,
                         parent->Constants()) {}
+  explicit ConstantBuilder(BasicBlock* parent)
+      : IRObjectBuilder(parent->GetGlobalContext(), parent,
+                        parent->Constants()) {}
 
   /// Create a new constant and append it to the end of current function or
   /// module using default data layout.
@@ -258,4 +263,4 @@ class IRBuilder final
 
 } // namespace halo
 
-#endif // HALO_LIB_IR_IR_BUILDER_H_
+#endif // HALO_LIB_IR_IR_BUILDER_H_

Diff for: lib/ir/basic_block.cc

@@ -20,10 +20,22 @@
 namespace halo {
 
 void BasicBlock::Print(std::ostream& os) const {
-  os << "BasicBlock: " << GetName() << "\n";
+  os << "BasicBlock: " << GetName() << "(";
+  int arg_idx = 0;
+  for (auto& arg : Args()) {
+    if (arg_idx++ != 0) {
+      os << ", ";
+    }
+    arg->Print(os);
+  }
+  os << ")\n";
+
+  for (auto& c : Constants()) {
+    c->Print(os);
+  }
   for (const auto& inst : *this) {
     inst->Print(os);
   }
 }
 
-} // namespace halo
+} // namespace halo

Diff for: lib/ir/function.cc

@@ -22,11 +22,10 @@
 namespace halo {
 
 ReturnInst* Function::GetReturnInst() const {
-  for (auto& bb : basic_blocks_) {
-    for (auto& inst : *bb) {
-      if (inst->GetOpCode() == OpCode::RETURN) {
-        return DynCast<ReturnInst>(inst.get());
-      }
+  HLCHECK(!basic_blocks_.empty());
+  for (auto& inst : *basic_blocks_.front()) {
+    if (inst->GetOpCode() == OpCode::RETURN) {
+      return DynCast<ReturnInst>(inst.get());
     }
   }
   return nullptr;
@@ -52,4 +51,4 @@ void Function::Print(std::ostream& os) const {
   }
 }
 
-} // namespace halo
+} // namespace halo

Diff for: lib/ir/ir_builder.cc

@@ -33,7 +33,7 @@ Function* FunctionBuilder::CreateFunction(const std::string& name) {
 
 BasicBlock* BasicBlockBuilder::CreateBasicBlock(const std::string& name) {
   auto bb = std::make_unique<BasicBlock>(GetContext(), name);
-  bb->parent_function_ = GetParent();
+  bb->parent_ = GetParent();
   return Insert(std::move(bb));
 }
 
@@ -44,7 +44,7 @@ Argument* ArgumentBuilder::CreateArgument(const std::string& name) {
 Argument* ArgumentBuilder::CreateArgument(const std::string& name,
                                           const Type& type) {
   auto arg = std::make_unique<Argument>(GetContext(), name, type);
-  arg->parent_function_ = GetParent();
+  arg->parent_ = GetParent();
   return Insert(std::move(arg));
 }
 
@@ -221,4 +221,4 @@ DummyInst* IRBuilder::CreateDummy(const std::string& name,
 
 #include "halo/lib/ir/ir_builder.cc.inc"
 
-} // namespace halo
+} // namespace halo