python jit代码示例

jit.py

# coding=utf-8
import dis
import mmap
import struct
from ctypes import *
from functools import wraps
from typing import Union

import numpy as np


class RID(object):
    RAX = np.uint8(0)
    RCX = np.uint8(1)
    RDX = np.uint8(2)
    RBX = np.uint8(3)
    RSP = np.uint8(4)
    RBP = np.uint8(5)
    RSI = np.uint8(6)
    RDI = np.uint8(7)
    R8 = np.uint8(8)
    R9 = np.uint8(9)
    R10 = np.uint8(10)
    R11 = np.uint8(11)
    R12 = np.uint8(12)
    R13 = np.uint8(13)
    R14 = np.uint8(14)
    R15 = np.uint8(15)


REX_64 = np.uint8(0b01001000)
MOV_RR = np.uint8(0x89)


class CodeBuffer(object):
    def __init__(self):
        self.data = bytearray()

    def write(self, val):
        self.data.append(np.uint8(val))


def bin_rr(code: CodeBuffer, op: Union[np.uint8, int], src: Union[np.uint8, int], target: Union[np.uint8, int],
           ext: int):
    op = np.uint8(op)
    src = np.uint8(src)
    target = np.uint8(target)
    rex = REX_64

    if src & 0x8:
        rex |= 1 << 2
        src &= ~np.uint8(0x8)

    if target & 0x8:
        rex |= 1
        target &= ~np.uint8(0x8)

    code.write(rex)
    if ext:
        code.write(0x0F)
    code.write(op)
    code.write((0b11 << 6) | (src << 3) | target)


def bin_mr(code: CodeBuffer, op: Union[np.uint8, int], src: Union[np.uint8, int], target: Union[np.uint8, int],
           offset: int,
           ext: int):
    op = np.uint8(op)
    src = np.uint8(src)
    target = np.uint8(target)
    rex = REX_64

    if src & 0x8:
        rex |= 1 << 2
        src &= ~np.uint8(0x8)

    if target & 0x8:
        rex |= 1
        target &= ~np.uint8(0x8)

    code.write(rex)
    if ext:
        code.write(0x0F)

    code.write(op)
    if 128 > offset >= -128:
        code.write((0b01 << 6) | (src << 3) | target)
        code.write(offset)
    else:
        code.write((0b10 << 6) | (src << 3) | target)
        for x in struct.pack("=i", offset):
            code.write(x)
    # code.write(0XFF)


def movq_ir(code: CodeBuffer, val: int, target: np.uint8):
    data = struct.pack("=q", val)
    rex = REX_64
    if target & 0x8:
        rex |= 1
        target &= ~np.uint8(0x8)
    code.write(rex)
    code.write(0xB8 + target)
    for i in range(8):
        code.write(data[i])


REG_STACK = ["rbx", "rcx", "rdi", "rsi", "r8", "r9",
             "r10", "r11", "r12", "r13", "r14", "r15"]

REG_STACKI = [RID.RBX, RID.RCX, RID.RDI, RID.RSI, RID.R8,
              RID.R9, RID.R10, RID.R11, RID.R11, RID.R12,
              RID.R13, RID.R14, RID.R15]


