From 199fb3017bec7a497ccceacecd21c5765c27b5c0 Mon Sep 17 00:00:00 2001 From: Alex Malyshev Date: Fri, 3 May 2024 10:00:09 -0700 Subject: [PATCH] Better track code buffers for compiled functions Summary: After NativeGenerator compiles code, it stores pointers to the function's entry points and the total size of the code, but it doesn't have a pointer to the start of the code. We need this for ELF writing because the goal is to serialize the entire code buffer, and not just the code that starts at the normal entry point. NativeGenerator and CompiledFunction will now expose the code buffer as a span of (const) bytes. Differential Revision: D56897617 fbshipit-source-id: 6d6668b7b0b737405c8fff85ff31525ce0614174 --- cinderx/Jit/codegen/gen_asm.cpp | 24 ++++++++--------- cinderx/Jit/codegen/gen_asm.h | 22 ++++++++++++++-- cinderx/Jit/compiler.cpp | 29 ++++++++++----------- cinderx/Jit/compiler.h | 24 ++++++++++++----- cinderx/Jit/elf.cpp | 2 +- cinderx/Jit/elf.h | 3 ++- cinderx/Jit/pyjit.cpp | 5 +--- cinderx/RuntimeTests/deopt_patcher_test.cpp | 4 +-- cinderx/RuntimeTests/elf_test.cpp | 2 +- 9 files changed, 71 insertions(+), 44 deletions(-) diff --git a/cinderx/Jit/codegen/gen_asm.cpp b/cinderx/Jit/codegen/gen_asm.cpp index c01ac3d8cbd..e200772f6b8 100644 --- a/cinderx/Jit/codegen/gen_asm.cpp +++ b/cinderx/Jit/codegen/gen_asm.cpp @@ -208,6 +208,11 @@ PhyLocation get_arg_location_phy_location(int arg) { return 0; } +std::span NativeGenerator::getCodeBuffer() const { + return std::span{ + reinterpret_cast(code_start_), compiled_size_}; +} + void* NativeGenerator::getVectorcallEntry() { if (vectorcall_entry_ != nullptr) { // already compiled @@ -397,7 +402,7 @@ void* NativeGenerator::getVectorcallEntry() { */ JIT_DCHECK(code.codeSize() < INT_MAX, "Code size is larger than INT_MAX"); - compiled_size_ = static_cast(code.codeSize()); + compiled_size_ = code.codeSize(); env_.code_rt->set_frame_size(env_.stack_frame_size); return vectorcall_entry_; } @@ -413,10 +418,6 @@ void* NativeGenerator::getStaticEntry() { JITRT_STATIC_ENTRY_OFFSET); } -int NativeGenerator::GetCompiledFunctionSize() const { - return compiled_size_; -} - int NativeGenerator::GetCompiledFunctionStackSize() const { return env_.stack_frame_size; } @@ -1392,10 +1393,9 @@ void NativeGenerator::generateCode(CodeHolder& codeholder) { generateDeoptExits(codeholder); ASM_CHECK_THROW(as_->finalize()); - void* code_top; - ASM_CHECK_THROW(CodeAllocator::get()->addCode(&code_top, &codeholder)); + ASM_CHECK_THROW(CodeAllocator::get()->addCode(&code_start_, &codeholder)); - // ------------- code_top + // ------------- code_start_ // ^ // | JITRT_STATIC_ENTRY_OFFSET (2 bytes, optional) // | JITRT_CALL_REENTRY_OFFSET (6 bytes) @@ -1420,7 +1420,7 @@ void NativeGenerator::generateCode(CodeHolder& codeholder) { env_.code_rt->debug_info()->resolvePending( env_.pending_debug_locs, *GetFunction(), codeholder); - vectorcall_entry_ = static_cast(code_top) + + vectorcall_entry_ = static_cast(code_start_) + codeholder.labelOffsetFromBase(vectorcall_entry_label); for (auto& entry : env_.unresolved_gen_entry_labels) { @@ -1437,14 +1437,14 @@ void NativeGenerator::generateCode(CodeHolder& codeholder) { compiled_size_ = codeholder.codeSize(); if (!g_dump_hir_passes_json.empty()) { - env_.annotations.disassembleJSON(*json, code_top, codeholder); + env_.annotations.disassembleJSON(*json, code_start_, codeholder); } JIT_LOGIF( g_dump_asm, "Disassembly for {}\n{}", GetFunction()->fullname, - env_.annotations.disassemble(code_top, codeholder)); + env_.annotations.disassemble(code_start_, codeholder)); { ThreadedCompileSerialize guard; for (auto& x : env_.function_indirections) { @@ -1468,7 +1468,7 @@ void NativeGenerator::generateCode(CodeHolder& codeholder) { // For perf, we want only the size of the code, so we get that directly from // the text sections. std::vector> code_sections; - populateCodeSections(code_sections, codeholder, code_top); + populateCodeSections(code_sections, codeholder, code_start_); perf::registerFunction(code_sections, func->fullname, prefix); } diff --git a/cinderx/Jit/codegen/gen_asm.h b/cinderx/Jit/codegen/gen_asm.h index 38d2e996157..f25a73e55ee 100644 --- a/cinderx/Jit/codegen/gen_asm.h +++ b/cinderx/Jit/codegen/gen_asm.h @@ -16,7 +16,9 @@ #include "cinderx/ThirdParty/asmjit/src/asmjit/asmjit.h" #include +#include #include +#include #include #include #include @@ -68,9 +70,24 @@ class NativeGenerator { } std::string GetFunctionName() const; + + // Get the buffer containing the compiled machine code. The start of this + // buffer is not guaranteed to be a valid entry point. + // + // Note: getVectorcallEntry() **must** be called before this is called. + std::span getCodeBuffer() const; + + // Get the entry point of the compiled function if it is called via a + // vectorcall. + // + // Note: This is where the function is actually compiled, it is done the first + // time this method is called. void* getVectorcallEntry(); + + // Get the entry point of the compiled function if it is called via a Static + // Python call. void* getStaticEntry(); - int GetCompiledFunctionSize() const; + int GetCompiledFunctionStackSize() const; int GetCompiledFunctionSpillStackSize() const; const hir::Function* GetFunction() const { @@ -90,6 +107,7 @@ class NativeGenerator { #endif private: const hir::Function* func_; + void* code_start_{nullptr}; void* vectorcall_entry_{nullptr}; asmjit::x86::Builder* as_{nullptr}; CodeHolderMetadata metadata_{CodeSection::kHot}; @@ -97,7 +115,7 @@ class NativeGenerator { void* deopt_trampoline_generators_{nullptr}; void* const failed_deferred_compile_trampoline_; - int compiled_size_{-1}; + size_t compiled_size_{0}; int spill_stack_size_{-1}; int frame_header_size_; int max_inline_depth_; diff --git a/cinderx/Jit/compiler.cpp b/cinderx/Jit/compiler.cpp index 9d815023b79..267c5ab5148 100644 --- a/cinderx/Jit/compiler.cpp +++ b/cinderx/Jit/compiler.cpp @@ -266,11 +266,11 @@ std::unique_ptr Compiler::Compile( ngen->SetJSONOutput(json.get()); } - void* entry = nullptr; + vectorcallfunc entry = nullptr; COMPILE_TIMER( irfunc->compilation_phase_timer, "Native code Generation", - entry = ngen->getVectorcallEntry()) + entry = reinterpret_cast(ngen->getVectorcallEntry())) if (entry == nullptr) { JIT_DLOG("Generating native code for {} failed", fullname); return nullptr; @@ -282,7 +282,6 @@ std::unique_ptr Compiler::Compile( irfunc->setCompilationPhaseTimer(nullptr); } - int func_size = ngen->GetCompiledFunctionSize(); int stack_size = ngen->GetCompiledFunctionStackSize(); int spill_stack_size = ngen->GetCompiledFunctionSpillStackSize(); @@ -301,6 +300,7 @@ std::unique_ptr Compiler::Compile( // Grab some fields off of irfunc and ngen before moving them. hir::Function::InlineFunctionStats inline_stats = std::move(irfunc->inline_function_stats); + std::span code = ngen->getCodeBuffer(); void* static_entry = ngen->getStaticEntry(); CodeRuntime* code_runtime = ngen->codeRuntime(); @@ -309,25 +309,24 @@ std::unique_ptr Compiler::Compile( return std::make_unique( std::move(irfunc), std::move(ngen), - reinterpret_cast(entry), + code, + entry, static_entry, code_runtime, - func_size, - stack_size, - spill_stack_size, - std::move(inline_stats), - hir_opcode_counts); - } else { - return std::make_unique( - reinterpret_cast(entry), - static_entry, - code_runtime, - func_size, stack_size, spill_stack_size, std::move(inline_stats), hir_opcode_counts); } + return std::make_unique( + code, + entry, + static_entry, + code_runtime, + stack_size, + spill_stack_size, + std::move(inline_stats), + hir_opcode_counts); } } // namespace jit diff --git a/cinderx/Jit/compiler.h b/cinderx/Jit/compiler.h index c794bb2d38a..7518018533d 100644 --- a/cinderx/Jit/compiler.h +++ b/cinderx/Jit/compiler.h @@ -10,6 +10,8 @@ #include "cinderx/Jit/hir/preload.h" #include "cinderx/Jit/runtime.h" +#include +#include #include namespace jit { @@ -22,18 +24,18 @@ namespace jit { class CompiledFunction { public: CompiledFunction( + std::span code, vectorcallfunc vectorcall_entry, void* static_entry, CodeRuntime* code_runtime, - int func_size, int stack_size, int spill_stack_size, hir::Function::InlineFunctionStats inline_function_stats, const hir::OpcodeCounts& hir_opcode_counts) - : vectorcall_entry_(vectorcall_entry), + : code_(code), + vectorcall_entry_(vectorcall_entry), static_entry_(static_entry), code_runtime_(code_runtime), - code_size_(func_size), stack_size_(stack_size), spill_stack_size_(spill_stack_size), inline_function_stats_(std::move(inline_function_stats)), @@ -41,6 +43,12 @@ class CompiledFunction { virtual ~CompiledFunction() {} + // Get the buffer containing the compiled machine code. The start of this + // buffer is not guaranteed to be a valid entry point. + std::span codeBuffer() const { + return code_; + } + vectorcallfunc vectorcallEntry() const { return vectorcall_entry_; } @@ -60,18 +68,22 @@ class CompiledFunction { return code_runtime_; } - int codeSize() const { - return code_size_; + size_t codeSize() const { + return code_.size(); } + int stackSize() const { return stack_size_; } + int spillStackSize() const { return spill_stack_size_; } + const hir::Function::InlineFunctionStats& inlinedFunctionsStats() const { return inline_function_stats_; } + const hir::OpcodeCounts& hirOpcodeCounts() const { return hir_opcode_counts_; } @@ -79,10 +91,10 @@ class CompiledFunction { private: DISALLOW_COPY_AND_ASSIGN(CompiledFunction); + const std::span code_; vectorcallfunc const vectorcall_entry_; void* const static_entry_; CodeRuntime* const code_runtime_; - const int code_size_; const int stack_size_; const int spill_stack_size_; hir::Function::InlineFunctionStats inline_function_stats_; diff --git a/cinderx/Jit/elf.cpp b/cinderx/Jit/elf.cpp index b64f4d9c99c..f5cfd0127b8 100644 --- a/cinderx/Jit/elf.cpp +++ b/cinderx/Jit/elf.cpp @@ -290,7 +290,7 @@ void writeElf( // Write out the actual sections themselves. for (const CodeEntry& entry : entries) { - write(os, entry.code.data(), entry.code.size()); + write(os, entry.code); } pad(os, elf.text_padding); diff --git a/cinderx/Jit/elf.h b/cinderx/Jit/elf.h index e189334a03a..5f2aceeb027 100644 --- a/cinderx/Jit/elf.h +++ b/cinderx/Jit/elf.h @@ -5,6 +5,7 @@ #include "cinderx/Common/log.h" #include "cinderx/Common/util.h" +#include #include #include #include @@ -431,7 +432,7 @@ struct Object { // Code entry to add to an ELF file. struct CodeEntry { - std::span code; + std::span code; std::string func_name; std::string file_name; size_t lineno{0}; diff --git a/cinderx/Jit/pyjit.cpp b/cinderx/Jit/pyjit.cpp index 7661a8394ef..f754dd94ceb 100644 --- a/cinderx/Jit/pyjit.cpp +++ b/cinderx/Jit/pyjit.cpp @@ -1186,10 +1186,7 @@ static PyObject* dump_elf(PyObject* /* self */, PyObject* arg) { CompiledFunction* compiled_func = jit_ctx->lookupFunc(func); elf::CodeEntry entry; - // TODO: What about the staticEntry? - entry.code = { - reinterpret_cast(compiled_func->vectorcallEntry()), - static_cast(compiled_func->codeSize())}; + entry.code = compiled_func->codeBuffer(); entry.func_name = funcFullname(func); if (code->co_filename != nullptr && PyUnicode_Check(code->co_filename)) { entry.file_name = unicodeAsString(code->co_filename); diff --git a/cinderx/RuntimeTests/deopt_patcher_test.cpp b/cinderx/RuntimeTests/deopt_patcher_test.cpp index c13e6864981..0cff530cc52 100644 --- a/cinderx/RuntimeTests/deopt_patcher_test.cpp +++ b/cinderx/RuntimeTests/deopt_patcher_test.cpp @@ -19,14 +19,14 @@ class DeoptPatcherTest : public RuntimeTest { if (entry == nullptr) { return nullptr; } - int func_size = ngen.GetCompiledFunctionSize(); + std::span code = ngen.getCodeBuffer(); int stack_size = ngen.GetCompiledFunctionStackSize(); int spill_stack_size = ngen.GetCompiledFunctionSpillStackSize(); return std::make_unique( + code, reinterpret_cast(entry), ngen.getStaticEntry(), ngen.codeRuntime(), - func_size, stack_size, spill_stack_size, jit::hir::Function::InlineFunctionStats{}, diff --git a/cinderx/RuntimeTests/elf_test.cpp b/cinderx/RuntimeTests/elf_test.cpp index 96724a539ab..23b4c5618d5 100644 --- a/cinderx/RuntimeTests/elf_test.cpp +++ b/cinderx/RuntimeTests/elf_test.cpp @@ -27,7 +27,7 @@ TEST_F(ElfTest, OneEntry) { std::stringstream ss; elf::CodeEntry entry; - entry.code = {reinterpret_cast(elf::writeEntries), 0x40}; + entry.code = {reinterpret_cast(elf::writeEntries), 0x40}; entry.func_name = "funcABC"; entry.file_name = "spaghetti.exe"; entry.lineno = 15;