[Tarantool-patches] [PATCH luajit] Optionally return PC position in jit.util.tracesnap().

[Tarantool-patches] [PATCH luajit] Optionally return PC position in jit.util.tracesnap().

[Sergey Bronnikov via Tarantool-patches]

From: Mike Pall <mike>

Suggested by Sergey Bronnikov.

(cherry picked from commit 8e6520a7aecd0517e792b359afbbfd7274791f5f)

The optional `getpos` argument is needed for ljopt [1], a tool that
verifies the correctness of LuaJIT optimizations by translating the
recorded IR into SMT-LIB and checking the formulas with an SMT solver.
Traces often lack sufficient information for accurate modeling, so
ljopt extracts the address of the bytecode pointed to by the snapshot
directly from LuaJIT during execution.

The LuaJIT internal API includes the function `jit.util.tracesnap(tr,
sn)`. In principle, this function could be extended to return the
program counter (pc). This patch does exactly that: an optional third
argument is added, and when it is truthy the second value returned by
`jit.util.tracesnap()` is the PC position of the snapshot in the
bytecode of the function where the trace starts.

Also, the patch adds a test that records a root loop trace and
checks that the returned position is the `FORL` jump target (the
loop entry), that no position is returned without the optional
argument or with a false value, and that every snapshot of the
trace reports a valid bytecode offset.

1. https://github.com/ligurio/ljopt

Sergey Bronnikov:
* added the description and the test for the feature
---

Branch: https://github.com/tarantool/luajit/tree/ligurio/gh-xxxx-tracesnap-getpos

 src/lib_jit.c                                 |  9 +-
 .../lj-1472-tracesnap-getpos.test.lua         | 89 +++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua

diff --git a/src/lib_jit.c b/src/lib_jit.c
index 9f870f68f..5a0d38611 100644
--- a/src/lib_jit.c
+++ b/src/lib_jit.c
@@ -351,11 +351,12 @@ LJLIB_CF(jit_util_tracek)
   return 0;
 }
 
