llvm-journey

kaleidoscope_codegen.cpp (6500B)

---

 
 #include "kaleidoscope_codegen.hpp"
 
 #include "kaleidoscope_ast.hpp"
 
 #include <algorithm>
 #include <string>
 #include <utility>
 #include <vector>
 
 // clang-format off
 #include <llvm/IR/LLVMContext.h>
 #include <llvm/IR/IRBuilder.h>
 #include <llvm/IR/Module.h>
 // clang-format on
 
 #include <llvm/ADT/APFloat.h>
 #include <llvm/IR/Constants.h>
 #include <llvm/IR/Value.h>
 #include <llvm/IR/Verifier.h>
 
 #include <llvm/IR/LegacyPassManager.h>
 #include <llvm/Transforms/InstCombine/InstCombine.h>
 #include <llvm/Transforms/Scalar.h>
 #include <llvm/Transforms/Scalar/GVN.h>
 
 namespace {
 
 int
 assert_fail(const char* expr, const char* file, int line, const char* func)
 {
   fprintf(
     stderr, "Assertion failed: %s (%s: %s: %d)\n", expr, file, func, line);
   fflush(NULL);
   abort();
 
   return 0; // dummy
 }
 #define ASSERT_ALWAYS(x) ((x) || assert_fail(#x, __FILE__, __LINE__, __func__))
 #ifndef NDEBUG
 #define ASSERT(x)
 #else
 #define ASSERT(x) ASSERT_ALWAYS(x)
 #endif
 
 inline llvm::LLVMContext&
 context()
 {
   static llvm::LLVMContext ctx;
 
   return ctx;
 }
 
 inline llvm::IRBuilder<>&
 builder()
 {
   static llvm::IRBuilder<> b(context());
 
   return b;
 }
 
 inline llvm::Module*
 module()
 {
   static auto m = std::make_unique<llvm::Module>("Kaleidoscope JIT", context());
 
   return m.get();
 }
 
 inline llvm::legacy::FunctionPassManager*
 fpm()
 {
   static auto f = [] {
     auto f = std::make_unique<llvm::legacy::FunctionPassManager>(module());
 
     f->add(llvm::createInstructionCombiningPass());
     f->add(llvm::createReassociatePass());
     f->add(llvm::createGVNPass());
     f->add(llvm::createCFGSimplificationPass());
 
     f->doInitialization();
     return f;
   }();
 
   return f.get();
 }
 
 template<typename K, typename V>
 class Map
 {
 private:
   std::vector<std::pair<K, V>> d_; // data
 
   auto find_(const K& k)
   {
     const auto& f = d_.begin();
     const auto& l = d_.end();
     auto r = std::lower_bound(f, l, k, [](const auto& pair, const auto& key) {
       return pair.first < key;
     });
 
     return std::make_pair(r, l);
   }
 
 public:
   V find(const K& k)
   {
     auto r = find_(k);
 
     if (r.first == r.second)
       return nullptr;
     return k == r.first->first ? r.first->second : nullptr;
   }
 
   void insert(const K& k, const V& v)
   {
     auto r = find_(k);
 
     if (r.first != r.second && k == r.first->first)
       assert(v == r.first->second);
     else
       d_.insert(r.first, std::make_pair(k, v));
   }
 
   template<typename F>
   void for_each(F f)
   {
     for (const auto& e : d_)
       f(e.first, e.second);
   }
 
   void clear() { d_.clear(); }
 };
 
 Map<std::string, llvm::Function*> funcs;
 Map<std::string, llvm::Value*> names;
 
 } // anonymous namespace
 
 llvm::Value*
 kal::codegen(const Number& n)
 {
   return llvm::ConstantFP::get(context(), llvm::APFloat{ n.value });
 }
 
 llvm::Value*
 kal::codegen(const Variable& v)
 {
   return names.find(v.name);
 }
 
 llvm::Value*
 kal::codegen(const BinaryOp& b)
 {
   auto l = kal::codegen(b.lhs);
   auto r = kal::codegen(b.rhs);
 
   ASSERT_ALWAYS(l != nullptr);
   ASSERT_ALWAYS(r != nullptr);
 
   switch (b.op) {
     case '*':
       return builder().CreateFMul(l, r, "multmp");
 
     case '/':
       return builder().CreateFDiv(l, r, "divtmp");
 
     case '+':
       return builder().CreateFAdd(l, r, "addtmp");
 
     case '-':
       return builder().CreateFSub(l, r, "subtmp");
 
     case '<':
       l = builder().CreateFCmpULT(l, r, "cmptmp");
       return builder().CreateUIToFP(
         l, llvm::Type::getDoubleTy(context()), "booltmp");
 
     default:
       ASSERT_ALWAYS(false && "unknown operator");
   }
 
   return nullptr;
 }
 
 llvm::Value*
 kal::codegen(const Call& c)
 {
   auto fun = funcs.find(c.callee);
   auto n = c.args.size();
 
   ASSERT_ALWAYS(fun != nullptr);
   ASSERT_ALWAYS(n == fun->arg_size());
 
   std::vector<llvm::Value*> args(n);
 
   for (auto i = 0u; i < n; i++) {
     args[i] = kal::codegen(c.args[i]);
     ASSERT_ALWAYS(args[i] != nullptr);
   }
 
   return builder().CreateCall(fun, args, "calltmp");
 }
 
 llvm::Value*
 kal::codegen(const kal::ASTNode& n)
 {
   using kal::node_type;
 
   switch (node_type(n)) {
     case kal::NodeType::Number: {
       kal::Number num;
 
       cast(n, num);
       return kal::codegen(num);
     }
 
     case kal::NodeType::Variable: {
       kal::Variable v;
 
       cast(n, v);
       return kal::codegen(v);
     }
 
     case kal::NodeType::BinaryOp: {
       kal::BinaryOp op;
 
       cast(n, op);
       return kal::codegen(op);
     }
 
     case kal::NodeType::Call: {
       kal::Call c;
 
       cast(n, c);
       return kal::codegen(c);
     }
 
     case kal::NodeType::None:
     case kal::NodeType::Prototype:
     case kal::NodeType::Function:
       break;
   }
 
   ASSERT_ALWAYS(false && "unknown NodeType");
   return nullptr;
 }
 
 llvm::Function*
 kal::codegen(const Prototype& p)
 {
   std::vector<llvm::Type*> dbls(p.params.size(),
                                 llvm::Type::getDoubleTy(context()));
   auto ftype =
     llvm::FunctionType::get(llvm::Type::getDoubleTy(context()), dbls, false);
   auto func = llvm::Function::Create(
     ftype, llvm::Function::ExternalLinkage, p.name, module());
 
   {
     auto i = 0;
 
     for (auto& arg : func->args())
       arg.setName(p.params[i++]);
   }
 
   funcs.insert(p.name, func);
   return func;
 }
 
 llvm::Function*
 kal::codegen(const Function& f, bool optimize)
 {
   auto func = funcs.find(f.proto.name);
 
   if (func == nullptr) {
     func = kal::codegen(f.proto);
     if (func != nullptr)
       funcs.insert(f.proto.name, func);
   }
 
   ASSERT_ALWAYS(func != nullptr);
 
   auto bb = llvm::BasicBlock::Create(context(), "entry", func);
 
   builder().SetInsertPoint(bb);
 
   names.clear();
   for (auto& arg : func->args())
     names.insert(arg.getName(), &arg);
 
   auto body = kal::codegen(f.body);
 
   ASSERT_ALWAYS(body != nullptr);
 
   builder().CreateRet(body);
   llvm::verifyFunction(*func);
   if (optimize)
     fpm()->run(*func);
 
   return func;
 }
 
 llvm::Function*
 kal::mkfunc(const std::string& name, ASTNode* f, ASTNode* l, bool optimize)
 {
   if (f == l)
     return nullptr;
 
   auto func = codegen(kal::Prototype{ name, {} });
 
   ASSERT_ALWAYS(func != nullptr);
 
   auto bb = llvm::BasicBlock::Create(context(), "entry", func);
 
   builder().SetInsertPoint(bb);
   while (f != l)
     kal::codegen(*f++);
   builder().CreateRet(kal::codegen(kal::Number{ 0 }));
 
   llvm::verifyFunction(*func);
   if (optimize)
     fpm()->run(*func);
   return func;
 }
 
 void
 kal::codegen_reset(void)
 {
   funcs.for_each([](const auto&, const auto& v) { v->eraseFromParent(); });
   funcs.clear();
   names.clear();
 }