Skip to content

Commit c0821d2

Browse files
ryantimwilsondanobi
ryantimwilson
authored and
danobi
committed
Migrate map keys to scratch buffers
1 parent 6959c4b commit c0821d2

22 files changed

+1176
-164
lines changed

‎.github/include/aot_skip.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ aot.call.str_big_tuple
4545
aot.call.tuple_scratch_buf
4646
# Same problem as #3392
4747
aot.call.variable_for_loop_scratch_buf
48+
# Same problem as #3392
49+
aot.call.map_key_scratch_buf
4850
aot.call.str_truncated
4951
aot.call.str_truncated_custom
5052
aot.call.strftime

‎src/ast/async_ids.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ namespace ast {
1212
DO(bpf_print) \
1313
DO(non_map_print) \
1414
DO(printf) \
15+
DO(map_key) \
1516
DO(read_map_value) \
1617
DO(skb_output) \
1718
DO(strftime) \

‎src/ast/codegen_helper.cpp‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,17 @@ bool needAssignMapStatementAllocation(const AssignMapStatement &assignment)
3030
return true;
3131
}
3232

33+
bool needMapKeyAllocation(const Map &map)
34+
{
35+
return needMapKeyAllocation(map, map.key_expr);
36+
}
37+
38+
bool needMapKeyAllocation(const Map &map, Expression *key_expr)
39+
{
40+
if (key_expr && inBpfMemory(key_expr->type)) {
41+
return !key_expr->type.IsSameSizeRecursive(map.key_type);
42+
}
43+
return true;
44+
}
45+
3346
} // namespace bpftrace::ast

‎src/ast/codegen_helper.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,5 +35,8 @@ inline AddrSpace find_addrspace_stack(const SizedType &ty)
3535

3636
bool needAssignMapStatementAllocation(const AssignMapStatement &assignment);
3737

38+
bool needMapKeyAllocation(const Map &map);
39+
bool needMapKeyAllocation(const Map &map, Expression *key_expr);
40+
3841
} // namespace ast
3942
} // namespace bpftrace

‎src/ast/irbuilderbpf.cpp‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -580,6 +580,19 @@ Value *IRBuilderBPF::CreateVariableAllocationInit(const SizedType &value_type,
580580
return alloc;
581581
}
582582

