Support limit, offset, top-k serial and parallel versions in codegen (cmu-db#1435)

* Initial limit commit

* Cleanup OrderByPlan

* Don't use SQL types for the result of the comparison when sorting

* Sorting for Top-K works. Probably slow.

* Optimize pop-then-push into single operation

* Fix lang::If to allow injection of custom then and else blocks

* Added early exit block to consumer context to allow operators to quick exit

* More comments

* Fast-exit limits

* Cleanup comments. Enabled parallel sorting.

* Fixed bytecode compilation. Added explicit call handlers for sorter to bytecode builder. Fixed parallel sorting top-k

* Set limit and offset to 0 when no limit/offset provided

* Fix offset

* Fix sort test in optimize to have well-defined order (some sort values were equal producing non-deterministic outputs). Fixed order by translator test to push a limit plan on top. This wasn't there before because we didn't need a limit

* Cleanup

* Fix limit without order-by clause

Author: pmenon

src/codegen/aggregation.cpp

@@ -417,7 +417,7 @@ void Aggregation::DoNullCheck(
         default: { break; }
       }
     }
-    agg_is_null.ElseBlock("Agg.IfAggIsNotNull");
+    agg_is_null.ElseBlock();
     {
       // (1)
       DoAdvanceValue(codegen, space, type, storage_index, update);

src/codegen/code_context.cpp

@@ -43,9 +43,8 @@ namespace {
 
 class PelotonMemoryManager : public llvm::SectionMemoryManager {
  public:
-  explicit PelotonMemoryManager(
-      const std::unordered_map<std::string,
-                               std::pair<llvm::Function *, CodeContext::FuncPtr>> &builtins)
+  explicit PelotonMemoryManager(const std::unordered_map<
+      std::string, std::pair<llvm::Function *, CodeContext::FuncPtr>> &builtins)
       : builtins_(builtins) {}
 
 #if LLVM_VERSION_GE(4, 0)
@@ -93,8 +92,8 @@ class PelotonMemoryManager : public llvm::SectionMemoryManager {
  private:
   // The code context
   const std::unordered_map<std::string,
-                           std::pair<llvm::Function *, CodeContext::FuncPtr>>
-      &builtins_;
+                           std::pair<llvm::Function *, CodeContext::FuncPtr>> &
+      builtins_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -143,9 +142,9 @@ class InstructionCounts : public llvm::ModulePass {
 
   void DumpStats() const {
     LOG_INFO("# functions: %" PRId64 " (%" PRId64
-              " external), # blocks: %" PRId64 ", # instructions: %" PRId64,
-              func_count_, external_func_count_, basic_block_count_,
-              total_inst_counts_);
+             " external), # blocks: %" PRId64 ", # instructions: %" PRId64,
+             func_count_, external_func_count_, basic_block_count_,
+             total_inst_counts_);
     for (const auto iter : counts_) {
       const char *inst_name = llvm::Instruction::getOpcodeName(iter.first);
       LOG_INFO("↳ %s: %" PRId64, inst_name, iter.second);
@@ -198,13 +197,13 @@ CodeContext::CodeContext()
   // references etc.
   std::unique_ptr<llvm::Module> m{module_};
   module_ = m.get();
-  engine_.reset(
-      llvm::EngineBuilder(std::move(m))
-          .setEngineKind(llvm::EngineKind::JIT)
-          .setMCJITMemoryManager(llvm::make_unique<PelotonMemoryManager>(builtins_))
-          .setMCPU(llvm::sys::getHostCPUName())
-          .setErrorStr(&err_str_)
-          .create());
+  engine_.reset(llvm::EngineBuilder(std::move(m))
+                    .setEngineKind(llvm::EngineKind::JIT)
+                    .setMCJITMemoryManager(
+                         llvm::make_unique<PelotonMemoryManager>(builtins_))
+                    .setMCPU(llvm::sys::getHostCPUName())
+                    .setErrorStr(&err_str_)
+                    .create());
   PELOTON_ASSERT(engine_ != nullptr);
 
   // The set of optimization passes we include
@@ -272,9 +271,13 @@ void CodeContext::RegisterBuiltin(llvm::Function *func_decl,
   builtins_[name] = std::make_pair(func_decl, func_impl);
 }
 
-std::pair<llvm::Function *, CodeContext::FuncPtr> CodeContext::LookupBuiltin(const std::string &name) const {
+std::pair<llvm::Function *, CodeContext::FuncPtr> CodeContext::LookupBuiltin(
+    const std::string &name) const {
   auto iter = builtins_.find(name);
-  return (iter == builtins_.end() ? std::make_pair<llvm::Function *, CodeContext::FuncPtr>(nullptr, nullptr) : iter->second);
+  return (iter == builtins_.end()
+              ? std::make_pair<llvm::Function *, CodeContext::FuncPtr>(nullptr,
+                                                                       nullptr)
+              : iter->second);
 }
 
 /// Verify all the functions that were created in this context
@@ -328,7 +331,6 @@ void CodeContext::Compile() {
   }
 
   // Log the module
-  LOG_TRACE("%s\n", GetIR().c_str());
   if (settings::SettingsManager::GetBool(settings::SettingId::dump_ir)) {
     LOG_DEBUG("%s\n", GetIR().c_str());
   }
@@ -351,7 +353,8 @@ size_t CodeContext::GetTypeAllocSizeInBits(llvm::Type *type) const {
   return GetDataLayout().getTypeAllocSizeInBits(type);
 }
 
-size_t CodeContext::GetStructElementOffset(llvm::StructType *type, size_t index) const {
+size_t CodeContext::GetStructElementOffset(llvm::StructType *type,
+                                           size_t index) const {
   return GetDataLayout().getStructLayout(type)->getElementOffset(index);
 }
 

src/codegen/consumer_context.cpp

@@ -20,14 +20,16 @@ namespace codegen {
 // Constructor
 ConsumerContext::ConsumerContext(CompilationContext &compilation_context,
                                  Pipeline &pipeline)
-    : ConsumerContext(compilation_context, pipeline, nullptr) {}
+    : ConsumerContext(compilation_context, pipeline, nullptr, nullptr) {}
 
 ConsumerContext::ConsumerContext(CompilationContext &compilation_context,
                                  Pipeline &pipeline,
-                                 PipelineContext *pipeline_context)
+                                 PipelineContext *pipeline_context,
+                                 llvm::BasicBlock *exit_block)
     : compilation_context_(compilation_context),
       pipeline_(pipeline),
-      pipeline_context_(pipeline_context) {}
+      pipeline_context_(pipeline_context),
+      exit_block_(exit_block) {}
 
 // Pass the row batch to the next operator in the pipeline
 void ConsumerContext::Consume(RowBatch &batch) {

src/codegen/expression/case_translator.cpp

@@ -6,7 +6,7 @@
 //
 // Identification: src/codegen/expression/case_translator.cpp
 //
-// Copyright (c) 2015-2017, Carnegie Mellon University Database Group
+// Copyright (c) 2015-2018, Carnegie Mellon University Database Group
 //
 //===----------------------------------------------------------------------===//
 
@@ -47,13 +47,20 @@ codegen::Value CaseTranslator::DeriveValue(CodeGen &codegen,
   codegen::Value ret;
   for (uint32_t i = 0; i < expr.GetWhenClauseSize(); i++) {
     codegen::Value cond = row.DeriveValue(codegen, *expr.GetWhenClauseCond(i));
-    lang::If when{codegen, cond.GetValue(), "case" + std::to_string(i)};
+    lang::If when(codegen, cond.GetValue(), "case" + std::to_string(i));
     {
+      // Derive the value of this case expression
       ret = row.DeriveValue(codegen, *expr.GetWhenClauseResult(i));
+
+      // Collect the value in order to merge at the very end
       branch_vals.emplace_back(ret, codegen->GetInsertBlock());
+
+      // Branch (unconditionally) to the merging block
+      codegen->CreateBr(merge_bb);
     }
-    when.EndIf(merge_bb);
+    when.EndIf();
   }
+
   // Compute the default clause
   // default_ret will have the same type as one of the ret's from above
   codegen::Value default_ret =

src/codegen/function_builder.cpp

@@ -61,11 +61,11 @@ llvm::Function *ConstructFunction(
 
 }  // namespace
 
-//===----------------------------------------------------------------------===//
-//
-// FunctionDeclaration
-//
-//===----------------------------------------------------------------------===//
+////////////////////////////////////////////////////////////////////////////////
+///
+/// FunctionDeclaration
+///
+////////////////////////////////////////////////////////////////////////////////
 
 FunctionDeclaration::FunctionDeclaration(
     CodeContext &cc, const std::string &name,
@@ -84,11 +84,11 @@ FunctionDeclaration FunctionDeclaration::MakeDeclaration(
   return FunctionDeclaration(cc, name, visibility, ret_type, args);
 }
 
-//===----------------------------------------------------------------------===//
-//
-// FunctionBuilder
-//
-//===----------------------------------------------------------------------===//
+////////////////////////////////////////////////////////////////////////////////
+///
+/// FunctionBuilder
+///
+////////////////////////////////////////////////////////////////////////////////
 
 // We preserve the state of any ongoing function construction in order to be
 // able to restore it after this function has been fully completed. Thus,
@@ -101,7 +101,8 @@ FunctionBuilder::FunctionBuilder(CodeContext &cc, llvm::Function *func_decl)
       previous_insert_point_(cc.GetBuilder().GetInsertBlock()),
       func_(func_decl),
       overflow_bb_(nullptr),
-      divide_by_zero_bb_(nullptr) {
+      divide_by_zero_bb_(nullptr),
+      return_bb_(nullptr) {
   // At this point, we've saved the current position during code generation and
   // we have a function declaration. Now:
   //  1. We define the "entry" block and attach it to the function. At this
@@ -175,10 +176,11 @@ llvm::BasicBlock *FunctionBuilder::GetOverflowBB() {
     return overflow_bb_;
   }
 
-  CodeGen codegen{code_context_};
+  CodeGen codegen(code_context_);
 
   // Save the current position so we can restore after we're done
-  auto *curr_position = codegen->GetInsertBlock();
+  auto *curr_block = codegen->GetInsertBlock();
+  auto curr_pos = codegen->GetInsertPoint();
 
   // Create the overflow block now, but don't attach to function just yet.
   overflow_bb_ = llvm::BasicBlock::Create(codegen.GetContext(), "overflow");
@@ -189,7 +191,7 @@ llvm::BasicBlock *FunctionBuilder::GetOverflowBB() {
   codegen->CreateUnreachable();
 
   // Restore position
-  codegen->SetInsertPoint(curr_position);
+  codegen->SetInsertPoint(curr_block, curr_pos);
 
   return overflow_bb_;
 }
@@ -204,10 +206,11 @@ llvm::BasicBlock *FunctionBuilder::GetDivideByZeroBB() {
     return divide_by_zero_bb_;
   }
 
-  CodeGen codegen{code_context_};
+  CodeGen codegen(code_context_);
 
   // Save the current position so we can restore after we're done
-  auto *curr_position = codegen->GetInsertBlock();
+  auto *curr_block = codegen->GetInsertBlock();
+  auto curr_pos = codegen->GetInsertPoint();
 
   // Create the overflow block now, but don't attach to function just yet.
   divide_by_zero_bb_ =
@@ -219,18 +222,42 @@ llvm::BasicBlock *FunctionBuilder::GetDivideByZeroBB() {
   codegen->CreateUnreachable();
 
   // Restore position
-  codegen->SetInsertPoint(curr_position);
+  codegen->SetInsertPoint(curr_block, curr_pos);
 
   return divide_by_zero_bb_;
 }
 
+llvm::BasicBlock *FunctionBuilder::GetExitBlock() {
+  if (return_bb_ != nullptr) {
+    return return_bb_;
+  }
+
+  CodeGen codegen(code_context_);
+
+  // Save the current position so we can restore after we're done
+  auto *curr_block = codegen->GetInsertBlock();
+  auto curr_pos = codegen->GetInsertPoint();
+
+  // Create the overflow block now, but don't attach to function just yet.
+  return_bb_ = llvm::BasicBlock::Create(codegen.GetContext(), "return");
+
+  // Quickly switch into the return block
+  codegen->SetInsertPoint(return_bb_);
+  codegen->CreateRetVoid();
+
+  // Restore position
+  codegen->SetInsertPoint(curr_block, curr_pos);
+
+  return return_bb_;
+}
+
 // Return the given value from the function and finish it
 void FunctionBuilder::ReturnAndFinish(llvm::Value *ret) {
   if (finished_) {
     return;
   }
 
-  CodeGen codegen{code_context_};
+  CodeGen codegen(code_context_);
 
   if (ret != nullptr) {
     codegen->CreateRet(ret);
@@ -239,12 +266,17 @@ void FunctionBuilder::ReturnAndFinish(llvm::Value *ret) {
     codegen->CreateRetVoid();
   }
 
-  // Add the overflow error block if it exists
+  // Add the return block, if it exists
+  if (return_bb_ != nullptr) {
+    return_bb_->insertInto(func_);
+  }
+
+  // Add the overflow error block, if it exists
   if (overflow_bb_ != nullptr) {
     overflow_bb_->insertInto(func_);
   }
 
-  // Add the divide-by-zero error block if it exists
+  // Add the divide-by-zero error block, if it exists
   if (divide_by_zero_bb_ != nullptr) {
     divide_by_zero_bb_->insertInto(func_);
   }
@@ -261,23 +293,33 @@ void FunctionBuilder::ReturnAndFinish(llvm::Value *ret) {
 }
 
 llvm::Value *FunctionBuilder::GetOrCacheVariable(
-    const std::string &name, const std::function<llvm::Value *()> &func) {
-  auto iter = cached_vars_.find(name);
-  if (iter != cached_vars_.end()) {
-    return iter->second;
-  } else {
-    CodeGen codegen{code_context_};
-    auto pos = codegen->GetInsertPoint();
-    if (entry_bb_->empty()) {
-      codegen->SetInsertPoint(entry_bb_);
-    } else {
-      codegen->SetInsertPoint(&entry_bb_->back());
-    }
-    auto *val = func();
-    codegen->SetInsertPoint(&*pos);
-    cached_vars_.emplace(name, val);
+    const std::string &name, const std::function<llvm::Value *()> &load_func) {
+  llvm::Value *&val = cached_vars_[name];
+  if (val != nullptr) {
     return val;
   }
+
+  CodeGen codegen(code_context_);
+
+  // Save current position
+  auto *curr_block = codegen->GetInsertBlock();
+  auto curr_pos = codegen->GetInsertPoint();
+
+  // Move to start of function
+  if (entry_bb_->empty()) {
+    codegen->SetInsertPoint(entry_bb_);
+  } else {
+    codegen->SetInsertPoint(&entry_bb_->back());
+  }
+
+  // Generate loading code and save in cache
+  val = load_func();
+
+  // Move back
+  codegen->SetInsertPoint(curr_block, curr_pos);
+
+  // Return value
+  return val;
 }
 
 }  // namespace codegen