wazevo(frontend): store module context for only imported calls (#2103)

internal/engine/wazevo/backend/backend_test.go

@@ -615,12 +615,10 @@ L5 (SSA Block: blk3):
 L1 (SSA Block: blk0):
 	mov x128?, x0
 	mov x129?, x1
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	bl f1
 	mov x130?, x0
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	mov x2, x130?
@@ -628,7 +626,6 @@ L1 (SSA Block: blk0):
 	mov x3, x131?
 	bl f2
 	mov x132?, x0
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	mov x2, x132?
@@ -645,23 +642,20 @@ L1 (SSA Block: blk0):
 	sub sp, sp, #0x10
 	orr x27, xzr, #0x10
 	str x27, [sp, #-0x10]!
-	str x0, [sp, #0x18]
-	str x1, [sp, #0x10]
-	str x1, [x0, #0x8]
+	str x0, [sp, #0x10]
+	str x1, [sp, #0x18]
 	bl f1
 	mov x2, x0
 	ldr x8, [sp, #0x10]
+	mov x0, x8
 	ldr x9, [sp, #0x18]
-	str x8, [x9, #0x8]
-	mov x0, x9
-	mov x1, x8
+	mov x1, x9
 	movz w3, #0x5, lsl 0
 	bl f2
 	mov x2, x0
 	ldr x8, [sp, #0x10]
-	ldr x9, [sp, #0x18]
-	str x8, [x9, #0x8]
-	mov x0, x9
+	mov x0, x8
+	ldr x8, [sp, #0x18]
 	mov x1, x8
 	bl f3
 	add sp, sp, #0x10
@@ -680,7 +674,6 @@ L1 (SSA Block: blk0):
 	mov x131?, x3
 	mov v132?.8b, v0.8b
 	mov v133?.8b, v1.8b
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	mov x2, x130?
@@ -736,7 +729,6 @@ L1 (SSA Block: blk0):
 	str x3, [sp, #0x14]
 	str s0, [sp, #0x1c]
 	str d1, [sp, #0x20]
-	str x1, [x0, #0x8]
 	ldr w8, [sp, #0x10]
 	mov x4, x8
 	ldr x9, [sp, #0x14]
@@ -794,7 +786,6 @@ L1 (SSA Block: blk0):
 	mov x131?, x3
 	mov v132?.8b, v0.8b
 	mov v133?.8b, v1.8b
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	mov x2, x130?
@@ -895,7 +886,6 @@ L1 (SSA Block: blk0):
 	str q19, [sp, #-0x10]!
 	orr x27, xzr, #0x40
 	str x27, [sp, #-0x10]!
-	str x1, [x0, #0x8]
 	bl f1
 	ldr w8, [sp, #-0xc0]
 	ldr x9, [sp, #-0xb8]
@@ -1884,7 +1874,6 @@ L1 (SSA Block: blk0):
 	ldr x131?, [x129?, #0x20]
 	ldr s132?, [x129?, #0x30]
 	ldr d133?, [x129?, #0x40]
-	str x129?, [x128?, #0x8]
 	mov x0, x128?
 	mov x1, x129?
 	bl f1
@@ -1917,7 +1906,6 @@ L1 (SSA Block: blk0):
 	str s0, [sp, #0x20]
 	ldr d1, [x1, #0x40]
 	str d1, [sp, #0x18]
-	str x1, [x0, #0x8]
 	bl f1
 	ldr x8, [sp, #0x10]
 	ldr w2, [x8, #0x10]

internal/engine/wazevo/frontend/frontend_test.go

@@ -551,12 +551,9 @@ signatures:
 	sig3: i64i64i32_i32i32
 
 blk0: (exec_ctx:i64, module_ctx:i64)
-	Store module_ctx, exec_ctx, 0x8
 	v2:i32 = Call f1:sig1, exec_ctx, module_ctx
 	v3:i32 = Iconst_32 0x5
-	Store module_ctx, exec_ctx, 0x8
 	v4:i32 = Call f2:sig2, exec_ctx, module_ctx, v2, v3
-	Store module_ctx, exec_ctx, 0x8
 	v5:i32, v6:i32 = Call f3:sig3, exec_ctx, module_ctx, v4
 	Jump blk_ret, v5, v6
 `,
@@ -569,7 +566,6 @@ signatures:
 	sig1: i64i64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64_v
 
 blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64, v4:f32, v5:f64)
-	Store module_ctx, exec_ctx, 0x8
 	Call f1:sig1, exec_ctx, module_ctx, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5, v2, v3, v4, v5
 	Jump blk_ret
 `,
@@ -582,7 +578,6 @@ signatures:
 	sig0: i64i64i32i64f32f64_i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64i32i64f32f64
 
 blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64, v4:f32, v5:f64)
-	Store module_ctx, exec_ctx, 0x8
 	v6:i32, v7:i64, v8:f32, v9:f64, v10:i32, v11:i64, v12:f32, v13:f64, v14:i32, v15:i64, v16:f32, v17:f64, v18:i32, v19:i64, v20:f32, v21:f64, v22:i32, v23:i64, v24:f32, v25:f64, v26:i32, v27:i64, v28:f32, v29:f64, v30:i32, v31:i64, v32:f32, v33:f64, v34:i32, v35:i64, v36:f32, v37:f64, v38:i32, v39:i64, v40:f32, v41:f64, v42:i32, v43:i64, v44:f32, v45:f64 = Call f1:sig0, exec_ctx, module_ctx, v2, v3, v4, v5
 	Jump blk_ret, v6, v7, v8, v9, v10, v11, v12, v13, v14, v15, v16, v17, v18, v19, v20, v21, v22, v23, v24, v25, v26, v27, v28, v29, v30, v31, v32, v33, v34, v35, v36, v37, v38, v39, v40, v41, v42, v43, v44, v45
 `,
@@ -982,11 +977,9 @@ blk2: () <-- (blk0)
 blk3: () <-- (blk2)
 	v5:i32 = Iconst_32 0x1
 	v6:i32 = Isub v2, v5
-	Store module_ctx, exec_ctx, 0x8
 	v7:i32 = Call f0:sig0, exec_ctx, module_ctx, v6
 	v8:i32 = Iconst_32 0x2
 	v9:i32 = Isub v2, v8
-	Store module_ctx, exec_ctx, 0x8
 	v10:i32 = Call f0:sig0, exec_ctx, module_ctx, v9
 	v11:i32 = Iadd v7, v10
 	Jump blk_ret, v11
@@ -1010,11 +1003,9 @@ blk2: () <-- (blk0)
 blk3: () <-- (blk2)
 	v5:i32 = Iconst_32 0x1
 	v6:i32 = Isub v2, v5
-	Store module_ctx, exec_ctx, 0x8
 	v7:i32 = Call f0:sig0, exec_ctx, module_ctx, v6
 	v8:i32 = Iconst_32 0x2
 	v9:i32 = Isub v2, v8
-	Store module_ctx, exec_ctx, 0x8
 	v10:i32 = Call f0:sig0, exec_ctx, module_ctx, v9
 	v11:i32 = Iadd v7, v10
 	Jump blk_ret, v11
@@ -1066,7 +1057,6 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
 	Jump blk1
 
 blk1: () <-- (blk0)
-	Store module_ctx, exec_ctx, 0x8
 	Call f1:sig1, exec_ctx, module_ctx
 	v5:i64 = Load module_ctx, 0x8
 	v6:i64 = Uload32 module_ctx, 0x10
@@ -1098,7 +1088,6 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
 	Jump fallthrough
 
 blk1: () <-- (blk0)
-	Store module_ctx, exec_ctx, 0x8
 	Call f1:sig1, exec_ctx, module_ctx
 	Jump blk3
 
@@ -1233,7 +1222,6 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 	v3:i64 = Load module_ctx, 0x20
 	v4:f32 = Load module_ctx, 0x30
 	v5:f64 = Load module_ctx, 0x40
-	Store module_ctx, exec_ctx, 0x8
 	Call f1:sig1, exec_ctx, module_ctx
 	v6:i32 = Load module_ctx, 0x10
 	v7:i64 = Load module_ctx, 0x20
@@ -1250,7 +1238,6 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 	v3:i64 = Load module_ctx, 0x20
 	v4:f32 = Load module_ctx, 0x30
 	v5:f64 = Load module_ctx, 0x40
-	Store module_ctx, exec_ctx, 0x8
 	Call f1:sig1, exec_ctx, module_ctx
 	v6:i32 = Load module_ctx, 0x10
 	v7:i64 = Load module_ctx, 0x20

internal/engine/wazevo/frontend/lower.go

@@ -1598,14 +1598,11 @@ func (c *Compiler) lowerCurrentOpcode() {
 			break
 		}
 
-		// Before transfer the control to the callee, we have to store the current module's moduleContextPtr
-		// into execContext.callerModuleContextPtr in case when the callee is a Go function.
-		//
-		// TODO: maybe this can be optimized out if this is in-module function calls. Investigate later.
-		c.storeCallerModuleContext()
-
 		var typIndex wasm.Index
 		if fnIndex < c.m.ImportFunctionCount {
+			// Before transfer the control to the callee, we have to store the current module's moduleContextPtr
+			// into execContext.callerModuleContextPtr in case when the callee is a Go function.
+			c.storeCallerModuleContext()
 			var fi int
 			for i := range c.m.ImportSection {
 				imp := &c.m.ImportSection[i]