The millionth register allocator. This time with scoped spills

5 files changed, 360 insertions, 21 deletions

diff --git a/Makefile b/Makefile
index 1886df7..215ac9e 100644
--- a/Makefile
+++ b/Makefile
@@ -18,11 +18,14 @@ COMMON_HEADERS =	\
 	crc32.h		\
 	stack-man.c
 
-all: genrender blitter testsuite iterjit
+all: genrender blitter testsuite regscope
 
 genrender: $(COMMON) main.c
 	$(CC) $(CFLAGS) -o genrender $(COMMON_SOURCES) main.c simple-reg.c $(LDFLAGS)
 
+regscope: regscope.c
+	$(CC) $(CFLAGS) -o regscope regscope.c $(COMMON_SOURCES) $(LDFLAGS)
+
 blitter: $(COMMON) blitter.c
 	$(CC) $(CFLAGS) -o blitter $(COMMON_SOURCES) blitter.c simple-reg.c $(LDFLAGS)
 
diff --git a/iterjit.c b/iterjit.c
index 57bcc42..76bd147 100644
--- a/iterjit.c
+++ b/iterjit.c
@@ -292,8 +292,17 @@ jit_alloc_gp (jit_t *jit)
 }
 
 reg_t jit_alloc_xmm (jit_t *jit);
-void jit_preserve_gp (jit_t *jit, reg_t reg);
-void jit_free_gp (jit_t *jit, reg_t reg);
+
+reg_t jit_preserve_gp (jit_t *jit, reg_t reg)
+{
+    return reg_alloc_preserve (&jit->gp_allocator, reg);
+}
+
+void jit_free_gp (jit_t *jit, reg_t reg)
+{
+    reg_alloc_free (&jit->gp_allocator, reg);
+}
+
 void jit_reload_gp (jit_t *jit, reg_t reg);
 void jit_reload_xmm (jit_t *jit, reg_t reg);
 void jit_free_gp (jit_t *jit, reg_t reg);
