Abstract DeoptPatcher into CodePatcher

Summary:
Making a more general abstraction that can patch arbitrary sections of code.  Or
at least up to 7 bytes :).

The CodePatcher stores a small buffer of data to patch in plus the destination
address.  When a patch is triggered, it will swap what it has in its buffer with
what is at the destination address.  This allows the CodePatcher to undo its
patch in the future with an "unpatch".

A "deopt patcher" is a specific case of this, and there's now a new subclass
called JumpPatcher that covers this case, but it's still more general in that
it's focused on patching in jumps rather than anything specific to deopts.

The 7-byte limit here is to keep the CodePatcher to a small size, currently 24
bytes.  If in the future we'll want to patch in 8-byte values we can replace its
vtable with a smaller enum value and shrink it further.

Reviewed By: jbower-fb

Differential Revision: D83591551

fbshipit-source-id: 160023a047397fb9ef21452e6914bcc49a1ee95d

Author: alexmalyshev

cinderx/Jit/code_patcher.cpp

@@ -0,0 +1,139 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include "cinderx/Jit/code_patcher.h"
+
+#include "cinderx/Common/log.h"
+#include "cinderx/Common/util.h"
+
+#include <array>
+#include <cstring>
+
+namespace jit {
+
+namespace {
+
+static_assert(
+    sizeof(CodePatcher) == 24,
+    "CodePatcher should be kept small as there could be many per function");
+
+// 5-byte nop - https://www.felixcloutier.com/x86/nop
+//
+// Asmjit supports multi-byte nops but for whatever reason we can't get it to
+// emit the 5-byte version.
+constexpr auto kJmpNopBytes =
+    std::to_array<uint8_t>({0x0f, 0x1f, 0x44, 0x00, 0x00});
+
+// Compute an x86-64 jump displacement operand.
+uint32_t jumpDisplacement(uintptr_t from, uintptr_t to) {
+  auto disp = to - (from + kJmpNopBytes.size());
+  JIT_CHECK(
+      fitsInt32(disp),
+      "Can't encode jump from {:#x} to {:#x} as relative",
+      from,
+      to);
+  return static_cast<uint32_t>(disp);
+}
+
+// Given the starting address and displacement operand of a jump instruction,
+// resolve it to a target address.
+uintptr_t resolveDisplacement(uintptr_t from, uint32_t displacement) {
+  return from + displacement + kJmpNopBytes.size();
+}
+
+} // namespace
+
+void CodePatcher::link(uintptr_t patchpoint, std::span<const uint8_t> data) {
+  JIT_CHECK(!isLinked(), "Trying to re-link a patcher");
+
+  patchpoint_ = reinterpret_cast<uint8_t*>(patchpoint);
+
+  JIT_CHECK(
+      data.size() <= data_.size(),
+      "Trying to link a patch point with {} bytes of data but only {} are "
+      "supported",
+      data.size(),
+      data_.size());
+
+  std::memcpy(data_.data(), data.data(), data.size());
+  data_len_ = data.size();
+
+  onLink();
+}
+
+void CodePatcher::patch() {
+  JIT_CHECK(isLinked(), "Trying to patch a patcher that isn't linked");
+  JIT_DLOG("Patching DeoptPatchPoint at {}", static_cast<void*>(patchpoint_));
+
+  swap();
+
+  is_patched_ = true;
+  onPatch();
+}
+
+void CodePatcher::unpatch() {
+  JIT_CHECK(isLinked(), "Trying to unpatch a patcher that isn't linked");
+  JIT_DLOG("Unpatching DeoptPatchPoint at {}", static_cast<void*>(patchpoint_));
+
+  swap();
+
+  is_patched_ = false;
+  onUnpatch();
+}
+
+bool CodePatcher::isLinked() const {
+  return patchpoint_ != nullptr;
+}
+
+bool CodePatcher::isPatched() const {
+  return is_patched_;
+}
+
+uint8_t* CodePatcher::patchpoint() const {
+  return patchpoint_;
+}
+
+std::span<const uint8_t> CodePatcher::storedBytes() const {
+  return std::span{data_.data(), data_len_};
+}
+
+void CodePatcher::swap() {
+  decltype(data_) temp;
+  std::memcpy(temp.data(), patchpoint_, data_len_);
+  std::memcpy(patchpoint_, data_.data(), data_len_);
+  std::memcpy(data_.data(), temp.data(), data_len_);
+}
+
+JumpPatcher::JumpPatcher() {
+  // Initializes to a nop.
+  std::memcpy(data_.data(), kJmpNopBytes.data(), kJmpNopBytes.size());
+  data_len_ = kJmpNopBytes.size();
+}
+
+void JumpPatcher::linkJump(uintptr_t patchpoint, uintptr_t jump_target) {
+  auto disp = jumpDisplacement(patchpoint, jump_target);
+
+  // 32 bit relative jump - https://www.felixcloutier.com/x86/jmp
+  std::array<uint8_t, kJmpNopBytes.size()> buf{};
+  buf[0] = 0xe9;
+  std::memcpy(buf.data() + 1, &disp, sizeof(uint32_t));
+
+  link(patchpoint, buf);
+}
+
+uint8_t* JumpPatcher::jumpTarget() const {
+  JIT_CHECK(
+      isLinked(), "Can't compute jump target before JumpPatcher is linked");
+
+  std::span<const uint8_t> bytes = storedBytes();
+  JIT_CHECK(
+      bytes.size() == 5,
+      "Must have linked a 5-byte 'jmp $DISP' instruction into a JumpPatcher");
+
+  uint32_t disp = 0;
+  std::memcpy(&disp, bytes.data() + 1, bytes.size() - 1);
+
+  return reinterpret_cast<uint8_t*>(
+      resolveDisplacement(reinterpret_cast<uintptr_t>(patchpoint_), disp));
+}
+
+} // namespace jit

