Skip to content

Commit

Permalink
Add implicit cast for 'sum' and 'count'
Browse files Browse the repository at this point in the history 
This utilizes 'map_lookup_percpu_elem' to iterate over
the per-cpu maps that 'sum' and 'count' use to aggregate
the values for each cpu so these maps can be used in kernel
space for conditionals and arithmetic.

Additionally, add explicit casting for printing or storing
just the value for these types of maps.

Issue: bpftrace#3126

Note: We can do the same to support to 'avg', 'max', 'min'
but those require the same user-space logic written in llvm ir.
  • Loading branch information
Jordan Rome committed May 21, 2024
1 parent 6d260ed commit f9e22d1
Show file tree
Hide file tree
Showing 6 changed files with 223 additions and 13 deletions.
142 changes: 142 additions & 0 deletions src/ast/irbuilderbpf.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include <iostream>
#include <sstream>
#include <thread>

#include <llvm/IR/DataLayout.h>
#include <llvm/IR/Module.h>
Expand Down Expand Up @@ -327,6 +328,35 @@ CallInst *IRBuilderBPF::createMapLookup(const std::string &map_name,
return createCall(lookup_func_type, lookup_func, { map_ptr, key }, name);
}

CallInst *IRBuilderBPF::createPerCpuMapLookup(const std::string &map_name,
Value *key,
Value *cpu,
const std::string &name)
{
return createPerCpuMapLookup(map_name, key, cpu, GET_PTR_TY(), name);
}

CallInst *IRBuilderBPF::createPerCpuMapLookup(const std::string &map_name,
Value *key,
Value *cpu,
PointerType *val_ptr_ty,
const std::string &name)
{
Value *map_ptr = GetMapVar(map_name);
// void *map_lookup_percpu_elem(struct bpf_map * map, void * key, u32 cpu)
// Return: Map value or NULL

assert(key->getType()->isPointerTy());
FunctionType *lookup_func_type = FunctionType::get(
val_ptr_ty, { map_ptr->getType(), key->getType(), getInt32Ty() }, false);
PointerType *lookup_func_ptr_type = PointerType::get(lookup_func_type, 0);
Constant *lookup_func = ConstantExpr::getCast(
Instruction::IntToPtr,
getInt64(libbpf::BPF_FUNC_map_lookup_percpu_elem),
lookup_func_ptr_type);
return createCall(lookup_func_type, lookup_func, { map_ptr, key, cpu }, name);
}

