diff --git a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
index 6ddd805..f4ed0e9 100644
--- a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
+++ b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
@@ -504,7 +504,7 @@ param int SHAKE256_RATE = 136;
 param int SHA3_256_RATE = 136;
 param int SHA3_512_RATE = 72;
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _sha3_256(reg ptr u8[32] out, reg u64 in inlen) -> reg ptr u8[32]
 {
   reg u256[7] state;
@@ -537,7 +537,7 @@ fn _sha3_256(reg ptr u8[32] out, reg u64 in inlen) -> reg ptr u8[32]
 }
 
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _sha3_256_32(reg ptr u8[32] out, reg ptr u8[KYBER_SYMBYTES] in) -> reg ptr u8[32]
 {
   reg u256[7] state;
@@ -585,7 +585,7 @@ fn _sha3_256_32(reg ptr u8[32] out, reg ptr u8[KYBER_SYMBYTES] in) -> reg ptr u8
 }
 
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _shake256_64(reg u64 out outlen, reg const ptr u8[64] in)
 {
   reg u256[7] state;
@@ -625,7 +625,7 @@ fn _shake256_64(reg u64 out outlen, reg const ptr u8[64] in)
 }
 
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _shake256_128_33(reg ptr u8[128] out, reg const ptr u8[33] in) -> stack u8[128]
 {
   reg u256[7] state;
@@ -679,7 +679,7 @@ fn _shake256_128_33(reg ptr u8[128] out, reg const ptr u8[33] in) -> stack u8[12
   return out;
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _sha3_512_64(reg ptr u8[64] out, reg const ptr u8[64] in) -> stack u8[64]
 {
   reg u256[7] state;
@@ -728,7 +728,7 @@ fn _sha3_512_64(reg ptr u8[64] out, reg const ptr u8[64] in) -> stack u8[64]
   return out;
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _sha3_512_32(reg ptr u8[64] out, reg const ptr u8[32] in) -> stack u8[64]
 {
   reg u256[7] state;
@@ -1237,7 +1237,7 @@ inline fn __keccakf1600_4x_avx2(reg ptr u256[25] a) -> reg ptr u256[25]
   return a;
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _keccakf1600_4x_avx2(reg ptr u256[25] a) -> reg ptr u256[25]
 {
   a = __keccakf1600_4x_avx2(a);
@@ -1256,7 +1256,7 @@ inline fn _keccakf1600_4x_avx2_(reg ptr u256[25] a) -> reg ptr u256[25]
 
 u256 SHAKE_SEP = (4u64)[0x8000000000000000, 0x8000000000000000, 0x8000000000000000, 0x8000000000000000];
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _shake128_absorb4x_34(reg ptr u256[25] s, reg ptr u8[34] m0 m1 m2 m3) -> reg ptr u256[25]
 {
 	inline int i;
@@ -1307,7 +1307,7 @@ fn _shake128_absorb4x_34(reg ptr u256[25] s, reg ptr u8[34] m0 m1 m2 m3) -> reg
 }
 
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _shake256_absorb4x_33(reg ptr u256[25] s, reg ptr u8[33] m0 m1 m2 m3) -> reg ptr u256[25]
 {
 	inline int i;
@@ -2459,7 +2459,7 @@ fn __poly_cbd_eta2(reg ptr u16[KYBER_N] rp, reg ptr u8[KYBER_ETA2*KYBER_N/4] buf
 }
 
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _poly_getnoise(reg ptr u16[KYBER_N] rp, reg ptr u8[KYBER_SYMBYTES] seed, reg u8 nonce) -> reg ptr u16[KYBER_N]
 {
   inline int i;
@@ -2549,7 +2549,7 @@ fn __shake256_squeezenblocks4x(reg ptr u256[25] state, reg ptr u8[NOISE_NBLOCKS
   return state, buf0, buf1, buf2, buf3;
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _poly_getnoise_eta1_4x(
   reg ptr u16[KYBER_N] r0 r1 r2 r3,
   reg ptr u8[KYBER_SYMBYTES] seed,
@@ -2593,7 +2593,7 @@ fn _poly_getnoise_eta1_4x(
   return r0, r1, r2, r3;
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn _poly_getnoise_eta1122_4x(
   reg ptr u16[KYBER_N] r0 r1 r2 r3,
   reg ptr u8[KYBER_SYMBYTES] seed,
@@ -4629,7 +4629,7 @@ fn __cmov(reg ptr u8[KYBER_SYMBYTES] dst, reg u64 src cnd) -> reg ptr u8[KYBER_S
 // Note2: due to the integration of hakyber implementations into libjade, this file is no longer
 //        required by kyber768/amd64/ref/
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn __crypto_kem_keypair_derand_jazz(reg u64 pkp, reg u64 skp, reg ptr u8[2*KYBER_SYMBYTES] coins)
 {
   stack u8[32] h_pk;
@@ -4668,7 +4668,7 @@ fn __crypto_kem_keypair_derand_jazz(reg u64 pkp, reg u64 skp, reg ptr u8[2*KYBER
   __fromstack32u8(skp, coins[32:32]);
 }
 
-#[returnaddress="stack"]
+#[returnaddress="mmx"]
 fn __crypto_kem_enc_derand_jazz(reg u64 ctp, reg u64 shkp, reg u64 pkp, reg ptr u8[KYBER_SYMBYTES] coins)
 {
   stack u8[KYBER_SYMBYTES * 2] buf kr;