@@ -307,7 +316,7 @@ src_a8r8g8b8_begin (jit_src_iter_t *src,
     reg_t image, tmp;
 
     jit_switch_group (jit, "src");
-    jit_preserve_gp (jit, info);
+    info = jit_preserve_gp (jit, info);
     src->line = jit_alloc_gp (jit);
     src->stride = jit_alloc_gp (jit);
     image = jit_alloc_gp (jit);
@@ -435,7 +444,7 @@ dest_a8r8g8b8_begin (jit_dest_iter_t *	dest,
     reg_t image, tmp;
 
     jit_switch_group (jit, "dest");
-    jit_preserve_gp (jit, info);
+    info = jit_preserve_gp (jit, info);
     dest->line = jit_alloc_gp (jit);
     dest->stride = jit_alloc_gp (jit);
     image = jit_alloc_gp (jit);
@@ -663,10 +672,12 @@ combine_over (jit_combiner_t *combiner, jit_t *jit,
 
     s = src->load_pixels (src, jit, n_pixels);
 
+    jit_switch_group (jit, "combiner");
+    s = jit_preserve_gp (jit, s);
+
     src->advance_pixels (src, jit, n_pixels);
 
     jit_switch_group (jit, "combiner");
-    jit_preserve_gp (jit, s);
 
     m00ff = zero = jit_alloc_xmm (jit);
     m_hi = jit_alloc_xmm (jit);
@@ -698,7 +709,7 @@ combine_over (jit_combiner_t *combiner, jit_t *jit,
     d = dest->load_pixels (dest, jit, n_pixels);
 
     jit_switch_group (jit, "combiner");
-    jit_preserve_gp (jit, d);
+    d = jit_preserve_gp (jit, d);
 
     m0080 = zero = jit_alloc_xmm (jit);
     m0101 = jit_alloc_xmm (jit);
@@ -768,25 +779,26 @@ generate_kernel (jit_t *jit,
 		 jit_combiner_t *combiner)
 {
     reg_t h, composite_info;
+    reg_t r_rbx, r_r12, r_r13, r_r14, r_r15;
 
     jit_switch_group (jit, "outer");
 
     composite_info = rsi;
     
-    jit_preserve_gp (jit, composite_info);
+    composite_info = jit_preserve_gp (jit, composite_info);
 
     /* Preserve callee-save registers */
-    jit_preserve_gp (jit, rbx);
-    jit_preserve_gp (jit, r12);
-    jit_preserve_gp (jit, r13);
-    jit_preserve_gp (jit, r14);
-    jit_preserve_gp (jit, r15);
+    r_rbx = jit_preserve_gp (jit, rbx);
+    r_r12 = jit_preserve_gp (jit, r12);
+    r_r13 = jit_preserve_gp (jit, r13);
+    r_r14 = jit_preserve_gp (jit, r14);
+    r_r15 = jit_preserve_gp (jit, r15);
 
     BEGIN_ASM (jit->fragment)
 	I_push, rbp,
 	I_mov, rbp, rsp,
     END_ASM();
-	
+
     h = jit_alloc_gp (jit);
 
     BEGIN_ASM (jit->fragment)
@@ -796,7 +808,7 @@ generate_kernel (jit_t *jit,
     END_ASM ();
 
     jit_free_gp (jit, composite_info);
-    
+
     /* begin */
     src_iter->begin (src_iter, jit, composite_info);
     if (mask_iter)
@@ -827,14 +839,28 @@ generate_kernel (jit_t *jit,
     src_iter->end (src_iter, jit);
 
     jit_switch_group (jit, "outer");
+
     /* Restore callee-save registers */
-    jit_reload_gp (jit, rbx);
-    jit_reload_gp (jit, r12);
-    jit_reload_gp (jit, r13);
-    jit_reload_gp (jit, r14);
-    jit_reload_gp (jit, r15);
+#define RESTORE(reg)							\
+    r_##reg = jit_reload_gp (jit, r_##reg);				\
+    if (r_##reg != reg)							\
+    {									\
+	BEGIN_ASM (jit->fragment)					\
+	    I_mov, reg, r_##reg,					\
+	END_ASM ();							\
+    }
+
     BEGIN_ASM (jit->fragment)
 	DEFINE_LABEL ("done"),
+    END_ASM();
+
+    RESTORE(rbx);
+    RESTORE(r12);
+    RESTORE(r13);
+    RESTORE(r14);
+    RESTORE(r15);
+
+    BEGIN_ASM (jit->fragment)
 	I_pop, rbp,
     END_ASM();
 }
diff --git a/reggroups.h b/reggroups.h
index 0813aa7..0984e82 100644
--- a/reggroups.h
+++ b/reggroups.h
@@ -31,7 +31,9 @@ struct reg_alloc_t
 };
 
 void
-reg_alloc_init (reg_alloc_t *ra, stack_man_t *stack, const reg_pool_t *reg_pool);
+reg_alloc_init (reg_alloc_t *ra,
+		stack_man_t *stack,
+		const reg_pool_t *reg_pool);
 
 void
 reg_alloc_switch_group (reg_alloc_t *ra, const char *name);
diff --git a/regscope.c b/regscope.c
new file mode 100644
index 0000000..ee7a001
--- /dev/null
+++ b/regscope.c
@@ -0,0 +1,234 @@
+#include <stdarg.h>
+#include <string.h>
+#include <stdio.h>
+#include "regscope.h"
+
+void
+reg_alloc_init (reg_alloc_t *ra,
+		fragment_t *frag,
+		stack_man_t *stack,
+		const reg_pool_t *reg_pool)
+{
+    int i;
+
+    ra->pool = reg_pool;
+    ra->stack = stack;
+    ra->fragment = frag;
+
+    for (i = 0; i < ra->pool->n_registers; ++i) 
+    {
+	int j;
+	ra->status[i].allocated = FALSE;
+	ra->status[i].spillable = FALSE;
+	ra->status[i].n_spills = 0;
+	
+	for (j = 0; j < MAX_SPILLS; ++j)
+	    ra->status[i].spill_offsets[j] = -1;
+    }
+}
+
+/* Allocates a register. The allocation is non-local, so it
+ * is allowed to survive the current begin/end_spill group.
+ */
+reg_t
+reg_alloc_alloc (reg_alloc_t *ra)
+{
+    int i;
+
+    for (i = 0; i < ra->pool->n_registers; ++i)
+    {
+	reg_status_t *status = &ra->status[i];
+
+	if (!status->allocated)
+	{
+	    status->allocated = 1;
+	    return ra->pool->registers[i];
+	}
+    }
+
+    return (reg_t)0;
+}
+
+static void
+one_register_begin_spill (reg_alloc_t *ra, reg_status_t *status, reg_t reg)
+{
+    if (!status->allocated)
+	fprintf (stderr, "Attempting to spill a register that isn't allocated\n");
+
+    if (status->spillable)
+	fprintf (stderr, "Attempting to spill a register that is already spillable\n");
+
+    status->spillable = TRUE;
+}
+
+static void
+one_register_end_spill (reg_alloc_t *ra, reg_status_t *status, reg_t reg)
+{
+    if (!status->allocated)
+	fprintf (stderr, "Attempting to unspill an unallocated register\n");
+
+    if (!status->spillable)
+    {
+	int offset;
+
+	if (status->n_spills < 0)
+	{
+	    fprintf (stderr, "No spill offset available for reloading\n");
+	    return;
+	}
+
+	offset = status->spill_offsets[--status->n_spills];
+
+	BEGIN_ASM (ra->fragment)
+	    ra->pool->move_instruction, reg, BASE(rsp, offset),
+	END_ASM ();
+    }
+
+    status->spillable = FALSE;
+}
+
+static int
+find_index (reg_alloc_t *ra, reg_t reg)
+{
+    int i;
+
+    for (i = 0; i < ra->pool->n_registers; ++i)
+    {
+	if (ra->pool->registers[i] == reg)
+	    return i;
+    }
+
+    return -1;
+}
+
+/* Calling this indicates that the registers given will not be
+ * needed until the corresponding reg_alloc_end_spill().
+ */
+void
+reg_alloc_begin_spill (reg_alloc_t *ra, reg_t reg, ...)
+{
+    va_list list;
+
+    va_start (list, reg);
+    while (reg != (reg_t)0)
+    {
+	int idx = find_index (ra, reg);
+
+	one_register_begin_spill (ra, &ra->status[idx], reg);
+
+	reg = va_arg (list, reg_t);
+    }
+    va_end (list);
+}
+
+void
+reg_alloc_end_spill (reg_alloc_t *ra, reg_t reg, ...)
+{
+    va_list list;
+
+    va_start (list, reg);
+    while (reg != (reg_t)0)
+    {
+	int idx = find_index (ra, reg);
+
+	one_register_end_spill (ra, &ra->status[idx], reg);
+
+	reg = va_arg (list, reg_t);
+    }
+    va_end (list);
+}
+
+/* This allocates a register. Calliing this function indicates that
+ * the register in question will be freed before the current 
+ * begin/end_spill group ends.
+ */
+reg_t
+reg_alloc_alloc_local (reg_alloc_t *ra)
+{
+    reg_t reg;
+    int i;
+
+    reg = reg_alloc_alloc (ra);
+    if (reg)
+	return reg;
+
+    /* All registers are in use, so find a spillable one */
+    for (i = 0; i < ra->pool->n_registers; ++i)
+    {
+	reg_status_t *status = &ra->status[i];
+	int offset;
+
+	if (status->spillable)
+	{
+	    reg_t reg = ra->pool->registers[i];
+
+	    if (!stack_manager_alloc (
+		    ra->stack, ra->pool->bytes_per_reg, &offset))
+	    {
+		fprintf (stderr, "OOM\n");
+	    }
+	    
+	    BEGIN_ASM (ra->fragment)
+		ra->pool->move_instruction, BASE(rsp, offset), reg,
+	    END_ASM ();
+
+	    if (status->n_spills >= MAX_SPILLS)
+	    {
+		fprintf (stderr, "Out of spill levels\n");
+	    }
+
+	    status->spill_offsets[status->n_spills++] = offset;
+	    status->allocated = TRUE;
+	    status->spillable = 0;
+
+	    return reg;
+	}
+    }
+
+    return (reg_t)0;
+}
+
+/* This resurrects a register that has been freed, but is known
+ * to still contain a useful value. The allocation is non-local,
+ * so it is allowed to survive the current begin/end_spill group.
+ *
+ * The returned register may or may not be the same as the
+ * input register.
+ */
+reg_t
+reg_alloc_alloc_preserve (reg_alloc_t *ra, reg_t reg)
+{
+    int idx = find_index (ra, reg);
+    reg_t new_loc;
+
+    if (!ra->status[i].allocated)
+    {
+	ra->status[i].allocated = TRUE;
+	return reg;
+    }
+
+    new_loc = reg_alloc_alloc (ra);
+    if (new_loc && new_loc != reg)
+    {
+	BEGIN_ASM (ra->fragment)
+	    ra->pool->move_instruction, new_loc, reg,
+	END_ASM (ra->fragment);
+    }
+
+    return new_loc;
+}
+
+void
+reg_alloc_free (reg_alloc_t *ra, reg_t reg)
+{
+    int idx = find_index (ra, reg);
+    reg_status_t *status = &ra->status[idx];
+
+    status->allocated = FALSE;
+}
+
+int
+main ()
+{
+    return 0;
+}
diff --git a/regscope.h b/regscope.h
new file mode 100644
index 0000000..4276064
--- /dev/null
+++ b/regscope.h
@@ -0,0 +1,74 @@
+#include "stack-man.h"
+#include "simplex86.h"
+
+typedef struct reg_alloc_t reg_alloc_t;
+typedef struct reg_status_t reg_status_t;
+typedef struct reg_pool_t reg_pool_t;
+
+#define MAX_REGISTERS (16)
+#define MAX_SPILLS    (16)
+
+struct reg_pool_t
+{
+    int         n_registers;
+    int         bytes_per_reg;
+    op_t	move_instruction;
+    reg_t       registers[MAX_REGISTERS];
+};
+
+struct reg_status_t
+{
+    int		allocated;
+    int		spillable;
+    int		n_spills;
+    int		spill_offsets[MAX_SPILLS];
+};
+
+struct reg_alloc_t
+{
+    reg_status_t	status[MAX_REGISTERS];
+    const reg_pool_t *	pool;
+    stack_man_t *	stack;
+    fragment_t *	fragment;
+};
+
+void
+reg_alloc_init (reg_alloc_t *ra,
+		fragment_t *frag,
+		stack_man_t *stack,
+		const reg_pool_t *reg_pool);
+
+/* Calling this indicates that the registers given will not be
+ * needed until the corresponding reg_alloc_end_spill().
+ */
+void
+reg_alloc_begin_spill (reg_alloc_t *ra, reg_t reg, ...);
+
+void
+reg_alloc_end_spill (reg_alloc_t *ra, reg_t reg, ...);
+
+/* This allocates a register. Calliing this function indicates that
+ * the register in question will be freed before the current 
+ * begin/end_spill group ends.
+ */
+reg_t
+reg_alloc_alloc_local (reg_alloc_t *ra);
+
+/* This resurrects a register that has been freed, but is known
+ * to still contain a useful value. The allocation is non-local,
+ * so it is allowed to survive the current begin/end_spill group.
+ *
+ * The returned register may or may not be the same as the
+ * input register.
+ */
+reg_t
+reg_alloc_alloc_preserve (reg_alloc_t *ra, reg_t reg);
+
+/* Allocates a register. The allocation is non-local, so it
+ * is allowed to survive the current begin/end_spill group.
+ */
+reg_t
+reg_alloc_alloc (reg_alloc_t *ra);
+
+void
+reg_alloc_free (reg_alloc_t *ra, reg_t reg);