py/emitnative: Simplify layout and loading of native function prelude.

dpgeorge · dpgeorge · commit 6d403eb6972b · 2024-02-20T10:56:24.000+11:00
Now native functions and native generators have similar behaviour: the
first machine-word of their code is an index to get to the prelude.  This
simplifies the handling of these types of functions, and also reduces the
size of the emitted native machine code by no longer requiring special code
at the start of the function to load a pointer to the prelude.

Signed-off-by: Damien George &lt;damien@micropython.org&gt;
diff --git a/py/bc.c b/py/bc.c
@@ -336,9 +336,9 @@ void mp_setup_code_state(mp_code_state_t *code_state, size_t n_args, size_t n_kw
 // On entry code_state should be allocated somewhere (stack/heap) and
 // contain the following valid entries:
 //    - code_state->fun_bc should contain a pointer to the function object
-//    - code_state->ip should contain a pointer to the beginning of the prelude
 //    - code_state->n_state should be the number of objects in the local state
 void mp_setup_code_state_native(mp_code_state_native_t *code_state, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    code_state->ip = mp_obj_fun_native_get_prelude_ptr(code_state->fun_bc);
     code_state->sp = &code_state->state[0] - 1;
     mp_setup_code_state_helper((mp_code_state_t *)code_state, n_args, n_kw, args);
 }
diff --git a/py/emitnative.c b/py/emitnative.c
@@ -253,7 +253,6 @@ struct _emit_t {
     int pass;
 
     bool do_viper_types;
-    bool prelude_offset_uses_u16_encoding;
 
     mp_uint_t local_vtype_alloc;
     vtype_kind_t *local_vtype;
@@ -519,8 +518,11 @@ STATIC void emit_native_start_pass(emit_t *emit, pass_kind_t pass, scope_t *scop
         // work out size of state (locals plus stack)
         emit->n_state = scope->num_locals + scope->stack_size;
 
+        // Store in the first machine-word an index used to the function's prelude.
+        // This is used at runtime by mp_obj_fun_native_get_prelude_ptr().
+        mp_asm_base_data(&emit->as->base, ASM_WORD_SIZE, (uintptr_t)emit->prelude_ptr_index);
+
         if (emit->scope->scope_flags & MP_SCOPE_FLAG_GENERATOR) {
-            mp_asm_base_data(&emit->as->base, ASM_WORD_SIZE, (uintptr_t)emit->prelude_ptr_index);
             mp_asm_base_data(&emit->as->base, ASM_WORD_SIZE, (uintptr_t)emit->start_offset);
             ASM_ENTRY(emit->as, emit->code_state_start);
 
@@ -576,15 +578,6 @@ STATIC void emit_native_start_pass(emit_t *emit, pass_kind_t pass, scope_t *scop
             // Set code_state.fun_bc
             ASM_MOV_LOCAL_REG(emit->as, LOCAL_IDX_FUN_OBJ(emit), REG_PARENT_ARG_1);
 
-            // Set code_state.ip, a pointer to the beginning of the prelude.  This pointer is found
-            // either directly in mp_obj_fun_bc_t.child_table (if there are no children), or in
-            // mp_obj_fun_bc_t.child_table[num_children] (if num_children > 0).
-            ASM_LOAD_REG_REG_OFFSET(emit->as, REG_PARENT_ARG_1, REG_PARENT_ARG_1, OFFSETOF_OBJ_FUN_BC_CHILD_TABLE);
-            if (emit->prelude_ptr_index != 0) {
-                ASM_LOAD_REG_REG_OFFSET(emit->as, REG_PARENT_ARG_1, REG_PARENT_ARG_1, emit->prelude_ptr_index);
-            }
-            emit_native_mov_state_reg(emit, emit->code_state_start + OFFSETOF_CODE_STATE_IP, REG_PARENT_ARG_1);
-
             // Set code_state.n_state (only works on little endian targets due to n_state being uint16_t)
             emit_native_mov_state_imm_via(emit, emit->code_state_start + OFFSETOF_CODE_STATE_N_STATE, emit->n_state, REG_ARG_1);
 
diff --git a/py/objfun.c b/py/objfun.c
@@ -401,7 +401,7 @@ mp_obj_t mp_obj_new_fun_bc(const mp_obj_t *def_args, const byte *code, const mp_
 STATIC mp_obj_t fun_native_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
     MP_STACK_CHECK();
     mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
-    mp_call_fun_t fun = MICROPY_MAKE_POINTER_CALLABLE((void *)self->bytecode);
+    mp_call_fun_t fun = mp_obj_fun_native_get_function_start(self);
     return fun(self_in, n_args, n_kw, args);
 }
 
diff --git a/py/objfun.h b/py/objfun.h
@@ -69,6 +69,34 @@ static inline mp_obj_t mp_obj_new_fun_viper(const void *fun_data, const mp_modul
     return MP_OBJ_FROM_PTR(o);
 }
 
+static inline const uint8_t *mp_obj_fun_native_get_prelude_ptr(const mp_obj_fun_bc_t *fun_native) {
+    // Obtain a pointer to the start of the function prelude, based on prelude_ptr_index.
+    uintptr_t prelude_ptr_index = ((uintptr_t *)fun_native->bytecode)[0];
+    const uint8_t *prelude_ptr;
+    if (prelude_ptr_index == 0) {
+        prelude_ptr = (const uint8_t *)fun_native->child_table;
+    } else {
+        prelude_ptr = (const uint8_t *)fun_native->child_table[prelude_ptr_index];
+    }
+    return prelude_ptr;
+}
+
+static inline void *mp_obj_fun_native_get_function_start(const mp_obj_fun_bc_t *fun_native) {
+    // Obtain a pointer to the start of the function executable machine code.
+    return MICROPY_MAKE_POINTER_CALLABLE((void *)(fun_native->bytecode + sizeof(uintptr_t)));
+}
+
+static inline void *mp_obj_fun_native_get_generator_start(const mp_obj_fun_bc_t *fun_native) {
+    // Obtain a pointer to the start of the generator executable machine code.
+    uintptr_t start_offset = ((uintptr_t *)fun_native->bytecode)[1];
+    return MICROPY_MAKE_POINTER_CALLABLE((void *)(fun_native->bytecode + start_offset));
+}
+
+static inline void *mp_obj_fun_native_get_generator_resume(const mp_obj_fun_bc_t *fun_native) {
+    // Obtain a pointer to the resume location of the generator executable machine code.
+    return MICROPY_MAKE_POINTER_CALLABLE((void *)&((uintptr_t *)fun_native->bytecode)[2]);
+}
+
 #endif
 
 #if MICROPY_EMIT_INLINE_ASM
diff --git a/py/objgenerator.c b/py/objgenerator.c
@@ -101,13 +101,7 @@ STATIC mp_obj_t native_gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_k
     mp_obj_fun_bc_t *self_fun = MP_OBJ_TO_PTR(self_in);
 
     // Determine start of prelude.
-    uintptr_t prelude_ptr_index = ((uintptr_t *)self_fun->bytecode)[0];
-    const uint8_t *prelude_ptr;
-    if (prelude_ptr_index == 0) {
-        prelude_ptr = (void *)self_fun->child_table;
-    } else {
-        prelude_ptr = (void *)self_fun->child_table[prelude_ptr_index];
-    }
+    const uint8_t *prelude_ptr = mp_obj_fun_native_get_prelude_ptr(self_fun);
 
     // Extract n_state from the prelude.
     const uint8_t *ip = prelude_ptr;
@@ -119,17 +113,14 @@ STATIC mp_obj_t native_gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_k
     // Parse the input arguments and set up the code state
     o->pend_exc = mp_const_none;
     o->code_state.fun_bc = self_fun;
-    o->code_state.ip = prelude_ptr;
     o->code_state.n_state = n_state;
-    o->code_state.sp = &o->code_state.state[0] - 1;
     mp_setup_code_state_native(&o->code_state, n_args, n_kw, args);
 
     // Indicate we are a native function, which doesn't use this variable
     o->code_state.exc_sp_idx = MP_CODE_STATE_EXC_SP_IDX_SENTINEL;
 
     // Prepare the generator instance for execution
-    uintptr_t start_offset = ((uintptr_t *)self_fun->bytecode)[1];
-    o->code_state.ip = MICROPY_MAKE_POINTER_CALLABLE((void *)(self_fun->bytecode + start_offset));
+    o->code_state.ip = mp_obj_fun_native_get_generator_start(self_fun);
 
     return MP_OBJ_FROM_PTR(o);
 }
@@ -208,9 +199,9 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_
 
     #if MICROPY_EMIT_NATIVE
     if (self->code_state.exc_sp_idx == MP_CODE_STATE_EXC_SP_IDX_SENTINEL) {
-        // A native generator, with entry point 2 words into the "bytecode" pointer
+        // A native generator.
         typedef uintptr_t (*mp_fun_native_gen_t)(void *, mp_obj_t);
-        mp_fun_native_gen_t fun = MICROPY_MAKE_POINTER_CALLABLE((const void *)(self->code_state.fun_bc->bytecode + 2 * sizeof(uintptr_t)));
+        mp_fun_native_gen_t fun = mp_obj_fun_native_get_generator_resume(self->code_state.fun_bc);
         ret_kind = fun((void *)&self->code_state, throw_value);
     } else
     #endif

Original file line number	Diff line number	Diff line change
`@@ -401,7 +401,7 @@ mp_obj_t mp_obj_new_fun_bc(const mp_obj_t def_args, const byte code, const mp_`
`401`	`401`	`STATIC mp_obj_t fun_native_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {`
`402`	`402`	`MP_STACK_CHECK();`
`403`	`403`	`mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);`
`404`		`- mp_call_fun_t fun = MICROPY_MAKE_POINTER_CALLABLE((void *)self->bytecode);`
	`404`	`+ mp_call_fun_t fun = mp_obj_fun_native_get_function_start(self);`
`405`	`405`	`return fun(self_in, n_args, n_kw, args);`
`406`	`406`	`}`
`407`	`407`