Use global scratch buffers for variables

This commit migrates variables to use global scratch variables if size > on_stack_limit.
This will allow users to store large strings/tuples/objects in local variables.

Author: ryantimwilson

.github/include/aot_skip.txt

@@ -43,6 +43,8 @@ aot.call.str_big_for
 aot.call.str_big_tuple
 # Same problem as #3392
 aot.call.tuple_scratch_buf
+# Same problem as #3392
+aot.call.variable_for_loop_scratch_buf
 aot.call.str_truncated
 aot.call.str_truncated_custom
 aot.call.strftime

src/ast/async_ids.h

@@ -19,6 +19,7 @@ namespace ast {
   DO(system)                                                                   \
   DO(time)                                                                     \
   DO(tuple)                                                                    \
+  DO(variable)                                                                 \
   DO(watchpoint)
 
 #define DEFINE_MEMBER_VAR(id, ...) int _##id = 0;

src/ast/irbuilderbpf.cpp

@@ -183,6 +183,15 @@ AllocaInst *IRBuilderBPF::CreateAllocaBPF(const SizedType &stype,
   return CreateAllocaBPF(ty, name);
 }
 
+void IRBuilderBPF::CreateAllocationInit(const SizedType &stype, Value *alloc)
+{
+  if (needMemcpy(stype)) {
+    CreateMemsetBPF(alloc, getInt8(0), stype.GetSize());
+  } else {
+    CreateStore(ConstantInt::get(GetType(stype), 0), alloc);
+  }
+}
+
 AllocaInst *IRBuilderBPF::CreateAllocaBPFInit(const SizedType &stype,
                                               const std::string &name)
 {
@@ -194,11 +203,7 @@ AllocaInst *IRBuilderBPF::CreateAllocaBPFInit(const SizedType &stype,
     llvm::Type *ty = GetType(stype);
     alloca = CreateAlloca(ty, nullptr, name);
     CreateLifetimeStart(alloca);
-    if (needMemcpy(stype)) {
-      CreateMemsetBPF(alloca, getInt8(0), stype.GetSize());
-    } else {
-      CreateStore(ConstantInt::get(ty, 0), alloca);
-    }
+    CreateAllocationInit(stype, alloca);
   });
   return alloca;
 }
@@ -553,6 +558,28 @@ Value *IRBuilderBPF::CreateWriteMapValueAllocation(const SizedType &value_type,
       loc);
 }
 