-/* local snap = jit.util.tracesnap(tr, sn) */
+/* local snap = jit.util.tracesnap(tr, sn[, getpos]) */
 LJLIB_CF(jit_util_tracesnap)
 {
   GCtrace *T = jit_checktrace(L);
   SnapNo sn = (SnapNo)lj_lib_checkint(L, 2);
+  int getpos = (L->base+2 < L->top && tvistruecond(L->base+2));
   if (T && sn < T->nsnap) {
     SnapShot *snap = &T->snap[sn];
     SnapEntry *map = &T->snapmap[snap->mapofs];
@@ -368,6 +369,12 @@ LJLIB_CF(jit_util_tracesnap)
     for (n = 0; n < nent; n++)
       setintV(lj_tab_setint(L, t, (int32_t)(n+2)), (int32_t)map[n]);
     setintV(lj_tab_setint(L, t, (int32_t)(nent+2)), (int32_t)SNAP(255, 0, 0));
+    if (getpos) {
+      const BCIns *pc = snap_pc(&map[nent]), *startpc = pc;
+      while (bc_op(*startpc) < BC_FUNCF) startpc--;
+      setintV(L->top++, (int)(pc - startpc));
+      return 2;
+    }
     return 1;
   }
   return 0;
diff --git a/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua b/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua
new file mode 100644
index 000000000..a56a7a679
--- /dev/null
+++ b/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua
@@ -0,0 +1,89 @@
+local tap = require('tap')
+local bit = require('bit')
+local jutil = require('jit.util')
+local vmdef = require('jit.vmdef')
+
+-- The test checks the optional `getpos` argument of
+-- `jit.util.tracesnap()` that makes the function return the PC
+-- position of the snapshot in the function bytecode.
+-- See also https://github.com/LuaJIT/LuaJIT/issues/1472.
+
+local test = tap.test('lj-1472-tracesnap-getpos'):skipcond({
+  ['Test requires JIT enabled'] = not jit.status(),
+})
+
+test:plan(8)
+
+-- Entries in `vmdef.bcnames` are 6 characters long.
+local BC_NAME_LENGTH = 6
+
+-- Return the name of the bytecode operation for the given
+-- instruction, that is encoded in the low byte.
+local function opname(ins)
+  local oidx = BC_NAME_LENGTH * bit.band(ins, 0xff)
+  local name = vmdef.bcnames:sub(oidx + 1, oidx + BC_NAME_LENGTH)
+  return (name:gsub('%s+$', ''))
+end
+
+local function payload()
+  local sum = 0
+  for i = 1, 100 do
+    sum = sum + i
+  end
+  return sum
+end
+
+-- Find the loop back-edge (FORL) in the function bytecode and
+-- compute its jump target, that is the first instruction of the
+-- loop body. `pc == 0` is the JCproto header, so real
+-- instructions start from `pc == 1`. JIT compilation is disabled
+-- while scanning, so this loop is not recorded as a trace.
+jit.off()
+local nbc = jutil.funcinfo(payload).bytecodes
+local loop_entry
+for pc = 1, nbc - 1 do
+  local ins = jutil.funcbc(payload, pc)
+  if opname(ins) == 'FORL' then
+    -- The jump offset is relative to the next instruction and is
+    -- biased by `BCBIAS_J`.
+    local jmp = bit.rshift(ins, 16) - 0x8000
+    loop_entry = pc + 1 + jmp
+    break
+  end
+end
+test:isnt(loop_entry, nil, 'FORL instruction is found')
+
+-- Record a root loop trace in `payload`. Snapshot #0 points to
+-- the loop entry (the FORL jump target).
+jit.on()
+jit.flush()
+jit.opt.start('hotloop=1')
+payload()
+
+local info = jutil.traceinfo(1)
+test:ok(info, 'the root loop trace is recorded')
+
+local snap, pos = jutil.tracesnap(1, 0, true)
+test:istable(snap, 'the snapshot table is returned')
+test:is(pos, loop_entry, 'the PC position points to the loop entry')
+
+local _, pos_without = jutil.tracesnap(1, 0)
+test:is(pos_without, nil, 'no PC position without getpos')
+
+local _, pos_false = jutil.tracesnap(1, 0, false)
+test:is(pos_false, nil, 'no PC position with getpos=false')
+
+local all_valid = true
+for sn = 0, info.nexit - 1 do
+  local _, p = jutil.tracesnap(1, sn, true)
+  if type(p) ~= 'number' or p < 0 or p >= nbc or
+     jutil.funcbc(payload, p) == nil then
+    all_valid = false
+  end
+end
+test:ok(all_valid, 'positions of all snapshots are valid offsets')
+
+test:is(select('#', jutil.tracesnap(1, info.nexit, true)), 0,
+        'out-of-range snapshot returns nothing')
+
+test:done(true)
-- 
2.51.0

Re: [Tarantool-patches] [PATCH luajit] Optionally return PC position in jit.util.tracesnap().

[Sergey Kaplun via Tarantool-patches]

Hi, Sergey!
Thanks for the patch!
LGTM, after fixing my comments below.

On 08.10.26, Sergey Bronnikov wrote:
> From: Mike Pall <mike>
> 
> Suggested by Sergey Bronnikov.
> 
> (cherry picked from commit 8e6520a7aecd0517e792b359afbbfd7274791f5f)
> 
> The optional `getpos` argument is needed for ljopt [1], a tool that
> verifies the correctness of LuaJIT optimizations by translating the
> recorded IR into SMT-LIB and checking the formulas with an SMT solver.
> Traces often lack sufficient information for accurate modeling, so
> ljopt extracts the address of the bytecode pointed to by the snapshot
> directly from LuaJIT during execution.
> 
> The LuaJIT internal API includes the function `jit.util.tracesnap(tr,
> sn)`. In principle, this function could be extended to return the
> program counter (pc). This patch does exactly that: an optional third
> argument is added, and when it is truthy the second value returned by
> `jit.util.tracesnap()` is the PC position of the snapshot in the
> bytecode of the function where the trace starts.
> 
> Also, the patch adds a test that records a root loop trace and
> checks that the returned position is the `FORL` jump target (the
> loop entry), that no position is returned without the optional
> argument or with a false value, and that every snapshot of the
> trace reports a valid bytecode offset.

I don't understand this paragraph. Usually we don't describe added tests
for the features (especially in the way "what was done" instead of "why
it was done").

> 
> 1. https://github.com/ligurio/ljopt
> 

Please, mention the backporting issues as well.

> Sergey Bronnikov:
> * added the description and the test for the feature
> ---
> 
> Branch: https://github.com/tarantool/luajit/tree/ligurio/gh-xxxx-tracesnap-getpos
> 
>  src/lib_jit.c                                 |  9 +-
>  .../lj-1472-tracesnap-getpos.test.lua         | 89 +++++++++++++++++++
>  2 files changed, 97 insertions(+), 1 deletion(-)
>  create mode 100644 test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua
> 
> diff --git a/src/lib_jit.c b/src/lib_jit.c
> index 9f870f68f..5a0d38611 100644
> --- a/src/lib_jit.c
> +++ b/src/lib_jit.c

<snipped>

> diff --git a/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua b/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua
> new file mode 100644
> index 000000000..a56a7a679
> --- /dev/null
> +++ b/test/tarantool-tests/lj-1472-tracesnap-getpos.test.lua
> @@ -0,0 +1,89 @@
> +local tap = require('tap')
> +local bit = require('bit')
> +local jutil = require('jit.util')
> +local vmdef = require('jit.vmdef')
> +
> +-- The test checks the optional `getpos` argument of
> +-- `jit.util.tracesnap()` that makes the function return the PC
> +-- position of the snapshot in the function bytecode.
> +-- See also https://github.com/LuaJIT/LuaJIT/issues/1472.
> +
> +local test = tap.test('lj-1472-tracesnap-getpos'):skipcond({
> +  ['Test requires JIT enabled'] = not jit.status(),
> +})
> +
> +test:plan(8)
> +
> +-- Entries in `vmdef.bcnames` are 6 characters long.
> +local BC_NAME_LENGTH = 6
> +
> +-- Return the name of the bytecode operation for the given
> +-- instruction, that is encoded in the low byte.
> +local function opname(ins)
> +  local oidx = BC_NAME_LENGTH * bit.band(ins, 0xff)
> +  local name = vmdef.bcnames:sub(oidx + 1, oidx + BC_NAME_LENGTH)
> +  return (name:gsub('%s+$', ''))

Minor: I suggest the more readable:
| return name:match('%a+')
Feel free to ignore.

> +end
> +
> +local function payload()
> +  local sum = 0
> +  for i = 1, 100 do

Minor: Typo: s/100/4/

> +    sum = sum + i
> +  end
> +  return sum
> +end
> +
> +-- Find the loop back-edge (FORL) in the function bytecode and
> +-- compute its jump target, that is the first instruction of the

Typo: s/that/which/

> +-- loop body. `pc == 0` is the JCproto header, so real
> +-- instructions start from `pc == 1`. JIT compilation is disabled
> +-- while scanning, so this loop is not recorded as a trace.
> +jit.off()
> +local nbc = jutil.funcinfo(payload).bytecodes
> +local loop_entry
> +for pc = 1, nbc - 1 do
> +  local ins = jutil.funcbc(payload, pc)
> +  if opname(ins) == 'FORL' then
> +    -- The jump offset is relative to the next instruction and is
> +    -- biased by `BCBIAS_J`.
> +    local jmp = bit.rshift(ins, 16) - 0x8000

Let's name the magic number:
| local BCBIAS_J = 0x8000

> +    loop_entry = pc + 1 + jmp
> +    break
> +  end
> +end
> +test:isnt(loop_entry, nil, 'FORL instruction is found')

It should be an assert instead, since this is the test precondition.

> +
> +-- Record a root loop trace in `payload`. Snapshot #0 points to
> +-- the loop entry (the FORL jump target).
> +jit.on()
> +jit.flush()
> +jit.opt.start('hotloop=1')
> +payload()
> +
> +local info = jutil.traceinfo(1)
> +test:ok(info, 'the root loop trace is recorded')

It should be an assert instead, since this is the test precondition.

> +local snap, pos = jutil.tracesnap(1, 0, true)
> +test:istable(snap, 'the snapshot table is returned')
> +test:is(pos, loop_entry, 'the PC position points to the loop entry')
> +
> +local _, pos_without = jutil.tracesnap(1, 0)
> +test:is(pos_without, nil, 'no PC position without getpos')
> +
> +local _, pos_false = jutil.tracesnap(1, 0, false)
> +test:is(pos_false, nil, 'no PC position with getpos=false')
> +
> +local all_valid = true
> +for sn = 0, info.nexit - 1 do
> +  local _, p = jutil.tracesnap(1, sn, true)
> +  if type(p) ~= 'number' or p < 0 or p >= nbc or
> +     jutil.funcbc(payload, p) == nil then
> +    all_valid = false
> +  end
> +end
> +test:ok(all_valid, 'positions of all snapshots are valid offsets')
> +
> +test:is(select('#', jutil.tracesnap(1, info.nexit, true)), 0,
> +        'out-of-range snapshot returns nothing')
> +

It would be nice to add any correctness test.
The most trivial in mind is the following:

| src/luajit -Ohotloop=1 -e '
|   local jutil = require"jit.util"
|   local function trivial()
|     local a = 0
|     return a
|   end
|   trivial()
|   trivial()
|
|   local _, b = jutil.tracesnap(1, 0, true)
|   print(b)
|   local _, b = jutil.tracesnap(1, 1, true)
|   print(b)
| '
| 1
| 2

The snapshot is emitted at the function entry and before returning the
result which makes it trivial to detect:

| ---- TRACE 1 start (command line):3
| 0001  KSHORT   0   0
| 0002  RET1     0   2
| ---- TRACE 1 IR
| ....        SNAP   #0   [ ---- ]
| ....        SNAP   #1   [ ---- +0   ]

Don't forget to reset the hotcounters during the recording of this
example.

> +test:done(true)
> -- 
> 2.51.0
> 

-- 
Best regards,
Sergey Kaplun