cinderx/Jit/code_patcher.h

@@ -0,0 +1,114 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#pragma once
+
+#include <cstdint>
+#include <span>
+
+namespace jit {
+
+// A CodePatcher is used by the runtime to overwrite parts of compiled code.
+// Often times this is used to patch in a jump to a deopt exit when an invariant
+// that the compiled code relies on is invalidated. It is intended to be used in
+// conjunction with the DeoptPatchpoint HIR instruction.
+//
+// Using a CodePatcher looks roughly like:
+//   1. Allocate a CodePatcher.
+//
+//   2. Allocate a DeoptPatchpoint HIR instruction linked to the CodePatcher
+//      from (1) and insert it into the appropriate point in the HIR
+//      instruction stream.
+//
+//   3. Link the CodePatcher from (1) to the appropriate address in the
+//      generated code after code generation is complete.
+//
+// A CodePatcher is only valid for as long as the compiled code to which it is
+// linked is alive, so care must be taken not to call `patch()` after the code
+// has been destroyed.
+class CodePatcher {
+ public:
+  virtual ~CodePatcher() = default;
+
+  // Link the patcher to a specific location in generated code. This is
+  // intended to be called by the JIT after code has been generated but before
+  // it is active.
+  //
+  // `patchpoint` contains the address of the first byte of the patchpoint.
+  // `data` contains the bytes that will be written on patching.
+  void link(uintptr_t patchpoint, std::span<const uint8_t> data);
+
+  // Overwrite the patchpoint.
+  //
+  // The patcher must be linked before this can be called.
+  void patch();
+
+  // Revert the patchpoint back to a nop.
+  //
+  // The patcher must be linked before this can be called.
+  void unpatch();
+
+  // Check if the patcher has been linked.
+  bool isLinked() const;
+
+  // Check if the patcher is currently patched.
+  bool isPatched() const;
+
+  // Get where in the code to patch.  Will be nullptr before the patcher is
+  // linked.
+  uint8_t* patchpoint() const;
+
+  // Get the bytes that are stored within the patcher right now.
+  //
+  // This either contains the bytes that will be patched in, or the bytes that
+  // were there originally.  The former is injected with patch(), the latter can
+  // be put back in with unpatch().
+  std::span<const uint8_t> storedBytes() const;
+
+ protected:
+  // Callback to execute after linking (e.g. subscribing to changes).
+  virtual void onLink() {}
+
+  // Callback to execute after patching (e.g. cleaning up the patcher).
+  virtual void onPatch() {}
+
+  // Callback to execute after unpatching.
+  virtual void onUnpatch() {}
+
+  // Swap data between this object and the actual patchpoint.
+  void swap();
+
+  // Where in the code we should patch.
+  uint8_t* patchpoint_{nullptr};
+
+  // Data that's written into the patch point.  This is swapped with what's
+  // already there, so that this can continuously patch and unpatch.
+  //
+  // The size of the array here is the total capacity, not necessarily all of it
+  // will be patched.
+  std::array<uint8_t, 7> data_{};
+
+  // Actual length of the data buffer above, from 0 to 7 bytes.
+  uint8_t data_len_ : 7 {0};
+
+  // Whether patch() has been called and a corresponding unpatch() has not yet
+  // been called..
+  bool is_patched_ : 1 {false};
+};
+
+// Subclass of a CodePatcher that is intended for patching in jumps.
+class JumpPatcher : public CodePatcher {
+ public:
+  JumpPatcher();
+  ~JumpPatcher() override = default;
+
+  // Specific form of link() for handling jumps.
+  //
+  // NB: The distance between the patchpoint and the jump target must fit into a
+  // signed 32-bit int.
+  void linkJump(uintptr_t patchpoint, uintptr_t jump_target);
+
+  // Get the jump target of this patcher.
+  uint8_t* jumpTarget() const;
+};
+
+} // namespace jit

