merge allocations in same basic block up to a size bound (#1125)

Previously, when we were allocating a large term, we would issue a
number of repeated calls to kore_alloc. This is inefficient because it
has to check whether there is enough space and execute a branch
instruction on each allocation. It is more efficient if we do one,
larger call to kore_alloc for all the memory that we need.

In order to efficiently implement this, we modify the way we generate
calls to kore_alloc in the code generator so that it will place the
allocation function at the beginning of the basic block and then, if
further allocations are required in the same basic block, it will modify
the call to the allocation function so it requests a larger amount of
memory. This only happens up to a certain bound, in order to prevent
memory from being wasted due to requesting particularly large chunks of
memory. In practice, this leads to most apply_rule functions only
calling kore_alloc once.

Author: Dwight Guth

include/kllvm/codegen/CreateTerm.h

@@ -1,13 +1,18 @@
 #ifndef CREATE_TERM_H
 #define CREATE_TERM_H
 
+#include <concepts>
+
 #include "kllvm/ast/AST.h"
 #include "kllvm/codegen/DecisionParser.h"
 
 #include "llvm/ADT/StringMap.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Value.h"
 
+// size up to which calls to the same allocation function will be merged.
+static constexpr int max_block_merge_size = 4096;
+
 namespace kllvm {
 
 class create_term {
@@ -143,12 +148,15 @@ bool is_injection_symbol(kore_pattern *p, kore_symbol *sym);
 
 void add_abort(llvm::BasicBlock *block, llvm::Module *module);
 
+template <typename T>
+  requires std::same_as<T, llvm::BasicBlock>
+           || std::same_as<T, llvm::Instruction>
 llvm::Value *allocate_term(
-    llvm::Type *alloc_type, llvm::BasicBlock *block,
+    llvm::Type *alloc_type, llvm::Value *len, T *insert_point,
     char const *alloc_fn = "kore_alloc");
 llvm::Value *allocate_term(
-    llvm::Type *alloc_type, llvm::Value *len, llvm::BasicBlock *block,
-    char const *alloc_fn = "kore_alloc");
+    llvm::Type *alloc_type, llvm::BasicBlock *block,
+    char const *alloc_fn = "kore_alloc", bool mergeable = false);
 } // namespace kllvm
 
 #endif // CREATE_TERM_H

include/kllvm/codegen/Util.h

@@ -23,6 +23,8 @@ llvm::Function *kore_heap_alloc(std::string const &name, llvm::Module *module);
 
 llvm::Instruction *create_malloc(
     llvm::BasicBlock *block, llvm::Value *alloc_size, llvm::Function *malloc_f);
+llvm::Instruction *create_malloc(
+    llvm::Instruction *inst, llvm::Value *alloc_size, llvm::Function *malloc_f);
 
 // getOrInsertFunction on module, aborting on failure
 template <class... Ts>

lib/codegen/CreateTerm.cpp

@@ -168,22 +168,58 @@ llvm::Value *get_block_header(
           llvm::Type::getInt64Ty(module->getContext()), header_val));
 }
 
+template <typename T>
+  requires std::same_as<T, llvm::BasicBlock>
+           || std::same_as<T, llvm::Instruction>
 llvm::Value *allocate_term(
-    llvm::Type *alloc_type, llvm::BasicBlock *block, char const *alloc_fn) {
-  return allocate_term(
-      alloc_type, llvm::ConstantExpr::getSizeOf(alloc_type), block, alloc_fn);
-}
-
-llvm::Value *allocate_term(
-    llvm::Type *alloc_type, llvm::Value *len, llvm::BasicBlock *block,
+    llvm::Type *alloc_type, llvm::Value *len, T *insert_point,
     char const *alloc_fn) {
   auto *malloc = create_malloc(
-      block, len, kore_heap_alloc(alloc_fn, block->getModule()));
+      insert_point, len, kore_heap_alloc(alloc_fn, insert_point->getModule()));
 
   set_debug_loc(malloc);
   return malloc;
 }
 