+Value *IRBuilderBPF::CreateVariableAllocationInit(const SizedType &value_type,
+                                                  const std::string &name,
+                                                  const location &loc)
+{
+  /* Hoist variable declaration and initialization to entry point of
+   * probe/subprogram. While we technically do not need this as variables
+   * are properly scoped, it eases debugging and is consistent with previous
+   * stack-only variable implementation. */
+  Value *alloc;
+  hoist([this, &value_type, &name, &loc, &alloc] {
+    alloc = createAllocation(bpftrace::globalvars::GlobalVar::VARIABLE_BUFFER,
+                             GetType(value_type),
+                             name,
+                             loc,
+                             [](AsyncIds &async_ids) {
+                               return async_ids.variable();
+                             });
+    CreateAllocationInit(value_type, alloc);
+  });
+  return alloc;
+}
+
 Value *IRBuilderBPF::createAllocation(
     bpftrace::globalvars::GlobalVar globalvar,
     llvm::Type *obj_type,

src/ast/irbuilderbpf.h

@@ -175,6 +175,9 @@ class IRBuilderBPF : public IRBuilder<> {
   Value *CreateWriteMapValueAllocation(const SizedType &value_type,
                                        const std::string &name,
                                        const location &loc);
+  Value *CreateVariableAllocationInit(const SizedType &value_type,
+                                      const std::string &name,
+                                      const location &loc);
   void CreateCheckSetRecursion(const location &loc, int early_exit_ret);
   void CreateUnSetRecursion(const location &loc);
   CallInst *CreateHelperCall(libbpf::bpf_func_id func_id,
@@ -314,6 +317,7 @@ class IRBuilderBPF : public IRBuilder<> {
                           const location &loc,
                           std::optional<std::function<size_t(AsyncIds &)>>
                               gen_async_id_cb = std::nullopt);
+  void CreateAllocationInit(const SizedType &stype, Value *alloc);
   Value *createScratchBuffer(globalvars::GlobalVar globalvar,
                              const location &loc,
                              size_t key);

src/ast/passes/codegen_llvm.cpp

@@ -2428,7 +2428,8 @@ void CodegenLLVM::visit(AssignMapStatement &assignment)
 }
 
 void CodegenLLVM::maybeAllocVariable(const std::string &var_ident,
-                                     const SizedType &var_type)
+                                     const SizedType &var_type,
+                                     const location &loc)
 {
   if (maybeGetVariable(var_ident) != nullptr) {
     // Already been allocated
@@ -2445,9 +2446,9 @@ void CodegenLLVM::maybeAllocVariable(const std::string &var_ident,
     alloca_type = CreatePointer(pointee_type, var_type.GetAS());
   }
 
-  AllocaInst *val = b_.CreateAllocaBPFInit(alloca_type, var_ident);
+  auto val = b_.CreateVariableAllocationInit(alloca_type, var_ident, loc);
   variables_[scope_stack_.back()][var_ident] = VariableLLVM{
-    val, val->getAllocatedType()
+    val, b_.GetType(alloca_type)
   };
 }
 
@@ -2479,7 +2480,7 @@ void CodegenLLVM::visit(AssignVarStatement &assignment)
 
   auto scoped_del = accept(assignment.expr);
 
-  maybeAllocVariable(var.ident, var.type);
+  maybeAllocVariable(var.ident, var.type, var.loc);
 
   if (var.type.IsArrayTy() || var.type.IsRecordTy()) {
     // For arrays and structs, only the pointer is stored
@@ -2512,7 +2513,7 @@ void CodegenLLVM::visit(VarDeclStatement &decl)
     // unused and has no type
     return;
   }
-  maybeAllocVariable(var.ident, var.type);
+  maybeAllocVariable(var.ident, var.type, var.loc);
 }
 
 void CodegenLLVM::visit(If &if_node)

src/ast/passes/codegen_llvm.h

@@ -265,7 +265,8 @@ class CodegenLLVM : public Visitor {
                             const SizedType &key_type);
 
   void maybeAllocVariable(const std::string &var_ident,
-                          const SizedType &var_type);
+                          const SizedType &var_type,
+                          const location &loc);
   VariableLLVM *maybeGetVariable(const std::string &);
   VariableLLVM &getVariable(const std::string &);
 

src/ast/passes/resource_analyser.cpp

@@ -86,6 +86,14 @@ std::optional<RequiredResources> ResourceAnalyser::analyse()
         bpftrace::globalvars::GlobalVar::MAX_CPU_ID);
   }
 
+  if (resources_.max_variable_size > 0) {
+    assert(resources_.variable_buffers > 0);
+    resources_.needed_global_vars.insert(
+        bpftrace::globalvars::GlobalVar::VARIABLE_BUFFER);
+    resources_.needed_global_vars.insert(
+        bpftrace::globalvars::GlobalVar::MAX_CPU_ID);
+  }
+
   return std::optional{ std::move(resources_) };
 }
 
