Un-revert D110114775: [CinderX] Relax AArch64 b/bl with branch stubs

Summary: Re-apply D109607577 by backing out its revert, D110114775. The linker issue that caused the revert has been fixed separately, so this restores the original branch-stub relaxation change exactly as it stood before the revert.

Reviewed By: mpage

Differential Revision: D110202421

fbshipit-source-id: d67afbd56b5be499f369110e785fec21be221f5c

Author: kddnewton

cinderx/Jit/code_allocator.cpp

@@ -335,48 +335,59 @@ void CodeAllocatorCinder::ensureSplitSpace(
 }
 
 AllocateResult CodeAllocatorCinder::addSplitCode(asmjit::CodeHolder* code) {
-  // Compute how much space each section type needs.
   size_t hot_size = 0;
   size_t cold_size = 0;
-  for (asmjit::Section* section : code->sections()) {
-    CodeSection cs = codeSectionFromName(section->name());
-    if (cs == CodeSection::kCold) {
-      cold_size += section->realSize();
-    } else {
-      hot_size += section->realSize();
+
+  for (;;) {
+    // Compute how much space each section type needs.
+    hot_size = 0;
+    cold_size = 0;
+    for (asmjit::Section* section : code->sections()) {
+      CodeSection cs = codeSectionFromName(section->name());
+      if (cs == CodeSection::kCold) {
+        cold_size += section->realSize();
+      } else {
+        hot_size += section->realSize();
+      }
     }
-  }
 
-  // Ensure we have enough space for both hot and cold code.
+    // Ensure we have enough space for both hot and cold code.
 #if defined(__aarch64__)
-  // On ARM64, branch displacements are limited (±128MB for B/BL, ±1MB for
-  // B.cond). Allocate hot and cold from a single contiguous region so
-  // cross-section jumps are always in range.
-  ensureSplitSpace(hot_size, cold_size);
+    // On ARM64, branch displacements are limited (±128MB for B/BL, ±1MB for
+    // B.cond). Allocate hot and cold from a single contiguous region so
+    // cross-section jumps are always in range.
+    ensureSplitSpace(hot_size, cold_size);
 #else
-  // On x86-64, RIP-relative addressing has a ±2GB range which is large enough
-  // that independent allocations are unlikely to exceed it in practice.
-  ensureSpace(hot_alloc_, hot_alloc_free_, hot_size, true);
-  ensureSpace(
-      cold_alloc_,
-      cold_alloc_free_,
-      cold_size,
-      getConfig().cold_code_huge_pages);
+    // On x86-64, RIP-relative addressing has a ±2GB range which is large enough
+    // that independent allocations are unlikely to exceed it in practice.
+    ensureSpace(hot_alloc_, hot_alloc_free_, hot_size, true);
+    ensureSpace(
+        cold_alloc_,
+        cold_alloc_free_,
+        cold_size,
+        getConfig().cold_code_huge_pages);
 #endif
 
-  // Fix up offsets for each code section before resolving links.
-  // All offsets are relative to the hot allocation base so that asmjit can
-  // resolve cross-section jumps correctly.
-  size_t hot_offset = 0;
-  size_t cold_offset = static_cast<size_t>(cold_alloc_ - hot_alloc_);
-  for (asmjit::Section* section : code->sections()) {
-    CodeSection cs = codeSectionFromName(section->name());
-    if (cs == CodeSection::kCold) {
-      section->setOffset(cold_offset);
-      cold_offset += section->realSize();
-    } else {
-      section->setOffset(hot_offset);
-      hot_offset += section->realSize();
+    // Fix up offsets for each code section before resolving links.
+    // All offsets are relative to the hot allocation base so that asmjit can
+    // resolve cross-section jumps correctly.
+    size_t hot_offset = 0;
+    size_t cold_offset = static_cast<size_t>(cold_alloc_ - hot_alloc_);
+    for (asmjit::Section* section : code->sections()) {
+      CodeSection cs = codeSectionFromName(section->name());
+      if (cs == CodeSection::kCold) {
+        section->setOffset(cold_offset);
+        cold_offset += section->realSize();
+      } else {
+        section->setOffset(hot_offset);
+        hot_offset += section->realSize();
+      }
+    }
+
+    bool changed = false;
+    PROPAGATE_ERROR(code->ensureBranchStubIslands(&changed));
+    if (!changed) {
+      break;
     }
   }
 

cinderx/Jit/codegen/autogen.cpp

@@ -1708,8 +1708,8 @@ void translateCall(Environ* env, const Instruction* instr) {
   auto input = instr->getInput(0);
 
   if (input->isImm()) {
-    // Use bl(imm) which leverages asmjit's relaxation to pick the optimal
-    // encoding: direct bl if within ±128MB, or ldr+blr via address table.
+    // Use bl(imm) so asmjit can pick the final encoding at relocation time:
+    // direct bl if within ±128MB, or a branch to an out-of-line stub.
     as->bl(static_cast<uint64_t>(input->getConstant()));
   } else if (input->isReg()) {
     as->blr(AT::getGp(input));

cinderx/Jit/codegen/code_section.cpp

@@ -16,7 +16,8 @@ const char* codeSectionName(CodeSection section) {
 }
 
 CodeSection codeSectionFromName(const char* name) {
-  if (strcmp(name, ".text") == 0 || strcmp(name, ".addrtab") == 0) {
+  if (strcmp(name, ".text") == 0 || strcmp(name, ".addrtab") == 0 ||
+      strcmp(name, ".a64stubs") == 0) {
     return CodeSection::kHot;
   }
   if (strcmp(name, ".coldtext") == 0) {

cinderx/RuntimeTests/branch_relaxation_test.cpp

@@ -15,6 +15,14 @@ using namespace cinderx::jit::codegen;
 
 namespace {
 
+extern "C" uint64_t branchRelaxationReturn42() {
+  return 42;
+}
+
+extern "C" uint64_t branchRelaxationAddOne(uint64_t value) {
+  return value + 1;
+}
+
 class BranchRelaxationTest : public ::testing::Test {
  public:
   void SetUp() override {
@@ -75,6 +83,40 @@ TEST_F(BranchRelaxationTest, InRangeCondBranch) {
   EXPECT_EQ(func(1), 42u);
 }
 
+TEST_F(BranchRelaxationTest, AbsoluteCallThroughCodeAllocator) {
+  asmjit::CodeHolder code;
+  code.init(code_allocator_->asmJitEnvironment());
+  arch::Builder as(&code);
+
+  as.stp(arch::fp, arch::lr, asmjit::a64::ptr_pre(asmjit::a64::sp, -16));
+  as.mov(arch::fp, asmjit::a64::sp);
+  as.bl(reinterpret_cast<uint64_t>(branchRelaxationAddOne));
+  as.add(asmjit::a64::x0, asmjit::a64::x0, 1);
+  as.mov(asmjit::a64::sp, arch::fp);
+  as.ldp(arch::fp, arch::lr, asmjit::a64::ptr_post(asmjit::a64::sp, 16));
+  as.ret(arch::lr);
+
+  void* fn = compileBuilder(as, code);
+  ASSERT_NE(fn, nullptr);
+
+  auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
+  EXPECT_EQ(func(41), 43u);
+}
+
+TEST_F(BranchRelaxationTest, AbsoluteBranchThroughCodeAllocator) {
+  asmjit::CodeHolder code;
+  code.init(code_allocator_->asmJitEnvironment());
+  arch::Builder as(&code);
+
+  as.b(reinterpret_cast<uint64_t>(branchRelaxationReturn42));
+
+  void* fn = compileBuilder(as, code);
+  ASSERT_NE(fn, nullptr);
+
+  auto func = reinterpret_cast<uint64_t (*)()>(fn);
+  EXPECT_EQ(func(), 42u);
+}
+
 // tbz has a 14-bit signed immediate (+/-32KB).
 TEST_F(BranchRelaxationTest, OutOfRangeTbzIsRelaxed) {
   asmjit::CodeHolder code;

cinderx/ThirdParty/asmjit/src/asmjit/arm/a64assembler.cpp

@@ -815,7 +815,7 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
   // is valid. All options that require special handling (including invalid
   // instruction) are handled by the next branch.
   // Check for < 8 bytes remaining because branch relaxation (cbz/cbnz,
-  // tbz/tbnz, b.cond, bl/b, adr) may emit two 32-bit words (8 bytes).
+  // tbz/tbnz, b.cond, adr) may emit two 32-bit words (8 bytes).
   InstOptions options = InstOptions(instId - 1 >= Inst::_kIdCount - 1) | InstOptions((size_t)(_bufferEnd - writer.cursor()) < 8) | instOptions() | forcedInstOptions();
 
   CondCode instCC = BaseInst::extractARMCondCode(instId);
@@ -854,7 +854,7 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
       goto InvalidInstruction;
 
     // Grow request, happens rarely. Reserve 8 bytes because branch relaxation
-    // (cbz/cbnz, tbz/tbnz, b.cond, bl/b, adr) may emit two 32-bit words.
+    // (cbz/cbnz, tbz/tbnz, b.cond, adr) may emit two 32-bit words.
     err = writer.ensureSpace(this, 8);
     if (ASMJIT_UNLIKELY(err))
       goto Failed;
@@ -5207,9 +5207,9 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
 
     size_t codeOffset = writer.offsetFrom(_bufferData);
 
-    // For b(imm) and bl(imm), use address table mechanism to handle out-of-range
-    // targets. At emit time we reserve 8 bytes, and during relocation we choose
-    // the optimal encoding based on the actual displacement.
+    // For b(imm) and bl(imm), use a branch stub to handle out-of-range
+    // targets. At emit time we reserve one branch instruction, and during
+    // relocation we choose either the target or the stub as its destination.
     if ((instId == Inst::kIdBl || instId == Inst::kIdB) &&
         offsetFormat.immBitCount() == 26 &&
         (baseAddress == Globals::kNoBaseAddress || _section->id() != 0))
@@ -5235,7 +5235,7 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
         pc &= ~uint64_t(4096 - 1);
 
       // For b(imm) and bl(imm) when base address is known, try the direct
-      // encoding first. If it doesn't fit, use the address table mechanism.
+      // encoding first. If it doesn't fit, use the branch stub mechanism.
       if ((instId == Inst::kIdBl || instId == Inst::kIdB) && offsetFormat.immBitCount() == 26) {
         int64_t displacement = int64_t(targetOffset - pc);
         int64_t dispImm = displacement >> 2;
@@ -5281,7 +5281,7 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
   }
 
   // --------------------------------------------------------------------------
-  // [EmitOp - b/bl Relaxation via Address Table]
+  // [EmitOp - b/bl Relaxation via Branch Stub]
   // --------------------------------------------------------------------------
 
 EmitOp_BranchReloc:
@@ -5295,21 +5295,21 @@ Error Assembler::_emit(InstId instId, const Operand_& o0, const Operand_& o1, co
     if (err)
       goto Failed;
 
+    size_t codeOffset = writer.offsetFrom(_bufferData);
     uint64_t targetOffset = rmRel->as<Imm>().valueAs<uint64_t>();
-    err = _code->addAddressToAddressTable(targetOffset);
+    err = _code->addAddressToA64BranchStubTable(targetOffset);
     if (err)
       goto Failed;
 
-    size_t codeOffset = writer.offsetFrom(_bufferData);
     re->_sourceSectionId = _section->id();
     re->_sourceOffset = codeOffset;
     re->_format = offsetFormat;
     re->_payload = targetOffset;
 
-    // Emit two 32-bit words (placeholder: nop; nop). These will be patched
-    // during relocation to either `b/bl target; nop` or `ldr x16, [pc+off]; br/blr x16`.
-    writer.emit32uLE(0xD503201F); // NOP
-    writer.emit32uLE(0xD503201F); // NOP
+    // Emit a 32-bit branch placeholder. During relocation this is patched to
+    // either `b/bl target` or `b/bl stub`, where the stub loads the full target
+    // address and branches to it.
+    writer.emit32uLE(opcode.get());
     goto EmitDone;
   }
 

cinderx/ThirdParty/asmjit/src/asmjit/arm/a64builder.cpp

@@ -96,8 +96,6 @@ static uint32_t estimateNodeSize(BaseNode* node, uint32_t currentOffset) noexcep
         return 8;
       if (realId == Inst::kIdLdr && op.isMem() && op.as<Mem>().hasBaseLabel())
         return 8;
-      if ((realId == Inst::kIdB || realId == Inst::kIdBl) && op.isImm())
-        return 8;
       if (realId == Inst::kIdMov && op.isImm() && op.as<Imm>().valueAs<uint64_t>() > 0xFFFF)
         return 16;
     }
@@ -149,9 +147,10 @@ Error Builder::relaxBranches() {
   ZoneVector<uint32_t> labelSections;
   ASMJIT_PROPAGATE(labelSections.resize(&_allocator, labelCount));
 
-  for (;;) {
+  auto updateNodePositions = [&]() noexcept {
     uint32_t sectionId = 0;
     uint32_t offset = 0;
+
     memset(labelSections.data(), 0, labelCount * sizeof(uint32_t));
 
     for (BaseNode* node = firstNode(); node; node = node->next()) {
@@ -171,9 +170,13 @@ Error Builder::relaxBranches() {
 
       offset += size;
     }
+  };
+
+  for (;;) {
+    updateNodePositions();
 
     bool changed = false;
-    sectionId = 0;
+    uint32_t sectionId = 0;
 
     for (BaseNode* node = firstNode(); node; node = node->next()) {
       if (node->isSection()) {