Use direct inc/dec for ++/-- op

This will save a call to bpf_map_update_elem Issue: bpftrace#3175
jordalgo · May 17, 2024 · 4848530 · 4848530
1 parent 6a1f3f0
commit 4848530
Show file tree

Hide file tree

Showing 9 changed files with 338 additions and 305 deletions.
diff --git a/src/ast/irbuilderbpf.cpp b/src/ast/irbuilderbpf.cpp
@@ -1519,7 +1519,9 @@ void IRBuilderBPF::CreateMapElemAdd(Value *ctx,
                                     Map &map,
                                     Value *key,
                                     Value *val,
-                                    const location &loc)
+                                    const location &loc,
+                                    AllocaInst *pre_post_val,
+                                    bool is_post_op)
 {
   CallInst *call = CreateMapLookup(map, key);
   SizedType &type = map.type;
@@ -1546,18 +1548,34 @@ void IRBuilderBPF::CreateMapElemAdd(Value *ctx,
 
   // createMapLookup  returns an u8*
   auto *cast = CreatePointerCast(call, value->getType(), "cast");
+
+  if (pre_post_val && is_post_op) {
+    CreateStore(CreateLoad(getInt64Ty(), cast), pre_post_val);
+  }
+
   CreateStore(CreateAdd(CreateLoad(getInt64Ty(), cast), val), cast);
 
+  if (pre_post_val && !is_post_op) {
+    CreateStore(CreateLoad(getInt64Ty(), cast), pre_post_val);
+  }
+
   CreateBr(lookup_merge_block);
 
   SetInsertPoint(lookup_failure_block);
 
   CreateMapElemInit(ctx, map, key, val, loc);
 
+  if (pre_post_val) {
+    if (is_post_op) {
+      CreateStore(getInt64(0), pre_post_val);
+    } else {
+      CreateStore(val, pre_post_val);
+    }
+  }
+
   CreateBr(lookup_merge_block);
   SetInsertPoint(lookup_merge_block);
   CreateLifetimeEnd(value);
-  return;
 }
 
 void IRBuilderBPF::CreatePerfEventOutput(Value *ctx,

diff --git a/src/ast/irbuilderbpf.h b/src/ast/irbuilderbpf.h
@@ -222,7 +222,9 @@ class IRBuilderBPF : public IRBuilder<> {
                         Map &map,
                         Value *key,
                         Value *val,
-                        const location &loc);
+                        const location &loc,
+                        AllocaInst *pre_post_val = nullptr,
+                        bool is_post_op = false);
   void CreateDebugOutput(std::string fmt_str,
                          const std::vector<Value *> &values,
                          const location &loc);

diff --git a/src/ast/passes/codegen_llvm.cpp b/src/ast/passes/codegen_llvm.cpp
@@ -3893,34 +3893,31 @@ void CodegenLLVM::createIncDec(Unop &unop)
   bool is_increment = unop.op == Operator::INCREMENT;
   SizedType &type = unop.expr->type;
   uint64_t step = type.IsPtrTy() ? type.GetPointeeTy()->GetSize() : 1;
+  if (!is_increment) {
+    step = -step;
+  }
 
   if (unop.expr->is_map) {
     Map &map = static_cast<Map &>(*unop.expr);
     auto [key, scoped_key_deleter] = getMapKey(map);
-    Value *oldval = b_.CreateMapLookupElem(ctx_, map, key, unop.loc);
-    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_newval");
-    if (is_increment)
-      b_.CreateStore(b_.CreateAdd(oldval, b_.GetIntSameSize(step, oldval)),
-                     newval);
-    else
-      b_.CreateStore(b_.CreateSub(oldval, b_.GetIntSameSize(step, oldval)),
-                     newval);
-    b_.CreateMapUpdateElem(ctx_, map.ident, key, newval, unop.loc);
-
-    if (unop.is_post_op)
-      expr_ = oldval;
-    else
-      expr_ = b_.CreateLoad(b_.GetType(map.type), newval);
-    b_.CreateLifetimeEnd(newval);
+    AllocaInst *pre_post_val = b_.CreateAllocaBPF(map.type,
+                                                  map.ident + "pre_post_val");
+    b_.CreateMapElemAdd(ctx_,
+                        map,
+                        key,
+                        b_.GetIntSameSize(step,
+                                          b_.getInt64(is_increment ? 1 : -1)),
+                        unop.loc,
+                        pre_post_val,
+                        unop.is_post_op);
+    expr_ = b_.CreateLoad(b_.GetType(map.type), pre_post_val);
+    b_.CreateLifetimeEnd(pre_post_val);
   } else if (unop.expr->is_variable) {
     Variable &var = static_cast<Variable &>(*unop.expr);
     Value *oldval = b_.CreateLoad(variables_[var.ident]->getAllocatedType(),
                                   variables_[var.ident]);
     Value *newval;
-    if (is_increment)
-      newval = b_.CreateAdd(oldval, b_.GetIntSameSize(step, oldval));
-    else
-      newval = b_.CreateSub(oldval, b_.GetIntSameSize(step, oldval));
+    newval = b_.CreateAdd(oldval, b_.GetIntSameSize(step, oldval));
     b_.CreateStore(newval, variables_[var.ident]);
 
     if (unop.is_post_op)