AArch64: Fix bug in remaining reductions handling

Signed-off-by: Paul Guyot <[email protected]>

Author: pguyot

libs/jit/src/jit_aarch64.erl

@@ -1836,11 +1836,11 @@ decrement_reductions_and_maybe_schedule_next(
     #state{stream_module = StreamModule, stream = Stream0, available_regs = [Temp | _]} = State0
 ) ->
     % Load reduction count
-    I1 = jit_aarch64_asm:ldr(Temp, ?JITSTATE_REDUCTIONCOUNT),
+    I1 = jit_aarch64_asm:ldr_w(Temp, ?JITSTATE_REDUCTIONCOUNT),
     % Decrement reduction count
     I2 = jit_aarch64_asm:subs(Temp, Temp, 1),
     % Store back the decremented value
-    I3 = jit_aarch64_asm:str(Temp, ?JITSTATE_REDUCTIONCOUNT),
+    I3 = jit_aarch64_asm:str_w(Temp, ?JITSTATE_REDUCTIONCOUNT),
     Stream1 = StreamModule:append(Stream0, <<I1/binary, I2/binary, I3/binary>>),
     BNEOffset = StreamModule:offset(Stream1),
     % Branch if reduction count is not zero
@@ -1879,11 +1879,11 @@ call_only_or_schedule_next(
     Label
 ) ->
     % Load reduction count
-    I1 = jit_aarch64_asm:ldr(Temp, ?JITSTATE_REDUCTIONCOUNT),
+    I1 = jit_aarch64_asm:ldr_w(Temp, ?JITSTATE_REDUCTIONCOUNT),
     % Decrement reduction count
     I2 = jit_aarch64_asm:subs(Temp, Temp, 1),
     % Store back the decremented value
-    I3 = jit_aarch64_asm:str(Temp, ?JITSTATE_REDUCTIONCOUNT),
+    I3 = jit_aarch64_asm:str_w(Temp, ?JITSTATE_REDUCTIONCOUNT),
     Stream1 = StreamModule:append(Stream0, <<I1/binary, I2/binary, I3/binary>>),
     BNEOffset = StreamModule:offset(Stream1),
     % Branch to label if reduction count is not zero

libs/jit/src/jit_aarch64_asm.erl

@@ -46,6 +46,7 @@
     ret/0,
     nop/0,
     str/2,
+    str_w/2,
     str/3,
     tst/2,
     tst_w/2,
@@ -225,13 +226,11 @@ ldr_w(Dst, {BaseReg, Offset}) when
     is_atom(BaseReg),
     is_integer(Offset),
     Offset >= 0,
-    Offset =< 32760,
-    (Offset rem 8) =:= 0
+    Offset =< 16380,
+    (Offset rem 4) =:= 0
 ->
     DstNum = reg_to_num(Dst),
     BaseRegNum = reg_to_num(BaseReg),
-    %% AArch64 LDR (immediate) encoding for 64-bit: 11111001010iiiiiiiiiiibbbbbttttt
-    %% 0xf9400000 | (Offset div 8) << 10 | BaseReg << 5 | Dst
     <<
         (16#B9400000 bor ((Offset div 4) bsl 10) bor (BaseRegNum bsl 5) bor DstNum):32/little
     >>.
@@ -588,6 +587,22 @@ str(Reg, {Base}, Imm) when
     BaseNum = reg_to_num(Base),
     <<(16#F8000400 bor ((Imm band 16#1FF) bsl 12) bor (BaseNum bsl 5) bor RegNum):32/little>>.
 
+%% Emit a store register (STR) instruction for 32-bit store to memory
+-spec str_w(aarch64_gpr_register(), {aarch64_gpr_register(), integer()}) -> binary().
+str_w(Src, {BaseReg, Offset}) when
+    is_atom(Src),
+    is_atom(BaseReg),
+    is_integer(Offset),
+    Offset >= 0,
+    Offset =< 16380,
+    (Offset rem 4) =:= 0
+->
+    SrcNum = reg_to_num(Src),
+    BaseRegNum = reg_to_num(BaseReg),
+    <<
+        (16#B9000000 bor ((Offset div 4) bsl 10) bor (BaseRegNum bsl 5) bor SrcNum):32/little
+    >>.
+
 %% Emit a load register (LDR) instruction for 64-bit store to memory, with store-update (writeback)
 -spec ldr
     (aarch64_gpr_register(), {aarch64_gpr_register(), integer()}, '!') -> binary();

tests/libs/jit/jit_aarch64_asm_tests.erl

@@ -139,8 +139,30 @@ ldr_test_() ->
 
 ldr_w_test_() ->
     [
-        ?_assertEqual(<<16#b9400821:32/little>>, jit_aarch64_asm:ldr_w(r1, {r1, 8})),
-        ?_assertEqual(<<16#b9406042:32/little>>, jit_aarch64_asm:ldr_w(r2, {r2, 96}))
+        ?_assertEqual(
+            asm(<<16#b9400821:32/little>>, "ldr w1, [x1, 8]"), jit_aarch64_asm:ldr_w(r1, {r1, 8})
+        ),
+        ?_assertEqual(
+            asm(<<16#b9406042:32/little>>, "ldr w2, [x2, 96]"), jit_aarch64_asm:ldr_w(r2, {r2, 96})
+        ),
+        ?_assertEqual(
+            asm(<<16#b97ffc60:32/little>>, "ldr w0, [x3, 16380]"),
+            jit_aarch64_asm:ldr_w(r0, {r3, 16380})
+        )
+    ].
+
+str_w_test_() ->
+    [
+        ?_assertEqual(
+            asm(<<16#b9000821:32/little>>, "str w1, [x1, 8]"), jit_aarch64_asm:str_w(r1, {r1, 8})
+        ),
+        ?_assertEqual(
+            asm(<<16#b9006042:32/little>>, "str w2, [x2, 96]"), jit_aarch64_asm:str_w(r2, {r2, 96})
+        ),
+        ?_assertEqual(
+            asm(<<16#b93ffc60:32/little>>, "str w0, [x3, 16380]"),
+            jit_aarch64_asm:str_w(r0, {r3, 16380})
+        )
     ].
 
 ldr_d_test_() ->