take loops into account in register allocatorHEADmaster

1 files changed, 143 insertions, 0 deletions

diff --git a/tests/z_matrix_mult.c b/tests/z_matrix_mult.c
new file mode 100644
index 0000000..8cfba69
--- /dev/null
+++ b/tests/z_matrix_mult.c
@@ -0,0 +1,143 @@
+#include <ejit/ejit.h>
+#include <assert.h>
+#include "do_jit.h"
+
+/* register allocator had a bug where registers whose last use was midway
+ * through a loop might get overwritten, here's one case that exposed the bug */
+
+#define X 10
+int A[X][X];
+int B[X][X];
+int C[X][X];
+
+bool do_jit = false;
+
+static void init_matrices(int A[X][X], int B[X][X])
+{
+	int counter = 1;
+	for (size_t i = 0; i < X; ++i)
+	for (size_t j = 0; j < X; ++j) {
+		A[i][j] = counter;
+		B[i][j] = counter;
+		C[i][j] = 0;
+
+		counter++;
+	}
+}
+
+static int hash(int C[X][X])
+{
+	int h = 0;
+	for (size_t i = 0; i < X; ++i)
+	for (size_t j = 0; j < X; ++j) {
+		h += C[i][j];
+	}
+
+	return h;
+}
+
+static struct ejit_func *compile()
+{
+#define IR EJIT_GPR(0)
+#define JR EJIT_GPR(1)
+#define KR EJIT_GPR(2)
+#define AR EJIT_GPR(3)
+#define BR EJIT_GPR(4)
+#define CR EJIT_GPR(5)
+#define AX EJIT_GPR(6)
+#define BX EJIT_GPR(7)
+#define CX EJIT_GPR(8)
+#define RR EJIT_GPR(9)
+#define TMP EJIT_GPR(10)
+
+	struct ejit_operand args[3] = {
+		EJIT_OPERAND_GPR(3, EJIT_TYPE(int *)), /* A */
+		EJIT_OPERAND_GPR(4, EJIT_TYPE(int *)), /* B */
+		EJIT_OPERAND_GPR(5, EJIT_TYPE(int *)), /* C */
+	};
+	struct ejit_func *f = ejit_create_func(EJIT_VOID, 3, args);
+
+	ejit_movi(f, IR, 0);
+	struct ejit_label out = ejit_label(f);
+	struct ejit_reloc outer = ejit_bgei(f, IR, X);
+
+	ejit_movi(f, JR, 0);
+	struct ejit_label mid = ejit_label(f);
+	struct ejit_reloc midder = ejit_bgei(f, JR, X);
+
+	ejit_movi(f, RR, 0);
+
+	ejit_movi(f, KR, 0);
+	struct ejit_label in = ejit_label(f);
+	struct ejit_reloc inner = ejit_bgei(f, KR, X);
+
+	/* A[i][k] at addr 4 * (i * 400 + k) */
+	ejit_movi(f, TMP, X);
+	ejit_mulr(f, AX, IR, TMP);
+	ejit_addr(f, AX, AX, KR);
+	ejit_movi(f, TMP, 4);
+	ejit_lshi(f, AX, AX, 2);
+	EJIT_LDXR(f, int, AX, AR, AX);
+
+	/* B[k][j] at addr 4 * (k * 400 + j) */
+	ejit_movi(f, TMP, X);
+	ejit_mulr(f, BX, KR, TMP);
+	ejit_addr(f, BX, BX, JR);
+	ejit_movi(f, TMP, 4);
+	ejit_lshi(f, BX, BX, 2);
+	EJIT_LDXR(f, int, BX, BR, BX);
+
+	/* R += A[i][k] * B[k][j] */
+	ejit_mulr(f, TMP, AX, BX);
+	ejit_addr(f, RR, RR, TMP);
+
+	/* increment inner */
+	ejit_addi(f, KR, KR, 1);
+	ejit_patch(f, ejit_jmp(f), in);
+	/* end of inner */
+	ejit_patch(f, inner, ejit_label(f));
+
+	/* C[i][j] at addr 4 * (i * 400 + j) */
+	ejit_movi(f, TMP, X);
+	ejit_mulr(f, CX, IR, TMP);
+	ejit_addr(f, CX, CX, JR);
+	ejit_lshi(f, CX, CX, 2);
+
+	/* C[i][j] = R */
+	EJIT_STXR(f, int, RR, CX, CR);
+
+	/* increment midder */
+	ejit_addi(f, JR, JR, 1);
+	ejit_patch(f, ejit_jmp(f), mid);
+	/* end of midder */
+	ejit_patch(f, midder, ejit_label(f));
+
+	/* increment outer */
+	ejit_addi(f, IR, IR, 1);
+	ejit_patch(f, ejit_jmp(f), out);
+	/* end of outer */
+	ejit_patch(f, outer, ejit_label(f));
+
+	/* return */
+	ejit_ret(f);
+	ejit_select_compile_func(f, 11, 0, EJIT_USE64(long), do_jit, true);
+	return f;
+}
+
+int main(int argc, char *argv[])
+{
+	(void)argv;
+	do_jit = argc > 1;
+
+	init_matrices(A, B);
+	struct ejit_func *f = compile();
+	struct ejit_arg args[3] = {
+		EJIT_ARG(A, int *),
+		EJIT_ARG(B, int *),
+		EJIT_ARG(C, int *)
+	};
+	ejit_run_func(f, 3, args);
+	assert(hash(C) == 2632750);
+	ejit_destroy_func(f);
+	return 0;
+}