@@ -114,8 +122,6 @@ void ResourceAnalyser::visit(Call &call)
 {
   Visitor::visit(call);
 
-  const auto on_stack_limit = bpftrace_.config_.get(
-      ConfigKeyInt::on_stack_limit);
   if (call.func == "printf" || call.func == "system" || call.func == "cat" ||
       call.func == "debugf") {
     // Implicit first field is the 64bit printf ID.
@@ -152,7 +158,7 @@ void ResourceAnalyser::visit(Call &call)
     // will use this information to create a percpu array map of large
     // enough size for all fmt string calls to use.
     const auto tuple_size = static_cast<uint64_t>(tuple->size);
-    if (tuple_size > on_stack_limit) {
+    if (exceeds_stack_limit(tuple_size)) {
       resources_.max_fmtstring_args_size = std::max(
           resources_.max_fmtstring_args_size,
           static_cast<uint64_t>(tuple_size));
@@ -227,7 +233,7 @@ void ResourceAnalyser::visit(Call &call)
       resources_.non_map_print_args.push_back(arg.type);
 
       const auto fmtstring_args_size = nonmap_headroom + arg.type.GetSize();
-      if (fmtstring_args_size > on_stack_limit) {
+      if (exceeds_stack_limit(fmtstring_args_size)) {
         resources_.max_fmtstring_args_size = std::max(
             resources_.max_fmtstring_args_size, fmtstring_args_size);
       }
@@ -237,7 +243,7 @@ void ResourceAnalyser::visit(Call &call)
         resources_.non_map_print_args.push_back(map.type);
 
         const auto fmtstring_args_size = nonmap_headroom + map.type.GetSize();
-        if (fmtstring_args_size > on_stack_limit) {
+        if (exceeds_stack_limit(fmtstring_args_size)) {
           resources_.max_fmtstring_args_size = std::max(
               resources_.max_fmtstring_args_size, fmtstring_args_size);
         }
@@ -261,7 +267,7 @@ void ResourceAnalyser::visit(Call &call)
   } else if (call.func == "delete") {
     auto &arg0 = *call.vargs.at(0);
     auto &map = static_cast<Map &>(arg0);
-    if (map.type.GetSize() > on_stack_limit) {
+    if (exceeds_stack_limit(map.type.GetSize())) {
       resources_.max_write_map_value_size = std::max(
           resources_.max_write_map_value_size, map.type.GetSize());
     }
@@ -279,7 +285,7 @@ void ResourceAnalyser::visit(Call &call)
 
   if (call.func == "str" || call.func == "buf" || call.func == "path") {
     const auto max_strlen = bpftrace_.config_.get(ConfigKeyInt::max_strlen);
-    if (max_strlen > on_stack_limit)
+    if (exceeds_stack_limit(max_strlen))
       resources_.str_buffers++;
   }
 
@@ -296,9 +302,7 @@ void ResourceAnalyser::visit(Map &map)
 
   update_map_info(map);
 
-  const auto on_stack_limit = bpftrace_.config_.get(
-      ConfigKeyInt::on_stack_limit);
-  if (map.type.GetSize() > on_stack_limit) {
+  if (exceeds_stack_limit(map.type.GetSize())) {
     resources_.read_map_value_buffers++;
     resources_.max_read_map_value_size = std::max(
         resources_.max_read_map_value_size, map.type.GetSize());
@@ -309,9 +313,7 @@ void ResourceAnalyser::visit(Tuple &tuple)
 {
   Visitor::visit(tuple);
 
-  const auto on_stack_limit = bpftrace_.config_.get(
-      ConfigKeyInt::on_stack_limit);
-  if (tuple.type.GetSize() > on_stack_limit) {
+  if (exceeds_stack_limit(tuple.type.GetSize())) {
     resources_.tuple_buffers++;
     resources_.max_tuple_size = std::max(resources_.max_tuple_size,
                                          tuple.type.GetSize());
@@ -323,9 +325,7 @@ void ResourceAnalyser::visit(For &f)
   Visitor::visit(f);
 
   // Need tuple per for loop to store key and value
-  const auto on_stack_limit = bpftrace_.config_.get(
-      ConfigKeyInt::on_stack_limit);
-  if (f.decl->type.GetSize() > on_stack_limit) {
+  if (exceeds_stack_limit(f.decl->type.GetSize())) {
     resources_.tuple_buffers++;
     resources_.max_tuple_size = std::max(resources_.max_tuple_size,
                                          f.decl->type.GetSize());
@@ -360,9 +360,7 @@ void ResourceAnalyser::visit(AssignMapStatement &assignment)
   update_map_info(*assignment.map);
 
   if (needAssignMapStatementAllocation(assignment)) {
-    const auto on_stack_limit = bpftrace_.config_.get(
-        ConfigKeyInt::on_stack_limit);
-    if (assignment.map->type.GetSize() > on_stack_limit) {
+    if (exceeds_stack_limit(assignment.map->type.GetSize())) {
       resources_.max_write_map_value_size = std::max(
           resources_.max_write_map_value_size, assignment.map->type.GetSize());
     }
@@ -379,14 +377,46 @@ void ResourceAnalyser::visit(Ternary &ternary)
   // blow it up. So we need a scratch buffer for it.
 
   if (ternary.type.IsStringTy()) {
-    const auto on_stack_limit = bpftrace_.config_.get(
-        ConfigKeyInt::on_stack_limit);
     const auto max_strlen = bpftrace_.config_.get(ConfigKeyInt::max_strlen);
-    if (max_strlen > on_stack_limit)
+    if (exceeds_stack_limit(max_strlen))
       resources_.str_buffers++;
   }
 }
 
+void ResourceAnalyser::update_variable_info(Variable &var)
+{
+  /* Note we don't check if a variable has been declared/assigned before.
+   * We do this to simplify the code and make it more robust to changes
+   * in other modules at the expense of memory over-allocation. Otherwise,
+   * we would need to track scopes like SemanticAnalyser and CodegenLLVM
+   * and duplicate scope tracking in a third module.
+   */
+  if (exceeds_stack_limit(var.type.GetSize())) {
+    resources_.variable_buffers++;
+    resources_.max_variable_size = std::max(resources_.max_variable_size,
+                                            var.type.GetSize());
+  }
+}
+
+void ResourceAnalyser::visit(AssignVarStatement &assignment)
+{
+  Visitor::visit(assignment);
+
+  update_variable_info(*assignment.var);
+}
+
+void ResourceAnalyser::visit(VarDeclStatement &decl)
+{
+  Visitor::visit(decl);
+
+  update_variable_info(*decl.var);
+}
+
+bool ResourceAnalyser::exceeds_stack_limit(size_t size)
+{
+  return size > bpftrace_.config_.get(ConfigKeyInt::on_stack_limit);
+}
+
 bool ResourceAnalyser::uses_usym_table(const std::string &fun)
 {
   return fun == "usym" || fun == "func" || fun == "ustack";

src/ast/passes/resource_analyser.h

@@ -37,12 +37,17 @@ class ResourceAnalyser : public Visitor {
   void visit(For &f) override;
   void visit(Ternary &ternary) override;
   void visit(AssignMapStatement &assignment) override;
+  void visit(AssignVarStatement &assignment) override;
+  void visit(VarDeclStatement &decl) override;
 
   // Determines whether the given function uses userspace symbol resolution.
   // This is used later for loading the symbol table into memory.
   bool uses_usym_table(const std::string &fun);
 
+  bool exceeds_stack_limit(size_t size);
+
   void update_map_info(Map &map);
+  void update_variable_info(Variable &var);
 
   RequiredResources resources_;
   Node *root_;

src/ast/passes/semantic_analyser.cpp

@@ -3004,22 +3004,6 @@ void SemanticAnalyser::visit(AssignVarStatement &assignment)
     if (storedTy.IsNoneTy())
       LOG(ERROR, assignment.expr->loc, err_)
           << "Invalid expression for assignment: " << storedTy;
-
-    // Scratch variables are on the BPF stack, which is only 512 bytes at time
-    // of writing, so large values do not fit on there. 200 is a ballpark of
-    // what probably won't work.
-    auto expr_size = storedTy.GetSize();
-    if (expr_size > 200) {
-      LOG(ERROR, assignment.loc, err_)
-          << "Value is too big "
-          << "(" << expr_size << " bytes) for the stack. "
-          << "Try reducing its size, storing it in a map, or creating it in "
-             "argument position to a helper call.\n\n"
-          << "Examples:\n"
-          << "    `$s = str(..);` => `$s = str(.., 32);`\n"
-          << "    `$s = str(..);` => `@s = str(..);`\n"
-          << "    `$s = str(..); print($s);` => `print(str(..));`\n\n";
-    }
   }
 }
 

src/globalvars.cpp

@@ -142,6 +142,7 @@ static void update_global_vars_rodata(
       case GlobalVar::GET_STR_BUFFER:
       case GlobalVar::READ_MAP_VALUE_BUFFER:
       case GlobalVar::WRITE_MAP_VALUE_BUFFER:
+      case GlobalVar::VARIABLE_BUFFER:
         break;
     }
   }
@@ -287,6 +288,12 @@ SizedType get_type(bpftrace::globalvars::GlobalVar global_var,
       assert(resources.max_write_map_value_size > 0);
       return make_rw_type(
           1, CreateArray(resources.max_write_map_value_size, CreateInt8()));
+    case GlobalVar::VARIABLE_BUFFER:
+      assert(resources.variable_buffers > 0);
+      assert(resources.max_variable_size > 0);
+      return make_rw_type(resources.variable_buffers,
+                          CreateArray(resources.max_variable_size,
+                                      CreateInt8()));
   }
   return {}; // unreachable
 }