CallInst *IRBuilderBPF::CreateGetJoinMap(BasicBlock *failure_callback,
const location &loc)
{
Expand Down Expand Up @@ -451,6 +481,118 @@ Value *IRBuilderBPF::CreateMapLookupElem(Value *ctx,
return ret;
}

Value *IRBuilderBPF::CreatePerCpuMapLookupElem(Value *ctx,
const std::string &map_name,
Value *key,
Value *cpu,
SizedType &type,
const location &loc)
{
assert(ctx && ctx->getType() == GET_PTR_TY());
CallInst *call = createPerCpuMapLookup(map_name, key, cpu);

// Check if result == 0
Function *parent = GetInsertBlock()->getParent();
BasicBlock *lookup_success_block = BasicBlock::Create(module_.getContext(),
"lookup_success",
parent);
BasicBlock *lookup_failure_block = BasicBlock::Create(module_.getContext(),
"lookup_failure",
parent);
BasicBlock *lookup_merge_block = BasicBlock::Create(module_.getContext(),
"lookup_merge",
parent);

AllocaInst *value = CreateAllocaBPF(type, "lookup_elem_val");
Value *condition = CreateICmpNE(
CreateIntCast(call, GET_PTR_TY(), true),
ConstantExpr::getCast(Instruction::IntToPtr, getInt64(0), GET_PTR_TY()),
"map_lookup_cond");
CreateCondBr(condition, lookup_success_block, lookup_failure_block);

SetInsertPoint(lookup_success_block);
// per cpu maps used by bpftrace only hold integer types for aggregation
// so we don't need to memcpy
assert(value->getAllocatedType() == getInt64Ty());
// createMapLookup returns an u8*
auto *cast = CreatePointerCast(call, value->getType(), "cast");
CreateStore(CreateLoad(getInt64Ty(), cast), value);
CreateBr(lookup_merge_block);

SetInsertPoint(lookup_failure_block);

CreateStore(getInt64(0), value);
CreateHelperError(
ctx, getInt32(0), libbpf::BPF_FUNC_map_lookup_percpu_elem, loc);
CreateBr(lookup_merge_block);

SetInsertPoint(lookup_merge_block);

// value is a pointer to i64
Value *ret = CreateLoad(getInt64Ty(), value);
CreateLifetimeEnd(value);
return ret;
}

Value *IRBuilderBPF::CreatePerCpuMapSumElems(Value *ctx,
Map &map,
Value *key,
const location &loc)
{
assert(ctx && ctx->getType() == GET_PTR_TY());
return CreatePerCpuMapSumElems(ctx, map.ident, key, map.type, loc);
}
Value *IRBuilderBPF::CreatePerCpuMapSumElems(Value *ctx,
const std::string &map_name,
Value *key,
SizedType &type,
const location &loc)
{
AllocaInst *sum = CreateAllocaBPF(getInt64Ty(), "sum");
AllocaInst *i = CreateAllocaBPF(getInt32Ty(), "i");

const auto processor_count = std::thread::hardware_concurrency();

CreateStore(getInt32(0), i);
CreateStore(getInt64(0), sum);

Function *parent = GetInsertBlock()->getParent();
BasicBlock *while_cond = BasicBlock::Create(module_.getContext(),
"while_cond",
parent);
BasicBlock *while_body = BasicBlock::Create(module_.getContext(),
"while_body",
parent);
BasicBlock *while_end = BasicBlock::Create(module_.getContext(),
"while_end",
parent);
CreateBr(while_cond);
SetInsertPoint(while_cond);
auto *cond = CreateICmp(CmpInst::ICMP_ULT,
CreateLoad(getInt32Ty(), i),
getInt32(processor_count),
"num_cpu.cmp");
CreateCondBr(cond, while_body, while_end);

SetInsertPoint(while_body);

Value *cpu_value = CreatePerCpuMapLookupElem(
ctx, map_name, key, CreateLoad(getInt32Ty(), i), type, loc);
// sum += cpu_value;
CreateStore(CreateAdd(cpu_value, CreateLoad(getInt64Ty(), sum)), sum);

// ++i;
CreateStore(CreateAdd(CreateLoad(getInt32Ty(), i), getInt32(1)), i);

CreateBr(while_cond);
SetInsertPoint(while_end);

CreateLifetimeEnd(i);
Value *ret = CreateLoad(getInt64Ty(), sum);
CreateLifetimeEnd(sum);
return ret;
}

void IRBuilderBPF::CreateMapUpdateElem(Value *ctx,
const std::string &map_ident,
Value *key,
Expand Down
27 changes: 27 additions & 0 deletions src/ast/irbuilderbpf.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,22 @@ class IRBuilderBPF : public IRBuilder<> {
Value *key,
SizedType &type,
const location &loc);

Value *CreatePerCpuMapLookupElem(Value *ctx,
const std::string &map_name,
Value *key,
Value *cpu,
SizedType &type,
const location &loc);
Value *CreatePerCpuMapSumElems(Value *ctx,
Map &map,
Value *key,
const location &loc);
Value *CreatePerCpuMapSumElems(Value *ctx,
const std::string &map_name,
Value *key,
SizedType &type,
const location &loc);
void CreateMapUpdateElem(Value *ctx,
const std::string &map_ident,
Value *key,
Expand Down Expand Up @@ -303,6 +319,17 @@ class IRBuilderBPF : public IRBuilder<> {
Value *key,
PointerType *val_ptr_ty,
const std::string &name = "lookup_elem");
CallInst *createPerCpuMapLookup(
const std::string &map_name,
Value *key,
Value *cpu,
const std::string &name = "lookup_percpu_elem");
CallInst *createPerCpuMapLookup(
const std::string &map_name,
Value *key,
Value *cpu,
PointerType *val_ptr_ty,
const std::string &name = "lookup_percpu_elem");
CallInst *createGetScratchMap(const std::string &map_name,
const std::string &name,
PointerType *val_ptr_ty,
Expand Down
44 changes: 33 additions & 11 deletions src/ast/passes/codegen_llvm.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1413,7 +1413,22 @@ void CodegenLLVM::visit(Offsetof &ofof)
void CodegenLLVM::visit(Map &map)
{
auto [key, scoped_key_deleter] = getMapKey(map);
Value *value = b_.CreateMapLookupElem(ctx_, map, key, map.loc);

auto map_info = bpftrace_.resources.maps_info.find(map.ident);
if (map_info == bpftrace_.resources.maps_info.end()) {
LOG(BUG) << "map name: \"" << map.ident << "\" not found";
}

const auto &val_type = map_info->second.value_type;
auto map_type = getBpfMapType(val_type, map_info->second.key);
Value *value;
if (map_type == libbpf::BPF_MAP_TYPE_PERCPU_ARRAY ||
map_type == libbpf::BPF_MAP_TYPE_PERCPU_HASH) {
value = b_.CreatePerCpuMapSumElems(ctx_, map, key, map.loc);
} else {
value = b_.CreateMapLookupElem(ctx_, map, key, map.loc);
}

expr_ = value;

if (dyn_cast<AllocaInst>(value))
Expand Down Expand Up @@ -3458,6 +3473,22 @@ void CodegenLLVM::generate_ir()
state_ = State::IR;
}

libbpf::bpf_map_type CodegenLLVM::getBpfMapType(const SizedType &val_type,
const MapKey &key)
{
auto map_type = libbpf::BPF_MAP_TYPE_HASH;
if (val_type.IsCountTy() && key.args_.empty()) {
map_type = libbpf::BPF_MAP_TYPE_PERCPU_ARRAY;
} else if (bpftrace_.feature_->has_map_percpu_hash() &&
(val_type.IsHistTy() || val_type.IsLhistTy() ||
val_type.IsCountTy() || val_type.IsSumTy() ||
val_type.IsMinTy() || val_type.IsMaxTy() || val_type.IsAvgTy() ||
val_type.IsStatsTy())) {
map_type = libbpf::BPF_MAP_TYPE_PERCPU_HASH;
}
return map_type;
}

void CodegenLLVM::createMapDefinition(const std::string &name,
libbpf::bpf_map_type map_type,
uint64_t max_entries,
Expand Down Expand Up @@ -3520,18 +3551,9 @@ void CodegenLLVM::generate_maps(const RequiredResources &resources)
const auto &key = info.key;

auto max_entries = bpftrace_.config_.get(ConfigKeyInt::max_map_keys);
auto map_type = libbpf::BPF_MAP_TYPE_UNSPEC;
auto map_type = getBpfMapType(val_type, key);
if (val_type.IsCountTy() && key.args_.empty()) {
max_entries = 1;
map_type = libbpf::BPF_MAP_TYPE_PERCPU_ARRAY;
} else if (bpftrace_.feature_->has_map_percpu_hash() &&
(val_type.IsHistTy() || val_type.IsLhistTy() ||
val_type.IsCountTy() || val_type.IsSumTy() ||
val_type.IsMinTy() || val_type.IsMaxTy() ||
val_type.IsAvgTy() || val_type.IsStatsTy())) {
map_type = libbpf::BPF_MAP_TYPE_PERCPU_HASH;
} else {
map_type = libbpf::BPF_MAP_TYPE_HASH;
}

createMapDefinition(name, map_type, max_entries, key, val_type);
Expand Down
2 changes: 2 additions & 0 deletions src/ast/passes/codegen_llvm.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,8 @@ class CodegenLLVM : public Visitor {
void createPrintMapCall(Call &call);
void createPrintNonMapCall(Call &call, int &id);

libbpf::bpf_map_type getBpfMapType(const SizedType &val_type,
const MapKey &key);
void createMapDefinition(const std::string &name,
libbpf::bpf_map_type map_type,
uint64_t max_entries,
Expand Down
6 changes: 4 additions & 2 deletions src/ast/passes/semantic_analyser.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1790,7 +1790,8 @@ void SemanticAnalyser::visit(Binop &binop)
return;
}

if (lht.IsIntTy() && rht.IsIntTy()) {
if ((lht.IsCountTy() || lht.IsSumTy() || lht.IsIntTy()) &&
(rht.IsCountTy() || rht.IsSumTy() || rht.IsIntTy())) {
binop_int(binop);
} else if (lht.IsArrayTy() && rht.IsArrayTy()) {
binop_array(binop);
Expand Down Expand Up @@ -2324,7 +2325,8 @@ void SemanticAnalyser::visit(Cast &cast)
}

if ((cast.type.IsIntTy() && !rhs.IsIntTy() && !rhs.IsPtrTy() &&
!rhs.IsCtxAccess() && !(rhs.IsArrayTy())) ||
!rhs.IsCtxAccess() && !rhs.IsArrayTy() && !rhs.IsCountTy() &&
!rhs.IsSumTy()) ||
// casting from/to int arrays must respect the size
(cast.type.IsArrayTy() &&
(!rhs.IsIntTy() || cast.type.GetSize() != rhs.GetSize())) ||
Expand Down
15 changes: 15 additions & 0 deletions tests/runtime/other
View file
Original file line number Diff line number Diff line change
Expand Up @@ -279,3 +279,18 @@ EXPECT stdin:1:37-41: WARNING: Can't lookup map element because it does not exis
i:ms:100 { @[1] = 1; printf("%d\n", @[2]); exit(); }
~~~~
TIMEOUT 1

NAME per_cpu_map_if_lt
PROG i:ms:1 { @ = count(); if (@ > 5) { printf("done\n"); exit(); }}
EXPECT done
TIMEOUT 5

NAME per_cpu_map_arithmetic
PROG BEGIN { @a = sum(10); @b = count(); $c = @a + @b; if ($c == 11) { printf("done\n"); } exit();}
EXPECT done
TIMEOUT 5

NAME per_cpu_map_cast
PROG BEGIN { @a = count(); @b = sum(10); printf("%d-%d\n", (uint64)@a, (int64)@b); exit();}
EXPECT 1-10
TIMEOUT 5

0 comments on commit f9e22d1

Please sign in to comment.