initial full matrix multiplication

4 files changed, 636 insertions, 51 deletions

diff --git a/deps/ejit b/deps/ejit
-Subproject b0a6350b96f6ef3f8b2af5a385474d885a41759
+Subproject a371bc41f4d208bcb2f8b56677504e1746e3bb8
diff --git a/include/berg/vm.h b/include/berg/vm.h
index 0f87452..66a7f06 100644
--- a/include/berg/vm.h
+++ b/include/berg/vm.h
@@ -15,13 +15,22 @@ enum berg_envcall {
 
 enum berg_op {
 	BERG_LABEL,
-	BERG_LD32,
-	BERG_ST32,
+	BERG_LDXR_I32,
+	BERG_STXR_32,
 	BERG_CALL,
 	BERG_RET,
 	BERG_ECALL,
 	BERG_MOVI,
 	BERG_MOVR,
+
+	BERG_ADDR,
+	BERG_ADDI,
+	BERG_LSHI,
+
+	BERG_MULR,
+
+	BERG_BGER,
+	BERG_JMP,
 };
 
 const char *berg_op_str(enum berg_op op);
@@ -55,6 +64,14 @@ struct berg_operand {
 	struct berg_gpr gpr;
 };
 
+struct berg_reloc {
+	size_t r;
+};
+
+struct berg_label {
+	size_t l;
+};
+
 static inline struct berg_operand berg_operand_gpr(struct berg_gpr gpr, enum berg_type type)
 {
 	return (struct berg_operand){.type = type, .gpr = gpr};
@@ -99,6 +116,48 @@ long berg_movr(struct berg_func *f,
 		struct berg_gpr ra,
 		struct berg_gpr rb);
 
+long berg_mulr(struct berg_func *f,
+		struct berg_gpr ra,
+		struct berg_gpr rb,
+		struct berg_gpr rc);
+
+long berg_addr(struct berg_func *f,
+		struct berg_gpr ra,
+		struct berg_gpr rb,
+		struct berg_gpr rc);
+
+long berg_stxr_32(struct berg_func *f,
+		struct berg_gpr rb,
+		struct berg_gpr rc,
+		struct berg_gpr rd);
+
+long berg_ldxr_i32(struct berg_func *f,
+		struct berg_gpr rb,
+		struct berg_gpr rc,
+		struct berg_gpr rd);
+
+long berg_addi(struct berg_func *f,
+		struct berg_gpr ra,
+		struct berg_gpr rb,
+		int64_t imm);
+
+long berg_lshi(struct berg_func *f,
+		struct berg_gpr ra,
+		struct berg_gpr rb,
+		int64_t imm);
+
+void berg_patch(struct berg_func *f,
+		struct berg_reloc r,
+		struct berg_label l);
+
+struct berg_reloc berg_bger(struct berg_func *f,
+		struct berg_gpr rb,
+		struct berg_gpr rc);
+
+struct berg_reloc berg_jmp(struct berg_func *f);
+
+struct berg_label berg_label(struct berg_func *f);
+
 long berg_ret(struct berg_func *f);
 
 #endif /* BERG_VM_H */
diff --git a/src/main.c b/src/main.c
index d09a87b..b344c11 100644
--- a/src/main.c
+++ b/src/main.c
@@ -5,8 +5,12 @@
 
 static struct berg_func *compile_init(struct berg_vm *vm)
 {
-	struct berg_gpr A = berg_gpr(0);
-	struct berg_gpr B = berg_gpr(1);
+	struct berg_gpr A = berg_gpr_tmp(0);
+	struct berg_gpr B = berg_gpr_tmp(1);
+	struct berg_gpr I = berg_gpr_tmp(2);
+	struct berg_gpr J = berg_gpr_tmp(3);
+	struct berg_gpr T = berg_gpr_tmp(4);
+
 	struct berg_operand operands[] = {
 		/** @todo pointer bounds in type signature? */
 		berg_operand_gpr(A, BERG_POINTER),
@@ -14,6 +18,42 @@ static struct berg_func *compile_init(struct berg_vm *vm)
 	};
 	struct berg_func *f = create_berg_func(vm, BERG_VOID, 2, operands);
 
+	/* outer loop */
+	berg_movi(f, I, 0);
+	struct berg_label outer = berg_label(f);
+
+	/* stop if I >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc outer_jmp = berg_bger(f, I, T);
+
+	/* inner loop */
+	berg_movi(f, J, 0);
+	struct berg_label inner = berg_label(f);
+
+	/* stop if J >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc inner_jmp = berg_bger(f, J, T);
+
+	/* &[I][J] = 4 * (J + MATRIX_SIZE * I) */
+	berg_movi(f, T, MATRIX_SIZE);
+	berg_mulr(f, T, T, I);
+	berg_addr(f, T, T, J);
+	berg_lshi(f, T, T, 2);
+
+	/* A[I][J] = I */
+	berg_stxr_32(f, I, A, T);
+
+	/* B[I][J] = J */
+	berg_stxr_32(f, J, B, T);
+
+	berg_addi(f, J, J, 1);
+	berg_patch(f, berg_jmp(f), inner);
+	berg_patch(f, inner_jmp, berg_label(f));
+
+	berg_addi(f, I, I, 1);
+	berg_patch(f, berg_jmp(f), outer);
+	berg_patch(f, outer_jmp, berg_label(f));
+
 	berg_ret(f);
 	compile_berg_func(f);
 	return f;
@@ -26,6 +66,14 @@ static struct berg_func *compile_mult(struct berg_vm *vm)
 	struct berg_gpr A = berg_gpr_tmp(0);
 	struct berg_gpr B = berg_gpr_tmp(1);
 	struct berg_gpr C = berg_gpr_tmp(2);
+	struct berg_gpr I = berg_gpr_tmp(3);
+	struct berg_gpr J = berg_gpr_tmp(4);
+	struct berg_gpr K = berg_gpr_tmp(5);
+	struct berg_gpr R = berg_gpr_tmp(6);
+	struct berg_gpr T = berg_gpr_tmp(7);
+	struct berg_gpr AT = berg_gpr_tmp(8);
+	struct berg_gpr BT = berg_gpr_tmp(9);
+
 	struct berg_operand operands[] = {
 		berg_operand_gpr(A, BERG_POINTER),
 		berg_operand_gpr(B, BERG_POINTER),
@@ -33,11 +81,73 @@ static struct berg_func *compile_mult(struct berg_vm *vm)
 	};
 	struct berg_func *f = create_berg_func(vm, BERG_VOID, 3, operands);
 
-	/* should be fixed in ejit itself but for now, just pretend we're doing
-	 * something with the registers */
-	berg_movr(f, B, A);
-	berg_movr(f, C, B);
-	berg_movr(f, A, C);
+	/* outer loop */
+	berg_movi(f, I, 0);
+	struct berg_label outer = berg_label(f);
+
+	/* stop if I >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc outer_jmp = berg_bger(f, I, T);
+
+	/* inner loop */
+	berg_movi(f, J, 0);
+	struct berg_label inner = berg_label(f);
+
+	/* stop if J >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc inner_jmp = berg_bger(f, J, T);
+
+	berg_movi(f, R, 0);
+
+	/* innermost loop */
+	berg_movi(f, K, 0);
+	struct berg_label innermost = berg_label(f);
+
+	/* stop if K >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc innermost_jmp = berg_bger(f, K, T);
+
+	/* A[I][K] */
+	/* &[I][K] = 4 * (K + MATRIX_SIZE * I) */
+	berg_movi(f, T, MATRIX_SIZE);
+	berg_mulr(f, T, T, I);
+	berg_addr(f, T, T, K);
+	berg_lshi(f, T, T, 2);
+
+	berg_ldxr_i32(f, AT, A, T);
+
+	/* B[K][J] */
+	berg_movi(f, T, MATRIX_SIZE);
+	berg_mulr(f, T, T, K);
+	berg_addr(f, T, T, J);
+	berg_lshi(f, T, T, 2);
+
+	berg_ldxr_i32(f, BT, B, T);
+
+	/* R += A[I][K] * B[K][J] */
+	berg_mulr(f, T, AT, BT);
+	berg_addr(f, R, R, T);
+
+	berg_addi(f, K, K, 1);
+	berg_patch(f, berg_jmp(f), innermost);
+	berg_patch(f, innermost_jmp, berg_label(f));
+
+	/* &[I][J] = 4 * (J + MATRIX_SIZE * I) */
+	berg_movi(f, T, MATRIX_SIZE);
+	berg_mulr(f, T, T, I);
+	berg_addr(f, T, T, J);
+	berg_lshi(f, T, T, 2);
+
+	/* C[I][J] = R */
+	berg_stxr_32(f, R, C, T);
+
+	berg_addi(f, J, J, 1);
+	berg_patch(f, berg_jmp(f), inner);
+	berg_patch(f, inner_jmp, berg_label(f));
+
+	berg_addi(f, I, I, 1);
+	berg_patch(f, berg_jmp(f), outer);
+	berg_patch(f, outer_jmp, berg_label(f));
 
 	berg_ret(f);
 	compile_berg_func(f);
@@ -46,12 +156,59 @@ static struct berg_func *compile_mult(struct berg_vm *vm)
 
 static struct berg_func *compile_dump(struct berg_vm *vm)
 {
+	struct berg_gpr A0 = berg_gpr_arg(0);
+
 	struct berg_gpr C = berg_gpr_tmp(0);
+	struct berg_gpr I = berg_gpr_tmp(1);
+	struct berg_gpr J = berg_gpr_tmp(2);
+	struct berg_gpr T = berg_gpr_tmp(3);
+	struct berg_gpr S = berg_gpr_tmp(4);
+
 	struct berg_operand operands[] = {
 		berg_operand_gpr(C, BERG_POINTER),
 	};
 	struct berg_func *f = create_berg_func(vm, BERG_VOID, 1, operands);
 
+	berg_movi(f, S, 0);
+
+	/* outer loop */
+	berg_movi(f, I, 0);
+	struct berg_label outer = berg_label(f);
+
+	/* stop if I >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc outer_jmp = berg_bger(f, I, T);
+
+	/* inner loop */
+	berg_movi(f, J, 0);
+	struct berg_label inner = berg_label(f);
+
+	/* stop if J >= MATRIX_SIZE */
+	berg_movi(f, T, MATRIX_SIZE);
+	struct berg_reloc inner_jmp = berg_bger(f, J, T);
+
+	/* &[I][J] = 4 * (J + MATRIX_SIZE * I) */
+	berg_movi(f, T, MATRIX_SIZE);
+	berg_mulr(f, T, T, I);
+	berg_addr(f, T, T, J);
+	berg_lshi(f, T, T, 2);
+
+	/* S += C[I][J] = I */
+	berg_ldxr_i32(f, T, C, T);
+	berg_addr(f, S, S, T);
+
+	berg_addi(f, J, J, 1);
+	berg_patch(f, berg_jmp(f), inner);
+	berg_patch(f, inner_jmp, berg_label(f));
+
+	berg_addi(f, I, I, 1);
+	berg_patch(f, berg_jmp(f), outer);
+	berg_patch(f, outer_jmp, berg_label(f));
+
+	/* print out 'hash' of matrix */
+	berg_movr(f, A0, S);
+	berg_ecall(f, T, BERG_PUTI32);
+
 	berg_ret(f);
 	compile_berg_func(f);
 	return f;
diff --git a/src/vm.c b/src/vm.c
index fdcd71b..3e98d2c 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -15,6 +15,19 @@
 #define VEC_TYPE struct berg_operand
 #include <conts/vec.h>
 
+#define VEC_NAME labels
+#define VEC_TYPE struct ejit_label
+#include <conts/vec.h>
+
+struct reloc_helper {
+	struct ejit_reloc reloc;
+	size_t target;
+};
+
+#define VEC_NAME relocs
+#define VEC_TYPE struct reloc_helper
+#include <conts/vec.h>
+
 struct alloc {
 	void *base, *top;
 };
@@ -37,6 +50,12 @@ struct berg_func {
 	/* vm instructions */
 	struct insns insns;
 
+	/* labels */
+	struct labels labels;
+
+	/* relocs */
+	struct relocs relocs;
+
 	/* arguments */
 	struct operands operands;
 
@@ -153,13 +172,22 @@ const uint8_t NOGPR = 255;
  * looking */
 const struct ejit_gpr ALLOC_MASK = EJIT_GPR(0);
 const struct ejit_gpr ALLOC_BUF = EJIT_GPR(1);
-const struct ejit_gpr VM_PTR = EJIT_GPR(2);
+const struct ejit_gpr TMP0 = EJIT_GPR(2);
+const struct ejit_gpr TMP1 = EJIT_GPR(3);
 
 static struct ejit_gpr ejit_gpr_from(uint8_t r)
 {
 	assert(r != NOGPR);
-	/* three reserved registers */
-	return EJIT_GPR(r + 3);
+	/* four reserved registers (seems like a lot?) */
+	return EJIT_GPR(r + 4);
+}
+
+static long reload_alloc_status(struct berg_func *f)
+{
+	EJIT_LDI(f->func, size_t, ALLOC_MASK, &f->vm->allocs.size);
+	EJIT_LDI(f->func, void *, ALLOC_BUF, &f->vm->allocs.buf);
+	ejit_subi(f->func, ALLOC_MASK, ALLOC_MASK, 1);
+	return 0;
 }
 
 struct berg_func *create_berg_func(struct berg_vm *vm, enum berg_type rtype, size_t argc, const struct berg_operand args[argc])
@@ -167,6 +195,7 @@ struct berg_func *create_berg_func(struct berg_vm *vm, enum berg_type rtype, siz
 	struct berg_func *f = calloc(1, sizeof(struct berg_func));
 	f->vm = vm;
 	f->insns = insns_create(0);
+	f->labels = labels_create(0);
 	f->operands = operands_create(argc);
 
 	struct ejit_operand ejit_operands[argc];
@@ -179,6 +208,8 @@ struct berg_func *create_berg_func(struct berg_vm *vm, enum berg_type rtype, siz
 	}
 
 	f->func = ejit_create_func(ejit_type_from(rtype), argc, ejit_operands);
+	reload_alloc_status(f);
+
 	funcs_append(&vm->funcs, f);
 	return f;
 }
@@ -208,6 +239,54 @@ static struct berg_insn insn_orr(enum berg_op op, struct berg_gpr ra, struct ber
 	};
 }
 
+static struct berg_insn insn_orrr(enum berg_op op, struct berg_gpr ra, struct berg_gpr rb, struct berg_gpr rc)
+{
+	return (struct berg_insn){
+		.op = op,
+		.ra = ra.r,
+		.rb = rb.r,
+		.rc = rc.r,
+		.rd = NOGPR,
+		.imm = 0
+	};
+}
+
+static struct berg_insn insn_orri(enum berg_op op, struct berg_gpr ra, struct berg_gpr rb, int64_t imm)
+{
+	return (struct berg_insn){
+		.op = op,
+		.ra = ra.r,
+		.rb = rb.r,
+		.rc = NOGPR,
+		.rd = NOGPR,
+		.imm = imm
+	};
+}
+
+static struct berg_insn insn_oxrri(enum berg_op op, struct berg_gpr rb, struct berg_gpr rc, int64_t imm)
+{
+	return (struct berg_insn){
+		.op = op,
+		.ra = NOGPR,
+		.rb = rb.r,
+		.rc = rc.r,
+		.rd = NOGPR,
+		.imm = imm
+	};
+}
+
+static struct berg_insn insn_oxrrr(enum berg_op op, struct berg_gpr rb, struct berg_gpr rc, struct berg_gpr rd)
+{
+	return (struct berg_insn){
+		.op = op,
+		.ra = NOGPR,
+		.rb = rb.r,
+		.rc = rc.r,
+		.rd = rd.r,
+		.imm = 0
+	};
+}
+
 static struct berg_insn insn_o(enum berg_op op)
 {
 	return (struct berg_insn){
@@ -232,6 +311,18 @@ static struct berg_insn insn_op(enum berg_op op, void *p)
 	};
 }
 
+static struct berg_insn insn_oi(enum berg_op op, int64_t imm)
+{
+	return (struct berg_insn){
+		.op = op,
+		.ra = NOGPR,
+		.rb = NOGPR,
+		.rc = NOGPR,
+		.rd = NOGPR,
+		.imm = imm
+	};
+}
+
 long berg_movi(struct berg_func *f, struct berg_gpr ra, int64_t imm)
 {
 	insns_append(&f->insns, insn_ori(BERG_MOVI, ra, imm));
@@ -264,18 +355,88 @@ long berg_call(struct berg_func *f, struct berg_func *c)
 	return 0;
 }
 
+long berg_addr(struct berg_func *f, struct berg_gpr ra, struct berg_gpr rb, struct berg_gpr rc)
+{
+	insns_append(&f->insns, insn_orrr(BERG_ADDR, ra, rb, rc));
+	return 0;
+}
+
+long berg_addi(struct berg_func *f, struct berg_gpr ra, struct berg_gpr rb, int64_t imm)
+{
+	insns_append(&f->insns, insn_orri(BERG_ADDI, ra, rb, imm));
+	return 0;
+}
+
+long berg_mulr(struct berg_func *f, struct berg_gpr ra, struct berg_gpr rb, struct berg_gpr rc)
+{
+	insns_append(&f->insns, insn_orrr(BERG_MULR, ra, rb, rc));
+	return 0;
+}
+
+long berg_lshi(struct berg_func *f, struct berg_gpr ra, struct berg_gpr rb, int64_t imm)
+{
+	assert(imm < 64);
+	insns_append(&f->insns, insn_orri(BERG_LSHI, ra, rb, imm));
+	return 0;
+}
+
+long berg_stxr_32(struct berg_func *f, struct berg_gpr rb, struct berg_gpr rc, struct berg_gpr rd)
+{
+	insns_append(&f->insns, insn_oxrrr(BERG_STXR_32, rb, rc, rd));
+	return 0;
+}
+
+long berg_ldxr_i32(struct berg_func *f, struct berg_gpr rb, struct berg_gpr rc, struct berg_gpr rd)
+{
+	insns_append(&f->insns, insn_oxrrr(BERG_LDXR_I32, rb, rc, rd));
+	return 0;
+}
+
+struct berg_reloc berg_bger(struct berg_func *f,
+		struct berg_gpr rb,
+		struct berg_gpr rc)
+{
+	struct berg_reloc r = {.r = insns_len(&f->insns)};
+	insns_append(&f->insns, insn_oxrri(BERG_BGER, rb, rc, 0));
+	return r;
+}
+
+struct berg_reloc berg_jmp(struct berg_func *f)
+{
+	struct berg_reloc r = {.r = insns_len(&f->insns)};
+	insns_append(&f->insns, insn_oi(BERG_JMP, 0));
+	return r;
+}
+
+struct berg_label berg_label(struct berg_func *f)
+{
+	return (struct berg_label){.l = insns_len(&f->insns)};
+}
+
+void berg_patch(struct berg_func *f, struct berg_reloc r, struct berg_label l)
+{
+	struct berg_insn *insn = insns_at(&f->insns, r.r);
+	insn->imm = l.l;
+}
+
 const char *berg_op_str(enum berg_op op)
 {
 #define CASE(x) case x: return #x;
 	switch (op) {
 		CASE(BERG_LABEL);
-		CASE(BERG_LD32);
-		CASE(BERG_ST32);
+		CASE(BERG_LDXR_I32);
+		CASE(BERG_STXR_32);
 		CASE(BERG_CALL);
 		CASE(BERG_RET);
 		CASE(BERG_ECALL);
 		CASE(BERG_MOVI);
 		CASE(BERG_MOVR);
+		CASE(BERG_ADDR);
+		CASE(BERG_ADDI);
+		CASE(BERG_LSHI);
+		CASE(BERG_MULR);
+		CASE(BERG_BGER);
+		CASE(BERG_JMP);
 	}
 
 	return "???";
@@ -328,18 +489,16 @@ static long compile_malloc(struct berg_func *f, struct berg_insn *i)
 {
 	struct berg_gpr size = berg_gpr_arg(0);
 
-	ejit_movi(f->func, VM_PTR, (int64_t)f->vm);
+	ejit_movi(f->func, TMP0, (int64_t)f->vm);
 	struct ejit_operand args[] = {
-		EJIT_OPERAND_GPR(VM_PTR.r, EJIT_POINTER),
+		EJIT_OPERAND_GPR(TMP0.r, EJIT_POINTER),
 		EJIT_OPERAND_GPR(ejit_gpr_from(size.r).r, EJIT_TYPE(size_t)),
 	};
 	ejit_escapei_l(f->func, escape_malloc, 2, args);
 	ejit_retval(f->func, ejit_gpr_from(i->ra));
 
 	/* reload allocation context if the hashmap has changed */
-	EJIT_LDI(f->func, size_t, ALLOC_MASK, &f->vm->allocs.size);
-	EJIT_LDI(f->func, void *, ALLOC_BUF, &f->vm->allocs.buf);
-	ejit_subi(f->func, ALLOC_MASK, ALLOC_MASK, 1);
+	reload_alloc_status(f);
 	return 0;
 }
 
@@ -347,9 +506,9 @@ static long compile_free(struct berg_func *f, struct berg_insn *i)
 {
 	struct berg_gpr ptr = berg_gpr_arg(0);
 
-	ejit_movi(f->func, VM_PTR, (int64_t)f->vm);
+	ejit_movi(f->func, TMP0, (int64_t)f->vm);
 	struct ejit_operand args[] = {
-		EJIT_OPERAND_GPR(VM_PTR.r, EJIT_POINTER),
+		EJIT_OPERAND_GPR(TMP0.r, EJIT_POINTER),
 		EJIT_OPERAND_GPR(ejit_gpr_from(ptr.r).r, EJIT_POINTER),
 	};
 	ejit_escapei_l(f->func, escape_free, 2, args);
@@ -360,11 +519,32 @@ static long compile_free(struct berg_func *f, struct berg_insn *i)
 	return 0;
 }
 
+static long escape_puti32(size_t argc, const struct ejit_arg args[argc])
+{
+	int32_t i32 = EJIT_PARAM(argc, args, 0, int32_t);
+	printf("%lli\n", (long long)i32);
+	return 0;
+}
+
+static long compile_puti32(struct berg_func *f, struct berg_insn *i)
+{
+	struct berg_gpr i32 = berg_gpr_arg(0);
+
+	struct ejit_operand args[] = {
+		EJIT_OPERAND_GPR(ejit_gpr_from(i32.r).r, EJIT_INT32),
+	};
+	ejit_escapei_l(f->func, escape_puti32, 1, args);
+	ejit_retval(f->func, ejit_gpr_from(i->ra));
+
+	return 0;
+}
+
 static long compile_ecall(struct berg_func *f, struct berg_insn *i)
 {
 	switch (i->imm) {
 	case BERG_MALLOC: return compile_malloc(f, i);
 	case BERG_FREE: return compile_free(f, i);
+	case BERG_PUTI32: return compile_puti32(f, i);
 	}
 
 	fprintf(stderr, "unknown envcall: %lli\n", (long long)i->imm);
@@ -392,9 +572,7 @@ static long compile_call(struct berg_func *f, struct berg_insn *i)
 
 	/* reload allocation map since we don't know if something was maybe
 	 * freed */
-	EJIT_LDI(f->func, size_t, ALLOC_MASK, &f->vm->allocs.size);
-	EJIT_LDI(f->func, void *, ALLOC_BUF, &f->vm->allocs.buf);
-	ejit_subi(f->func, ALLOC_MASK, ALLOC_MASK, 1);
+	reload_alloc_status(f);
 	return 0;
 }
 
@@ -404,39 +582,228 @@ static long compile_movr(struct berg_func *f, struct berg_insn *i)
 	return 0;
 }
 
-long compile_berg_func(struct berg_func *f)
+static long compile_bger(struct berg_func *f, struct berg_insn *i)
 {
-	/** @todo eventually should make this be block-specific */
-	long ret = 0;
-	foreach(insns, i, &f->insns) switch (i->op) {
-	case BERG_RET:
-		if ((ret = compile_ret(f)))
-			return ret;
-		break;
+	struct ejit_reloc r = ejit_bger(f->func, ejit_gpr_from(i->rb), ejit_gpr_from(i->rc));
+	relocs_append(&f->relocs, (struct reloc_helper){.reloc = r, .target = i->imm});
+	return 0;
+}
 
-	case BERG_MOVI:
-		if ((ret = compile_movi(f, i)))
-			return ret;
-		break;
+static long compile_jmp(struct berg_func *f, struct berg_insn *i)
+{
+	struct ejit_reloc r = ejit_jmp(f->func);
+	relocs_append(&f->relocs, (struct reloc_helper){.reloc = r, .target = i->imm});
+	return 0;
+}
 
-	case BERG_MOVR:
-		if ((ret = compile_movr(f, i)))
-			return ret;
-		break;
+static long compile_mulr(struct berg_func *f, struct berg_insn *i)
+{
+	ejit_mulr(f->func, ejit_gpr_from(i->ra), ejit_gpr_from(i->rb), ejit_gpr_from(i->rc));
+	return 0;
+}
 
-	case BERG_CALL:
-		if ((ret = compile_call(f, i)))
-			return ret;
-		break;
+static long compile_addr(struct berg_func *f, struct berg_insn *i)
+{
+	ejit_addr(f->func, ejit_gpr_from(i->ra), ejit_gpr_from(i->rb), ejit_gpr_from(i->rc));
+	return 0;
+}
 
-	case BERG_ECALL:
-		if ((ret = compile_ecall(f, i)))
-			return ret;
-		break;
+static long compile_addi(struct berg_func *f, struct berg_insn *i)
+{
+	ejit_addi(f->func, ejit_gpr_from(i->ra), ejit_gpr_from(i->rb), i->imm);
+	return 0;
+}
 
-	default:
-		fprintf(stderr, "unhandled op %s\n", berg_op_str(i->op));
-		return -1;
+static long compile_lshi(struct berg_func *f, struct berg_insn *i)
+{
+	ejit_lshi(f->func, ejit_gpr_from(i->ra), ejit_gpr_from(i->rb), i->imm);
+	return 0;
+}
+
+static long compile_reg_bounds_check(struct berg_func *f, ejit_escape_i_t handler, uint8_t base, uint8_t offset, uint8_t data, enum ejit_type type)
+{
+	/* load base */
+	ejit_rshi(f->func, TMP0, ejit_gpr_from(base), 4);
+	ejit_andr(f->func, TMP0, TMP0, ALLOC_MASK);
+
+	static_assert(sizeof(struct alloc) == 1ULL << 4);
+	ejit_lshi(f->func, TMP0, TMP0, 4);
+
+	/** @todo would probably be faster to have inline boolean checks instead of
+	 * double jumps, should add support for something like that to ejit */
+
+	EJIT_LDXR(f->func, size_t, TMP1, ALLOC_BUF, TMP0);
+	/* check that base is larger than loaded value (value can be NULL at
+	 * this point) */
+	struct ejit_reloc base_check = ejit_bger_u(f->func, TMP1, ejit_gpr_from(base));
+
+	/* do slow fallback */
+	struct ejit_operand args[] = {
+		EJIT_OPERAND_GPR(TMP0.r, EJIT_POINTER),
+		EJIT_OPERAND_GPR(ejit_gpr_from(base).r, EJIT_POINTER),
+		EJIT_OPERAND_GPR(ejit_gpr_from(offset).r, EJIT_TYPE(intptr_t)),
+		EJIT_OPERAND_GPR(ejit_gpr_from(data).r, type)
+	};
+	ejit_movi(f->func, TMP0, (uintptr_t)f->vm);
+	ejit_escapei_l(f->func, handler, 4, args);
+
+	struct ejit_reloc base_continue = ejit_jmp(f->func);
+	ejit_patch(f->func, base_check, ejit_label(f->func));
+
+	/* load top */
+	ejit_addi(f->func, TMP1, TMP0, offsetof(struct alloc, top));
+	EJIT_LDXR(f->func, void *, TMP1, ALLOC_BUF, TMP1);
+	ejit_addr(f->func, TMP0, ejit_gpr_from(base), ejit_gpr_from(offset));
+	struct ejit_reloc top_check = ejit_bler_u(f->func, TMP0, TMP1);
+
+	/* do slow fallback */
+	ejit_movi(f->func, TMP0, (uintptr_t)f->vm);
+	ejit_escapei_l(f->func, handler, 4, args);
+
+	struct ejit_label ok = ejit_label(f->func);
+	ejit_patch(f->func, top_check, ok);
+	ejit_patch(f->func, base_continue, ok);
+	return 0;
+}
+
+static long escape_stxr_32_bounds_check(size_t argc, const struct ejit_arg args[argc])
+{
+	struct berg_vm *vm = EJIT_PARAM(argc, args, 0, struct berg_vm *);
+	void *base = EJIT_PARAM(argc, args, 1, void *);
+	intptr_t offset = EJIT_PARAM(argc, args, 2, intptr_t);
+	int32_t data = EJIT_PARAM(argc, args, 3, int32_t);
+
+	fprintf(stderr, "stxr32 fallback bounds check unimplemented\n");
+	abort();
+	return 0;
+}
+
+static long compile_stxr_32(struct berg_func *f, struct berg_insn *i)
+{
+	/* for now, always check if address is valid */
+	uint8_t base = i->rc;
+	uint8_t offset = i->rd;
+	uint8_t data = i->rb;
+
+	compile_reg_bounds_check(f, escape_stxr_32_bounds_check, base, offset, data, EJIT_INT32);
+
+	/* ok, do actual store */
+	ejit_stxr_32(f->func, ejit_gpr_from(data), ejit_gpr_from(base), ejit_gpr_from(offset));
+	return 0;
+}
+
+static long escape_ldxr_i32_bounds_check(size_t argc, const struct ejit_arg args[argc])
+{
+	struct berg_vm *vm = EJIT_PARAM(argc, args, 0, struct berg_vm *);
+	void *base = EJIT_PARAM(argc, args, 1, void *);
+	intptr_t offset = EJIT_PARAM(argc, args, 2, intptr_t);
+	int32_t data = EJIT_PARAM(argc, args, 3, int32_t);
+
+	fprintf(stderr, "ldxr32 fallback bounds check unimplemented\n");
+	abort();
+	return 0;
+}
+
+static long compile_ldxr_i32(struct berg_func *f, struct berg_insn *i)
+{
+	/* for now, always check if address is valid */
+	uint8_t base = i->rc;
+	uint8_t offset = i->rd;
+	uint8_t data = i->rb;
+
+	compile_reg_bounds_check(f, escape_ldxr_i32_bounds_check, base, offset, data, EJIT_INT32);
+
+	/* ok, do actual store */
+	ejit_stxr_32(f->func, ejit_gpr_from(data), ejit_gpr_from(base), ejit_gpr_from(offset));
+	return 0;
+}
+
+long compile_berg_func(struct berg_func *f)
+{
+	/** @todo eventually should make this be block-specific */
+	long ret = 0;
+	foreach(insns, i, &f->insns) {
+		labels_append(&f->labels, ejit_label(f->func));
+
+		switch (i->op) {
+		case BERG_RET:
+			if ((ret = compile_ret(f)))
+				return ret;
+			break;
+
+		case BERG_MOVI:
+			if ((ret = compile_movi(f, i)))
+				return ret;
+			break;
+
+		case BERG_MOVR:
+			if ((ret = compile_movr(f, i)))
+				return ret;
+			break;
+
+		case BERG_CALL:
+			if ((ret = compile_call(f, i)))
+				return ret;
+			break;
+
+		case BERG_ECALL:
+			if ((ret = compile_ecall(f, i)))
+				return ret;
+			break;
+
+		case BERG_MULR:
+			if ((ret = compile_mulr(f, i)))
+				return ret;
+			break;
+
+		case BERG_ADDR:
+			if ((ret = compile_addr(f, i)))
+				return ret;
+			break;
+
+		case BERG_ADDI:
+			if ((ret = compile_addi(f, i)))
+				return ret;
+			break;
+
+		case BERG_LSHI:
+			if ((ret = compile_lshi(f, i)))
+				return ret;
+			break;
+
+		case BERG_STXR_32:
+			if ((ret = compile_stxr_32(f, i)))
+				return ret;
+			break;
+
+		case BERG_LDXR_I32:
+			if ((ret = compile_ldxr_i32(f, i)))
+				return ret;
+			break;
+
+		case BERG_LABEL:
+			break;
+
+		case BERG_BGER:
+			if ((ret = compile_bger(f, i)))
+				return ret;
+			break;
+
+		case BERG_JMP:
+			if ((ret = compile_jmp(f, i)))
+				return ret;
+			break;
+
+		default:
+			fprintf(stderr, "unhandled op %s\n", berg_op_str(i->op));
+			return -1;
+		}
+	}
+
+	/* patch relocs */
+	foreach(relocs, r, &f->relocs) {
+		struct ejit_label *label = labels_at(&f->labels, r->target);
+		ejit_patch(f->func, r->reloc, *label);
 	}
 
 	ejit_compile_func(f->func);
@@ -455,6 +822,8 @@ static void destroy_berg_func(struct berg_func *f)
 {
 	insns_destroy(&f->insns);
 	operands_destroy(&f->operands);
+	labels_destroy(&f->labels);
+	relocs_destroy(&f->relocs);
 	ejit_destroy_func(f->func);
 	free(f);
 }