JIT: add register value caching across all backends

Add register value tracking (jit_regs) to all JIT backends (x86-64,
aarch64, armv6m, riscv32). This enables optimization of redundant loads
when a register already contains the needed value.

Signed-off-by: Paul Guyot <pguyot@kallisys.net>

Author: pguyot

libs/jit/src/CMakeLists.txt

@@ -25,6 +25,7 @@ include(BuildErlang)
 set(ERLANG_MODULES
     jit
     jit_precompile
+    jit_regs
     jit_stream_binary
     jit_stream_flash
     jit_stream_mmap

libs/jit/src/jit_aarch64.erl

@@ -0,0 +1,197 @@
+%
+% This file is part of AtomVM.
+%
+% Copyright 2025-2026 Paul Guyot <pguyot@kallisys.net>
+%
+% Licensed under the Apache License, Version 2.0 (the "License");
+% you may not use this file except in compliance with the License.
+% You may obtain a copy of the License at
+%
+%    http://www.apache.org/licenses/LICENSE-2.0
+%
+% Unless required by applicable law or agreed to in writing, software
+% distributed under the License is distributed on an "AS IS" BASIS,
+% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+% See the License for the specific language governing permissions and
+% limitations under the License.
+%
+% SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
+%
+
+%% @doc Track CPU register contents for the JIT backend.
+%%
+%% This module maintains knowledge about what each CPU register holds
+%% (VM x/y register values, immediates, pointers), enabling the backend to
+%% skip redundant loads.
+%%
+%% Tracking must be invalidated at:
+%% - Labels (any offset can be a branch target)
+%% - Function calls (ABI clobbers caller-saved registers)
+%% - Branches (the target has unknown incoming register state)
+%%
+%% The tracked information includes:
+%% - `contents`: maps cpu_reg -> what the register holds
+
+-module(jit_regs).
+
+-export([
+    new/0,
+    get_contents/2,
+    set_contents/3,
+    invalidate_reg/2,
+    invalidate_all/1,
+    invalidate_volatile/2,
+    invalidate_vm_loc/2,
+    find_reg_with_contents/2,
+    merge/2,
+    stack_push/2,
+    stack_pop/1,
+    stack_clear/1,
+    stack_contents/1,
+    value_to_contents/2,
+    vm_dest_to_contents/2,
+    regs_to_mask/2
+]).
+
+-export_type([regs/0, contents/0]).
+
+-type vm_loc() ::
+    {x_reg, non_neg_integer()}
+    | {y_reg, non_neg_integer()}.
+
+%% What a CPU register holds
+-type contents() ::
+    vm_loc()
+    | {ptr, vm_loc()}
+    %% Register holds a known immediate value
+    | {imm, integer()}
+    %% Register holds the address of CP
+    | cp
+    %% Register holds the module index
+    | module_index
+    %% Unknown / clobbered
+    | unknown.
+
+-record(regs, {
+    contents = #{} :: #{atom() => contents()},
+    stack = [] :: [atom() | contents()]
+}).
+
+-opaque regs() :: #regs{}.
+
+%% @doc Create a new empty register tracking state.
+-spec new() -> regs().
+new() ->
+    #regs{}.
+
+%% @doc Get what a CPU register currently holds.
+-spec get_contents(regs(), atom()) -> contents().
+get_contents(#regs{contents = C}, Reg) ->
+    maps:get(Reg, C, unknown).
+
+%% @doc Record that a CPU register now holds the given contents.
+-spec set_contents(regs(), atom(), contents()) -> regs().
+set_contents(#regs{contents = C} = Regs, Reg, Contents) ->
+    Regs#regs{contents = C#{Reg => Contents}}.
+
+%% @doc Invalidate tracking for a single CPU register (e.g. it was clobbered).
+-spec invalidate_reg(regs(), atom()) -> regs().
+invalidate_reg(#regs{contents = C} = Regs, Reg) ->
+    Regs#regs{contents = maps:remove(Reg, C)}.
+
+%% @doc Invalidate all register tracking (e.g. at a label or unknown branch target).
+-spec invalidate_all(regs()) -> regs().
+invalidate_all(Regs) ->
+    Regs#regs{contents = #{}, stack = []}.
+
+%% @doc Invalidate registers that are volatile across a C function call.
+%% On x86-64 System V ABI, all our scratch registers (rax, rcx, rdx, rsi, rdi,
+%% r8, r9, r10, r11) are caller-saved, so after a C call they're all clobbered.
+%% However, the special registers (rdi=ctx, rsi=jit_state, rdx=native_interface)
+%% are restored by the JIT after the call via push/pop, so we keep their tracking.
+-spec invalidate_volatile(regs(), [atom()]) -> regs().
+invalidate_volatile(#regs{contents = C0} = Regs, PreservedRegs) ->
+    C1 = maps:filter(fun(Reg, _) -> lists:member(Reg, PreservedRegs) end, C0),
+    Regs#regs{contents = C1}.
+
+%% @doc Invalidate all CPU registers that reference a given VM location.
+%% Call this when a VM register is written to, so that any CPU register
+%% that was caching its old value is invalidated.
+-spec invalidate_vm_loc(regs(), vm_loc()) -> regs().
+invalidate_vm_loc(#regs{contents = C} = Regs, VmLoc) ->
+    C1 = maps:filter(fun(_Reg, Val) -> Val =/= VmLoc end, C),
+    Regs#regs{contents = C1}.
+
+%% @doc Find a CPU register that holds the given contents.
+%% Returns `{ok, Reg}` or `none`.
+-spec find_reg_with_contents(regs(), contents()) -> {ok, atom()} | none.
+find_reg_with_contents(#regs{contents = C}, Contents) ->
+    find_in_map(maps:iterator(C), Contents).
+
+find_in_map(Iterator, Contents) ->
+    case maps:next(Iterator) of
+        {Reg, Contents, _Next} -> {ok, Reg};
+        {_Reg, _Other, Next} -> find_in_map(Next, Contents);
+        none -> none
+    end.
+
+%% @doc Merge two register tracking states (for control flow merge points).
+%% Only keeps information that is consistent in both states.
+-spec merge(regs(), regs()) -> regs().
+merge(#regs{contents = C1}, #regs{contents = C2}) ->
+    %% Keep only entries that match in both maps
+    MergedContents = maps:filter(
+        fun(Reg, Val) -> maps:get(Reg, C2, undefined) =:= Val end,
+        C1
+    ),
+    #regs{contents = MergedContents, stack = []}.
+
+%% @doc Record a push to the C stack.
+-spec stack_push(regs(), atom() | contents()) -> regs().
+stack_push(#regs{stack = S} = Regs, Value) ->
+    Regs#regs{stack = [Value | S]}.
+
+%% @doc Record a pop from the C stack.
+-spec stack_pop(regs()) -> {atom() | contents(), regs()}.
+stack_pop(#regs{stack = [Top | Rest]} = Regs) ->
+    {Top, Regs#regs{stack = Rest}};
+stack_pop(#regs{stack = []} = Regs) ->
+    {unknown, Regs}.
+
+%% @doc Clear the C stack tracking.
+-spec stack_clear(regs()) -> regs().
+stack_clear(Regs) ->
+    Regs#regs{stack = []}.
+
+%% @doc Get the current C stack contents.
+-spec stack_contents(regs()) -> [atom() | contents()].
+stack_contents(#regs{stack = S}) -> S.
+
+%% @doc Convert a backend value to a contents descriptor for tracking.
+%% MaxReg is the maximum number of x registers (typically ?MAX_REG from jit.hrl).
+-spec value_to_contents(term(), non_neg_integer()) -> contents().
+value_to_contents(cp, _MaxReg) -> cp;
+value_to_contents({x_reg, N}, _MaxReg) when is_integer(N) -> {x_reg, N};
+value_to_contents({x_reg, extra}, MaxReg) -> {x_reg, MaxReg};
+value_to_contents({y_reg, N}, _MaxReg) -> {y_reg, N};
+value_to_contents(Imm, _MaxReg) when is_integer(Imm) -> {imm, Imm};
+value_to_contents({ptr, _}, _MaxReg) -> unknown;
+value_to_contents(_, _MaxReg) -> unknown.
+
+%% @doc Convert a VM destination register to a contents descriptor for tracking.
+%% MaxReg is the maximum number of x registers (typically ?MAX_REG from jit.hrl).
+-spec vm_dest_to_contents(term(), non_neg_integer()) -> contents().
+vm_dest_to_contents({x_reg, X}, MaxReg) when is_integer(X), X < MaxReg -> {x_reg, X};
+vm_dest_to_contents({x_reg, extra}, MaxReg) -> {x_reg, MaxReg};
+vm_dest_to_contents({y_reg, Y}, _MaxReg) -> {y_reg, Y};
+vm_dest_to_contents(_, _MaxReg) -> unknown.
+
+%% @doc Convert a list of register atoms to a bitmask.
+%% Skips non-register entries like `imm`, `jit_state`, and `stack`.
+%% RegBitFn maps register atoms to their bit positions.
+-spec regs_to_mask([atom()], fun((atom()) -> non_neg_integer())) -> non_neg_integer().
+regs_to_mask([], _RegBitFn) -> 0;
+regs_to_mask([imm | T], RegBitFn) -> regs_to_mask(T, RegBitFn);
+regs_to_mask([jit_state | T], RegBitFn) -> regs_to_mask(T, RegBitFn);
+regs_to_mask([stack | T], RegBitFn) -> regs_to_mask(T, RegBitFn);
+regs_to_mask([Reg | T], RegBitFn) -> RegBitFn(Reg) bor regs_to_mask(T, RegBitFn).