+static bool is_basic_alloc(std::string const &alloc_fn) {
+  return alloc_fn == "kore_alloc" || alloc_fn == "kore_alloc_old"
+         || alloc_fn == "kore_alloc_always_gc";
+}
+
+llvm::Value *allocate_term(
+    llvm::Type *alloc_type, llvm::BasicBlock *block, char const *alloc_fn,
+    bool mergeable) {
+  llvm::DataLayout layout(block->getModule());
+  auto type_size = layout.getTypeAllocSize(alloc_type).getFixedValue();
+  auto *ty = llvm::Type::getInt64Ty(block->getContext());
+  if (mergeable) {
+    if (auto *first = block->getFirstNonPHI()) {
+      if (auto *call = llvm::dyn_cast<llvm::CallInst>(first)) {
+        if (auto *func = call->getCalledFunction()) {
+          if (auto *size
+              = llvm::dyn_cast<llvm::ConstantInt>(call->getOperand(0))) {
+            if (func->getName() == alloc_fn && is_basic_alloc(alloc_fn)
+                && size->getLimitedValue() + type_size < max_block_merge_size) {
+              call->setOperand(
+                  0, llvm::ConstantExpr::getAdd(
+                         size, llvm::ConstantInt::get(ty, type_size)));
+              auto *ret = llvm::GetElementPtrInst::Create(
+                  llvm::Type::getInt8Ty(block->getContext()), call, {size},
+                  "alloc_chunk", block);
+              set_debug_loc(ret);
+              return ret;
+            }
+          }
+        }
+      }
+    }
+    return allocate_term(
+        alloc_type, llvm::ConstantInt::get(ty, type_size), block, alloc_fn);
+  }
+  return allocate_term(
+      alloc_type, llvm::ConstantInt::get(ty, type_size), block, alloc_fn);
+}
+
 value_type term_type(
     kore_pattern *pattern, llvm::StringMap<value_type> &substitution,
     kore_definition *definition) {
@@ -686,7 +722,8 @@ llvm::Value *create_term::create_function_call(
     // we don't use alloca here because the tail call optimization pass for llvm
     // doesn't handle correctly functions with alloca
     alloc_sret = allocate_term(
-        return_type, current_block_, get_collection_alloc_fn(return_cat.cat));
+        return_type, current_block_, get_collection_alloc_fn(return_cat.cat),
+        true);
     sret_type = return_type;
     real_args.insert(real_args.begin(), alloc_sret);
     types.insert(types.begin(), alloc_sret->getType());
@@ -759,7 +796,8 @@ llvm::Value *create_term::not_injection_case(
     children.push_back(child_value);
     idx++;
   }
-  llvm::Value *block = allocate_term(block_type, current_block_);
+  llvm::Value *block
+      = allocate_term(block_type, current_block_, "kore_alloc", true);
   llvm::Value *block_header_ptr = llvm::GetElementPtrInst::CreateInBounds(
       block_type, block,
       {llvm::ConstantInt::get(llvm::Type::getInt64Ty(ctx_), 0),
@@ -1162,7 +1200,7 @@ std::string make_apply_rule_function(
       if (!arg->getType()->isPointerTy()) {
         auto *ptr = allocate_term(
             arg->getType(), creator.get_current_block(),
-            get_collection_alloc_fn(cat.cat));
+            get_collection_alloc_fn(cat.cat), true);
         new llvm::StoreInst(arg, ptr, creator.get_current_block());
         arg = ptr;
       }

lib/codegen/Util.cpp

@@ -29,6 +29,12 @@ llvm::Instruction *create_malloc(
   return llvm::CallInst::Create(malloc_f, {alloc_size}, "", block);
 }
 
+llvm::Instruction *create_malloc(
+    llvm::Instruction *inst, llvm::Value *alloc_size,
+    llvm::Function *malloc_f) {
+  return llvm::CallInst::Create(malloc_f, {alloc_size}, "", inst);
+}
+
 llvm::Constant *get_offset_of_member(
     [[maybe_unused]] llvm::Module *mod, llvm::StructType *struct_ty,
     int nth_member) {