583+
Value *IRBuilderBPF::CreateMapKeyAllocation(const SizedType &value_type,
584+
const std::string &name,
585+
const location &loc)
586+
{
587+
return createAllocation(bpftrace::globalvars::GlobalVar::MAP_KEY_BUFFER,
588+
GetType(value_type),
589+
name,
590+
loc,
591+
[](AsyncIds &async_ids) {
592+
return async_ids.map_key();
593+
});
594+
}
595+
583596
Value *IRBuilderBPF::createAllocation(
584597
bpftrace::globalvars::GlobalVar globalvar,
585598
llvm::Type *obj_type,

‎src/ast/irbuilderbpf.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,9 @@ class IRBuilderBPF : public IRBuilder<> {
178178
Value *CreateVariableAllocationInit(const SizedType &value_type,
179179
const std::string &name,
180180
const location &loc);
181+
Value *CreateMapKeyAllocation(const SizedType &value_type,
182+
const std::string &name,
183+
const location &loc);
181184
void CreateCheckSetRecursion(const location &loc, int early_exit_ret);
182185
void CreateUnSetRecursion(const location &loc);
183186
CallInst *CreateHelperCall(libbpf::bpf_func_id func_id,

‎src/ast/passes/codegen_llvm.cpp‎

Lines changed: 35 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -712,9 +712,10 @@ void CodegenLLVM::visit(Call &call)
712712
b_.getInt64Ty(),
713713
call.vargs.front()->type.IsSigned());
714714
Value *log2 = b_.CreateCall(log2_func_, { expr_, k }, "log2");
715-
AllocaInst *key = getHistMapKey(map, log2);
715+
auto key = getHistMapKey(map, log2, call.loc);
716716
b_.CreateMapElemAdd(ctx_, map, key, b_.getInt64(1), call.loc);
717-
b_.CreateLifetimeEnd(key);
717+
if (dyn_cast<AllocaInst>(key))
718+
b_.CreateLifetimeEnd(key);
718719
expr_ = nullptr;
719720
} else if (call.func == "lhist") {
720721
if (!linear_func_)
@@ -750,10 +751,10 @@ void CodegenLLVM::visit(Call &call)
750751
{ value, min, max, step },
751752
"linear");
752753

753-
AllocaInst *key = getHistMapKey(map, linear);
754+
auto key = getHistMapKey(map, linear, call.loc);
754755
b_.CreateMapElemAdd(ctx_, map, key, b_.getInt64(1), call.loc);
755-
756-
b_.CreateLifetimeEnd(key);
756+
if (dyn_cast<AllocaInst>(key))
757+
b_.CreateLifetimeEnd(key);
757758
expr_ = nullptr;
758759
} else if (call.func == "delete") {
759760
auto &arg0 = *call.vargs.at(0);
@@ -3018,13 +3019,19 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
30183019
Expression *key_expr)
30193020
{
30203021
Value *key;
3021-
bool alloca_created_here = true;
3022+
const auto alloca_created_here = needMapKeyAllocation(map, key_expr);
3023+
30223024
if (key_expr) {
30233025
auto scoped_del = accept(key_expr);
3026+
const auto &key_type = map.key_type;
3027+
// Allocation needs to be done after recursing via accept(key_expr)
3028+
// so expr_ will be set properly
3029+
key = alloca_created_here ? b_.CreateMapKeyAllocation(key_type,
3030+
map.ident + "_key",
3031+
key_expr->loc)
3032+
: expr_;
30243033
if (inBpfMemory(key_expr->type)) {
3025-
auto &key_type = map.key_type;
30263034
if (!key_expr->type.IsSameSizeRecursive(key_type)) {
3027-
key = b_.CreateAllocaBPF(key_type, map.ident + "_key");
30283035
b_.CreateMemsetBPF(key, b_.getInt8(0), key_type.GetSize());
30293036
if (key_expr->type.IsTupleTy()) {
30303037
createTupleCopy(key_expr->type, key_type, key, expr_);
@@ -3036,25 +3043,23 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
30363043
<< " Expression Type Size: " << key_expr->type.GetSize();
30373044
}
30383045
} else {
3039-
key = expr_;
30403046
// Call-ee freed
30413047
scoped_del.disarm();
3042-
3043-
// We don't have enough visibility into where key comes from to safely
3044-
// end its lifetime. It could be a variable, for example.
3045-
alloca_created_here = false;
30463048
}
30473049
} else if (map.key_type.IsIntTy()) {
3048-
key = b_.CreateAllocaBPF(map.key_type, map.ident + "_key");
30493050
// Integers are always stored as 64-bit in map keys
30503051
b_.CreateStore(
30513052
b_.CreateIntCast(expr_, b_.getInt64Ty(), key_expr->type.IsSigned()),
30523053
key);
30533054
} else {
3054-
key = b_.CreateAllocaBPF(key_expr->type, map.ident + "_key");
30553055
if (key_expr->type.IsArrayTy() || key_expr->type.IsRecordTy()) {
30563056
// We need to read the entire array/struct and save it
3057-
b_.CreateProbeRead(ctx_, key, key_expr->type, expr_, key_expr->loc);
3057+
b_.CreateProbeRead(ctx_,
3058+
b_.CreatePointerCast(
3059+
key, b_.GetType(key_expr->type)->getPointerTo()),
3060+
key_expr->type,
3061+
expr_,
3062+
key_expr->loc);
30583063
} else {
30593064
b_.CreateStore(
30603065
b_.CreateIntCast(expr_, b_.getInt64Ty(), key_expr->type.IsSigned()),
@@ -3063,19 +3068,23 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
30633068
}
30643069
} else {
30653070
// No map key (e.g., @ = 1;). Use 0 as a key.
3066-
key = b_.CreateAllocaBPF(CreateUInt64(), map.ident + "_key");
3071+
assert(alloca_created_here);
3072+
key = b_.CreateMapKeyAllocation(CreateUInt64(),
3073+
map.ident + "_key",
3074+
map.loc);
30673075
b_.CreateStore(b_.getInt64(0), key);
30683076
}
30693077

30703078
auto key_deleter = [this, key, alloca_created_here]() {
3071-
if (alloca_created_here)
3079+
if (alloca_created_here && dyn_cast<AllocaInst>(key))
30723080
b_.CreateLifetimeEnd(key);
30733081
};
30743082
return std::make_tuple(key, ScopedExprDeleter(std::move(key_deleter)));
30753083
}
30763084

3077-
AllocaInst *CodegenLLVM::getMultiMapKey(Map &map,
3078-
const std::vector<Value *> &extra_keys)
3085+
Value *CodegenLLVM::getMultiMapKey(Map &map,
3086+
const std::vector<Value *> &extra_keys,
3087+
const location &loc)
30793088
{
30803089
size_t size = map.key_type.GetSize();
30813090
for (auto *extra_key : extra_keys) {
@@ -3084,7 +3093,9 @@ AllocaInst *CodegenLLVM::getMultiMapKey(Map &map,
30843093

30853094
// If key ever changes to not be allocated here, be sure to update getMapKey()
30863095
// as well to take the new lifetime semantics into account.
3087-
AllocaInst *key = b_.CreateAllocaBPF(size, map.ident + "_key");
3096+
auto key = b_.CreateMapKeyAllocation(CreateArray(size, CreateInt8()),
3097+
map.ident + "_key",
3098+
loc);
30883099
auto *key_type = ArrayType::get(b_.getInt8Ty(), size);
30893100

30903101
int offset = 0;
@@ -3153,12 +3164,12 @@ AllocaInst *CodegenLLVM::getMultiMapKey(Map &map,
31533164
return key;
31543165
}
31553166

3156-
AllocaInst *CodegenLLVM::getHistMapKey(Map &map, Value *log2)
3167+
Value *CodegenLLVM::getHistMapKey(Map &map, Value *log2, const location &loc)
31573168
{
31583169
if (map.key_expr)
3159-
return getMultiMapKey(map, { log2 });
3170+
return getMultiMapKey(map, { log2 }, loc);
31603171

3161-
AllocaInst *key = b_.CreateAllocaBPF(CreateUInt64(), map.ident + "_key");
3172+
auto key = b_.CreateMapKeyAllocation(CreateUInt64(), map.ident + "_key", loc);
31623173
b_.CreateStore(log2, key);
31633174
return key;
31643175
}

‎src/ast/passes/codegen_llvm.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ class CodegenLLVM : public Visitor {
7676
void visit(Program &program) override;
7777
void visit(Block &block) override;
7878

79-
AllocaInst *getHistMapKey(Map &map, Value *log2);
79+
Value *getHistMapKey(Map &map, Value *log2, const location &loc);
8080
int getNextIndexForProbe();
8181
Value *createLogicalAnd(Binop &binop);
8282
Value *createLogicalOr(Binop &binop);
@@ -189,7 +189,9 @@ class CodegenLLVM : public Visitor {
189189
[[nodiscard]] std::tuple<Value *, ScopedExprDeleter> getMapKey(
190190
Map &map,
191191
Expression *key_expr);
192-
AllocaInst *getMultiMapKey(Map &map, const std::vector<Value *> &extra_keys);
192+
Value *getMultiMapKey(Map &map,
193+
const std::vector<Value *> &extra_keys,
194+
const location &loc);
193195

194196
void compareStructure(SizedType &our_type, llvm::Type *llvm_type);
195197

‎src/ast/passes/resource_analyser.cpp‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,14 @@ std::optional<RequiredResources> ResourceAnalyser::analyse()
9494
bpftrace::globalvars::GlobalVar::MAX_CPU_ID);
9595
}
9696

97+
if (resources_.max_map_key_size > 0) {
98+
assert(resources_.map_key_buffers > 0);
99+
resources_.needed_global_vars.insert(
100+
bpftrace::globalvars::GlobalVar::MAP_KEY_BUFFER);
101+
resources_.needed_global_vars.insert(
102+
bpftrace::globalvars::GlobalVar::MAX_CPU_ID);
103+
}
104+
97105
return std::optional{ std::move(resources_) };
98106
}
99107

@@ -289,6 +297,61 @@ void ResourceAnalyser::visit(Call &call)
289297
resources_.str_buffers++;
290298
}
291299

300+
/* Aggregation functions like count/sum/max are always called like:
301+
* @ = count()
302+
* Thus, we visit AssignMapStatement AST node which visits the map and
303+
* assigns a map key buffer. Thus, there is no need to assign another
304+
* buffer here.
305+
*
306+
* The exceptions are:
307+
* 1. lhist/hist because the map key buffer includes both the key itself
308+
* and the bucket ID from a call to linear/log2 functions.
309+
* 2. has_key/delete because the map key buffer allocation depends on
310+
* arguments to the function e.g.
311+
* delete(@, 2)
312+
* requires a map key buffer to hold arg1 = 2 but map.key_expr is null
313+
* so the map key buffer check in visit(Map &map) doesn't work as is.
314+
*/
315+
if (call.func == "lhist" || call.func == "hist") {
316+
Map &map = *call.map;
317+
// Allocation is always needed for lhist/hist. But we need to allocate
318+
// space for both map key and the bucket ID from a call to linear/log2
319+
// functions.
320+
const auto map_key_size = map.key_expr ? map.key_type.GetSize() +
321+
CreateUInt64().GetSize()
322+
: CreateUInt64().GetSize();
323+
if (exceeds_stack_limit(map_key_size)) {
324+
resources_.map_key_buffers++;
325+
resources_.max_map_key_size = std::max(resources_.max_map_key_size,
326+
map_key_size);
327+
}
328+
} else if (call.func == "has_key") {
329+
auto &arg0 = *call.vargs.at(0);
330+
auto &map = static_cast<Map &>(arg0);
331+
// has_key does not work on scalar maps (e.g. @a = 1), so we
332+
// don't need to check if map.key_expr is set
333+
if (needMapKeyAllocation(map, call.vargs.at(1)) &&
334+
exceeds_stack_limit(map.key_type.GetSize())) {
335+
resources_.map_key_buffers++;
336+
resources_.max_map_key_size = std::max(resources_.max_map_key_size,
337+
map.key_type.GetSize());
338+
}
339+
} else if (call.func == "delete") {
340+
auto &arg0 = *call.vargs.at(0);
341+
auto &map = static_cast<Map &>(arg0);
342+
const auto deleteNeedMapKeyAllocation =
343+
call.vargs.size() > 1 ? needMapKeyAllocation(map, call.vargs.at(1))
344+
: needMapKeyAllocation(map);
345+
// delete always expects a map and key, so we don't need to check if
346+
// map.key_expr is set
347+
if (deleteNeedMapKeyAllocation &&
348+
exceeds_stack_limit(map.key_type.GetSize())) {
349+
resources_.map_key_buffers++;
350+
resources_.max_map_key_size = std::max(resources_.max_map_key_size,
351+
map.key_type.GetSize());
352+
}
353+
}
354+
292355
if (uses_usym_table(call.func)) {
293356
// mark probe as using usym, so that the symbol table can be pre-loaded
294357
// and symbols resolved even when unavailable at resolution time
@@ -307,6 +370,7 @@ void ResourceAnalyser::visit(Map &map)
307370
resources_.max_read_map_value_size = std::max(
308371
resources_.max_read_map_value_size, map.type.GetSize());
309372
}
373+
maybe_allocate_map_key_buffer(map);
310374
}
311375

312376
void ResourceAnalyser::visit(Tuple &tuple)
@@ -365,6 +429,7 @@ void ResourceAnalyser::visit(AssignMapStatement &assignment)
365429
resources_.max_write_map_value_size, assignment.map->type.GetSize());
366430
}
367431
}
432+
maybe_allocate_map_key_buffer(*assignment.map);
368433
}
369434

370435
void ResourceAnalyser::visit(Ternary &ternary)
@@ -429,6 +494,17 @@ void ResourceAnalyser::update_map_info(Map &map)
429494
map_info.key_type = map.key_type;
430495
}
431496

497+
void ResourceAnalyser::maybe_allocate_map_key_buffer(const Map &map)
498+
{
499+
const auto map_key_size = map.key_expr ? map.key_type.GetSize()
500+
: CreateUInt64().GetSize();
501+
if (needMapKeyAllocation(map) && exceeds_stack_limit(map_key_size)) {
502+
resources_.map_key_buffers++;
503+
resources_.max_map_key_size = std::max(resources_.max_map_key_size,
504+
map_key_size);
505+
}
506+
}
507+
432508
Pass CreateResourcePass()
433509
{
434510
auto fn = [](Node &n, PassContext &ctx) {

‎src/ast/passes/resource_analyser.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ class ResourceAnalyser : public Visitor {
4646

4747
bool exceeds_stack_limit(size_t size);
4848

49+
void maybe_allocate_map_key_buffer(const Map &map);
50+
4951
void update_map_info(Map &map);
5052
void update_variable_info(Variable &var);
5153

0 commit comments

Comments
 (0)