Add builtin numaid to get NUMA Node ID

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to
 ## Unreleased
 
 #### Added
+- Add builtin: `numaid`
+  - [#2177](https://github.com/iovisor/bpftrace/pull/2177)
+
 #### Changed
 #### Deprecated
 #### Removed

docs/reference_guide.md

@@ -1831,6 +1831,7 @@ NetworkManager:1155 /var/lib/sss/mc/passwd (deleted)
 - `gid` - Group ID
 - `nsecs` - Nanosecond timestamp
 - `elapsed` - Nanoseconds since bpftrace initialization
+- `numaid` - NUMA Node ID
 - `cpu` - Processor ID
 - `comm` - Process name
 - `kstack` - Kernel stack trace

man/adoc/bpftrace.adoc

@@ -936,6 +936,12 @@ For string arguments use the `str()` call to retrieve the value.
 | n/a
 | PID of the child process
 
+| numaid
+| uint32
+| 5.8
+| numa_node_id
+| ID of the NUMA node executing the BPF program
+
 | cpu
 | uint32
 | 4.1

src/ast/irbuilderbpf.cpp

@@ -906,6 +906,19 @@ CallInst *IRBuilderBPF::CreateGetUidGid()
   return createCall(getuidgid_func, {}, "get_uid_gid");
 }
 
+CallInst *IRBuilderBPF::CreateGetNumaId()
+{
+  // long bpf_get_numa_node_id(void)
+  // Return: NUMA Node ID
+  FunctionType *numaid_func_type = FunctionType::get(getInt32Ty(), false);
+  PointerType *numaid_func_ptr_type = PointerType::get(numaid_func_type, 0);
+  Constant *numaid_func = ConstantExpr::getCast(
+      Instruction::IntToPtr,
+      getInt64(libbpf::BPF_FUNC_get_numa_node_id),
+      numaid_func_ptr_type);
+  return createCall(numaid_func, {}, "get_numa_id");
+}
+
 CallInst *IRBuilderBPF::CreateGetCpuId()
 {
   // u32 bpf_raw_smp_processor_id(void)

src/ast/irbuilderbpf.h

@@ -120,6 +120,7 @@ class IRBuilderBPF : public IRBuilder<>
   CallInst   *CreateGetPidTgid();
   CallInst   *CreateGetCurrentCgroupId();
   CallInst   *CreateGetUidGid();
+  CallInst   *CreateGetNumaId();
   CallInst   *CreateGetCpuId();
   CallInst   *CreateGetCurrentTask();
   CallInst   *CreateGetRandom();

src/ast/passes/codegen_llvm.cpp

@@ -172,6 +172,11 @@ void CodegenLLVM::visit(Builtin &builtin)
       expr_ = b_.CreateLShr(uidgid, 32);
     }
   }
