AArch64: increase coverage and use cbz/tbz/tbnz

pguyot · pguyot · commit de9c67130e77 · 2025-08-23T13:18:11.000+02:00
Also simplify {RegOrTuple, '&amp;', Val, '!=', 0} when possible

Signed-off-by: Paul Guyot &lt;pguyot@kallisys.net&gt;
diff --git a/libs/jit/src/jit_aarch64.erl b/libs/jit/src/jit_aarch64.erl
@@ -567,7 +567,19 @@ if_block(
     OffsetAfter = StreamModule:offset(Stream2),
     %% Patch the conditional branch instruction to jump to the end of the block
     BranchOffset = OffsetAfter - (Offset + BranchInstrOffset),
-    NewBranchInstr = jit_aarch64_asm:bcc(CC, BranchOffset),
+    NewBranchInstr =
+        case CC of
+            {cbz, Reg} ->
+                jit_aarch64_asm:cbz(Reg, BranchOffset);
+            {cbz_w, Reg} ->
+                jit_aarch64_asm:cbz_w(Reg, BranchOffset);
+            {tbz, Reg, Bit} ->
+                jit_aarch64_asm:tbz(Reg, Bit, BranchOffset);
+            {tbnz, Reg, Bit} ->
+                jit_aarch64_asm:tbnz(Reg, Bit, BranchOffset);
+            CC when is_atom(CC) ->
+                jit_aarch64_asm:bcc(CC, BranchOffset)
+        end,
     Stream3 = StreamModule:replace(Stream2, Offset + BranchInstrOffset, NewBranchInstr),
     merge_used_regs(State2#state{stream = Stream3}, State1#state.used_regs).
 
@@ -601,7 +613,19 @@ if_else_block(
     OffsetAfter = StreamModule:offset(Stream3),
     %% Patch the conditional branch to jump to the else block
     ElseBranchOffset = OffsetAfter - (Offset + BranchInstrOffset),
-    NewBranchInstr = jit_aarch64_asm:bcc(CC, ElseBranchOffset),
+    NewBranchInstr =
+        case CC of
+            {cbz, Reg} ->
+                jit_aarch64_asm:cbz(Reg, ElseBranchOffset);
+            {cbz_w, Reg} ->
+                jit_aarch64_asm:cbz_w(Reg, ElseBranchOffset);
+            {tbz, Reg, Bit} ->
+                jit_aarch64_asm:tbz(Reg, Bit, ElseBranchOffset);
+            {tbnz, Reg, Bit} ->
+                jit_aarch64_asm:tbnz(Reg, Bit, ElseBranchOffset);
+            CC when is_atom(CC) ->
+                jit_aarch64_asm:bcc(CC, ElseBranchOffset)
+        end,
     Stream4 = StreamModule:replace(Stream3, Offset + BranchInstrOffset, NewBranchInstr),
     %% Build the else block
     StateElse = State2#state{
@@ -621,16 +645,10 @@ if_else_block(
 
 -spec if_block_cond(state(), condition()) -> {state(), jit_aarch64_asm:cc(), non_neg_integer()}.
 if_block_cond(#state{stream_module = StreamModule, stream = Stream0} = State0, {Reg, '<', 0}) ->
-    I1 = jit_aarch64_asm:tst(Reg, Reg),
-    % pl = positive or zero (>=0)
-    I2 = jit_aarch64_asm:bcc(pl, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
+    I = jit_aarch64_asm:tbz(Reg, 63, 0),
+    Stream1 = StreamModule:append(Stream0, I),
     State1 = State0#state{stream = Stream1},
-    {State1, pl, byte_size(I1)};
+    {State1, {tbz, Reg, 63}, 0};
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0,
     {RegA, '<', RegB}
@@ -653,17 +671,11 @@ if_block_cond(
             {free, Reg0} -> Reg0;
             RegOrTuple -> RegOrTuple
         end,
-    I1 = jit_aarch64_asm:tst(Reg, Reg),
-    % ne = not equal
-    I2 = jit_aarch64_asm:bcc(ne, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
+    I = jit_aarch64_asm:cbz(Reg, 0),
+    Stream1 = StreamModule:append(Stream0, I),
     State1 = if_block_free_reg(RegOrTuple, State0),
     State2 = State1#state{stream = Stream1},
-    {State2, ne, byte_size(I1)};
+    {State2, {cbz, Reg}, 0};
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0, {'(int)', RegOrTuple, '==', 0}
 ) ->
@@ -672,16 +684,11 @@ if_block_cond(
             {free, Reg0} -> Reg0;
             RegOrTuple -> RegOrTuple
         end,
-    I1 = jit_aarch64_asm:tst_w(Reg, Reg),
-    I2 = jit_aarch64_asm:bcc(ne, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
+    I = jit_aarch64_asm:cbz_w(Reg, 0),
+    Stream1 = StreamModule:append(Stream0, I),
     State1 = if_block_free_reg(RegOrTuple, State0),
     State2 = State1#state{stream = Stream1},
-    {State2, ne, byte_size(I1)};
+    {State2, {cbz_w, Reg}, 0};
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0,
     {RegOrTuple, '!=', Val}
@@ -783,17 +790,12 @@ if_block_cond(
             {free, Reg0} -> Reg0;
             RegOrTuple -> RegOrTuple
         end,
-    % Test low 8 bits
-    I1 = jit_aarch64_asm:tst_w(Reg, 16#FF),
-    I2 = jit_aarch64_asm:bcc(ne, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
+    % Test lowest bit
+    I = jit_aarch64_asm:tbnz(Reg, 0, 0),
+    Stream1 = StreamModule:append(Stream0, I),
     State1 = if_block_free_reg(RegOrTuple, State0),
     State2 = State1#state{stream = Stream1},
-    {State2, ne, byte_size(I1)};
+    {State2, {tbnz, Reg, 0}, 0};
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0,
     {'(bool)', RegOrTuple, '!=', false}
@@ -803,17 +805,44 @@ if_block_cond(
             {free, Reg0} -> Reg0;
             RegOrTuple -> RegOrTuple
         end,
-    % Test low 8 bits
-    I1 = jit_aarch64_asm:tst_w(Reg, 16#FF),
+    % Test lowest bit
+    I = jit_aarch64_asm:tbz(Reg, 0, 0),
+    Stream1 = StreamModule:append(Stream0, I),
+    State1 = if_block_free_reg(RegOrTuple, State0),
+    State2 = State1#state{stream = Stream1},
+    {State2, {tbz, Reg, 0}, 0};
+if_block_cond(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        available_regs = [Temp | _]
+    } = State0,
+    {RegOrTuple, '&', Val, '!=', 0}
+) ->
+    Reg =
+        case RegOrTuple of
+            {free, Reg0} -> Reg0;
+            RegOrTuple -> RegOrTuple
+        end,
+    % Test bits
+    TestCode =
+        try
+            jit_aarch64_asm:tst(Reg, Val)
+        catch
+            error:{unencodable_immediate, Val} ->
+                TestCode0 = jit_aarch64_asm:mov(Temp, Val),
+                TestCode1 = jit_aarch64_asm:tst(Reg, Temp),
+                <<TestCode0/binary, TestCode1/binary>>
+        end,
     I2 = jit_aarch64_asm:bcc(eq, 0),
     Code = <<
-        I1/binary,
+        TestCode/binary,
         I2/binary
     >>,
     Stream1 = StreamModule:append(Stream0, Code),
     State1 = if_block_free_reg(RegOrTuple, State0),
     State2 = State1#state{stream = Stream1},
-    {State2, eq, byte_size(I1)};
+    {State2, eq, byte_size(TestCode)};
 if_block_cond(
     #state{
         stream_module = StreamModule,
@@ -853,53 +882,7 @@ if_block_cond(
     Stream3 = StreamModule:append(Stream2, I3),
     State3 = State1#state{stream = Stream3},
     State4 = if_block_free_reg(RegTuple, State3),
-    {State4, eq, OffsetAfter - OffsetBefore};
-if_block_cond(
-    #state{
-        stream_module = StreamModule,
-        stream = Stream0
-    } = State0,
-    {RegOrTuple, '&', Val}
-) ->
-    Reg =
-        case RegOrTuple of
-            {free, Reg0} -> Reg0;
-            RegOrTuple -> RegOrTuple
-        end,
-    % Test bits
-    I1 = jit_aarch64_asm:tst(Reg, Val),
-    I2 = jit_aarch64_asm:bcc(eq, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
-    State1 = if_block_free_reg(RegOrTuple, State0),
-    State2 = State1#state{stream = Stream1},
-    {State2, eq, byte_size(I1)};
-if_block_cond(
-    #state{
-        stream_module = StreamModule,
-        stream = Stream0
-    } = State0,
-    {'(bool)', RegOrTuple, '&', Val}
-) ->
-    Reg =
-        case RegOrTuple of
-            {free, Reg0} -> Reg0;
-            RegOrTuple -> RegOrTuple
-        end,
-    % Test 8-bit value
-    I1 = jit_aarch64_asm:tst_w(Reg, Val),
-    I2 = jit_aarch64_asm:bcc(eq, 0),
-    Code = <<
-        I1/binary,
-        I2/binary
-    >>,
-    Stream1 = StreamModule:append(Stream0, Code),
-    State1 = if_block_free_reg(RegOrTuple, State0),
-    State2 = State1#state{stream = Stream1},
-    {State2, eq, byte_size(I1)}.
+    {State4, eq, OffsetAfter - OffsetBefore}.
 
 -spec if_block_free_reg(aarch64_register() | {free, aarch64_register()}, state()) -> state().
 if_block_free_reg({free, Reg}, State0) ->
diff --git a/libs/jit/src/jit_aarch64_asm.erl b/libs/jit/src/jit_aarch64_asm.erl
@@ -30,6 +30,10 @@
     blr/1,
     br/1,
     brk/1,
+    cbz/2,
+    cbz_w/2,
+    tbz/3,
+    tbnz/3,
     cmp/2,
     cmp32/2,
     and_/3,
@@ -779,6 +783,39 @@ bcc(Cond, Offset) when is_atom(Cond), is_integer(Offset) ->
     Offset19 = Offset div 4,
     <<(16#54000000 bor ((Offset19 band 16#7FFFF) bsl 5) bor CondNum):32/little>>.
 
+%% Emit a compare and branch on zero
+-spec cbz(aarch64_gpr_register(), integer()) -> binary().
+cbz(Rt, Offset) when is_integer(Offset) ->
+    RtNum = reg_to_num(Rt),
+    Offset19 = Offset div 4,
+    <<(16#B4000000 bor ((Offset19 band 16#7FFFF) bsl 5) bor RtNum):32/little>>.
+
+-spec cbz_w(aarch64_gpr_register(), integer()) -> binary().
+cbz_w(Rt, Offset) when is_integer(Offset) ->
+    RtNum = reg_to_num(Rt),
+    Offset19 = Offset div 4,
+    <<(16#34000000 bor ((Offset19 band 16#7FFFF) bsl 5) bor RtNum):32/little>>.
+
+%% Emit a test bit and branch if zero
+-spec tbz(aarch64_gpr_register(), 0..63, integer()) -> binary().
+tbz(Rt, Bit, Offset) when Offset >= -32768 andalso Offset < 32768 ->
+    RtNum = reg_to_num(Rt),
+    Offset14 = Offset div 4,
+    <<
+        ((Bit band 32 bsl 26) bor 16#36000000 bor (Bit band 31 bsl 19) bor
+            ((Offset14 band 16#3FFF) bsl 5) bor RtNum):32/little
+    >>.
+
+%% Emit a test bit and branch if not zero
+-spec tbnz(aarch64_gpr_register(), 0..63, integer()) -> binary().
+tbnz(Rt, Bit, Offset) when Offset >= -32768 andalso Offset < 32768 ->
+    RtNum = reg_to_num(Rt),
+    Offset14 = Offset div 4,
+    <<
+        ((Bit band 32 bsl 26) bor 16#37000000 bor (Bit band 31 bsl 19) bor
+            ((Offset14 band 16#3FFF) bsl 5) bor RtNum):32/little
+    >>.
+
 %% Emit a compare instruction
 -spec cmp(aarch64_gpr_register(), aarch64_gpr_register() | integer()) -> binary().
 cmp(Rn, Rm) when is_atom(Rn), is_atom(Rm) ->
diff --git a/tests/libs/jit/jit_aarch64_asm_tests.erl b/tests/libs/jit/jit_aarch64_asm_tests.erl
@@ -389,6 +389,54 @@ bcc_test_() ->
         )
     ].
 
+cbz_test_() ->
+    [
+        ?_assertEqual(
+            asm(<<16#b4000401:32/little>>, "cbz x1, 128"),
+            jit_aarch64_asm:cbz(r1, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#34000402:32/little>>, "cbz w2, 128"),
+            jit_aarch64_asm:cbz_w(r2, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#b4fffc03:32/little>>, "cbz x3, -128"),
+            jit_aarch64_asm:cbz(r3, -128)
+        )
+    ].
+
+tbz_test_() ->
+    [
+        ?_assertEqual(
+            asm(<<16#b6f80400:32/little>>, "tbz x0, #63, 128"),
+            jit_aarch64_asm:tbz(r0, 63, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#36180400:32/little>>, "tbz x0, #3, 128"),
+            jit_aarch64_asm:tbz(r0, 3, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#363ffc03:32/little>>, "tbz x3, #7, -128"),
+            jit_aarch64_asm:tbz(r3, 7, -128)
+        )
+    ].
+
+tbnz_test_() ->
+    [
+        ?_assertEqual(
+            asm(<<16#37000400:32/little>>, "tbnz x0, #0, 128"),
+            jit_aarch64_asm:tbnz(r0, 0, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#37180400:32/little>>, "tbnz x0, #3, 128"),
+            jit_aarch64_asm:tbnz(r0, 3, 128)
+        ),
+        ?_assertEqual(
+            asm(<<16#373ffc03:32/little>>, "tbnz x3, #7, -128"),
+            jit_aarch64_asm:tbnz(r3, 7, -128)
+        )
+    ].
+
 stp_test_() ->
     [
         ?_assertEqual(
diff --git a/tests/libs/jit/jit_aarch64_tests.erl b/tests/libs/jit/jit_aarch64_tests.erl
