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

@@ -208,6 +208,11 @@ PhyLocation get_arg_location_phy_location(int arg) {
   return 0;
 }
 
+std::span<const std::byte> NativeGenerator::getCodeBuffer() const {
+  return std::span{
+      reinterpret_cast<const std::byte*>(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<int>(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<char*>(code_top) +
+  vectorcall_entry_ = static_cast<char*>(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<std::pair<void*, std::size_t>> code_sections;
-  populateCodeSections(code_sections, codeholder, code_top);
+  populateCodeSections(code_sections, codeholder, code_start_);
   perf::registerFunction(code_sections, func->fullname, prefix);
 }
 

cinderx/Jit/codegen/gen_asm.h

@@ -16,7 +16,9 @@
 #include "cinderx/ThirdParty/asmjit/src/asmjit/asmjit.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <list>
+#include <span>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -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<const std::byte> 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,14 +107,15 @@ 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};
   void* deopt_trampoline_{nullptr};
   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_;

cinderx/Jit/compiler.cpp

@@ -266,11 +266,11 @@ std::unique_ptr<CompiledFunction> 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<vectorcallfunc>(ngen->getVectorcallEntry()))
   if (entry == nullptr) {
     JIT_DLOG("Generating native code for {} failed", fullname);
     return nullptr;
@@ -282,7 +282,6 @@ std::unique_ptr<CompiledFunction> 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<CompiledFunction> 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<const std::byte> code = ngen->getCodeBuffer();
   void* static_entry = ngen->getStaticEntry();
   CodeRuntime* code_runtime = ngen->codeRuntime();
 
@@ -309,25 +309,24 @@ std::unique_ptr<CompiledFunction> Compiler::Compile(
     return std::make_unique<CompiledFunctionDebug>(
         std::move(irfunc),
         std::move(ngen),
-        reinterpret_cast<vectorcallfunc>(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<CompiledFunction>(
-        reinterpret_cast<vectorcallfunc>(entry),
-        static_entry,
-        code_runtime,
-        func_size,
         stack_size,
         spill_stack_size,
         std::move(inline_stats),
         hir_opcode_counts);
   }
+  return std::make_unique<CompiledFunction>(
+      code,
+      entry,
+      static_entry,
+      code_runtime,
+      stack_size,
+      spill_stack_size,
+      std::move(inline_stats),
+      hir_opcode_counts);
 }
 
 } // namespace jit

cinderx/Jit/compiler.h

@@ -10,6 +10,8 @@
 #include "cinderx/Jit/hir/preload.h"
 #include "cinderx/Jit/runtime.h"
 
+#include <cstddef>
+#include <span>
 #include <utility>
 
 namespace jit {
@@ -22,25 +24,31 @@ namespace jit {
 class CompiledFunction {
  public:
   CompiledFunction(
+      std::span<const std::byte> 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)),
         hir_opcode_counts_(hir_opcode_counts) {}
 
   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<const std::byte> codeBuffer() const {
+    return code_;
+  }
+
   vectorcallfunc vectorcallEntry() const {
     return vectorcall_entry_;
   }
@@ -60,29 +68,33 @@ 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_;
   }
 
  private:
   DISALLOW_COPY_AND_ASSIGN(CompiledFunction);
 
+  const std::span<const std::byte> 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_;

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);
 

cinderx/Jit/elf.h

@@ -5,6 +5,7 @@
 #include "cinderx/Common/log.h"
 #include "cinderx/Common/util.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <iosfwd>
@@ -431,7 +432,7 @@ struct Object {
 
 // Code entry to add to an ELF file.
 struct CodeEntry {
-  std::span<uint8_t> code;
+  std::span<const std::byte> code;
   std::string func_name;
   std::string file_name;
   size_t lineno{0};

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<uint8_t*>(compiled_func->vectorcallEntry()),
-        static_cast<size_t>(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);

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<const std::byte> code = ngen.getCodeBuffer();
     int stack_size = ngen.GetCompiledFunctionStackSize();
     int spill_stack_size = ngen.GetCompiledFunctionSpillStackSize();
     return std::make_unique<jit::CompiledFunction>(
+        code,
         reinterpret_cast<vectorcallfunc>(entry),
         ngen.getStaticEntry(),
         ngen.codeRuntime(),
-        func_size,
         stack_size,
         spill_stack_size,
         jit::hir::Function::InlineFunctionStats{},

cinderx/RuntimeTests/elf_test.cpp

@@ -27,7 +27,7 @@ TEST_F(ElfTest, OneEntry) {
   std::stringstream ss;
 
   elf::CodeEntry entry;
-  entry.code = {reinterpret_cast<uint8_t*>(elf::writeEntries), 0x40};
+  entry.code = {reinterpret_cast<const std::byte*>(elf::writeEntries), 0x40};
   entry.func_name = "funcABC";
   entry.file_name = "spaghetti.exe";
   entry.lineno = 15;