Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

For loops: Allow sharing variables with main program #3014

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

ajor
Copy link
Member

@ajor ajor commented Feb 21, 2024

This lets programs such as this run:

BEGIN
{
  @map[0] = 1; @map[1] = 2;
  $cnt = 0;
  for ($kv : @map) {
    $cnt++; // Previously would have errored with "Variables defined outside of a for-loop can not be accessed in the loop's scope"
  }
  print($cnt);
}

Implementation Details

  1. Determine which variables need to be shared with the loop callback
  2. Pack pointers to them into a context struct
  3. Pass pointer to the context struct to the callback function
  4. In the callback, override the shared variables so that they read and write through the context pointers instead of directly from their original addresses

Pseudo code for the transformation we apply:

Before:

PROBE {
  $str = "hello";
  $not_shared = 2;
  $len = 0;
  @map[11, 12] = "c";
  for ($kv : @map) {
    print($str);
    $len++;
  }
  print($len);
  print($not_shared);
}

After:

struct ctx_t {
  string *str;
  uint64 *len;
};
PROBE {
  $str = "hello";
  $not_shared = 2;
  $len = 0;
  @map[11, 12] = "c";

  ctx_t ctx { .str = &$str, .len = &$len };
  bpf_for_each_map_elem(@map, &map_for_each_cb, &ctx, 0);

  print($len);
  print($not_shared);
}
long map_for_each_cb(bpf_map *map,
                     const void *key,
                     void *value,
                     void *ctx) {
  $kv = (((uint64, uint64))key, (string)value);
  $str = ((ctx_t*)ctx)->str;
  $len = ((ctx_t*)ctx)->len;

  print($str);
  $len++;
}
Checklist

N/A: For-loops haven't been released yet and this is basic functionality expected of for-loops so is covered by the existing docs/changelog.
- [ ] Language changes are updated in man/adoc/bpftrace.adoc and if needed in docs/reference_guide.md
- [ ] User-visible and non-trivial changes updated in CHANGELOG.md

  • The new behaviour is covered by tests

@ajor ajor force-pushed the loop-maps-ctx branch 2 times, most recently from 915ca3b to f00dedd Compare March 22, 2024 17:32
@ajor ajor force-pushed the loop-maps-ctx branch 2 times, most recently from 14cc73e to 70e0798 Compare May 13, 2024 15:02
1. Determine which variables need to be shared with the loop callback
2. Pack pointers to them into a context struct
3. Pass pointer to the context struct to the callback function
4. In the callback, override the shared variables so that they read and
   write through the context pointers instead of directly from their
   original addresses

See the comment in semantic_analyser.cpp for pseudo code of this
transformation.
@ajor ajor marked this pull request as ready for review May 13, 2024 15:54
@@ -351,19 +352,27 @@ void Printer::visit(While &while_block)
}
}

void Printer::visit(For &for_loop)
void Printer::visit(For &f)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I know f is more terse but I much prefer the variable name for_loop.

@@ -176,7 +182,7 @@ class StructManager {

private:
std::map<std::string, std::shared_ptr<Struct>> struct_map_;
std::unordered_set<std::shared_ptr<Struct>> tuples_;
std::unordered_set<std::shared_ptr<Struct>> anonymous_types_;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

naming nit: maybe make this consistent with the method e.g. anon_types_ or AddAnonymousStruct

@@ -259,7 +258,11 @@ class CodegenLLVM : public Visitor {
int current_usdt_location_index_{ 0 };
bool inside_subprog_ = false;

std::map<std::string, AllocaInst *> variables_;
struct VariableLLVM {
llvm::Value *value;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why make the value generic, it seems like all the places where this is used are creating a AllocaInst?

llvm::Value *value;
llvm::Type *type;
};
std::map<std::string, VariableLLVM> variables_;
Copy link
Contributor

@jordalgo jordalgo May 15, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since this is class level is there a case of possible collision where two (or more probes) use the same scratch variable name and reference them in a loop?

Or no, because each prog has it's own ELF file...

ctx_t = b_.GetStructType("ctx_t", ctx_field_types);
ctx = b_.CreateAllocaBPF(ctx_t, "ctx");

for (size_t i = 0; i < ctx_fields.size(); i++) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

super duper nit

Suggested change
for (size_t i = 0; i < ctx_fields.size(); i++) {
for (size_t i = 0; i < ctx_fields.size(); ++i) {

Comment on lines +4252 to +4253
auto *field_expr = variables_[field.name].value;
auto *field_ty = field_expr->getType();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

super nit (makes it slightly easier to grok IMHO)

Suggested change
auto *field_expr = variables_[field.name].value;
auto *field_ty = field_expr->getType();
auto *field_ty = variables_[field.name].value->getType();

* write through the context pointers instead of directly from their
* original addresses
*
* Example transformation with context:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like the example!

// function as the context parameter.
CollectNodes<Variable> vars_referenced;
std::unordered_set<std::string> var_set;
for (auto *stmt : *f.stmts) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: looks like you could combine these two loops to get referenced and new vars.

Copy link
Contributor

@jordalgo jordalgo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Few small questions and nits but otherwise LGTM.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants