ARM64: Fix exit stub patching.

Contributed by Javier Guerra Giraldez.

(cherry picked from commit 9da0653)

When the side trace is assembled, it is linked to its parent trace. For
this purpose, JIT runs through the parent trace mcode and updates jump
instruction targeted to the corresponding exitno. Prior to this patch,
these instructions were patched unconditionally, that leads to errors if
the jump target address is out of the value ranges specified in ARM64
references[1][2][3][4][5][6].

As a result of the patch <lj_asm_patchexit> considers value ranges of
the jump targets and updates directly only those instructions fitting
the particular jump range. Moreover, the corresponding jump in the pad
leading to <lj_vm_exit_handler> is also patched, so those instructions,
that are not updated before, targets to the linked side trace too.

Additionally, there is some refactoring of jump targets assembling in
scope of this patch.

Igor Munkin:
* added the description and the test for the problem

[1]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/B
[2]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/B-cond
[3]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/CBZ
[4]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/CBNZ
[5]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/TBZ
[6]: https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/TBNZ

Resolves tarantool/tarantool#6098
Part of tarantool/tarantool#5629

Reviewed-by: Sergey Kaplun <skaplun@tarantool.org>
Reviewed-by: Kirill Yukhin <kyukhin@tarantool.org>
Signed-off-by: Igor Munkin <imun@tarantool.org>

Author: Mike Pall

src/lj_asm_arm64.h

@@ -56,11 +56,11 @@ static void asm_exitstub_setup(ASMState *as, ExitNo nexits)
     asm_mclimit(as);
   /* 1: str lr,[sp]; bl ->vm_exit_handler; movz w0,traceno; bl <1; bl <1; ... */
   for (i = nexits-1; (int32_t)i >= 0; i--)
-    *--mxp = A64I_LE(A64I_BL|((-3-i)&0x03ffffffu));
-  *--mxp = A64I_LE(A64I_MOVZw|A64F_U16(as->T->traceno));
+    *--mxp = A64I_LE(A64I_BL | A64F_S26(-3-i));
+  *--mxp = A64I_LE(A64I_MOVZw | A64F_U16(as->T->traceno));
   mxp--;
-  *mxp = A64I_LE(A64I_BL|(((MCode *)(void *)lj_vm_exit_handler-mxp)&0x03ffffffu));
-  *--mxp = A64I_LE(A64I_STRx|A64F_D(RID_LR)|A64F_N(RID_SP));
+  *mxp = A64I_LE(A64I_BL | A64F_S26(((MCode *)(void *)lj_vm_exit_handler-mxp)));
+  *--mxp = A64I_LE(A64I_STRx | A64F_D(RID_LR) | A64F_N(RID_SP));
   as->mctop = mxp;
 }
 
@@ -77,7 +77,7 @@ static void asm_guardcc(ASMState *as, A64CC cc)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_cond_branch(as, cc^1, p-1);
     return;
   }
@@ -91,7 +91,7 @@ static void asm_guardtnb(ASMState *as, A64Ins ai, Reg r, uint32_t bit)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_tnb(as, ai^0x01000000u, r, bit, p-1);
     return;
   }
@@ -105,7 +105,7 @@ static void asm_guardcnb(ASMState *as, A64Ins ai, Reg r)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_cnb(as, ai^0x01000000u, r, p-1);
     return;
   }
@@ -1858,7 +1858,7 @@ static void asm_loop_fixup(ASMState *as)
     p[-2] |= ((uint32_t)delta & mask) << 5;
   } else {
     ptrdiff_t delta = target - (p - 1);
-    p[-1] = A64I_B | ((uint32_t)(delta) & 0x03ffffffu);
+    p[-1] = A64I_B | A64F_S26(delta);
   }
 }
 
@@ -1927,7 +1927,7 @@ static void asm_tail_fixup(ASMState *as, TraceNo lnk)
   }
   /* Patch exit branch. */
   target = lnk ? traceref(as->J, lnk)->mcode : (MCode *)lj_vm_exit_interp;
