Skip to content

Commit 2b3ede9

Browse files
pguyotpguyot
committed
JIT: Optimize memory usage by patching jump table asap
Signed-off-by: Paul Guyot <pguyot@kallisys.net>
1 parent 593b3ee commit 2b3ede9

File tree

7 files changed

+468
-307
lines changed

7 files changed

+468
-307
lines changed

‎libs/estdlib/src/code_server.erl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ set_native_code(_Module, _LabelsCount, _Stream) ->
152152
load(Module) ->
153153
case erlang:system_info(emu_flavor) of
154154
jit ->
155-
% atomvm_heap_growth, fibonacci divides compilation time by two
155+
% atomvm_heap_growth, fibonacci reduces compilation time
156156
{Pid, Ref} = spawn_opt(
157157
fun() ->
158158
try

‎libs/jit/src/jit_aarch64.erl‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@
134134
stream :: stream(),
135135
offset :: non_neg_integer(),
136136
branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
137+
jump_table_start :: non_neg_integer(),
137138
available_regs :: [aarch64_register()],
138139
used_regs :: [aarch64_register()],
139140
labels :: [{integer() | reference(), integer()}],
@@ -233,6 +234,7 @@ new(Variant, StreamModule, Stream) ->
233234
stream_module = StreamModule,
234235
stream = Stream,
235236
branches = [],
237+
jump_table_start = 0,
236238
offset = StreamModule:offset(Stream),
237239
available_regs = ?AVAILABLE_REGS,
238240
used_regs = [],
@@ -355,22 +357,21 @@ assert_all_native_free(#state{
355357
%% @return Updated backend state
356358
%%-----------------------------------------------------------------------------
357359
-spec jump_table(state(), pos_integer()) -> state().
358-
jump_table(State, LabelsCount) ->
359-
jump_table0(State, 0, LabelsCount).
360+
jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
361+
JumpTableStart = StreamModule:offset(Stream0),
362+
jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
360363

361364
-spec jump_table0(state(), non_neg_integer(), pos_integer()) -> state().
362365
jump_table0(State, N, LabelsCount) when N > LabelsCount ->
363366
State;
364367
jump_table0(
365-
#state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
368+
#state{stream_module = StreamModule, stream = Stream0} = State,
366369
N,
367370
LabelsCount
368371
) ->
369-
Offset = StreamModule:offset(Stream0),
370372
BranchInstr = jit_aarch64_asm:b(0),
371-
Reloc = {N, Offset, b},
372373
Stream1 = StreamModule:append(Stream0, BranchInstr),
373-
jump_table0(State#state{stream = Stream1, branches = [Reloc | Branches]}, N + 1, LabelsCount).
374+
jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
374375

375376
%%-----------------------------------------------------------------------------
376377
%% @doc Rewrite stream to update all branches for labels.
@@ -2343,5 +2344,22 @@ add_label(#state{stream_module = StreamModule, stream = Stream} = State, Label)
23432344
%% @return Updated backend state
23442345
%%-----------------------------------------------------------------------------
23452346
-spec add_label(state(), integer() | reference(), integer()) -> state().
2347+
add_label(
2348+
#state{
2349+
stream_module = StreamModule,
2350+
stream = Stream0,
2351+
jump_table_start = JumpTableStart,
2352+
labels = Labels
2353+
} = State,
2354+
Label,
2355+
LabelOffset
2356+
) when is_integer(Label) ->
2357+
% Patch the jump table entry immediately
2358+
% Each b instruction is 4 bytes
2359+
JumpTableEntryOffset = JumpTableStart + Label * 4,
2360+
RelativeOffset = LabelOffset - JumpTableEntryOffset,
2361+
BranchInstr = jit_aarch64_asm:b(RelativeOffset),
2362+
Stream1 = StreamModule:replace(Stream0, JumpTableEntryOffset, BranchInstr),
2363+
State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
23462364
add_label(#state{labels = Labels} = State, Label, Offset) ->
23472365
State#state{labels = [{Label, Offset} | Labels]}.

‎libs/jit/src/jit_armv6m.erl‎

Lines changed: 37 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@
134134
stream :: stream(),
135135
offset :: non_neg_integer(),
136136
branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
137+
jump_table_start :: non_neg_integer(),
137138
available_regs :: [armv6m_register()],
138139
used_regs :: [armv6m_register()],
139140
labels :: [{integer() | reference(), integer()}],
@@ -247,6 +248,7 @@ new(Variant, StreamModule, Stream) ->
247248
stream_module = StreamModule,
248249
stream = Stream,
249250
branches = [],
251+
jump_table_start = 0,
250252
offset = StreamModule:offset(Stream),
251253
available_regs = ?AVAILABLE_REGS,
252254
used_regs = [],
@@ -380,13 +382,14 @@ assert_all_native_free(#state{
380382
%% @return Updated backend state
381383
%%-----------------------------------------------------------------------------
382384
-spec jump_table(state(), pos_integer()) -> state().
383-
jump_table(State, LabelsCount) ->
384-
jump_table0(State, 0, LabelsCount).
385+
jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
386+
JumpTableStart = StreamModule:offset(Stream0),
387+
jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
385388

386389
jump_table0(State, N, LabelsCount) when N > LabelsCount ->
387390
State;
388391
jump_table0(
389-
#state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
392+
#state{stream_module = StreamModule, stream = Stream0} = State,
390393
N,
391394
LabelsCount
392395
) ->
@@ -399,15 +402,7 @@ jump_table0(
399402
JumpEntry = <<I1/binary, I2/binary, I3/binary, I4/binary, 16#FFFFFFFF:32>>,
400403
Stream1 = StreamModule:append(Stream0, JumpEntry),
401404

402-
% Add relocation for the data entry so update_branches/2 can patch the jump target
403-
DataOffset = StreamModule:offset(Stream1) - 4,
404-
% Calculate the offset of the add instruction (3rd instruction, at offset 4 from entry start)
405-
EntryStartOffset = StreamModule:offset(Stream1) - 12,
406-
AddInstrOffset = EntryStartOffset + 4,
407-
DataReloc = {N, DataOffset, {jump_table_data, AddInstrOffset}},
408-
UpdatedState = State#state{stream = Stream1, branches = [DataReloc | Branches]},
409-
410-
jump_table0(UpdatedState, N + 1, LabelsCount).
405+
jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
411406

412407
%%-----------------------------------------------------------------------------
413408
%% @doc Rewrite stream to update all branches for labels.
@@ -500,13 +495,7 @@ update_branches(
500495
I4 = <<RelativeOffset:32/little>>,
501496
<<I1/binary, I2/binary, I3/binary, I4/binary>>
502497
end
503-
end;
504-
{jump_table_data, AddInstrOffset} ->
505-
% Calculate offset from 'add pc, pc, r3' instruction + 4 to target label
506-
% PC when add instruction executes
507-
AddPC = AddInstrOffset + 4,
508-
RelativeOffset = LabelOffset - AddPC,
509-
<<RelativeOffset:32/little>>
498+
end
510499
end,
511500
Stream1 = StreamModule:replace(Stream0, Offset, NewInstr),
512501
update_branches(State#state{stream = Stream1, branches = BranchesT}).
@@ -3288,5 +3277,34 @@ add_label(#state{stream_module = StreamModule, stream = Stream0} = State0, Label
32883277
%% @return Updated backend state
32893278
%%-----------------------------------------------------------------------------
32903279
-spec add_label(state(), integer() | reference(), integer()) -> state().
3280+
add_label(
3281+
#state{
3282+
stream_module = StreamModule,
3283+
stream = Stream0,
3284+
jump_table_start = JumpTableStart,
3285+
labels = Labels
3286+
} = State,
3287+
Label,
3288+
LabelOffset
3289+
) when is_integer(Label) ->
3290+
% Patch the jump table entry immediately
3291+
% Each jump table entry is 12 bytes:
3292+
% - ldr r3, [pc, 4] (2 bytes) at offset 0
3293+
% - push {...} (2 bytes) at offset 2
3294+
% - add pc, r3 (2 bytes) at offset 4
3295+
% - nop (2 bytes) at offset 6
3296+
% - data (4 bytes) at offset 8
3297+
JumpTableEntryStart = JumpTableStart + Label * 12,
3298+
DataOffset = JumpTableEntryStart + 8,
3299+
AddInstrOffset = JumpTableEntryStart + 4,
3300+
3301+
% Calculate offset from 'add pc, pc, r3' instruction + 4 to target label
3302+
% PC when add instruction executes
3303+
AddPC = AddInstrOffset + 4,
3304+
RelativeOffset = LabelOffset - AddPC,
3305+
DataBytes = <<RelativeOffset:32/little>>,
3306+
3307+
Stream1 = StreamModule:replace(Stream0, DataOffset, DataBytes),
3308+
State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
32913309
add_label(#state{labels = Labels} = State, Label, Offset) ->
32923310
State#state{labels = [{Label, Offset} | Labels]}.

‎libs/jit/src/jit_x86_64.erl‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@
115115
stream :: stream(),
116116
offset :: non_neg_integer(),
117117
branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
118+
jump_table_start :: non_neg_integer(),
118119
available_regs :: [x86_64_register()],
119120
used_regs :: [x86_64_register()],
120121
labels :: [{integer() | reference(), integer()}],
@@ -218,6 +219,7 @@ new(Variant, StreamModule, Stream) ->
218219
stream_module = StreamModule,
219220
stream = Stream,
220221
branches = [],
222+
jump_table_start = 0,
221223
offset = StreamModule:offset(Stream),
222224
available_regs = ?AVAILABLE_REGS,
223225
used_regs = [],
@@ -340,21 +342,21 @@ assert_all_native_free(State) ->
340342
%% @return Updated backend state
341343
%%-----------------------------------------------------------------------------
342344
-spec jump_table(state(), pos_integer()) -> state().
343-
jump_table(State, LabelsCount) ->
344-
jump_table0(State, 0, LabelsCount).
345+
jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
346+
JumpTableStart = StreamModule:offset(Stream0),
347+
jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
345348

346349
jump_table0(State, N, LabelsCount) when N > LabelsCount ->
347350
State;
348351
jump_table0(
349-
#state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
352+
#state{stream_module = StreamModule, stream = Stream0} = State,
350353
N,
351354
LabelsCount
352355
) ->
353-
Offset = StreamModule:offset(Stream0),
354-
{RelocOffset, I1} = jit_x86_64_asm:jmp_rel32(1),
355-
Reloc = {N, Offset + RelocOffset, 32},
356+
% Placeholder, encodes with 0xffffffff
357+
{_RelocOffset, I1} = jit_x86_64_asm:jmp_rel32(4),
356358
Stream1 = StreamModule:append(Stream0, I1),
357-
jump_table0(State#state{stream = Stream1, branches = [Reloc | Branches]}, N + 1, LabelsCount).
359+
jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
358360

359361
%%-----------------------------------------------------------------------------
360362
%% @doc Rewrite stream to update all branches for labels.
@@ -2086,5 +2088,22 @@ add_label(#state{stream_module = StreamModule, stream = Stream} = State, Label)
20862088
add_label(State, Label, Offset).
20872089

20882090
-spec add_label(state(), integer() | reference(), integer()) -> state().
2091+
add_label(
2092+
#state{
2093+
stream_module = StreamModule,
2094+
stream = Stream0,
2095+
jump_table_start = JumpTableStart,
2096+
labels = Labels
2097+
} = State,
2098+
Label,
2099+
LabelOffset
2100+
) when is_integer(Label) ->
2101+
% Patch the jump table entry immediately
2102+
% Each jmp_rel32 instruction is 5 bytes
2103+
JumpTableEntryOffset = JumpTableStart + Label * 5,
2104+
RelativeOffset = LabelOffset - JumpTableEntryOffset,
2105+
{_RelocOffset, JmpInstruction} = jit_x86_64_asm:jmp_rel32(RelativeOffset),
2106+
Stream1 = StreamModule:replace(Stream0, JumpTableEntryOffset, JmpInstruction),
2107+
State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
20892108
add_label(#state{labels = Labels} = State, Label, Offset) ->
20902109
State#state{labels = [{Label, Offset} | Labels]}.

0 commit comments

Comments
 (0)