def jit_compile(bcp, cb: CodeBuffer):
    cb.write(0x50 + RID.RBP)

    bin_rr(cb, MOV_RR, RID.RSP, RID.RBP, 0)
    bin_mr(cb, MOV_RR, RID.RDI, RID.RBP, -8, 0)
    bin_mr(cb, MOV_RR, RID.RSI, RID.RBP, -16, 0)
    bin_mr(cb, MOV_RR, RID.RDX, RID.RBP, -24, 0)
    bin_mr(cb, MOV_RR, RID.RCX, RID.RBP, -32, 0)
    bin_mr(cb, MOV_RR, RID.R8, RID.RBP, -40, 0)
    bin_mr(cb, MOV_RR, RID.R9, RID.RBP, -48, 0)

    sp = 0

    for ins in bcp:
        op = ins.split(" ")

        if op[0] == "NOP":
            continue

        elif op[0] == "RET":
            bin_rr(cb, MOV_RR, RID.RBX, RID.RAX, 0)
            cb.write(0x58 + RID.RBP)
            cb.write(0xc3)

        elif op[0] == "PUSHV":
            offset = (int(op[1]) + 1) * -8
            bin_mr(cb, 0x8B, REG_STACKI[sp], RID.RBP, offset, 0)
            sp += 1

        elif op[0] == "PUSHC":
            movq_ir(cb, int(op[1]), REG_STACKI[sp])
            sp += 1

        elif op[0] == "ADD":
            bin_rr(cb, 0x01, REG_STACKI[sp - 1], REG_STACKI[sp - 2], 0)
            sp -= 1

        elif op[0] == "SUB":
            bin_rr(cb, 0x29, REG_STACKI[sp - 1], REG_STACKI[sp - 2], 0)
            sp -= 1

        elif op[0] == "MUL":
            bin_rr(cb, 0xAF, REG_STACKI[sp - 2], REG_STACKI[sp - 1], 1)
            sp -= 1

        elif op[0] == "DIV":
            bin_rr(cb, MOV_RR, RID.RAX, REG_STACKI[sp], 0)
            bin_rr(cb, MOV_RR, RID.RDX, REG_STACKI[sp + 1], 0)

            bin_rr(cb, 0x31, RID.RDX, RID.RDX, 0)

            bin_rr(cb, MOV_RR, REG_STACKI[sp - 2], RID.RAX, 0)

            bin_mr(cb, MOV_RR, REG_STACKI[sp - 1], RID.RBP, -56, 0)

            cb.write(REX_64)
            cb.write(0x99)

            cb.write(REX_64)
            cb.write(0xF7)
            cb.write((0b01 << 6) | (7 << 3) | RID.RBP)
            cb.write(-56)

            bin_rr(cb, MOV_RR, RID.RAX, REG_STACKI[sp - 2], 0)
            bin_rr(cb, MOV_RR, REG_STACKI[sp], RID.RAX, 0)
            bin_rr(cb, MOV_RR, REG_STACKI[sp + 1], RID.RAX, 0)
            sp -= 1


def flush_to_exec(code: bytearray, *args):
    codeblock = mmap.mmap(-1, len(code) + 1, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS,
                          mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC,
                          0, 0)
    codeblock.write(code)
    codeblock.seek(0)

    p = addressof(c_uint.from_buffer(codeblock))
    fptr = cast(c_char_p(p), CFUNCTYPE(*args))
    setattr(fptr, "__mmap__", codeblock)
    return fptr


def with_jit(fn):
    opcodes = []

    for ins in dis.get_instructions(fn):
        if ins.opname == "LOAD_FAST":
            opcodes.append("PUSHV {}".format(ins.arg))
        elif ins.opname == "LOAD_CONST":
            opcodes.append("PUSHC {}".format(ins.argval))
        elif ins.opname == "BINARY_MULTIPLY":
            opcodes.append("MUL")
        elif ins.opname == "BINARY_FLOOR_DIVIDE":
            opcodes.append("DIV")
        elif ins.opname == "BINARY_ADD":
            opcodes.append("ADD")
        elif ins.opname == "BINARY_SUBTRACT":
            opcodes.append("SUB")
        elif ins.opname == "RETURN_VALUE":
            opcodes.append("RET")

    # print("\n".join(opcodes))

    s = CodeBuffer()
    jit_compile(opcodes, s)

    func_ptr = flush_to_exec(s.data, c_ulonglong)
    print(" ".join(["%02x" % x for x in s.data]))

    @wraps(fn)
    def wrapper(*args):
        return func_ptr(*args)

    return wrapper


@with_jit
def my_func(a, b, c, d):
    return c + a + a // b + d + + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d * c + a + a // b + d + + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d


def my_func2(a, b, c, d):
    return c + a + a // b + d + + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d * c + a + a // b + d + + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d + c + a + a // b + d


import timeit

print(timeit.timeit("my_func(100, 202, 302, 41)", globals=globals()))
print(timeit.timeit("my_func2(100, 202, 302, 41)", globals=globals()))

print(my_func(100, 202, 302, 41))
print(my_func2(100, 202, 302, 41))

大神，有微信吗？加我微信 18511963308
有偿帮忙改装我家的面板，也是大金的面板，苦于不懂任何代码，辗转到您github下求救