Fix atomic operation alignment causing uprobe failures

.github/workflows/bpf_conformance.yml

@@ -10,14 +10,14 @@ concurrency:
   cancel-in-progress: true
 jobs:
   build:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
     - uses: actions/checkout@v2
       with:
         submodules: 'recursive'
 
     - name: Install lcov (Ubuntu)
-      run: sudo apt install -y lcov llvm-15-dev libzstd-dev libboost-all-dev
+      run: sudo apt-get update && sudo apt-get install -y lcov llvm-15-dev libzstd-dev libboost-all-dev
 
     - name: build
       run: 

.github/workflows/unit-test.yml

@@ -7,14 +7,14 @@ on:
     branches: "main"
 jobs:
   build:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
     - uses: actions/checkout@v2
       with:
         submodules: 'recursive'
 
     - name: Install Ninja and lcov (Ubuntu)
-      run: sudo apt install -y ninja-build lcov llvm-15-dev libzstd-dev
+      run: sudo apt-get update && sudo apt-get install -y ninja-build lcov llvm-15-dev libzstd-dev
 
     - name: build
       run:

.gitignore

@@ -23,4 +23,7 @@ dump*.txt
 /out.ptx
 /test.*
 /bpf_program.spv
-/bpf_program.spvasm
+/bpf_program.spvasm
+__pycache__/
+_codeql_build_dir/
+_codeql_detected_source_root

src/compiler.cpp

@@ -1331,7 +1331,7 @@ this conversion.
 						is64 ? builder.getInt64Ty() :
 						       builder.getInt32Ty(),
 						regs[inst.src]),
-					MaybeAlign(0),
+					MaybeAlign(is64 ? 8 : 4),
 					AtomicOrdering::Monotonic,
 					AtomicOrdering::Monotonic);
 				builder.CreateStore(

src/compiler_utils.cpp

@@ -359,9 +359,18 @@ void emitAtomicBinOp(llvm::IRBuilder<> &builder, llvm::Value **regs,
 			       builder.CreateLoad(builder.getInt64Ty(),
 						  regs[inst.src]),
 			       builder.getInt32Ty()),
-		llvm::MaybeAlign(32), llvm::AtomicOrdering::Monotonic);
+		llvm::MaybeAlign(is64 ? 8 : 4), llvm::AtomicOrdering::Monotonic);
 	if (is_fetch) {
-		builder.CreateStore(oldValue, regs[inst.src]);
+		if (is64) {
+			builder.CreateStore(oldValue, regs[inst.src]);
+		} else {
+			// For 32-bit atomics, zero-extend the fetched value to 64 bits
+			// before storing into the 64-bit register slot to preserve
+			// eBPF register semantics.
+			auto zextOldValue =
+				builder.CreateZExt(oldValue, builder.getInt64Ty());
+			builder.CreateStore(zextOldValue, regs[inst.src]);
+		}
 	}
 }
 

test/unit-test/bpf_prog_test.cpp

@@ -22,3 +22,27 @@ TEST_CASE("Test aot compilation for extenal function")
 	REQUIRE(vm.exec(&mem, sizeof(mem), ret) == 0);
 	REQUIRE(ret == 4);
 }
+
+TEST_CASE("Test atomic add operation 64-bit")
+{
+	bpftime::llvmbpf_vm vm;
+	REQUIRE(vm.load_code((const void *)bpf_atomic_add_64,
+			     sizeof(bpf_atomic_add_64) - 1) == 0);
+	uint64_t ret = 0;
+	uint64_t mem = 0;
+
+	REQUIRE(vm.exec(&mem, sizeof(mem), ret) == 0);
+	REQUIRE(ret == 1);  // counter should be 1 after atomic add
+}
+
+TEST_CASE("Test atomic add with fetch 32-bit")
+{
+	bpftime::llvmbpf_vm vm;
+	REQUIRE(vm.load_code((const void *)bpf_atomic_add_fetch_32,
+			     sizeof(bpf_atomic_add_fetch_32) - 1) == 0);
+	uint64_t ret = 0;
+	uint64_t mem = 0;
+
+	REQUIRE(vm.exec(&mem, sizeof(mem), ret) == 0);
+	REQUIRE(ret == 0);  // should return old value (0) before add
+}

test/unit-test/bpf_progs.h

@@ -436,4 +436,54 @@ uint64_t ffi_print_integer(uint64_t a, uint64_t b, uint64_t _c, uint64_t _d,
 	return 0;
 }
 
+/*
+Atomic add test (uses 64-bit atomic instruction on 64-bit memory)
+Note: The instruction opcode 0xdb corresponds to a 64-bit (DW) atomic operation
+operating on 8-byte memory. This test validates correct handling of 64-bit
+atomic operations and alignment.
+
+uint64_t test_atomic_add() {
+    uint64_t counter = 0;  // 64-bit integer stored in 8 bytes
+    __sync_fetch_and_add(&counter, 1);  // uses a 64-bit atomic instruction
+    return counter;
+}
+Bytecode:
+0: mov r1, 0x0
+1: stxdw [r10-8], r1           // store 64-bit double word at 8-byte aligned slot
+2: mov r1, 0x1
+3: mov r2, r10
+4: add r2, -8                  // r2 = &counter (64-bit, aligned)
+5: atomic_add64 [r2], r1       // opcode 0xdb = 64-bit atomic add instruction
+6: ldxdw r0, [r10-8]           // load 64-bit double word
+7: exit
+*/
+const unsigned char bpf_atomic_add_64[] =
+	"\xb7\x01\x00\x00\x00\x00\x00\x00"  // mov r1, 0x0
+	"\x7b\x1a\xf8\xff\x00\x00\x00\x00"  // stxdw [r10-8], r1
+	"\xb7\x01\x00\x00\x01\x00\x00\x00"  // mov r1, 0x1
+	"\xbf\xa2\x00\x00\x00\x00\x00\x00"  // mov r2, r10
+	"\x07\x02\x00\x00\xf8\xff\xff\xff"  // add r2, -8
+	"\xdb\x12\x00\x00\x00\x00\x00\x00"  // atomic_add64 [r2], r1
+	"\x79\xa0\xf8\xff\x00\x00\x00\x00"  // ldxdw r0, [r10-8]
+	"\x95\x00\x00\x00\x00\x00\x00\x00"; // exit
+
+/*
+Atomic add 32-bit test with fetch (returns old value)
+int test_atomic_add_fetch() {
+    int counter = 0;
+    int old = __sync_fetch_and_add(&counter, 5);
+    return old;
+}
+Bytecode uses opcode 0xc3 = 32-bit atomic add with fetch (EBPF_ATOMIC_ADD | EBPF_ATOMIC_OP_FETCH)
+*/
+const unsigned char bpf_atomic_add_fetch_32[] =
+	"\xb7\x01\x00\x00\x00\x00\x00\x00"  // mov r1, 0x0
+	"\x63\x1a\xfc\xff\x00\x00\x00\x00"  // stxw [r10-4], r1
+	"\xb7\x01\x00\x00\x05\x00\x00\x00"  // mov r1, 0x5
+	"\xbf\xa2\x00\x00\x00\x00\x00\x00"  // mov r2, r10
+	"\x07\x02\x00\x00\xfc\xff\xff\xff"  // add r2, -4
+	"\xc3\x12\x00\x00\x01\x00\x00\x00"  // atomic_add32_fetch [r2], r1 (32-bit with fetch, returns old value in r1)
+	"\xbf\x10\x00\x00\x00\x00\x00\x00"  // mov r0, r1
+	"\x95\x00\x00\x00\x00\x00\x00\x00"; // exit
+
 #endif