luajit: experimentally add sink optimization

Author: RaphaelIT7

luajit/README.md

@@ -7,6 +7,7 @@ Included changes:
 \- [+] Added block debug functionality for function to prevent someone from getting functions they shouldn't have.<br>
 \- [+] Added read only table functionality to prevent someone from modifying a table while its being used by another thread<br>
 \- [+] Exposed `luaopen_jit_profile` & `lua_index2adr` for HolyLib.<br>
-\- [+] Implemented fix for FFI Sandwich/LUA VM re-entry through JIT trace (See https://github.com/LuaJIT/LuaJIT/pull/1165)<br> 
+\- [+] Implemented fix for FFI Sandwich/LUA VM re-entry through JIT trace (See https://github.com/LuaJIT/LuaJIT/pull/1165)<br>
+\- [+] Experimentally implemented Sink optimization (See https://github.com/LuaJIT/LuaJIT/pull/652)<br>
 \- [#] Made `cdata` return the type as `LUA_TUSERDATA` so that we can more easily allow FFI -> C calls without needing to hook 10 functions (& the TypeID also conflicted with gmod)<br>
 \- [#] Improved `GMODLUA_GetUserType` to directly do it's stuff without using the Lua stack<br>

luajit/src/jit/dump.lua

@@ -396,11 +396,15 @@ local function dump_snap(tr)
 end
 
 -- Return a register name or stack slot for a rid/sp location.
-local function ridsp_name(ridsp, ins)
+local function ridsp_name(ridsp, ins, op)
   if not disass then disass = require("jit.dis_"..jit.arch) end
   local rid, slot = band(ridsp, 0xff), shr(ridsp, 8)
   if rid == 253 or rid == 254 then
-    return (slot == 0 or slot == 255) and " {sink" or format(" {%04d", ins-slot)
+    if op == "TNEW  " or op == "TDUP  " or op == "CNEW  " then
+      return (slot == 0) and " {sink" or format(" {ri%02d", slot)
+    else
+      return (slot == 0 or slot == 255) and " {sink" or format(" {%04d", ins-slot)
+    end
   end
   if ridsp > 255 then return format("[%x]", slot*4) end
   if rid < 128 then return disass.regname(rid) end
@@ -485,7 +489,7 @@ local function dump_ir(tr, dumpsnap, dumpreg)
 	   (dumpreg or op ~= "RENAME") then
       local rid = band(ridsp, 255)
       if dumpreg then
-	out:write(format("%04d %-6s", ins, ridsp_name(ridsp, ins)))
+	out:write(format("%04d %-6s", ins, ridsp_name(ridsp, ins, op)))
       else
 	out:write(format("%04d ", ins))
       end

luajit/src/lj_asm.c

@@ -348,7 +348,7 @@ static void ra_dprintf(ASMState *as, const char *fmt, ...)
 #define ra_weak(as, r)		rset_set(as->weakset, (r))
 #define ra_noweak(as, r)	rset_clear(as->weakset, (r))
 
-#define ra_used(ir)		(ra_hasreg((ir)->r) || ra_hasspill((ir)->s))
+#define ra_used(ir)		(ra_hasreg((ir)->r) || ((ir)->r != RID_SUNK && (ir)->r != RID_SINK && ra_hasspill((ir)->s))
 
 /* Setup register allocator. */
 static void ra_setup(ASMState *as)
@@ -941,7 +941,8 @@ static int asm_sunk_store(ASMState *as, IRIns *ira, IRIns *irs)
     }
     return 0;
   } else {
-    return (ira + irs->s == irs);  /* Quick check. */
+    if (ira + irs->s != irs) return 0;
+    return irs->o == IR_ASTORE || irs->o == IR_HSTORE || irs->o == IR_FSTORE || irs->o == IR_XSTORE;
   }
 }
 

luajit/src/lj_opt_sink.c

@@ -14,6 +14,7 @@
 #include "lj_jit.h"
 #include "lj_iropt.h"
 #include "lj_target.h"
+#include "lj_dispatch.h"
 
 /* Some local macros to save typing. Undef'd at the end. */
 #define IR(ref)		(&J->cur.ir[(ref)])
@@ -54,6 +55,9 @@ static int sink_checkphi(jit_State *J, IRIns *ira, IRRef ref)
     IRIns *ir = IR(ref);
     if (irt_isphi(ir->t) || (ir->o == IR_CONV && ir->op2 == IRCONV_NUM_INT &&
 			     irt_isphi(IR(ir->op1)->t))) {
+      if ((ira->prev & 0x1FFF) == 0x1FFF)
+        return 0; /* This would cause an overflow, just force the allocation to not be sunken. */
+
       ira->prev++;
       return 1;  /* Sinkable PHI. */
     }
@@ -69,6 +73,14 @@ static int sink_checkphi(jit_State *J, IRIns *ira, IRRef ref)
   return 1;  /* Constant (non-PHI). */
 }
 
+/* Set prev of all instructions to 0. */
+static void sink_prepare(jit_State *J) {
+  IRIns *ir, *irlast = IR(J->cur.nins-1);
+  for (ir = irlast ; ir->o != IR_BASE; ir--) {
+    ir->prev = 0;
+  }
+}
+
 /* Mark non-sinkable allocations using single-pass backward propagation.
 **
 ** Roots for the marking process are:
@@ -79,7 +91,7 @@ static int sink_checkphi(jit_State *J, IRIns *ira, IRRef ref)
 ** - Stores with non-constant keys.
 ** - All stored values.
 */
-static void sink_mark_ins(jit_State *J)
+static int sink_mark_ins(jit_State *J, int lightsink)
 {
   IRIns *ir, *irlast = IR(J->cur.nins-1);
   for (ir = irlast ; ; ir--) {
@@ -152,26 +164,34 @@ static void sink_mark_snap(jit_State *J, SnapShot *snap)
 }
 
 /* Iteratively remark PHI refs with differing marks or PHI value counts. */
-static void sink_remark_phi(jit_State *J)
+static int sink_remark_phi(jit_State *J)
 {
   IRIns *ir;
   int remark;
+  int require_remark = 0;
   do {
     remark = 0;
     for (ir = IR(J->cur.nins-1); ir->o == IR_PHI; ir--) {
       IRIns *irl = IR(ir->op1), *irr = IR(ir->op2);
-      if (!((irl->t.irt ^ irr->t.irt) & IRT_MARK) && irl->prev == irr->prev)
+      if (!((irl->t.irt ^ irr->t.irt) & IRT_MARK) && (irl->prev & 0x1FFF) == (irr->prev & 0x1FFF))
 	continue;
       remark |= (~(irl->t.irt & irr->t.irt) & IRT_MARK);
+      if ((IR(ir->op1)->prev & 0x2000) || (IR(ir->op2)->prev & 0x2000)) {
+        IR(ir->op1)->prev &= ~0x2000;
+        IR(ir->op2)->prev &= ~0x2000;
+        require_remark |= (~(irl->t.irt & irr->t.irt) & IRT_MARK);
+      }
       irt_setmark(IR(ir->op1)->t);
       irt_setmark(IR(ir->op2)->t);
     }
   } while (remark);
+  return require_remark;
 }
 
 /* Sweep instructions and tag sunken allocations and stores. */
 static void sink_sweep_ins(jit_State *J)
 {
+  int index = 0;
   IRIns *ir, *irbase = IR(REF_BASE);
   for (ir = IR(J->cur.nins-1) ; ir >= irbase; ir--) {
     switch (ir->o) {
@@ -198,8 +218,14 @@ static void sink_sweep_ins(jit_State *J)
 #endif
     case IR_TNEW: case IR_TDUP:
       if (!irt_ismarked(ir->t)) {
+  if (ir->prev & 0x8000) {
+    index++;  /* A sunken store requires this for unsinking. */
+    lj_assertJ(index <= 0xFF, "Too many heavy sinks");
+    ir->prev = REGSP(RID_SINK, index);
+  } else {
+    ir->prev = REGSP(RID_SINK, 0);
+  }
 	ir->t.irt &= ~IRT_GUARD;
-	ir->prev = REGSP(RID_SINK, 0);
 	J->cur.sinktags = 1;  /* Signal present SINK tags to assembler. */
       } else {
 	irt_clearmark(ir->t);
@@ -244,15 +270,19 @@ void lj_opt_sink(jit_State *J)
   if ((J->flags & need) == need &&
       (J->chain[IR_TNEW] || J->chain[IR_TDUP] ||
        (LJ_HASFFI && (J->chain[IR_CNEW] || J->chain[IR_CNEWI])))) {
+    sink_prepare(J);
     if (!J->loopref)
       sink_mark_snap(J, &J->cur.snap[J->cur.nsnap-1]);
-    sink_mark_ins(J);
-    if (J->loopref)
-      sink_remark_phi(J);
+    int heavysinks;
+    int dolightsink = 0;
+    do {
+      heavysinks = sink_mark_ins(J, dolightsink);
+      dolightsink |= heavysinks >= 0xFF;
+    } while ((J->loopref && sink_remark_phi(J)) || heavysinks >= 0xFF);
     sink_sweep_ins(J);
   }
 }
 
 #undef IR
 
-#endif
+#endif