Skip to content

Commit 6b8742c

Browse files
kddnewtonmeta-codesync[bot]
kddnewton
authored and
meta-codesync[bot]
committed
Move aarch64 branch relaxation from Assembler to Builder
Summary: The vendored asmjit Assembler used pessimistic 8-byte encoding (instruction + NOP) for every forward-reference conditional branch, `adr`, and `ldr` literal on aarch64, wasting up to 50% of the encoded output on NOP padding. This eliminates all NOP padding for label-relative instructions by moving branch relaxation to an iterative Builder pass and switching `adr`/`ldr` literal to compact 4-byte label links. ## What changed **a64builder.cpp** — New `Builder::relaxBranches()` pass called from `finalize()` between `runPasses()` and `serializeTo()`. The algorithm: 1. Walks the node list computing per-section offsets and recording each label's section assignment. 2. Scans for conditional branches (`b.cond`, `cbz`/`cbnz`, `tbz`/`tbnz`) whose displacement exceeds the instruction's immediate range, or that cross section boundaries. 3. Expands out-of-range branches by replacing `b.cc target` with `b.inv_cc skip; b target; skip:`, using the existing `arm::negateCond()` utility for condition inversion. 4. Iterates until no more expansions are needed (guaranteed to converge since expansions are monotonic — each branch is expanded at most once). For the common case (all branches in range), this is a single O(N) walk that confirms nothing needs expansion. **a64assembler.cpp** — Three changes: 1. `emitCondBranchRelax()` replaced with `emitCondBranch()`. Forward references now emit 4 bytes (just the opcode) instead of 8 bytes (opcode + NOP). The `inversionMask` parameter is removed since the Builder handles condition inversion. 2. `adr Rd, label` forward references now emit 4 bytes with a `kAArch64_Adr` label link instead of 8 bytes (adr + NOP) with a RelocEntry. Backward references that don't fit are errors (the Builder should have caught these). 3. `ldr Rd, [label]` (literal pool loads) forward references now emit 4 bytes with a `kAArch64_LdrLiteral` label link instead of 8 bytes (ldr + NOP) with a RelocEntry. Note: `load_addr` (absolute address) and `b`/`bl` to absolute addresses still use 8-byte RelocEntry encoding because their displacements depend on the base address, which is unknown until `relocateToBase()`. **codeholder.h** — Added `OffsetType::kAArch64_LdrLiteral` for the new ldr literal label link format. **codeholder.cpp** — `bindLabel()` and `resolveUnresolvedLinks()` updated to handle the two new offset types (`kAArch64_Adr` with its split 21-bit immediate encoding, `kAArch64_LdrLiteral` with its 19-bit shifted immediate). The conditional branch path was already simplified in the initial change to just patch 4 bytes. ## AArch64 instruction range limits | Instruction | Immediate | Range | NOP eliminated? | |---|---|---|---| | `b.cond` | 19-bit signed | +/-1MB | Yes | | `cbz`/`cbnz` | 19-bit signed | +/-1MB | Yes | | `tbz`/`tbnz` | 14-bit signed | +/-32KB | Yes | | `adr` | 21-bit signed | +/-1MB | Yes | | `ldr` literal | 19-bit signed | +/-1MB | Yes | | `b`/`bl` (absolute) | 26-bit signed | +/-128MB | No (base address unknown at Builder time) | | `load_addr` (absolute) | — | — | No (same reason) | ## Performance Compilation speed benchmark (`fbcode//cinderx/benchmarks:compile-time`) shows 36ms total for both baseline and with changes across multiple runs — zero measurable regression. Reviewed By: alexmalyshev Differential Revision: D104062548 fbshipit-source-id: 88db5486d4526fd8d9b948993988e433a2949e31
1 parent 9475033 commit 6b8742c

6 files changed

Lines changed: 607 additions & 107 deletions

File tree

‎cinderx/RuntimeTests/branch_relaxation_test.cpp‎