-  p[-1] = A64I_B | (((target-p)+1)&0x03ffffffu);
+  p[-1] = A64I_B | A64F_S26((target-p)+1);
 }
 
 /* Prepare tail of code. */
@@ -1990,40 +1990,50 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
 {
   MCode *p = T->mcode;
   MCode *pe = (MCode *)((char *)p + T->szmcode);
-  MCode *cstart = NULL, *cend = p;
+  MCode *cstart = NULL;
   MCode *mcarea = lj_mcode_patch(J, p, 0);
   MCode *px = exitstub_trace_addr(T, exitno);
+  /* Note: this assumes a trace exit is only ever patched once. */
   for (; p < pe; p++) {
     /* Look for exitstub branch, replace with branch to target. */
+    ptrdiff_t delta = target - p;
     MCode ins = A64I_LE(*p);
     if ((ins & 0xff000000u) == 0x54000000u &&
 	((ins ^ ((px-p)<<5)) & 0x00ffffe0u) == 0) {
-      /* Patch bcc exitstub. */
-      *p = A64I_LE((ins & 0xff00001fu) | (((target-p)<<5) & 0x00ffffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch bcc, if within range. */
+      if (A64F_S_OK(delta, 19)) {
+	*p = A64I_LE((ins & 0xff00001fu) | A64F_S19(delta));
+	if (!cstart) cstart = p;
+      }
     } else if ((ins & 0xfc000000u) == 0x14000000u &&
 	       ((ins ^ (px-p)) & 0x03ffffffu) == 0) {
-      /* Patch b exitstub. */
-      *p = A64I_LE((ins & 0xfc000000u) | ((target-p) & 0x03ffffffu));
-      cend = p+1;
+      /* Patch b. */
+      lua_assert(A64F_S_OK(delta, 26));
+      *p = A64I_LE((ins & 0xfc000000u) | A64F_S26(delta));
       if (!cstart) cstart = p;
     } else if ((ins & 0x7e000000u) == 0x34000000u &&
 	       ((ins ^ ((px-p)<<5)) & 0x00ffffe0u) == 0) {
-      /* Patch cbz/cbnz exitstub. */
-      *p = A64I_LE((ins & 0xff00001f) | (((target-p)<<5) & 0x00ffffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch cbz/cbnz, if within range. */
+      if (A64F_S_OK(delta, 19)) {
+	*p = A64I_LE((ins & 0xff00001fu) | A64F_S19(delta));
+	if (!cstart) cstart = p;
+      }
     } else if ((ins & 0x7e000000u) == 0x36000000u &&
 	       ((ins ^ ((px-p)<<5)) & 0x0007ffe0u) == 0) {
-      /* Patch tbz/tbnz exitstub. */
-      *p = A64I_LE((ins & 0xfff8001fu) | (((target-p)<<5) & 0x0007ffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch tbz/tbnz, if within range. */
+      if (A64F_S_OK(delta, 14)) {
+	*p = A64I_LE((ins & 0xfff8001fu) | A64F_S14(delta));
+	if (!cstart) cstart = p;
+      }
     }
   }
-  lua_assert(cstart != NULL);
-  lj_mcode_sync(cstart, cend);
+  {  /* Always patch long-range branch in exit stub itself. */
+    ptrdiff_t delta = target - px;
+    lua_assert(A64F_S_OK(delta, 26));
+    *px = A64I_B | A64F_S26(delta);
+    if (!cstart) cstart = px;
+  }
+  lj_mcode_sync(cstart, px+1);
   lj_mcode_patch(J, mcarea, 1);
 }
 

src/lj_emit_arm64.h

@@ -241,7 +241,7 @@ static void emit_loadk(ASMState *as, Reg rd, uint64_t u64, int is64)
 #define mcpofs(as, k) \
   ((intptr_t)((uintptr_t)(k) - (uintptr_t)(as->mcp - 1)))
 #define checkmcpofs(as, k) \
-  ((((mcpofs(as, k)>>2) + 0x00040000) >> 19) == 0)
+  (A64F_S_OK(mcpofs(as, k)>>2, 19))
 
 static Reg ra_allock(ASMState *as, intptr_t k, RegSet allow);
 
@@ -312,32 +312,32 @@ static void emit_cond_branch(ASMState *as, A64CC cond, MCode *target)
 {
   MCode *p = --as->mcp;
   ptrdiff_t delta = target - p;
-  lua_assert(((delta + 0x40000) >> 19) == 0);
+  lua_assert(A64F_S_OK(delta, 19));
   *p = A64I_BCC | A64F_S19(delta) | cond;
 }
 
 static void emit_branch(ASMState *as, A64Ins ai, MCode *target)
 {
   MCode *p = --as->mcp;
   ptrdiff_t delta = target - p;
-  lua_assert(((delta + 0x02000000) >> 26) == 0);
-  *p = ai | ((uint32_t)delta & 0x03ffffffu);
+  lua_assert(A64F_S_OK(delta, 26));
+  *p = ai | A64F_S26(delta);
 }
 
 static void emit_tnb(ASMState *as, A64Ins ai, Reg r, uint32_t bit, MCode *target)
 {
   MCode *p = --as->mcp;
   ptrdiff_t delta = target - p;
-  lua_assert(bit < 63 && ((delta + 0x2000) >> 14) == 0);
+  lua_assert(bit < 63 && A64F_S_OK(delta, 14));
   if (bit > 31) ai |= A64I_X;
-  *p = ai | A64F_BIT(bit & 31) | A64F_S14((uint32_t)delta & 0x3fffu) | r;
+  *p = ai | A64F_BIT(bit & 31) | A64F_S14(delta) | r;
 }
 
 static void emit_cnb(ASMState *as, A64Ins ai, Reg r, MCode *target)
 {
   MCode *p = --as->mcp;
   ptrdiff_t delta = target - p;
-  lua_assert(((delta + 0x40000) >> 19) == 0);
+  lua_assert(A64F_S_OK(delta, 19));
   *p = ai | A64F_S19(delta) | r;
 }
 
@@ -347,8 +347,8 @@ static void emit_call(ASMState *as, void *target)
 {
   MCode *p = --as->mcp;
   ptrdiff_t delta = (char *)target - (char *)p;
-  if ((((delta>>2) + 0x02000000) >> 26) == 0) {
-    *p = A64I_BL | ((uint32_t)(delta>>2) & 0x03ffffffu);
+  if (A64F_S_OK(delta>>2, 26)) {
+    *p = A64I_BL | A64F_S26(delta>>2);
   } else {  /* Target out of range: need indirect call. But don't use R0-R7. */
     Reg r = ra_allock(as, i64ptr(target),
 		      RSET_RANGE(RID_X8, RID_MAX_GPR)-RSET_FIXED);

src/lj_target_arm64.h

@@ -132,9 +132,9 @@ static LJ_AINLINE uint32_t *exitstub_trace_addr_(uint32_t *p, uint32_t exitno)
 #define A64F_IMMR(x)	((x) << 16)
 #define A64F_U16(x)	((x) << 5)
 #define A64F_U12(x)	((x) << 10)
-#define A64F_S26(x)	(x)
+#define A64F_S26(x)	(((uint32_t)(x) & 0x03ffffffu))
 #define A64F_S19(x)	(((uint32_t)(x) & 0x7ffffu) << 5)
-#define A64F_S14(x)	((x) << 5)
+#define A64F_S14(x)	(((uint32_t)(x) & 0x3fffu) << 5)
 #define A64F_S9(x)	((x) << 12)
 #define A64F_BIT(x)	((x) << 19)
 #define A64F_SH(sh, x)	(((sh) << 22) | ((x) << 10))
@@ -145,6 +145,9 @@ static LJ_AINLINE uint32_t *exitstub_trace_addr_(uint32_t *p, uint32_t exitno)
 #define A64F_LSL16(x)	(((x) / 16) << 21)
 #define A64F_BSH(sh)	((sh) << 10)
 
+/* Check for valid field range. */
+#define A64F_S_OK(x, b)	((((x) + (1 << (b-1))) >> (b)) == 0)
+
 typedef enum A64Ins {
   A64I_S = 0x20000000,
   A64I_X = 0x80000000,

test/tarantool-tests/CMakeLists.txt

@@ -57,6 +57,7 @@ macro(BuildTestCLib lib sources)
 endmacro()
 
 add_subdirectory(gh-4427-ffi-sandwich)
+add_subdirectory(gh-6098-fix-side-exit-patching-on-arm64)
 add_subdirectory(gh-6189-cur_L)
 add_subdirectory(lj-49-bad-lightuserdata)
 add_subdirectory(lj-flush-on-trace)

test/tarantool-tests/gh-6098-fix-side-exit-patching-on-arm64.test.lua

@@ -0,0 +1,129 @@
+local tap = require('tap')
+local test = tap.test('gh-6098-fix-side-exit-patching-on-arm64')
+test:plan(1)
+
+-- The function to be tested for side exit patching:
+-- * At the beginning of the test case, the <if> branch is
+--   recorded as a root trace.
+-- * After <refuncs> (and some other hotspots) are recorded, the
+--   <else> branch is recorded as a side trace.
+-- When JIT is linking the side trace to the corresponding side
+-- exit, it patches the jump targets.
+local function cbool(cond)
+  if cond then
+    return 1
+  else
+    return 0
+  end
+end
+
+-- XXX: Function template below produces 8Kb mcode for ARM64, so
+-- we need to compile at least 128 traces to exceed 1Mb delta
+-- between <cbool> root trace side exit and <cbool> side trace.
+-- Unfortunately, we have no other option for extending this jump
+-- delta, since the base of the current mcode area (J->mcarea) is
+-- used as a hint for mcode allocator (see lj_mcode.c for info).
+local FUNCS = 128
+local recfuncs = { }
+for i = 1, FUNCS do
+  -- This is a quite heavy workload (though it doesn't look like
+  -- one at first). Each load from a table is type guarded. Each
+  -- table lookup (for both stores and loads) is guarded for table
+  -- <hmask> value and metatable presence. The code below results
+  -- to 8Kb of mcode for ARM64 in practice.
+  recfuncs[i] = assert(load(([[
+    return function(src)
+      local p = %d
+      local tmp = { }
+      local dst = { }
+      for i = 1, 3 do
+        tmp.a = src.a * p   tmp.j = src.j * p   tmp.s = src.s * p
+        tmp.b = src.b * p   tmp.k = src.k * p   tmp.t = src.t * p
+        tmp.c = src.c * p   tmp.l = src.l * p   tmp.u = src.u * p
+        tmp.d = src.d * p   tmp.m = src.m * p   tmp.v = src.v * p
+        tmp.e = src.e * p   tmp.n = src.n * p   tmp.w = src.w * p
+        tmp.f = src.f * p   tmp.o = src.o * p   tmp.x = src.x * p
+        tmp.g = src.g * p   tmp.p = src.p * p   tmp.y = src.y * p
+        tmp.h = src.h * p   tmp.q = src.q * p   tmp.z = src.z * p
+        tmp.i = src.i * p   tmp.r = src.r * p
+
+        dst.a = tmp.z + p   dst.j = tmp.q + p   dst.s = tmp.h + p
+        dst.b = tmp.y + p   dst.k = tmp.p + p   dst.t = tmp.g + p
+        dst.c = tmp.x + p   dst.l = tmp.o + p   dst.u = tmp.f + p
+        dst.d = tmp.w + p   dst.m = tmp.n + p   dst.v = tmp.e + p
+        dst.e = tmp.v + p   dst.n = tmp.m + p   dst.w = tmp.d + p
+        dst.f = tmp.u + p   dst.o = tmp.l + p   dst.x = tmp.c + p
+        dst.g = tmp.t + p   dst.p = tmp.k + p   dst.y = tmp.b + p
+        dst.h = tmp.s + p   dst.q = tmp.j + p   dst.z = tmp.a + p
+        dst.i = tmp.r + p   dst.r = tmp.i + p
+      end
+      dst.tmp = tmp
+      return dst
+    end
+  ]]):format(i)), ('Syntax error in function recfuncs[%d]'):format(i))()
+end
+
+-- Make compiler work hard:
+-- * No optimizations at all to produce more mcode.
+-- * Try to compile all compiled paths as early as JIT can.
+-- * Allow to compile 2Mb of mcode to be sure the issue occurs.
+jit.opt.start(0, 'hotloop=1', 'hotexit=1', 'maxmcode=2048')
+
+-- First call makes <cbool> hot enough to be recorded next time.
+cbool(true)
+-- Second call records <cbool> body (i.e. <if> branch). This is
+-- a root trace for <cbool>.
+cbool(true)
+
+for i = 1, FUNCS do
+  -- XXX: FNEW is NYI, hence loop recording fails at this point.
+  -- The recording is aborted on purpose: we are going to record
+  -- <FUNCS> number of traces for functions in <recfuncs>.
+  -- Otherwise, loop recording might lead to a very long trace
+  -- error (via return to a lower frame), or a trace with lots of
+  -- side traces. We need neither of this, but just bunch of
+  -- traces filling the available mcode area.
+  local function tnew(p)
+    return {
+      a = p + 1, f = p + 6,  k = p + 11, p = p + 16, u = p + 21, z = p + 26,
+      b = p + 2, g = p + 7,  l = p + 12, q = p + 17, v = p + 22,
+      c = p + 3, h = p + 8,  m = p + 13, r = p + 18, w = p + 23,
+      d = p + 4, i = p + 9,  n = p + 14, s = p + 19, x = p + 24,
+      e = p + 5, j = p + 10, o = p + 15, t = p + 20, y = p + 25,
+    }
+  end
+  -- Each function call produces a trace (see the template for the
+  -- function definition above).
+  recfuncs[i](tnew(i))
+end
+
+-- XXX: I tried to make the test in pure Lua, but I failed to
+-- implement the robust solution. As a result I've implemented a
+-- tiny Lua C API module to route the flow through C frames and
+-- make JIT work the way I need to reproduce the fail. See the
+-- usage below.
+-- <pxcall> is just a wrapper for <lua_call> with "multiargs" and
+-- "multiret" with the same signature as <pcall>.
+local pxcall = require('libproxy').proxycall
+
+-- XXX: Here is the dessert: JIT is aimed to work better for
+-- highly biased code. It means, the root trace should be the
+-- most popular flow. Furthermore, JIT also considers the fact,
+-- that frequently taken side exits are *also* popular, and
+-- compiles the side traces for such popular exits. However,
+-- to recoup his attempts JIT try to compile the flow as far
+-- as it can (see <lj_record_ret> in lj_record.c for more info).
+--
+-- Such "kind" behaviour occurs in our case: if one calls <cbool>
+-- the native way, JIT continues recording in a lower frame after
+-- returning from <cbool>. As a result, the second call is also
+-- recorded, but it has to trigger the far jump to the side trace.
+-- However, if the lower frame is not the Lua one, JIT doesn't
+-- proceed the further flow recording and assembles the trace. In
+-- this case, the second call jumps to <cbool> root trace, hits
+-- the assertion guard and jumps to <cbool> side trace.
+pxcall(cbool, false)
+cbool(false)
+
+test:ok(true)
+os.exit(test:check() and 0 or 1)

test/tarantool-tests/gh-6098-fix-side-exit-patching-on-arm64/CMakeLists.txt

@@ -0,0 +1 @@
+BuildTestCLib(libproxy libproxy.c)

test/tarantool-tests/gh-6098-fix-side-exit-patching-on-arm64/libproxy.c

@@ -0,0 +1,23 @@
+#include <lua.h>
+#include <lauxlib.h>
+
+/*
+ * Function with the signature similar to Lua <pcall> builtin,
+ * that routes the flow through C frame.
+ */
+static int proxycall(lua_State *L)
+{
+	lua_call(L, lua_gettop(L) - 1, LUA_MULTRET);
+	return lua_gettop(L);
+}
+
+static const struct luaL_Reg libproxy[] = {
+	{"proxycall", proxycall},
+	{NULL, NULL}
+};
+
+LUA_API int luaopen_libproxy(lua_State *L)
+{
+	luaL_register(L, "libproxy", libproxy);
+	return 1;
+}