diff --git a/jit/compiler.py b/jit/compiler.py index 9c239b5..02a90bd 100644 --- a/jit/compiler.py +++ b/jit/compiler.py @@ -588,6 +588,8 @@ def compile_instructions(self, mfunction, block, index=0): elif isinstance(instruction, model.LOAD_BUILD_CLASS): self.stub_handler.compile_class_stub(mfunction) + context.push_value("class_stub_address") + # Get the following class name which should be a LOAD_CONST instruction const_number = block.instructions[i + 2].argument name = block.function.consts[const_number] @@ -1029,15 +1031,20 @@ def compile_instructions(self, mfunction, block, index=0): # default arguments pass - # TODO : temporary, the return address will be after the call to the stub - address = stub_handler.lib.get_address(stub_handler.ffi.from_buffer(self.global_allocator.code_section), self.global_allocator.code_offset + 13) + stub_function = self.stub_handler.compile_function_stub(mfunction, nbargs) - stub_address = self.stub_handler.compile_function_stub(mfunction, nbargs, address) + self.compile_make_function(block, i, stub_function) - allocator.encode(asm.MOV(asm.r10, stub_address)) - allocator.encode(asm.CALL(asm.r10)) + # Clean the stack of the two TOS + allocator.encode(asm.ADD(asm.registers.rsp, 16)) + context.pop_value() + context.pop_value() - context.decrease_stack_size() + # Push the address of the stub on the stack + allocator.encode(asm.MOV(asm.r10, stub_function.first_address)) + allocator.encode(asm.PUSH(asm.r10)) + + # context.push_value("stub_address") elif isinstance(instruction, model.BUILD_SLICE): self.nyi() @@ -1176,6 +1183,38 @@ def compile_load_class(self, allocator, block, i, instruction): compare_operators = ('<', '<=', '==', '!=', '>', '>=', 'in', 'not in', 'is', 'is not', 'exception match', 'BAD') + # Compile the MAKE_FUNCTION instruction at given index in the block + # block : currently compiled basic block + # index : index on this instruction in the block + # stub_function : StubFunction instance of this stub + def compile_make_function(self, block, index, stub_function): + # This function try to analyze the IR + # In the normal case (no closure or default arguments) a Make_function has the following structure on the stack: + # | name of function | + # | code of function | + # + # The MAKE_FUNCTION is supposed to be after two LOAD_CONSTS with the code and the name of the function + if block.instructions[index].argument == 0 and index >= 2: + tos = block.instructions[index-1] + tos1 = block.instructions[index-2] + + if isinstance(tos, model.LOAD_CONST) and isinstance(tos1, model.LOAD_CONST): + function_name = block.function.consts[tos.argument] + function_code = block.function.consts[tos1.argument] + + stub_function.function_name = function_name + stub_function.code = function_code + + # Special case for primitive functions and addresses + if "standard_library.py" in block.function.filename: + short_name = function_name.replace("twopy_", "") + stub_handler.primitive_addresses[short_name] = stub_function.first_address + + return + + # We should always be in the previous case for now + self.nyi() + # Functions used to compile a comparison then a jump after (a if) # mfunction : Current compiled function # instruction : Current python Bytecode instruction @@ -1540,8 +1579,13 @@ def __call__(self, *args): def compile_prolog(self): offset_before = self.jitcompiler.global_allocator.code_offset - # Save rbp + # Save rbp and other registers self.encode(asm.PUSH(asm.rbp)) + self.encode(asm.PUSH(asm.rbx)) + self.encode(asm.PUSH(asm.r12)) + self.encode(asm.PUSH(asm.r13)) + self.encode(asm.PUSH(asm.r14)) + self.encode(asm.PUSH(asm.r15)) self.encode(asm.MOV(asm.rbp, asm.registers.rsp)) # Call the function just after this prolog @@ -1551,6 +1595,11 @@ def compile_prolog(self): # Restore the stack instructions.append(asm.MOV(asm.registers.rsp, asm.rbp)) + instructions.append(asm.POP(asm.r15)) + instructions.append(asm.POP(asm.r14)) + instructions.append(asm.POP(asm.r13)) + instructions.append(asm.POP(asm.r12)) + instructions.append(asm.POP(asm.rbx)) instructions.append(asm.POP(asm.rbp)) instructions.append(asm.RET()) diff --git a/jit/ffi_definitions.py b/jit/ffi_definitions.py index 838b0d0..06a1696 100644 --- a/jit/ffi_definitions.py +++ b/jit/ffi_definitions.py @@ -34,6 +34,9 @@ // Callback to trigger the compilation of a function extern "Python+C" void python_callback_function_stub(uint64_t, uint64_t, uint64_t, uint64_t); + + // Callback to trigger the compilation of a function with some information in the stub + extern "Python+C" void python_callback_function_stub_simplified(uint64_t, uint64_t); // Callback for type tests extern "Python+C" void python_callback_type_stub(uint64_t, int, int); @@ -81,6 +84,8 @@ static void python_callback_bb_stub(uint64_t rsp); static void python_callback_function_stub(uint64_t, uint64_t, uint64_t, uint64_t); + + static void python_callback_function_stub_simplified(uint64_t, uint64_t); static void python_callback_type_stub(uint64_t, int, int); @@ -107,9 +112,22 @@ { python_callback_bb_stub(rsp[-2]); } - + // Handle the compilation of a function's stub + // Case where the Stub is already filled with some information void function_stub(uint64_t* rsp) + { + uint64_t return_address = rsp[-2]; + + // Where to jump after the call to this stub + uint64_t address_after = rsp[0]; + + // Callback to python to trigger the compilation of the function + python_callback_function_stub_simplified(return_address, address_after); + } + + // Handle the compilation of a function's stub + /*void function_stub_old(uint64_t* rsp) { uint64_t* code_address = (uint64_t*)rsp[-2]; @@ -128,7 +146,7 @@ // Callback to python to trigger the compilation of the function python_callback_function_stub(name_id, code_id, return_address, rsp[3]); - } + }*/ // Handle compilation of a type-test stub void type_test_stub(uint64_t* rsp) @@ -145,7 +163,7 @@ python_callback_type_stub(return_address_aligned, id_variable, type_value); } - + void class_stub(uint64_t* rsp) { python_callback_class_stub(rsp[-1], rsp[0]); @@ -154,7 +172,7 @@ void print_stack(uint64_t* rsp) { printf("Print the stack\\n"); - for(int i=-1; i!=7; i++) + for(int i=-2; i!=7; i++) printf("\\t 0x%lx stack[%d] = 0x%lx\\n", (long int)&rsp[i], i, rsp[i]); } diff --git a/jit/stub_handler.py b/jit/stub_handler.py index 3af3888..5eeb36f 100644 --- a/jit/stub_handler.py +++ b/jit/stub_handler.py @@ -143,8 +143,7 @@ def compile_stub(self, mfunction, stub): # Compile a stub to a function # mfunction: The simple_interpreter.Function # nbargs : number of arguments in the registers, used by C function later - # address_after : where to jump after the stub - def compile_function_stub(self, mfunction, nbargs, address_after): + def compile_function_stub(self, mfunction, nbargs): # Now encode the stub stub_label = asm.Label("Stub_label_" + str(mfunction.name)) @@ -162,17 +161,22 @@ def compile_function_stub(self, mfunction, nbargs, address_after): # Align the stack on 16 bits mfunction.allocator.encode_stub(asm.MOV(asm.rax, asm.registers.rsp)) - mfunction.allocator.encode_stub(asm.AND(asm.registers.rsp, -16)) - mfunction.allocator.encode_stub(asm.PUSH(asm.registers.rsp)) + mfunction.allocator.encode_stub(asm.AND(asm.registers.rax, -16)) + mfunction.allocator.encode_stub(asm.PUSH(asm.registers.rax)) mfunction.allocator.encode_stub(asm.MOV(asm.r10, function_address)) mfunction.allocator.encode_stub(asm.CALL(asm.r10)) - # Now put additional informations for the stub + # Now put additional information for the stub # Force min 8 bits encoding for this value nbargs_bytes = encode_bytes(nbargs) stub_function = StubFunction() + stub_function.first_address = address + + address_after = lib.get_address(ffi.from_buffer(jitcompiler_instance.global_allocator.code_section), + jitcompiler_instance.global_allocator.stub_offset) + self.stub_dictionary[address_after] = stub_function self.data_addresses[address_after] = jitcompiler_instance.global_allocator.stub_offset @@ -186,7 +190,7 @@ def compile_function_stub(self, mfunction, nbargs, address_after): for i in range(5): mfunction.allocator.encode_stub(asm.NOP()) - return address + return stub_function # Compile a jump to an error # error_code The code of the error sent to C @@ -216,9 +220,6 @@ def compile_class_stub(self, mfunction): c_buffer = ffi.from_buffer(jitcompiler_instance.global_allocator.code_section) address_code = lib.get_address(c_buffer, jitcompiler_instance.global_allocator.code_offset) - # Be able to find them in the collection during the later callback - self.class_stub_addresses.append(address_stub) - # Call the stub function mfunction.allocator.encode_stub(asm.MOV(asm.rdi, asm.registers.rsp)) @@ -235,6 +236,9 @@ def compile_class_stub(self, mfunction): stub.data_address = offset stub.return_address = address_code + # Be able to find them in the collection during the later callback + self.class_stub_addresses.append(address_code) + # Indicate this offset correspond to the "return address" on the stub after the call to C returned self.data_addresses[return_address] = offset self.stub_dictionary[return_address] = stub @@ -291,22 +295,60 @@ def python_callback_bb_stub(rsp): # Delete the entry del stubhandler_instance.stub_dictionary[rsp] + # This function is called when a stub is executed, need to compile a function + # @ffi.def_extern() + # def python_callback_function_stub(name_id, code_id, return_address, canary_value): + # # Generate the Function object in the model + # name = jitcompiler_instance.consts[name_id] + # code = jitcompiler_instance.consts[code_id] + # + # #TODO: find a way to get this name and code from elsewhere + # # we could put them into the stub + # function = jitcompiler_instance.interpreter.generate_function(code, name, jitcompiler_instance.interpreter.mainmodule, False) + # + # # We may need to generate a class + # if canary_value in stubhandler_instance.class_stub_addresses: + # function.is_class = True + # function.mclass = objects.JITClass(function, name) + # + # # Add this class-function to the global collection + # jitcompiler_instance.class_functions.append(function) + # + # # Trigger the compilation of the given function + # jitcompiler_instance.compile_function(function) + # + # if jitcompiler_instance.interpreter.args.asm: + # jitcompiler_instance.global_allocator.disassemble_asm() + # + # stub = stubhandler_instance.stub_dictionary[return_address] + # stub.data_address = stubhandler_instance.data_addresses[return_address] + # + # stub.clean(return_address, function.allocator.code_address, canary_value) + # This function is called when a stub is executed, need to compile a function @ffi.def_extern() - def python_callback_function_stub(name_id, code_id, return_address, canary_value): + def python_callback_function_stub_simplified(return_address, address_after): + # Get the stub and its information + stub = stubhandler_instance.stub_dictionary[return_address] + stub.data_address = stubhandler_instance.data_addresses[return_address] + # Generate the Function object in the model - name = jitcompiler_instance.consts[name_id] - code = jitcompiler_instance.consts[code_id] + name = stub.function_name + code = stub.code - function = jitcompiler_instance.interpreter.generate_function(code, name, jitcompiler_instance.interpreter.mainmodule, False) + function = jitcompiler_instance.interpreter.generate_function(code, name, + jitcompiler_instance.interpreter.mainmodule, + False) + # TODO: handle the case where we are compiling a class # We may need to generate a class - if canary_value in stubhandler_instance.class_stub_addresses: - function.is_class = True - function.mclass = objects.JITClass(function, name) - - # Add this class-function to the global collection - jitcompiler_instance.class_functions.append(function) + # if address_after in stubhandler_instance.class_stub_addresses: + # print("Compiling a class") + # function.is_class = True + # function.mclass = objects.JITClass(function, name) + # + # # Add this class-function to the global collection + # jitcompiler_instance.class_functions.append(function) # Trigger the compilation of the given function jitcompiler_instance.compile_function(function) @@ -314,10 +356,7 @@ def python_callback_function_stub(name_id, code_id, return_address, canary_value if jitcompiler_instance.interpreter.args.asm: jitcompiler_instance.global_allocator.disassemble_asm() - stub = stubhandler_instance.stub_dictionary[return_address] - stub.data_address = stubhandler_instance.data_addresses[return_address] - - stub.clean(return_address, function.allocator.code_address, canary_value) + stub.clean(address_after, function.allocator.code_address) @ffi.def_extern() def python_callback_type_stub(return_address, id_variable, type_value): @@ -622,6 +661,15 @@ class StubFunction(Stub): def __init__(self): super().__init__() + # The address of the first encoded instruction in stub section for this stub + self.first_address = None + + # The name of the function + self.function_name = None + + # The CodeObject of this function + self.code = None + # Write instructions to restore the context before returning to asm # return_address : where to jump after this stub # function_address : address of the newly compiled function, to put on TOS after returning @@ -629,22 +677,19 @@ def __init__(self): def clean(self, return_address, function_address, canary_value=None): instructions = [] - # restore rsp - instructions.append(asm.POP(asm.registers.rsp).encode()) - if canary_value in stubhandler_instance.class_stub_addresses: instructions.append(asm.ADD(asm.registers.rsp, 32).encode()) else: - # Discard the three top values on the stack - instructions.append(asm.ADD(asm.registers.rsp, 24).encode()) + # Discard the return value on the stack + instructions.append(asm.ADD(asm.registers.rsp, 16).encode()) - # Now push the function address + # Now call the newly compiled function instructions.append(asm.MOV(asm.rax, function_address).encode()) - instructions.append(asm.PUSH(asm.rax).encode()) + instructions.append(asm.CALL(asm.rax).encode()) - # Finally, jump to the correct destination - instructions.append(asm.MOV(asm.rax, return_address).encode()) - instructions.append(asm.JMP(asm.rax).encode()) + # After the call, jump to the code section with the return_address + instructions.append(asm.MOV(asm.r10, return_address).encode()) + instructions.append(asm.JMP(asm.r10).encode()) offset = self.data_address for i in instructions: