llvm-jit: patch lddw helpers at compile time

vm/llvm-jit/include/compat_llvm.hpp

@@ -41,11 +41,11 @@ class bpftime_llvm_jit_vm : public bpftime::vm::compat::bpftime_vm_impl {
 	}
 
     private:
-	uint64_t (*map_by_fd)(uint32_t) = nullptr;
-	uint64_t (*map_by_idx)(uint32_t) = nullptr;
-	uint64_t (*map_val)(uint64_t) = nullptr;
-	uint64_t (*var_addr)(uint32_t) = nullptr;
-	uint64_t (*code_addr)(uint32_t) = nullptr;
+	uint64_t (*map_by_fd)(uint32_t);
+	uint64_t (*map_by_idx)(uint32_t);
+	uint64_t (*map_val)(uint64_t);
+	uint64_t (*var_addr)(uint32_t);
+	uint64_t (*code_addr)(uint32_t);
 	std::vector<ebpf_inst> instructions;
 	std::vector<std::optional<external_function> > ext_funcs;
 

vm/llvm-jit/src/compat/compat_llvm.cpp

@@ -1,10 +1,39 @@
 #include "spdlog/spdlog.h"
 #include <bpftime_vm_compat.hpp>
 #include <cerrno>
+#include <cstdint>
 #include <memory>
 #include <ebpf_inst.h>
 #include "compat_llvm.hpp"
 #include "../llvm/llvm_jit_context.hpp"
+
+static uint64_t default_map_by_fd(uint32_t)
+{
+	SPDLOG_WARN("map_by_fd not defined, but called");
+	return 0;
+}
+
+static uint64_t default_map_by_idx(uint32_t)
+{
+	SPDLOG_WARN("map_by_idx not defined, but called");
+	return 0;
+}
+
+static uint64_t default_map_val(uint64_t)
+{
+	SPDLOG_WARN("map_val not defined, but called");
+	return 0;
+}
+static uint64_t default_var_addr(uint32_t)
+{
+	SPDLOG_WARN("var_addr not defined, but called");
+	return 0;
+}
+static uint64_t default_code_addr(uint32_t)
+{
+	SPDLOG_WARN("code_addr not defined, but called");
+	return 0;
+}
 namespace bpftime::vm::compat
 {
 
@@ -21,6 +50,11 @@ bpftime_llvm_jit_vm::bpftime_llvm_jit_vm() : ext_funcs(MAX_EXT_FUNCS)
 
 {
 	this->jit_ctx = std::make_unique<llvm_bpf_jit_context>(this);
+	map_by_fd = &default_map_by_fd;
+	map_by_idx = &default_map_by_idx;
+	map_val = &default_map_val;
+	code_addr = &default_code_addr;
+	var_addr = &default_var_addr;
 }
 
 std::string bpftime_llvm_jit_vm::get_error_message()
@@ -98,11 +132,16 @@ void bpftime_llvm_jit_vm::set_lddw_helpers(uint64_t (*map_by_fd)(uint32_t),
 					   uint64_t (*var_addr)(uint32_t),
 					   uint64_t (*code_addr)(uint32_t))
 {
-	this->map_by_fd = map_by_fd;
-	this->map_by_idx = map_by_idx;
-	this->map_val = map_val;
-	this->var_addr = var_addr;
-	this->code_addr = code_addr;
+	if (map_by_fd)
+		this->map_by_fd = map_by_fd;
+	if (map_by_idx)
+		this->map_by_idx = map_by_idx;
+	if (map_val)
+		this->map_val = map_val;
+	if (var_addr)
+		this->var_addr = var_addr;
+	if (code_addr)
+		this->code_addr = code_addr;
 }
 
 std::vector<uint8_t> bpftime_llvm_jit_vm::do_aot_compile(bool print_ir)

vm/llvm-jit/src/llvm/compiler.cpp

@@ -86,8 +86,11 @@ const size_t MAX_LOCAL_FUNC_DEPTH = 32;
 */
 Expected<ThreadSafeModule> llvm_bpf_jit_context::generateModule(
 	const std::vector<std::string> &extFuncNames,
-	const std::vector<std::string> &lddwHelpers)
+	const std::vector<std::string> &lddwHelpers,
+	bool patch_map_val_at_compile_time)
 {
+	SPDLOG_DEBUG("Generating module: patch_map_val_at_compile_time={}",
+		     patch_map_val_at_compile_time);
 	auto context = std::make_unique<LLVMContext>();
 	auto jitModule = std::make_unique<Module>("bpf-jit", *context);
 	const auto &insts = vm->instructions;
@@ -155,17 +158,15 @@ Expected<ThreadSafeModule> llvm_bpf_jit_context::generateModule(
 	// The main function
 	Function *bpf_func = Function::Create(
 		FunctionType::get(Type::getInt64Ty(*context),
-				  { llvm::PointerType::getUnqual(llvm::Type::getInt8Ty(*context)),
+				  { llvm::PointerType::getUnqual(
+					    llvm::Type::getInt8Ty(*context)),
 				    Type::getInt64Ty(*context) },
 				  false),
 		Function::ExternalLinkage, "bpf_main", jitModule.get());
-
-	// Get args of uint64_t bpf_main(uint64_t, uint64_t)
-	llvm::Argument* mem = bpf_func->getArg(0);
-	llvm::Argument* mem_len = bpf_func->getArg(1);
-
-
 
+	// Get args of uint64_t bpf_main(uint64_t, uint64_t)
+	llvm::Argument *mem = bpf_func->getArg(0);
+	llvm::Argument *mem_len = bpf_func->getArg(1);
 
 	std::vector<Value *> regs;
 	std::vector<BasicBlock *> allBlocks;
@@ -631,173 +632,124 @@ Expected<ThreadSafeModule> llvm_bpf_jit_context::generateModule(
 				(((uint64_t)((uint32_t)nextInst.imm)) << 32);
 			pc++;
 
-			SPDLOG_TRACE("Load LDDW val= {} part1={:x} part2={:x}",
+			SPDLOG_DEBUG("Load LDDW val= {} part1={:x} part2={:x}",
 				     val, (uint64_t)inst.imm,
 				     (uint64_t)nextInst.imm);
 			if (inst.src_reg == 0) {
 				SPDLOG_DEBUG("Emit lddw helper 0 at pc {}", pc);
 				builder.CreateStore(builder.getInt64(val),
 						    regs[inst.dst_reg]);
 			} else if (inst.src_reg == 1) {
-				if (auto itr = lddwHelper.find(
-					    LDDW_HELPER_MAP_BY_FD);
-				    itr != lddwHelper.end())
-
-				{
+				SPDLOG_DEBUG(
+					"Emit lddw helper 1 (map_by_fd) at pc {}, imm={}, patched at compile time",
+					pc, inst.imm);
+				builder.CreateStore(
+					builder.getInt64(
+						vm->map_by_fd(inst.imm)),
+					regs[inst.dst_reg]);
+			} else if (inst.src_reg == 2) {
+				SPDLOG_DEBUG(
+					"Emit lddw helper 2 (map_by_fd + map_val) at pc {}, imm1={}, imm2={}",
+					pc, inst.imm, nextInst.imm);
+				auto mapPtr = vm->map_by_fd(inst.imm);
+				if (patch_map_val_at_compile_time) {
 					SPDLOG_DEBUG(
-						"Emit lddw helper 1 (map_by_fd) at pc {}, imm={}",
-						pc, inst.imm);
+						"map_val is required to be evaluated at compile time");
 					builder.CreateStore(
-						builder.CreateCall(
-							lddwHelperWithUint32,
-							itr->second,
-							{ builder.getInt32(
-								inst.imm) }),
+						builder.getInt64(
+							vm->map_val(mapPtr) +
+							nextInst.imm),
 						regs[inst.dst_reg]);
 				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 1, which requires map_by_fd",
-						llvm::inconvertibleErrorCode());
-				}
-			} else if (inst.src_reg == 2) {
-				if (auto itrMapByFd = lddwHelper.find(
-					    LDDW_HELPER_MAP_BY_FD);
-				    itrMapByFd != lddwHelper.end()) {
+					SPDLOG_DEBUG(
+						"map_val is required to be evaluated at runtime, emitting calling instructions");
+
 					if (auto itrMapVal = lddwHelper.find(
 						    LDDW_HELPER_MAP_VAL);
 					    itrMapVal != lddwHelper.end()) {
-						auto retMapByFd = builder.CreateCall(
-							lddwHelperWithUint32,
-							itrMapByFd->second,
-							{ builder.getInt32(
-								inst.imm) });
 						auto retMapVal = builder.CreateCall(
 							lddwHelperWithUint64,
 							itrMapVal->second,
-							{ retMapByFd });
+							{ builder.getInt64(
+								mapPtr) });
 						auto finalRet = builder.CreateAdd(
 							retMapVal,
 							builder.getInt64(
 								nextInst.imm));
 						builder.CreateStore(
 							finalRet,
 							regs[inst.dst_reg]);
-						SPDLOG_DEBUG(
-							"Emit lddw helper 2 (map_by_fd + map_val) at pc {}, imm1={}, imm2={}",
-							pc, inst.imm,
-							nextInst.imm);
+
 					} else {
 						return llvm::make_error<
 							llvm::StringError>(
-							"Using lddw helper 2, which requires map_val",
+							"Using lddw helper 2, which requires map_val to be defined. This should not happen",
 							llvm::inconvertibleErrorCode());
 					}
-
-				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 2, which requires map_by_fd",
-						llvm::inconvertibleErrorCode());
 				}
+
 			} else if (inst.src_reg == 3) {
-				if (auto itr = lddwHelper.find(
-					    LDDW_HELPER_VAR_ADDR);
-				    itr != lddwHelper.end()) {
-					builder.CreateStore(
-						builder.CreateCall(
-							lddwHelperWithUint32,
-							itr->second,
-							{ builder.getInt32(
-								inst.imm) }),
-						regs[inst.dst_reg]);
-					SPDLOG_DEBUG(
-						"Emit lddw helper 3 (var_addr) at pc {}, imm1={}",
-						pc, inst.imm);
-				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 3, which requires var_addr",
-						llvm::inconvertibleErrorCode());
-				}
+				SPDLOG_DEBUG(
+					"Emit lddw helper 3 (var_addr) at pc {}, imm1={}",
+					pc, inst.imm);
+				builder.CreateStore(
+					builder.getInt64(
+						vm->var_addr(inst.imm)),
+					regs[inst.dst_reg]);
 			} else if (inst.src_reg == 4) {
-				if (auto itr = lddwHelper.find(
-					    LDDW_HELPER_CODE_ADDR);
-				    itr != lddwHelper.end()) {
-					builder.CreateStore(
-						builder.CreateCall(
-							lddwHelperWithUint32,
-							itr->second,
-							{ builder.getInt32(
-								inst.imm) }),
-						regs[inst.dst_reg]);
-					SPDLOG_DEBUG(
-						"Emit lddw helper 4 (code_addr) at pc {}, imm1={}",
-						pc, inst.imm);
-				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 4, which requires code_addr",
-						llvm::inconvertibleErrorCode());
-				}
+				SPDLOG_DEBUG(
+					"Emit lddw helper 4 (code_addr) at pc {}, imm1={}",
+					pc, inst.imm);
+				builder.CreateStore(
+					builder.getInt64(
+						vm->code_addr(inst.imm)),
+					regs[inst.dst_reg]);
 			} else if (inst.src_reg == 5) {
-				if (auto itr = lddwHelper.find(
-					    LDDW_HELPER_MAP_BY_IDX);
-				    itr != lddwHelper.end()) {
+				SPDLOG_DEBUG(
+					"Emit lddw helper 4 (map_by_idx) at pc {}, imm1={}",
+					pc, inst.imm);
+				builder.CreateStore(
+					builder.getInt64(
+						vm->map_by_idx(inst.imm)),
+					regs[inst.dst_reg]);
+			} else if (inst.src_reg == 6) {
+				auto mapPtr = vm->map_by_idx(inst.imm);
+				SPDLOG_DEBUG(
+					"Emit lddw helper 6 (map_by_idx + map_val) at pc {}, imm1={}, imm2={}",
+					pc, inst.imm, nextInst.imm);
+				if (patch_map_val_at_compile_time) {
+					SPDLOG_DEBUG(
+						"Required to evaluate map_val at compile time");
 					builder.CreateStore(
-						builder.CreateCall(
-							lddwHelperWithUint32,
-							itr->second,
-							{ builder.getInt32(
-								inst.imm) }),
+						builder.getInt64(
+							vm->map_val(mapPtr) +
+							nextInst.imm),
 						regs[inst.dst_reg]);
-					SPDLOG_DEBUG(
-						"Emit lddw helper 4 (map_by_idx) at pc {}, imm1={}",
-						pc, inst.imm);
 				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 5, which requires map_by_idx",
-						llvm::inconvertibleErrorCode());
-				}
-			} else if (inst.src_reg == 6) {
-				if (auto itrMapByIdx = lddwHelper.find(
-					    LDDW_HELPER_MAP_BY_IDX);
-				    itrMapByIdx != lddwHelper.end()) {
+					SPDLOG_DEBUG(
+						"Required to evaluate map_val at runtime time");
 					if (auto itrMapVal = lddwHelper.find(
 						    LDDW_HELPER_MAP_VAL);
 					    itrMapVal != lddwHelper.end()) {
-						auto retMapByIdx = builder.CreateCall(
-							lddwHelperWithUint32,
-							itrMapByIdx->second,
-							{ builder.getInt32(
-								inst.imm) });
 						auto retMapVal = builder.CreateCall(
 							lddwHelperWithUint64,
 							itrMapVal->second,
-							{ retMapByIdx });
+							{ builder.getInt64(
+								mapPtr) });
 						auto finalRet = builder.CreateAdd(
 							retMapVal,
 							builder.getInt64(
 								nextInst.imm));
 						builder.CreateStore(
 							finalRet,
 							regs[inst.dst_reg]);
-						SPDLOG_DEBUG(
-							"Emit lddw helper 6 (map_by_idx + map_val) at pc {}, imm1={}, imm2={}",
-							pc, inst.imm,
-							nextInst.imm);
+
 					} else {
 						return llvm::make_error<
 							llvm::StringError>(
 							"Using lddw helper 6, which requires map_val",
 							llvm::inconvertibleErrorCode());
 					}
-				} else {
-					return llvm::make_error<
-						llvm::StringError>(
-						"Using lddw helper 6, which requires map_by_idx",
-						llvm::inconvertibleErrorCode());
 				}
 			}
 			break;