+  else if (builtin.ident == "numaid")
+  {
+    Value *tmp = b_.CreateGetNumaId();
+    expr_ = b_.CreateZExt(tmp, b_.getInt64Ty());
+  }
   else if (builtin.ident == "cpu")
   {
     expr_ = b_.CreateGetCpuId();

src/ast/passes/semantic_analyser.cpp

@@ -285,7 +285,7 @@ void SemanticAnalyser::visit(Builtin &builtin)
            builtin.ident == "pid" || builtin.ident == "tid" ||
            builtin.ident == "cgroup" || builtin.ident == "uid" ||
            builtin.ident == "gid" || builtin.ident == "cpu" ||
-           builtin.ident == "rand")
+           builtin.ident == "rand" || builtin.ident == "numaid")
   {
     builtin.type = CreateUInt64();
     if (builtin.ident == "cgroup" &&

src/lexer.l

@@ -37,7 +37,7 @@ hspace   [ \t]
 vspace   [\n\r]
 space    {hspace}|{vspace}
 path     :(\\.|[_\-\./a-zA-Z0-9#\*])+
-builtin  arg[0-9]|args|cgroup|comm|cpid|cpu|ctx|curtask|elapsed|func|gid|nsecs|pid|probe|rand|retval|sarg[0-9]|tid|uid|username
+builtin  arg[0-9]|args|cgroup|comm|cpid|numaid|cpu|ctx|curtask|elapsed|func|gid|nsecs|pid|probe|rand|retval|sarg[0-9]|tid|uid|username
 call     avg|buf|cat|cgroupid|clear|count|delete|exit|hist|join|kaddr|kptr|ksym|lhist|macaddr|max|min|ntop|override|print|printf|cgroup_path|reg|signal|sizeof|stats|str|strftime|strncmp|sum|system|time|uaddr|uptr|usym|zero|path|unwatch|bswap
 
 /* Don't add to this! Use builtin OR call not both */

tests/codegen/builtin_numaid.cpp

@@ -0,0 +1,16 @@
+#include "common.h"
+
+namespace bpftrace {
+namespace test {
+namespace codegen {
+
+TEST(codegen, builtin_numaid)
+{
+  test("kprobe:f { @x = numaid }",
+
+       NAME);
+}
+
+} // namespace codegen
+} // namespace test
+} // namespace bpftrace

tests/codegen/llvm/builtin_numaid.ll

@@ -0,0 +1,37 @@
+; ModuleID = 'bpftrace'
+source_filename = "bpftrace"
+target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"
+target triple = "bpf-pc-linux"
+
+; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64 %0, i64 %1) #0
+
+define i64 @"kprobe:f"(i8* %0) section "s_kprobe:f_1" {
+entry:
+  %"@x_val" = alloca i64, align 8
+  %"@x_key" = alloca i64, align 8
+  %get_numa_id = call i32 inttoptr (i64 42 to i32 ()*)()
+  %1 = zext i32 %get_numa_id to i64
+  %2 = bitcast i64* %"@x_key" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %2)
+  store i64 0, i64* %"@x_key", align 8
+  %3 = bitcast i64* %"@x_val" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %3)
+  store i64 %1, i64* %"@x_val", align 8
+  %pseudo = call i64 @llvm.bpf.pseudo(i64 1, i64 0)
+  %update_elem = call i64 inttoptr (i64 2 to i64 (i64, i64*, i64*, i64)*)(i64 %pseudo, i64* %"@x_key", i64* %"@x_val", i64 0)
+  %4 = bitcast i64* %"@x_val" to i8*
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %4)
+  %5 = bitcast i64* %"@x_key" to i8*
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %5)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg %0, i8* nocapture %1) #1
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.end.p0i8(i64 immarg %0, i8* nocapture %1) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nofree nosync nounwind willreturn }

tests/parser.cpp

@@ -53,6 +53,7 @@ TEST(Parser, builtin_variables)
   test("kprobe:f { gid }", "Program\n kprobe:f\n  builtin: gid\n");
   test("kprobe:f { nsecs }", "Program\n kprobe:f\n  builtin: nsecs\n");
   test("kprobe:f { elapsed }", "Program\n kprobe:f\n  builtin: elapsed\n");
+  test("kprobe:f { numaid }", "Program\n kprobe:f\n  builtin: numaid\n");
   test("kprobe:f { cpu }", "Program\n kprobe:f\n  builtin: cpu\n");
   test("kprobe:f { curtask }", "Program\n kprobe:f\n  builtin: curtask\n");
   test("kprobe:f { rand }", "Program\n kprobe:f\n  builtin: rand\n");

tests/runtime/builtin

@@ -28,6 +28,11 @@ PROG i:ms:1 { printf("SUCCESS %d\n", elapsed); exit(); }
 EXPECT SUCCESS -?[0-9][0-9]*
 TIMEOUT 5
 
+NAME numaid
+PROG i:ms:1 { printf("SUCCESS %d\n", numaid); exit(); }
+EXPECT SUCCESS -?[0-9][0-9]*
+TIMEOUT 5
+
 NAME cpu
 PROG i:ms:1 { printf("SUCCESS %d\n", cpu); exit(); }
 EXPECT SUCCESS -?[0-9][0-9]*

tests/semantic_analyser.cpp

@@ -146,6 +146,7 @@ TEST(semantic_analyser, builtin_variables)
   test("kprobe:f { gid }", 0);
   test("kprobe:f { nsecs }", 0);
   test("kprobe:f { elapsed }", 0);
+  test("kprobe:f { numaid }", 0);
   test("kprobe:f { cpu }", 0);
   test("kprobe:f { curtask }", 0);
   test("kprobe:f { rand }", 0);