cinderx/Jit/codegen/autogen.cpp

@@ -3,9 +3,9 @@
 #include "cinderx/Jit/codegen/autogen.h"
 
 #include "cinderx/Common/util.h"
+#include "cinderx/Jit/code_patcher.h"
 #include "cinderx/Jit/codegen/gen_asm_utils.h"
 #include "cinderx/Jit/codegen/x86_64.h"
-#include "cinderx/Jit/deopt_patcher.h"
 #include "cinderx/Jit/frame.h"
 #include "cinderx/Jit/generators_rt.h"
 #include "cinderx/Jit/jit_rt.h"
@@ -260,10 +260,15 @@ void TranslateGuard(Environ* env, const Instruction* instr) {
 void TranslateDeoptPatchpoint(Environ* env, const Instruction* instr) {
   auto as = env->as;
 
-  // Generate patchpoint
+  auto patcher =
+      reinterpret_cast<JumpPatcher*>(instr->getInput(0)->getMemoryAddress());
+
+  // Generate patchpoint by writing in an appropriately sized nop.  As a future
+  // optimization, we may be able to avoid reserving space for the patchpoint if
+  // we can prove that the following bytes are not the target of a jump.
   auto patchpoint_label = as->newLabel();
   as->bind(patchpoint_label);
-  for (uint8_t byte : kJmpNopBytes) {
+  for (uint8_t byte : patcher->storedBytes()) {
     as->db(byte);
   }
 
@@ -276,8 +281,6 @@ void TranslateDeoptPatchpoint(Environ* env, const Instruction* instr) {
 
   // The runtime will link the patcher to the appropriate point in the code
   // once code generation has completed.
-  auto patcher =
-      reinterpret_cast<DeoptPatcher*>(instr->getInput(0)->getMemoryAddress());
   env->pending_deopt_patchers.emplace_back(
       patcher, patchpoint_label, deopt_label);
 }

cinderx/Jit/codegen/environ.h

@@ -53,9 +53,9 @@ struct Environ {
   std::vector<DeoptExit> deopt_exits;
 
   struct PendingDeoptPatcher {
-    PendingDeoptPatcher(DeoptPatcher* p, asmjit::Label pp, asmjit::Label de)
+    PendingDeoptPatcher(JumpPatcher* p, asmjit::Label pp, asmjit::Label de)
         : patcher(p), patchpoint(pp), deopt_exit(de) {}
-    DeoptPatcher* patcher;
+    JumpPatcher* patcher;
 
     // Location of the patchpoint
     asmjit::Label patchpoint;

cinderx/Jit/codegen/gen_asm.cpp

@@ -1422,7 +1422,7 @@ void NativeGenerator::linkDeoptPatchers(const asmjit::CodeHolder& code) {
   for (const auto& udp : env_.pending_deopt_patchers) {
     uint64_t patchpoint = base + code.labelOffsetFromBase(udp.patchpoint);
     uint64_t deopt_exit = base + code.labelOffsetFromBase(udp.deopt_exit);
-    udp.patcher->link(patchpoint, deopt_exit);
+    udp.patcher->linkJump(patchpoint, deopt_exit);
 
     // Register patcher with the runtime if it is type-based.
     if (auto typed_patcher = dynamic_cast<TypeDeoptPatcher*>(udp.patcher)) {
@@ -1433,8 +1433,8 @@ void NativeGenerator::linkDeoptPatchers(const asmjit::CodeHolder& code) {
   // Any patchers that aren't linked at this point are pointing to patch points
   // that were optimized out.  It's safe to delete them.
   std::erase_if(
-      const_cast<hir::Function*>(func_)->deopt_patchers,
-      [](std::unique_ptr<DeoptPatcher>& patcher) {
+      const_cast<hir::Function*>(func_)->code_patchers,
+      [](std::unique_ptr<CodePatcher>& patcher) {
         return !patcher->isLinked();
       });
 }

cinderx/Jit/compiled_function.cpp

@@ -58,9 +58,9 @@ void CompiledFunction::setCompileTime(std::chrono::nanoseconds time) {
   compile_time_ = time;
 }
 
-void CompiledFunction::setDeoptPatchers(
-    std::vector<std::unique_ptr<DeoptPatcher>>&& deopt_patchers) {
-  deopt_patchers_ = std::move(deopt_patchers);
+void CompiledFunction::setCodePatchers(
+    std::vector<std::unique_ptr<CodePatcher>>&& code_patchers) {
+  code_patchers_ = std::move(code_patchers);
 }
 
 void CompiledFunction::setHirFunc(std::unique_ptr<hir::Function>&& irfunc) {

cinderx/Jit/compiled_function.h

@@ -54,8 +54,8 @@ bool isJitCompiled(const PyFunctionObject* func);
 #ifdef __cplusplus
 
 #include "cinderx/Common/util.h"
+#include "cinderx/Jit/code_patcher.h"
 #include "cinderx/Jit/code_runtime.h"
-#include "cinderx/Jit/deopt_patcher.h"
 #include "cinderx/Jit/hir/hir.h"
 
 #include <chrono>
@@ -130,8 +130,8 @@ class CompiledFunction {
   std::chrono::nanoseconds compileTime() const;
   void setCompileTime(std::chrono::nanoseconds time);
 
-  void setDeoptPatchers(
-      std::vector<std::unique_ptr<DeoptPatcher>>&& deopt_patchers);
+  void setCodePatchers(
+      std::vector<std::unique_ptr<CodePatcher>>&& code_patchers);
 
   void setHirFunc(std::unique_ptr<hir::Function>&& irfunc);
 
@@ -154,8 +154,8 @@ class CompiledFunction {
   std::chrono::nanoseconds compile_time_;
   hir::Function::InlineFunctionStats inline_function_stats_;
   hir::OpcodeCounts hir_opcode_counts_;
-  // All the deopt patchers pointing to patch points in this function.
-  std::vector<std::unique_ptr<DeoptPatcher>> deopt_patchers_;
+  // All the code patchers pointing to patch points in this function.
+  std::vector<std::unique_ptr<CodePatcher>> code_patchers_;
   std::unique_ptr<hir::Function> irfunc_;
   CodeRuntime* runtime_;
 };

cinderx/Jit/compiler.cpp

@@ -253,7 +253,7 @@ std::unique_ptr<CompiledFunction> Compiler::Compile(
       hir_opcode_counts,
       code_runtime);
   compiled_func->setCompileTime(compile_time);
-  compiled_func->setDeoptPatchers(std::move(irfunc->deopt_patchers));
+  compiled_func->setCodePatchers(std::move(irfunc->code_patchers));
   if (getConfig().log.debug) {
     irfunc->setCompilationPhaseTimer(nullptr);
     compiled_func->setHirFunc(std::move(irfunc));