Revert D109607577: Relax AArch64 b/bl with branch stubs

Differential Revision:
D109607577

Original commit changeset: e8c0f43980a9

Original Phabricator Diff: D109607577

fbshipit-source-id: 22aa51cd3f37b8e35a38e7370465d008fca6ac91

Author: amrshennawi

cinderx/Jit/code_allocator.cpp

@@ -335,59 +335,48 @@ 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 (;;) {
-    // 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();
-      }
+  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();
-      }
-    }
-
-    bool changed = false;
-    PROPAGATE_ERROR(code->ensureBranchStubIslands(&changed));
-    if (!changed) {
-      break;
+  // 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();
     }
   }
 

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) so asmjit can pick the final encoding at relocation time:
-    // direct bl if within ±128MB, or a branch to an out-of-line stub.
+    // Use bl(imm) which leverages asmjit's relaxation to pick the optimal
+    // encoding: direct bl if within ±128MB, or ldr+blr via address table.
     as->bl(static_cast<uint64_t>(input->getConstant()));
   } else if (input->isReg()) {
     as->blr(AT::getGp(input));

cinderx/Jit/codegen/code_section.cpp

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

cinderx/RuntimeTests/branch_relaxation_test.cpp

@@ -15,14 +15,6 @@ 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 {
@@ -83,40 +75,6 @@ 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, adr) may emit two 32-bit words (8 bytes).
+  // tbz/tbnz, b.cond, bl/b, 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, adr) may emit two 32-bit words.
+    // (cbz/cbnz, tbz/tbnz, b.cond, bl/b, 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 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.
+    // 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.
     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 branch stub mechanism.
+      // encoding first. If it doesn't fit, use the address table 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 Branch Stub]
+  // [EmitOp - b/bl Relaxation via Address Table]
   // --------------------------------------------------------------------------
 
 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->addAddressToA64BranchStubTable(targetOffset);
+    err = _code->addAddressToAddressTable(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 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());
+    // 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
     goto EmitDone;
   }
 

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

@@ -96,6 +96,8 @@ 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;
     }
@@ -147,10 +149,9 @@ Error Builder::relaxBranches() {
   ZoneVector<uint32_t> labelSections;
   ASMJIT_PROPAGATE(labelSections.resize(&_allocator, labelCount));
 
-  auto updateNodePositions = [&]() noexcept {
+  for (;;) {
     uint32_t sectionId = 0;
     uint32_t offset = 0;
-
     memset(labelSections.data(), 0, labelCount * sizeof(uint32_t));
 
     for (BaseNode* node = firstNode(); node; node = node->next()) {
@@ -170,13 +171,9 @@ Error Builder::relaxBranches() {
 
       offset += size;
     }
-  };
-
-  for (;;) {
-    updateNodePositions();
 
     bool changed = false;
-    uint32_t sectionId = 0;
+    sectionId = 0;
 
     for (BaseNode* node = firstNode(); node; node = node->next()) {
       if (node->isSection()) {