Lines changed: 370 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,370 @@
1+
// Copyright (c) Meta Platforms, Inc. and affiliates.
2+
3+
#include <gtest/gtest.h>
4+
5+
#include "cinderx/Jit/code_allocator.h"
6+
#include "cinderx/Jit/codegen/arch.h"
7+
8+
#include <cstdint>
9+
#include <memory>
10+
11+
#if defined(CINDER_AARCH64)
12+
13+
using namespace jit;
14+
using namespace jit::codegen;
15+
16+
namespace {
17+
18+
class BranchRelaxationTest : public ::testing::Test {
19+
public:
20+
void SetUp() override {
21+
code_allocator_ = std::unique_ptr<ICodeAllocator>(CodeAllocator::make());
22+
}
23+
24+
void TearDown() override {
25+
code_allocator_.reset();
26+
}
27+
28+
void* compileBuilder(arch::Builder& as, asmjit::CodeHolder& code) {
29+
EXPECT_EQ(as.finalize(), asmjit::kErrorOk);
30+
AllocateResult result = code_allocator_->addCode(&code);
31+
EXPECT_EQ(result.error, asmjit::kErrorOk);
32+
return result.addr;
33+
}
34+
35+
size_t finalizeAndGetSize(arch::Builder& as, asmjit::CodeHolder& code) {
36+
EXPECT_EQ(as.finalize(), asmjit::kErrorOk);
37+
return code.textSection()->bufferSize();
38+
}
39+
40+
size_t countNops(asmjit::CodeHolder& code) {
41+
constexpr uint32_t kAArch64Nop = 0xD503201Fu;
42+
asmjit::Section* text = code.textSection();
43+
const uint8_t* buf = text->data();
44+
size_t size = text->bufferSize();
45+
size_t count = 0;
46+
for (size_t i = 0; i + 3 < size; i += 4) {
47+
uint32_t inst;
48+
memcpy(&inst, buf + i, 4);
49+
if (inst == kAArch64Nop)
50+
count++;
51+
}
52+
return count;
53+
}
54+
55+
std::unique_ptr<ICodeAllocator> code_allocator_;
56+
};
57+
58+
TEST_F(BranchRelaxationTest, InRangeCondBranch) {
59+
asmjit::CodeHolder code;
60+
code.init(code_allocator_->asmJitEnvironment());
61+
arch::Builder as(&code);
62+
63+
asmjit::Label target = as.newLabel();
64+
65+
as.cbz(asmjit::a64::x0, target);
66+
as.mov(asmjit::a64::x0, 42);
67+
as.bind(target);
68+
as.ret(asmjit::a64::x30);
69+
70+
void* fn = compileBuilder(as, code);
71+
ASSERT_NE(fn, nullptr);
72+
73+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
74+
EXPECT_EQ(func(0), 0u);
75+
EXPECT_EQ(func(1), 42u);
76+
}
77+
78+
// tbz has a 14-bit signed immediate (+/-32KB).
79+
TEST_F(BranchRelaxationTest, OutOfRangeTbzIsRelaxed) {
80+
asmjit::CodeHolder code;
81+
code.init(code_allocator_->asmJitEnvironment());
82+
arch::Builder as(&code);
83+
84+
asmjit::Label target = as.newLabel();
85+
86+
as.tbz(asmjit::a64::x0, 0, target);
87+
88+
for (int i = 0; i < 9000; i++) {
89+
as.nop();
90+
}
91+
92+
as.mov(asmjit::a64::x0, 1);
93+
as.ret(asmjit::a64::x30);
94+
95+
as.bind(target);
96+
as.mov(asmjit::a64::x0, 0);
97+
as.ret(asmjit::a64::x30);
98+
99+
void* fn = compileBuilder(as, code);
100+
ASSERT_NE(fn, nullptr);
101+
102+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
103+
EXPECT_EQ(func(0), 0u);
104+
EXPECT_EQ(func(1), 1u);
105+
EXPECT_EQ(func(2), 0u);
106+
EXPECT_EQ(func(3), 1u);
107+
}
108+
109+
// cbz has a 19-bit signed immediate (+/-1MB).
110+
TEST_F(BranchRelaxationTest, OutOfRangeCbzIsRelaxed) {
111+
asmjit::CodeHolder code;
112+
code.init(code_allocator_->asmJitEnvironment());
113+
arch::Builder as(&code);
114+
115+
asmjit::Label target = as.newLabel();
116+
117+
as.cbz(asmjit::a64::x0, target);
118+
119+
for (int i = 0; i < 270000; i++) {
120+
as.nop();
121+
}
122+
123+
as.mov(asmjit::a64::x0, 1);
124+
as.ret(asmjit::a64::x30);
125+
126+
as.bind(target);
127+
as.mov(asmjit::a64::x0, 0);
128+
as.ret(asmjit::a64::x30);
129+
130+
void* fn = compileBuilder(as, code);
131+
ASSERT_NE(fn, nullptr);
132+
133+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
134+
EXPECT_EQ(func(0), 0u);
135+
EXPECT_EQ(func(1), 1u);
136+
}
137+
138+
TEST_F(BranchRelaxationTest, OutOfRangeCondBranchIsRelaxed) {
139+
asmjit::CodeHolder code;
140+
code.init(code_allocator_->asmJitEnvironment());
141+
arch::Builder as(&code);
142+
143+
asmjit::Label target = as.newLabel();
144+
145+
as.cmp(asmjit::a64::x0, 0);
146+
as.b_eq(target);
147+
148+
for (int i = 0; i < 270000; i++) {
149+
as.nop();
150+
}
151+
152+
as.mov(asmjit::a64::x0, 1);
153+
as.ret(asmjit::a64::x30);
154+
155+
as.bind(target);
156+
as.mov(asmjit::a64::x0, 0);
157+
as.ret(asmjit::a64::x30);
158+
159+
void* fn = compileBuilder(as, code);
160+
ASSERT_NE(fn, nullptr);
161+
162+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
163+
EXPECT_EQ(func(0), 0u);
164+
EXPECT_EQ(func(1), 1u);
165+
}
166+
167+
TEST_F(BranchRelaxationTest, OutOfRangeBackwardTbzIsRelaxed) {
168+
asmjit::CodeHolder code;
169+
code.init(code_allocator_->asmJitEnvironment());
170+
arch::Builder as(&code);
171+
172+
asmjit::Label loop_top = as.newLabel();
173+
asmjit::Label done = as.newLabel();
174+
175+
as.b(done);
176+
177+
as.bind(loop_top);
178+
as.nop();
179+
180+
for (int i = 0; i < 9000; i++) {
181+
as.nop();
182+
}
183+
184+
as.tbz(asmjit::a64::x0, 0, loop_top);
185+
186+
as.bind(done);
187+
as.ret(asmjit::a64::x30);
188+
189+
void* fn = compileBuilder(as, code);
190+
ASSERT_NE(fn, nullptr);
191+
192+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
193+
EXPECT_EQ(func(42), 42u);
194+
}
195+
196+
TEST_F(BranchRelaxationTest, NoNopPaddingForInRangeBranches) {
197+
asmjit::CodeHolder code;
198+
code.init(code_allocator_->asmJitEnvironment());
199+
arch::Builder as(&code);
200+
201+
asmjit::Label target = as.newLabel();
202+
203+
as.cbz(asmjit::a64::x0, target);
204+
as.mov(asmjit::a64::x0, 42);
205+
as.bind(target);
206+
as.ret(asmjit::a64::x30);
207+
208+
// cbz(4) + mov(4) + ret(4) = 12 bytes, no NOP padding.
209+
EXPECT_EQ(finalizeAndGetSize(as, code), 12u);
210+
EXPECT_EQ(countNops(code), 0u);
211+
}
212+
213+
TEST_F(BranchRelaxationTest, MultipleForwardBranchesNoNops) {
214+
asmjit::CodeHolder code;
215+
code.init(code_allocator_->asmJitEnvironment());
216+
arch::Builder as(&code);
217+
218+
asmjit::Label l1 = as.newLabel();
219+
asmjit::Label l2 = as.newLabel();
220+
asmjit::Label l3 = as.newLabel();
221+
222+
as.cbz(asmjit::a64::x0, l1);
223+
as.tbz(asmjit::a64::x0, 1, l2);
224+
as.cmp(asmjit::a64::x0, 0);
225+
as.b_eq(l3);
226+
as.mov(asmjit::a64::x0, 99);
227+
as.bind(l1);
228+
as.bind(l2);
229+
as.bind(l3);
230+
as.ret(asmjit::a64::x30);
231+
232+
// 6 instructions * 4 bytes = 24 bytes, no NOP padding.
233+
EXPECT_EQ(finalizeAndGetSize(as, code), 24u);
234+
EXPECT_EQ(countNops(code), 0u);
235+
}
236+
237+
// Unconditional b to a cross-section label resolves correctly.
238+
TEST_F(BranchRelaxationTest, UnconditionalBranchCrossSection) {
239+
asmjit::CodeHolder code;
240+
code.init(code_allocator_->asmJitEnvironment());
241+
242+
asmjit::Section* cold;
243+
ASSERT_EQ(
244+
code.newSection(
245+
&cold,
246+
".cold",
247+
SIZE_MAX,
248+
code.textSection()->flags(),
249+
code.textSection()->alignment()),
250+
asmjit::kErrorOk);
251+
252+
arch::Builder as(&code);
253+
254+
asmjit::Label cold_target = as.newLabel();
255+
256+
// Hot section: branch to cold, then return 1 (never reached).
257+
as.b(cold_target);
258+
as.mov(asmjit::a64::x0, 1);
259+
as.ret(asmjit::a64::x30);
260+
261+
// Cold section: return 42.
262+
as.section(cold);
263+
as.bind(cold_target);
264+
as.mov(asmjit::a64::x0, 42);
265+
as.ret(asmjit::a64::x30);
266+
267+
void* fn = compileBuilder(as, code);
268+
ASSERT_NE(fn, nullptr);
269+
270+
auto func = reinterpret_cast<uint64_t (*)()>(fn);
271+
EXPECT_EQ(func(), 42u);
272+
}
273+
274+
// Conditional branch relaxation works across sections: the expanded
275+
// unconditional b targets the cross-section label.
276+
TEST_F(BranchRelaxationTest, CondBranchCrossSectionIsRelaxed) {
277+
asmjit::CodeHolder code;
278+
code.init(code_allocator_->asmJitEnvironment());
279+
280+
asmjit::Section* cold;
281+
ASSERT_EQ(
282+
code.newSection(
283+
&cold,
284+
".cold",
285+
SIZE_MAX,
286+
code.textSection()->flags(),
287+
code.textSection()->alignment()),
288+
asmjit::kErrorOk);
289+
290+
arch::Builder as(&code);
291+
292+
asmjit::Label cold_target = as.newLabel();
293+
294+
// Hot section: conditional branch to cold section.
295+
as.cbz(asmjit::a64::x0, cold_target);
296+
as.mov(asmjit::a64::x0, 1);
297+
as.ret(asmjit::a64::x30);
298+
299+
// Cold section: return 0.
300+
as.section(cold);
301+
as.bind(cold_target);
302+
as.mov(asmjit::a64::x0, 0);
303+
as.ret(asmjit::a64::x30);
304+
305+
void* fn = compileBuilder(as, code);
306+
ASSERT_NE(fn, nullptr);
307+
308+
auto func = reinterpret_cast<uint64_t (*)(uint64_t)>(fn);
309+
EXPECT_EQ(func(0), 0u);
310+
EXPECT_EQ(func(1), 1u);
311+
}
312+
313+
// adr and ldr literal use 8-byte RelocEntry encoding to allow relaxation
314+
// to adrp+add / adrp+ldr for large displacements at relocation time.
315+
TEST_F(BranchRelaxationTest, AdrForwardRefUsesReloc) {
316+
asmjit::CodeHolder code;
317+
code.init(code_allocator_->asmJitEnvironment());
318+
arch::Builder as(&code);
319+
320+
asmjit::Label target = as.newLabel();
321+
322+
as.adr(asmjit::a64::x0, target);
323+
as.ret(asmjit::a64::x30);
324+
as.bind(target);
325+
as.nop();
326+
327+
// adr(4) + nop(4) + ret(4) + nop(4) = 16 bytes.
328+
EXPECT_EQ(finalizeAndGetSize(as, code), 16u);
329+
}
330+
331+
TEST_F(BranchRelaxationTest, LdrLiteralForwardRefUsesReloc) {
332+
asmjit::CodeHolder code;
333+
code.init(code_allocator_->asmJitEnvironment());
334+
arch::Builder as(&code);
335+
336+
asmjit::Label pool_entry = as.newLabel();
337+
338+
as.ldr(asmjit::a64::x0, asmjit::a64::ptr(pool_entry));
339+
as.ret(asmjit::a64::x30);
340+
as.bind(pool_entry);
341+
uint64_t value = 0xDEADBEEFCAFEBABEull;
342+
as.embed(&value, sizeof(value));
343+
344+
// ldr(4) + nop(4) + ret(4) + data(8) = 20 bytes.
345+
EXPECT_EQ(finalizeAndGetSize(as, code), 20u);
346+
}
347+
348+
TEST_F(BranchRelaxationTest, LdrLiteralLoadsCorrectValue) {
349+
asmjit::CodeHolder code;
350+
code.init(code_allocator_->asmJitEnvironment());
351+
arch::Builder as(&code);
352+
353+
asmjit::Label pool_entry = as.newLabel();
354+
355+
as.ldr(asmjit::a64::x0, asmjit::a64::ptr(pool_entry));
356+
as.ret(asmjit::a64::x30);
357+
as.bind(pool_entry);
358+
uint64_t value = 42;
359+
as.embed(&value, sizeof(value));
360+
361+
void* fn = compileBuilder(as, code);
362+
ASSERT_NE(fn, nullptr);
363+
364+
auto func = reinterpret_cast<uint64_t (*)()>(fn);
365+
EXPECT_EQ(func(), 42u);
366+
}
367+
368+
} // namespace
369+
370+
#endif // CINDER_AARCH64

0 commit comments

Comments
 (0)