├── .gitignore ├── .meta_version ├── .restrain ├── cython_rts │ ├── include │ │ └── typeint.h │ └── src │ │ └── typeint.c └── info.json ├── .style.yapf ├── LICENSE ├── README.md ├── datatype_gen.py ├── docs ├── DonationPrivacy.md ├── GetStarted.md └── What-is-CPython-Compatible.md ├── gen_instnames.py ├── mk_rts.py ├── requirements.txt ├── restrain_jit ├── __init__.py ├── abs_compiler │ ├── __init__.py │ ├── from_bc.py │ ├── instrnames.py │ └── py_apis.py ├── becython │ ├── __init__.py │ ├── cy_jit.py │ ├── cy_jit_common.py │ ├── cy_jit_ext_template.py │ ├── cy_jit_hotspot_comp.py │ ├── cy_loader.py │ ├── cy_method_codegen.py │ ├── cython_rts │ │ ├── Design1.md │ │ ├── Design2.md │ │ ├── RestrainJIT.pxd │ │ ├── RestrainJIT.pyx │ │ ├── __init__.py │ │ ├── hotspot.pxd │ │ └── hotspot.pyx │ ├── cython_vm.py │ ├── mono_vm.gen │ ├── mono_vm.py │ ├── phi_elim.py │ ├── phi_node_analysis.py │ ├── phi_vm.gen │ ├── phi_vm.py │ ├── relabel.py │ ├── representations.gen │ ├── representations.py │ ├── stack_vm_instructions.gen │ ├── stack_vm_instructions.py │ └── tools.py ├── bejulia │ ├── __init__.py │ ├── basics.py │ ├── functional.py │ ├── instructions │ ├── instructions.py │ ├── jl_init.py │ ├── jl_protocol.py │ ├── julia_vm.py │ ├── pragmas.py │ ├── representations │ ├── representations.py │ ├── simple_julia_py.py │ └── tools.py ├── config.py ├── cpy_compat.py ├── jit_info.py ├── utils.py └── vm │ ├── __init__.py │ └── am.py ├── setup.py ├── static ├── donate-weixin.png ├── p1.png ├── p2.png ├── p3.png └── procedure.png └── tests ├── __init__.py ├── becy ├── __init__.py ├── cfg_and_phi.py ├── design1_proof_of_concepts.py ├── design2_practice_1.py ├── design2_proof_of_concepts.py ├── dev.py ├── one_func_one_class.py ├── phi.py ├── test_datatype.py ├── test_if.py ├── test_loop.py └── typeid.py └── bejl ├── __init__.py ├── __main__.py ├── test_apis.py ├── test_array_op.py ├── test_from_bc.py ├── test_functional.py ├── test_instructions.py ├── test_load.py ├── test_load_pyjulia.py ├── test_simple_case.py ├── test_stage_count_import_star.py ├── test_stage_count_main.py └── test_stage_count_tested.py /.gitignore: -------------------------------------------------------------------------------- 1 | **__pycache__/ 2 | dev_test/ 3 | Digraph.gv* 4 | .idea/ 5 | .spyproject/ 6 | **.egg-info/ 7 | dist/ 8 | **~ 9 | build/ 10 | 11 | 12 | # Created by https://www.gitignore.io/api/macos,python 13 | # Edit at https://www.gitignore.io/?templates=macos,python 14 | 15 | ### macOS ### 16 | # General 17 | .DS_Store 18 | .AppleDouble 19 | .LSOverride 20 | 21 | # Icon must end with two \r 22 | Icon 23 | 24 | # Thumbnails 25 | ._* 26 | 27 | # Files that might appear in the root of a volume 28 | .DocumentRevisions-V100 29 | .fseventsd 30 | .Spotlight-V100 31 | .TemporaryItems 32 | .Trashes 33 | .VolumeIcon.icns 34 | .com.apple.timemachine.donotpresent 35 | 36 | # Directories potentially created on remote AFP share 37 | .AppleDB 38 | .AppleDesktop 39 | Network Trash Folder 40 | Temporary Items 41 | .apdisk 42 | 43 | ### Python ### 44 | # Byte-compiled / optimized / DLL files 45 | __pycache__/ 46 | *.py[cod] 47 | *$py.class 48 | 49 | # C extensions 50 | *.so 51 | 52 | # Distribution / packaging 53 | .Python 54 | build/ 55 | develop-eggs/ 56 | dist/ 57 | downloads/ 58 | eggs/ 59 | .eggs/ 60 | lib/ 61 | lib64/ 62 | parts/ 63 | sdist/ 64 | var/ 65 | wheels/ 66 | pip-wheel-metadata/ 67 | share/python-wheels/ 68 | *.egg-info/ 69 | .installed.cfg 70 | *.egg 71 | MANIFEST 72 | 73 | # PyInstaller 74 | # Usually these files are written by a python script from a template 75 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 76 | *.manifest 77 | *.spec 78 | 79 | # Installer logs 80 | pip-log.txt 81 | pip-delete-this-directory.txt 82 | 83 | # Unit test / coverage reports 84 | htmlcov/ 85 | .tox/ 86 | .nox/ 87 | .coverage 88 | .coverage.* 89 | .cache 90 | nosetests.xml 91 | coverage.xml 92 | *.cover 93 | .hypothesis/ 94 | .pytest_cache/ 95 | 96 | # Translations 97 | *.mo 98 | *.pot 99 | 100 | # Django stuff: 101 | *.log 102 | local_settings.py 103 | db.sqlite3 104 | db.sqlite3-journal 105 | 106 | # Flask stuff: 107 | instance/ 108 | .webassets-cache 109 | 110 | # Scrapy stuff: 111 | .scrapy 112 | 113 | # Sphinx documentation 114 | docs/_build/ 115 | 116 | # PyBuilder 117 | target/ 118 | 119 | # Jupyter Notebook 120 | .ipynb_checkpoints 121 | 122 | # IPython 123 | profile_default/ 124 | ipython_config.py 125 | 126 | # pyenv 127 | .python-version 128 | 129 | # pipenv 130 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 131 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 132 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 133 | # install all needed dependencies. 134 | #Pipfile.lock 135 | 136 | # celery beat schedule file 137 | celerybeat-schedule 138 | 139 | # SageMath parsed files 140 | *.sage.py 141 | 142 | # Environments 143 | .env 144 | .venv 145 | env/ 146 | venv/ 147 | ENV/ 148 | env.bak/ 149 | venv.bak/ 150 | 151 | # Spyder project settings 152 | .spyderproject 153 | .spyproject 154 | 155 | # Rope project settings 156 | .ropeproject 157 | 158 | # mkdocs documentation 159 | /site 160 | 161 | # mypy 162 | .mypy_cache/ 163 | .dmypy.json 164 | dmypy.json 165 | 166 | # Pyre type checker 167 | .pyre/ 168 | 169 | # End of https://www.gitignore.io/api/macos,python -------------------------------------------------------------------------------- /.meta_version: -------------------------------------------------------------------------------- 1 | method: autoinc 2 | current: 0.0.1 -------------------------------------------------------------------------------- /.restrain/cython_rts/include/typeint.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | int64_t ptrtoint(void* ptr); 4 | void* inttoptr(int64_t i); 5 | int check_ptr_eq(void*, void*); 6 | void* unsafe_cast(void*); -------------------------------------------------------------------------------- /.restrain/cython_rts/src/typeint.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | int64_t ptrtoint(void* ptr){ 4 | return (int64_t) ptr; 5 | } 6 | 7 | void* inttoptr(int64_t i){ 8 | return (void*) i; 9 | } 10 | 11 | int check_ptr_eq(void* a, void* b){ 12 | return ptrtoint(a) == ptrtoint(b); 13 | } 14 | 15 | void* unsafe_cast(void* f){ 16 | return f; 17 | } -------------------------------------------------------------------------------- /.restrain/info.json: -------------------------------------------------------------------------------- 1 | { 2 | "julia": { 3 | "image": "$JULIA_ROOT/julia-1.2.0/lib/julia/sys.so", 4 | "lib": "$JULIA_ROOT/julia-1.2.0/bin/../lib/libjulia.so.1", 5 | "bin": "$JULIA_ROOT/julia-1.2.0/bin" 6 | }, 7 | "cython": { 8 | "rts": "$RESTRAIN_ROOT/cython_rts/" 9 | } 10 | } 11 | -------------------------------------------------------------------------------- /.style.yapf: -------------------------------------------------------------------------------- 1 | [style] 2 | BASED_ON_STYLE: pep8 3 | COLUMN_LIMIT: 92 4 | SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED: true 5 | ALLOW_MULTILINE_LAMBDAS: false 6 | ALLOW_MULTILINE_DICTIONARY_KEYS: false 7 | BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF: true 8 | BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION: 2 9 | COALESCE_BRACKETS: true 10 | SPLIT_COMPLEX_COMPREHENSION: true 11 | SPLIT_BEFORE_DICT_SET_GENERATOR: true 12 | SPLIT_BEFORE_BITWISE_OPERATOR: true 13 | NO_SPACES_AROUND_SELECTED_BINARY_OPERATORS: false 14 | SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET: true -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | The MIT License (MIT) 3 | Copyright (c) 2019 thautwarm 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 18 | IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, 19 | DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR 20 | OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE 21 | OR OTHER DEALINGS IN THE SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | The developement of this project is suspended due to the complexity of managing multi-language dependencies and Python's verbose functional style codebase. We're now working on creating a static compiled programming language [RemuLang](https://github.com/RemuLang) targeting CPython bytecode, and after the available version of RemuLang released 3 | we'll resume the development of Restrain-JIT. 4 | 5 | 6 | [![What's CPython Compat](https://img.shields.io/badge/Hyping-What's "CPython Compatible"-Orange.svg?style=flat)](docs/What-is-CPython-Compatible.md) 7 | 8 | ## Python Restrain JIT 9 | 10 | [![赞助/Donotion](https://img.shields.io/badge/Donation-赞助-Teal.svg?style=flat)](docs/DonationPrivacy.md) 11 | [![开始参与开发/Dev Guide](https://img.shields.io/badge/Start Devel-开发参与指南-Purple.svg?style=flat)](docs/GetStarted.md) 12 | 13 | 14 | The first and yet the only "CPython compatible" Python JIT, over the world. 15 | 16 | This comes with my talk on PyConChina 2019. 17 | 18 | ## Restrain JIT: The Cython Back End 19 | 20 | For the previous suspending Julia back end, check the branch [julia-backend](https://github.com/thautwarm/restrain-jit/tree/julia-beckend). 21 | 22 | Cython is a widely used Python to C compiler, although it's lack of optimizations, it compiles fast, 23 | which greatly reduces the annoyance of JIT overhead. 24 | 25 | Currently the Cython back end works for many fundamental Python constructs. The constructs haven't suppported 26 | yet are exceptions(but you can wrap it as a function) and closures(ready to support today). 27 | 28 | The Cython back end did much more on the compilation stuffs, like Control Flow Analysis, many kinds of abstract interpretations, SSA conversions and 29 | corresponding passes. One of the most awesome pass is the Phi-Node analysis pass, 30 | which converts the semantics of Stack Virtual Machine to a Register-based Virtual Machine, and generates Phi nodes: 31 | [phi_node_analysis](https://github.com/thautwarm/restrain-jit/blob/master/restrain_jit/becython/phi_node_analysis.py) 32 | 33 | Besides, this is still based on Python bytecode, so Restrain JIT on Cython is still free of hack and available in any case. 34 | 35 | The developement of Cython back end is a joy. Yes, debugging is so fast that I can make faster developement iterations, 36 | and no need to wait half a minute when I want to re-run codes :) 37 | -------------------------------------------------------------------------------- /datatype_gen.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | 4 | def find_paths(p: Path): 5 | if not p.is_dir(): 6 | if p.suffix == '.gen': 7 | yield p 8 | else: 9 | for i in p.iterdir(): 10 | if i == p: 11 | continue 12 | yield from find_paths(i) 13 | 14 | 15 | for FROM, TO in [(path, path.with_suffix('.py')) 16 | for path in find_paths(Path('.'))]: 17 | with FROM.open() as f: 18 | text = f.read() 19 | 20 | defs = [[e.strip() for e in i.strip().split()] for i in text.split(';')] 21 | code = [ 22 | 'from enum import Enum, auto as _auto', 'import abc', 23 | 'import typing as t', 'from dataclasses import dataclass' 24 | ] 25 | defs = list(filter(None, defs)) 26 | 27 | for each in defs: 28 | code.append('') 29 | code.append('') 30 | 31 | head, *each = each 32 | if head == 'typevar': 33 | var, *args = each 34 | cov = "" 35 | if args: 36 | cov = "bound=" + "t.Union[" + ', '.join(map(repr, args)) + "]" 37 | code.append(f'{var} = t.TypeVar({var!r}, {cov})') 38 | if head == 'import': 39 | code.append(f'from {each[0]} import *') 40 | elif head == "enum": 41 | name, *variants = each 42 | code.append(f'class {name}(Enum):') 43 | for v in variants: 44 | code.append(f' {v} = _auto()') 45 | else: 46 | code.append(' pass') 47 | elif head == 'abc': 48 | abc = each[0] 49 | code.append(f'class {abc}:') 50 | code.append(' pass') 51 | continue 52 | elif head == 'data': 53 | name, *fields = each 54 | code.append('@dataclass(frozen=True, order=True)') 55 | code.append(f'class {name}:') 56 | for v in fields: 57 | code.append(' ' + v) 58 | else: 59 | code.append(' pass') 60 | code.append('') 61 | 62 | with TO.open('w') as f: 63 | f.write('\n'.join(code)) 64 | -------------------------------------------------------------------------------- /docs/DonationPrivacy.md: -------------------------------------------------------------------------------- 1 | # Donotation Policy 2 | 3 | 赞助方式可能逐渐增加. 4 | 5 | 对于所有赞助, **请从以下列表中选择至少一个用途**; 对于没有留言的赞助方式, 请用邮件告知twshere@outlook.com对应的业务单. 6 | 7 | - 用于赞助[thautwarm](https://github.com/thautwarm/)在研究生和博士阶段, 在**Programming Language**专业进行深造学习的费用(学费, 校内最低档次的住宿费等**必要的费用**) 8 | - 用于赞助Python Restrain JIT的开发和维护 9 | - 用于以上两种目的 10 | 11 | 赞助费使用明细公开. 12 | 13 | | Donation Means | Info | 14 | |:----------------|:-------------------------------------| 15 | | Weixin | [QRCode](https://raw.githubusercontent.com/thautwarm/restrain-jit/master/static/donate-weixin.png) | -------------------------------------------------------------------------------- /docs/GetStarted.md: -------------------------------------------------------------------------------- 1 | # Getting Started With Development 2 | 3 | 4 | ## Prerequisite 5 | - Python: CPython 3.6 or newer 6 | - Cython: recommend 0.29+(tested) 7 | 8 | ## `.restrain` Dotfiles 9 | 10 | Copy https://github.com/thautwarm/restrain-jit/tree/master/.restrain to 11 | your user directory, and edit `~/.restrain/info.json`, to specify 12 | `"cython.rts"`. 13 | 14 | ```json 15 | { 16 | "julia": {/*whatever*/}, 17 | "cython": { 18 | "rts": "~/.restrain/cython_rts" 19 | } 20 | } 21 | ``` 22 | 23 | Then use `g++`(a C++ compiler compatible to that compiles your Python) to generate necessary libraries. 24 | (Mac OSX user please see follow chapter "Fix Build Bug On Mac OSX" ) 25 | 26 | ``` 27 | ~/.restrain > cd cython_rts/src 28 | ~/.restrain/cython_rts/src > g++ -I../include -fPIC -c typeint.c 29 | ~/.restrain/cython_rts/src > mv typeint.o ../lib/typeint 30 | ``` 31 | 32 | 33 | Then check [tests/becy](https://github.com/thautwarm/restrain-jit/tree/master/tests/becy) for examples, e.g., 34 | 35 | - [Loop](https://github.com/thautwarm/restrain-jit/blob/master/tests/becy/test_loop.py) 36 | - [If](https://github.com/thautwarm/restrain-jit/blob/master/tests/becy/test_if.py) 37 | 38 | ### Fix Build Bug on Mac OSX 39 | Mac OSX user may encounter `g++` command error like: `ld: library not found for -l:typeint`. 40 | 41 | Use follow command instead: 42 | 43 | ``` 44 | ~/.restrain > cd cython_rts/src 45 | ~/.restrain/cython_rts/src > g++ -I../include -fPIC -c typeint.c -o libtypeint.so 46 | ~/.restrain/cython_rts/src > mv libtypeint.so ../lib/libtypeint.so 47 | ``` 48 | 49 | Change code in `becython/cy_loader.py` 50 | 51 | ``` 52 | -- libraries=[':typeint', *extra_libs]) 53 | ++ libraries=['typeint', *extra_libs]) 54 | ``` 55 | -------------------------------------------------------------------------------- /docs/What-is-CPython-Compatible.md: -------------------------------------------------------------------------------- 1 | # What's "CPython Compatible" 2 | 3 | "CPython Compatible" here might not be that proper after a simple glance, but I haven't found out a word for this. 4 | 5 | A "CPython Compatible" JIT roughly means: 6 | 7 | - Should be able to work with existing and prospective C-extensions of the CPython world. 8 | 9 | - Should be able to keep the same semantics bewteen the jitted code and original CPython code. 10 | 11 | - Should be able to support JITing most of Python objects(intentionally, everything other than objects from c-extensions, builtin/frozen modules), if specified. 12 | 13 | - Should be able to taken advantage as a regular Python 3-rd party library. 14 | 15 | Besides, I sometimes want to call it an **available** Python JIT. 16 | 17 | **If you get a better word to describe above properties, you could tell us at once, which will be specially appreciated.** 18 | -------------------------------------------------------------------------------- /gen_instnames.py: -------------------------------------------------------------------------------- 1 | import opcode 2 | 3 | with open("restrain_jit/ir/instrnames.py", 'w') as f: 4 | for k in opcode.opmap: 5 | f.write(f"{k} = {k!r}\n") 6 | -------------------------------------------------------------------------------- /mk_rts.py: -------------------------------------------------------------------------------- 1 | from Redy.Tools.PathLib import Path 2 | 3 | Path(".restrain/cython_rts").move_to("~/.restrain") -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | julia==0.4.1 2 | numpy>=1.6.0 3 | bytecode>=0.7 4 | -------------------------------------------------------------------------------- /restrain_jit/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /restrain_jit/abs_compiler/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/restrain_jit/abs_compiler/__init__.py -------------------------------------------------------------------------------- /restrain_jit/abs_compiler/from_bc.py: -------------------------------------------------------------------------------- 1 | import types 2 | import bytecode as bc 3 | import typing as t 4 | from restrain_jit.abs_compiler import instrnames as InstrNames 5 | from restrain_jit.abs_compiler import py_apis as RT 6 | from restrain_jit.vm import am 7 | u_map = { 8 | InstrNames.UNARY_POSITIVE: RT.py_pos, 9 | InstrNames.UNARY_NEGATIVE: RT.py_neg, 10 | InstrNames.UNARY_NOT: RT.py_not, 11 | InstrNames.UNARY_INVERT: RT.py_inv 12 | } 13 | 14 | cmp_map = { 15 | bc.Compare.EQ: RT.py_eq, 16 | bc.Compare.NE: RT.py_neq, 17 | bc.Compare.IS: RT.py_is, 18 | bc.Compare.IS_NOT: RT.py_is_not, 19 | bc.Compare.LT: RT.py_lt, 20 | bc.Compare.LE: RT.py_le, 21 | bc.Compare.GT: RT.py_gt, 22 | bc.Compare.GE: RT.py_ge, 23 | bc.Compare.IN: RT.py_in, 24 | bc.Compare.NOT_IN: RT.py_not_in, 25 | bc.Compare.EXC_MATCH: RT.py_exc_match 26 | } 27 | 28 | bin_map = { 29 | InstrNames.BINARY_POWER: RT.py_pow, 30 | InstrNames.BINARY_MULTIPLY: RT.py_mul, 31 | InstrNames.BINARY_MATRIX_MULTIPLY: RT.py_matmul, 32 | InstrNames.BINARY_FLOOR_DIVIDE: RT.py_floordiv, 33 | InstrNames.BINARY_TRUE_DIVIDE: RT.py_truediv, 34 | InstrNames.BINARY_MODULO: RT.py_mod, 35 | InstrNames.BINARY_ADD: RT.py_add, 36 | InstrNames.BINARY_SUBTRACT: RT.py_sub, 37 | InstrNames.BINARY_SUBSCR: RT.py_subscr, 38 | InstrNames.BINARY_LSHIFT: RT.py_lsh, 39 | InstrNames.BINARY_RSHIFT: RT.py_rsh, 40 | InstrNames.BINARY_AND: RT.py_and, 41 | InstrNames.BINARY_XOR: RT.py_xor, 42 | InstrNames.BINARY_OR: RT.py_or, 43 | } 44 | 45 | ibin_map = { 46 | InstrNames.INPLACE_POWER: RT.py_ipow, 47 | InstrNames.INPLACE_MULTIPLY: RT.py_imul, 48 | InstrNames.INPLACE_MATRIX_MULTIPLY: RT.py_imatmul, 49 | InstrNames.INPLACE_FLOOR_DIVIDE: RT.py_ifloordiv, 50 | InstrNames.INPLACE_TRUE_DIVIDE: RT.py_itruediv, 51 | InstrNames.INPLACE_MODULO: RT.py_imod, 52 | InstrNames.INPLACE_ADD: RT.py_iadd, 53 | InstrNames.INPLACE_SUBTRACT: RT.py_isub, 54 | InstrNames.INPLACE_LSHIFT: RT.py_ilsh, 55 | InstrNames.INPLACE_RSHIFT: RT.py_irsh, 56 | InstrNames.INPLACE_AND: RT.py_iand, 57 | InstrNames.INPLACE_XOR: RT.py_ixor, 58 | InstrNames.INPLACE_OR: RT.py_ior, 59 | } 60 | 61 | 62 | def label_to_name(label: bc.BasicBlock): 63 | return f"block-{id(label):x}" 64 | 65 | 66 | def pop_n(n: int): 67 | xs = [] 68 | for i in range(n): 69 | a = yield am.pop() 70 | xs.append(a) 71 | xs.reverse() 72 | return xs 73 | 74 | 75 | class Interpreter: 76 | """ 77 | From Python bytecode to finally tagless representations. 78 | Given such a finally tagless `result`, and a specified 79 | Virtual Machine `VMx` we have 80 | `instance_representation = run_machine(result, VMx)` 81 | where the `instance_representation` can be any 82 | - intermediate representation like LLVM IR 83 | - executable instructions like JVM bytecode, WASM, MIPS asm or Python bytecode 84 | - codes of high level language like Python, C/C++ 85 | """ 86 | 87 | def __init__(self, lineno_start: int): 88 | self.lineno = lineno_start 89 | 90 | def abs_i_cfg(self, b: bc.ControlFlowGraph): 91 | for each in b: 92 | assert isinstance(each, bc.BasicBlock) 93 | label_name = label_to_name(each) 94 | last_label = yield am.last_block_end() 95 | if last_label == label_name: 96 | meta = yield am.meta() 97 | unwind = meta[label_name] 98 | yield from unwind() 99 | yield am.label(label_name) 100 | for instr in each: 101 | yield from self.abs_i(instr) 102 | 103 | def abs_i(self, b: t.Union[bc.Instr, bc.SetLineno]): 104 | current_lineno = b.lineno 105 | if current_lineno != self.lineno: 106 | self.lineno = current_lineno 107 | yield am.set_lineno(current_lineno) 108 | 109 | if b.name == InstrNames.COMPARE_OP: 110 | yield from abs_i_cmp(b) 111 | elif "JUMP" in b.name: 112 | yield from abs_i_jump(b) 113 | elif b.name.startswith('UNARY'): 114 | yield from abs_i_unary(b) 115 | elif b.name.startswith('BINARY'): 116 | yield from abs_i_binary(b) 117 | elif b.name.startswith('INPLACE'): 118 | yield from abs_i_inplace_binary(b) 119 | elif b.name == InstrNames.ROT_TWO: 120 | a1 = yield am.pop() 121 | a2 = yield am.pop() 122 | yield am.push(a2) 123 | yield am.push(a1) 124 | elif b.name == InstrNames.ROT_THREE: 125 | a1 = yield am.pop() 126 | a2 = yield am.pop() 127 | a3 = yield am.pop() 128 | yield am.push(a3) 129 | yield am.push(a1) 130 | yield am.push(a2) 131 | elif b.name == InstrNames.NOP: 132 | pass 133 | elif b.name == InstrNames.DUP_TOP: 134 | a = yield am.pop() 135 | yield am.push(a) 136 | yield am.push(a) 137 | elif b.name == InstrNames.DUP_TOP_TWO: 138 | a = yield am.peek(1) 139 | yield am.push(a) 140 | yield am.push(a) 141 | 142 | elif b.name == InstrNames.GET_YIELD_FROM_ITER: 143 | a = yield am.pop() 144 | a = yield from RT.py_iter(a) 145 | yield am.push(a) 146 | elif b.name == InstrNames.STORE_SUBSCR: 147 | tos = yield am.pop() 148 | tos1 = yield am.pop() 149 | tos2 = yield am.pop() 150 | yield from RT.py_setitem(tos1, tos, tos2) 151 | elif b.name == InstrNames.DELETE_SUBSCR: 152 | tos = yield am.pop() 153 | tos1 = yield am.pop() 154 | yield from RT.py_delitem(tos1, tos) 155 | elif b.name == InstrNames.PRINT_EXPR: 156 | tos = yield am.pop() 157 | yield from RT.py_printexpr(tos) 158 | elif b.name == InstrNames.SET_ADD: 159 | tos = yield am.pop() 160 | subj = yield am.peek(b.arg) # TODO: check correctness 161 | yield from RT.py_set_add(subj, tos) 162 | elif b.name == InstrNames.LIST_APPEND: 163 | tos = yield am.pop() 164 | subj = yield am.peek(b.arg) 165 | yield from RT.py_list_append(subj, tos) 166 | elif b.name == InstrNames.MAP_ADD: 167 | tos = yield am.pop() 168 | tos1 = yield am.pop() 169 | subj = yield am.peek(b.arg) 170 | yield from RT.py_map_add(subj, tos, tos1) 171 | elif b.name == InstrNames.RETURN_VALUE: 172 | a = yield am.pop() 173 | yield am.ret(a) 174 | elif b.name == InstrNames.YIELD_VALUE: 175 | a = yield am.pop() 176 | yield am.yield_return(a) 177 | 178 | elif b.name == InstrNames.YIELD_FROM: 179 | a = yield am.pop() 180 | yield from RT.yield_from(a) 181 | elif b.name == InstrNames.STORE_NAME: 182 | assert isinstance(b.arg, str) 183 | raise NotImplemented 184 | elif b.name == InstrNames.DELETE_NAME: 185 | raise NotImplemented 186 | elif b.name == InstrNames.UNPACK_SEQUENCE: 187 | val = yield am.pop() 188 | val = yield from RT.py_to_tuple(val) 189 | for idx in range(b.arg): 190 | idx = yield am.const(idx) 191 | a = yield from RT.py_subscr(val, idx) 192 | yield am.push(a) 193 | elif b.name == InstrNames.UNPACK_EX: 194 | val = yield am.pop() 195 | val = yield from RT.py_to_tuple(val) 196 | n = yield from RT.py_len(val) 197 | tail = b.arg // 256 198 | init = b.arg % 256 199 | start_idx = None 200 | for start_idx in range(init): 201 | start_idx = am.const(start_idx) 202 | a = yield from RT.py_subscr(val, start_idx) 203 | yield am.push(a) 204 | 205 | start_idx = yield am.const(start_idx) 206 | 207 | tail = yield am.const(tail) 208 | _1 = yield am.const(1) 209 | end_idx = yield from RT.py_sub(n, tail) 210 | a = yield from RT.py_build_slice(start_idx, _1, end_idx) 211 | yield am.push(a) 212 | 213 | for end_idx in range(init): 214 | off = yield am.const(end_idx + 1) 215 | end_idx = yield from RT.py_sub(n, off) 216 | a = yield from RT.py_subscr(val, end_idx) 217 | yield am.push(a) 218 | 219 | elif b.name == InstrNames.FORMAT_VALUE: 220 | a = yield am.pop() 221 | if b.arg & 0x03 == 0x00: 222 | a = yield from RT.py_format(a) 223 | elif b.arg & 0x03 == 0x01: 224 | a = yield from RT.py_to_str(a) 225 | a = yield from RT.py_format(a) 226 | elif b.arg & 0x03 == 0x02: 227 | a = yield from RT.py_to_repr(a) 228 | a = yield from RT.py_format(a) 229 | elif b.arg & 0x03 == 0x03: 230 | a = yield from RT.py_to_ascii(a) 231 | a = yield from RT.py_format(a) 232 | elif b.arg & 0x04 == 0x04: 233 | a = yield from RT.py_to_ascii(a) 234 | a = yield from RT.py_format(a) 235 | else: 236 | raise ValueError(f"invalid format flag {b.arg}") 237 | yield am.push(a) 238 | elif b.name == InstrNames.BUILD_SLICE: 239 | tos = yield am.pop() 240 | tos1 = yield am.pop() 241 | tos2 = yield am.pop() 242 | a = yield from RT.py_build_slice(tos2, tos1, tos) 243 | yield am.push(a) 244 | elif b.name == InstrNames.LOAD_METHOD: 245 | a = yield am.pop() 246 | yield from RT.py_load_method_(a, b.arg) 247 | elif b.name == InstrNames.CALL_METHOD: 248 | xs = yield from pop_n(b.arg + 2) 249 | a = yield from RT.py_call_method(*xs) 250 | yield am.push(a) 251 | elif b.name == InstrNames.MAKE_FUNCTION: 252 | arg = b.arg 253 | if arg & 0x01: 254 | raise NotImplemented 255 | if arg & 0x02: 256 | raise NotImplemented 257 | if arg & 0x04: 258 | raise NotImplemented 259 | if arg & 0x08: 260 | name = yield am.pop() 261 | name = yield am.from_const(name) 262 | code = yield am.pop() 263 | code = yield am.from_const(code) 264 | assert isinstance(code, types.CodeType) 265 | closure = yield am.pop() 266 | fpr = yield from RT.py_mk_func(name, code) 267 | a = yield from RT.py_mk_closure(closure, fpr) 268 | 269 | else: 270 | name = yield am.pop() 271 | name = yield am.from_const(name) 272 | code = yield am.pop() 273 | code = yield am.from_const(code) 274 | assert isinstance(code, types.CodeType) 275 | a = yield from RT.py_mk_func(name, code) 276 | yield am.push(a) 277 | 278 | elif b.name == InstrNames.CALL_FUNCTION: 279 | c = b.arg 280 | xs = yield from pop_n(c) 281 | f = yield am.pop() 282 | a = yield from RT.py_call_func(f, *xs) 283 | yield am.push(a) 284 | 285 | elif b.name == InstrNames.LOAD_CLOSURE: 286 | arg = b.arg 287 | assert isinstance(arg, (bc.CellVar, bc.FreeVar)) 288 | a = yield am.reg_of(arg.name) 289 | yield am.push(a) 290 | 291 | elif b.name == InstrNames.LOAD_DEREF: 292 | arg = b.arg 293 | assert isinstance(arg, (bc.CellVar, bc.FreeVar)) 294 | a = yield am.load(arg.name) 295 | yield am.push(a) 296 | 297 | elif b.name == InstrNames.STORE_DEREF: 298 | arg = b.arg 299 | assert isinstance(arg, (bc.CellVar, bc.FreeVar)) 300 | 301 | a = yield am.pop() 302 | yield am.store(arg.name, a) 303 | 304 | elif b.name == InstrNames.LOAD_FAST: 305 | assert isinstance(b.arg, str) 306 | reg = yield am.reg_of(b.arg) 307 | yield am.push(reg) 308 | 309 | elif b.name == InstrNames.STORE_FAST: 310 | assert isinstance(b.arg, str) 311 | v = yield am.pop() 312 | yield am.assign(b.arg, v) 313 | 314 | elif b.name == InstrNames.LOAD_GLOBAL: 315 | arg = b.arg 316 | assert isinstance(arg, str) 317 | yield am.require_global(arg) 318 | a = yield am.from_higher("", arg) 319 | yield am.push(a) 320 | 321 | elif b.name == InstrNames.STORE_GLOBAL: 322 | yield am.require_global(b.arg) 323 | raise NotImplemented 324 | 325 | elif b.name == InstrNames.LOAD_CONST: 326 | arg = b.arg 327 | a = yield am.const(arg) 328 | yield am.push(a) 329 | 330 | elif b.name == InstrNames.STORE_ATTR: 331 | tos = yield am.pop() 332 | val = yield am.pop() 333 | yield from RT.py_store_attr(tos, val, b.arg) 334 | 335 | elif b.name == InstrNames.LOAD_ATTR: 336 | tos = yield am.pop() 337 | a = yield from RT.py_get_attr(tos, b.arg) 338 | yield am.push(a) 339 | 340 | elif b.name == InstrNames.DELETE_ATTR: 341 | tos = yield am.pop() 342 | yield from RT.py_del_attr(tos, b.arg) 343 | 344 | elif b.name == InstrNames.BUILD_TUPLE: 345 | xs = yield from pop_n(b.arg) 346 | a = yield from RT.py_mk_tuple(xs) 347 | yield am.push(a) 348 | 349 | elif b.name == InstrNames.BUILD_LIST: 350 | xs = yield from pop_n(b.arg) 351 | a = yield from RT.py_mk_list(xs) 352 | yield am.push(a) 353 | 354 | elif b.name == InstrNames.BUILD_SET: 355 | xs = yield from pop_n(b.arg) 356 | a = yield from RT.py_mk_set(xs) 357 | yield am.push(a) 358 | 359 | elif b.name == InstrNames.BUILD_MAP: 360 | xs = yield from pop_n(2 * b.arg) 361 | ks = xs[::2] 362 | vs = xs[1::2] 363 | ks = yield from RT.py_mk_tuple(ks) 364 | vs = yield from RT.py_mk_tuple(vs) 365 | a = yield from RT.py_mk_map(ks, vs) 366 | yield am.push(a) 367 | 368 | elif b.name == InstrNames.BUILD_CONST_KEY_MAP: 369 | ks = yield am.pop() 370 | vs = yield from pop_n(b.arg) 371 | vs = yield from RT.py_mk_tuple(vs) 372 | a = yield from RT.py_mk_map(ks, vs) 373 | yield a.push(a) 374 | 375 | elif b.name == InstrNames.BUILD_STRING: 376 | xs = yield from pop_n(b.arg) 377 | a = yield from RT.py_cat_strs(xs) 378 | yield am.push(a) 379 | 380 | elif b.name == InstrNames.FOR_ITER: 381 | f = yield am.pop() 382 | a = yield from RT.py_call_func(f) 383 | label = b.arg 384 | assert isinstance(label, bc.BasicBlock) 385 | check_if_nothing = yield from RT.py_is_none(a) 386 | yield am.jump_if(label_to_name(label), check_if_nothing) 387 | _0 = yield am.const(0) 388 | _1 = yield am.const(1) 389 | elt = yield from RT.py_subscr(a, _0) 390 | st = yield from RT.py_subscr(a, _1) 391 | yield am.push(st) 392 | yield am.push(elt) 393 | 394 | elif b.name == InstrNames.GET_ITER: 395 | a = yield am.pop() 396 | a = yield from RT.py_get_no_exception_iter(a) 397 | yield am.push(a) 398 | 399 | elif b.name == InstrNames.SETUP_LOOP: 400 | pass 401 | elif b.name == InstrNames.POP_BLOCK: 402 | pass 403 | elif b.name == InstrNames.POP_TOP: 404 | yield am.pop() 405 | elif b.name == InstrNames.SETUP_WITH: 406 | arg = b.arg 407 | assert isinstance(arg, bc.BasicBlock) 408 | end_label = label_to_name(arg) 409 | meta = yield am.meta() 410 | var = yield am.pop() 411 | 412 | def unwind(): 413 | # python 'with' block exits with a pushed 'None' 414 | yield am.pop() 415 | yield am.pop_block() 416 | yield from RT.py_exit(var) 417 | 418 | meta[end_label] = unwind 419 | entered = yield from RT.py_enter(var) 420 | yield am.push(entered) 421 | yield am.push_block(end_label) 422 | elif b.name == InstrNames.END_FINALLY: 423 | exc = yield am.pop_exception() 424 | label_no_ext = yield am.alloc() 425 | py_none = yield from RT.py_none() 426 | exc_is_none = yield from RT.py_is(py_none, exc) 427 | yield am.jump_if(label_no_ext, exc_is_none) 428 | yield from RT.py_throw(exc) 429 | yield am.label(label_no_ext) 430 | 431 | elif b.name == InstrNames.SETUP_FINALLY: 432 | arg = b.arg 433 | assert isinstance(arg, bc.BasicBlock) 434 | end_label = label_to_name(arg) 435 | meta = yield am.meta() 436 | 437 | def unwind(): 438 | yield am.pop_block() 439 | 440 | meta[end_label] = unwind 441 | yield am.push_block(end_label) 442 | elif b.name == InstrNames.SETUP_EXCEPT: 443 | arg = b.arg 444 | assert isinstance(arg, bc.BasicBlock) 445 | end_label = label_to_name(arg) 446 | meta = yield am.meta() 447 | 448 | def unwind(): 449 | yield am.pop_block() 450 | exc = yield am.pop_exception(must=True) 451 | py_none = yield from RT.py_none() 452 | yield am.push(py_none) 453 | yield am.push(exc) 454 | yield am.push(py_none) 455 | 456 | meta[end_label] = unwind 457 | yield am.push_block(end_label) 458 | elif b.name in (InstrNames.WITH_CLEANUP_FINISH, 459 | InstrNames.WITH_CLEANUP_START, 460 | InstrNames.POP_EXCEPT): 461 | pass 462 | elif b.name == InstrNames.RAISE_VARARGS: 463 | c = b.arg 464 | if c is not 1: 465 | raise ValueError( 466 | "Raise statement must take 1 argument due to the limitations of current implementation." 467 | ) 468 | err = yield am.pop() 469 | yield from RT.py_throw(err) 470 | else: 471 | raise NotImplementedError( 472 | f"instruction {b} not supported yet") 473 | 474 | 475 | def abs_i_unary(b: bc.Instr): 476 | a = yield am.pop() 477 | f = u_map.get(b.name, None) 478 | if f is None: 479 | raise ValueError(f"unknown unary instruction {b}") 480 | a = yield from f(a) 481 | yield am.push(a) 482 | 483 | 484 | def abs_i_jump(b: bc.Instr): 485 | label_name = label_to_name(b.arg) 486 | if b.name in (InstrNames.JUMP_FORWARD, InstrNames.JUMP_ABSOLUTE): 487 | yield am.jump(label_name) 488 | 489 | elif b.name == InstrNames.POP_JUMP_IF_FALSE: 490 | a = yield am.pop() 491 | a = yield from RT.py_not(a) 492 | yield am.jump_if(label_name, a) 493 | elif b.name == InstrNames.POP_JUMP_IF_TRUE: 494 | a = yield am.pop() 495 | yield am.jump_if(label_name, a) 496 | elif b.name == InstrNames.JUMP_IF_FALSE_OR_POP: 497 | a = yield am.pop() 498 | o = yield from RT.py_not(a) 499 | yield am.jump_if_push(label_name, o, a) 500 | elif b.name == InstrNames.JUMP_IF_TRUE_OR_POP: 501 | a = yield am.pop() 502 | yield am.jump_if_push(label_name, a, a) 503 | else: 504 | raise ValueError(f"unknown jump instruction {b}") 505 | 506 | 507 | def abs_i_binary(b: bc.Instr): 508 | a2 = yield am.pop() 509 | a1 = yield am.pop() 510 | f = bin_map.get(b.name, None) 511 | if f is None: 512 | raise ValueError(f"unknown binary instruction {b}") 513 | c = yield from f(a1, a2) 514 | yield am.push(c) 515 | 516 | 517 | def abs_i_inplace_binary(b: bc.Instr): 518 | a2 = yield am.pop() 519 | a1 = yield am.pop() 520 | f = ibin_map.get(b.name, None) 521 | if f is None: 522 | raise ValueError(f"unknown binary instruction {b}") 523 | c = yield from f(a1, a2) 524 | yield am.push(c) 525 | 526 | 527 | def abs_i_cmp(b: bc.Instr): 528 | arg: bc.Compare = b.arg 529 | f = cmp_map[arg] 530 | a2 = yield am.pop() 531 | a1 = yield am.pop() 532 | a = yield from f(a1, a2) 533 | yield am.push(a) 534 | -------------------------------------------------------------------------------- /restrain_jit/abs_compiler/instrnames.py: -------------------------------------------------------------------------------- 1 | POP_TOP = 'POP_TOP' 2 | ROT_TWO = 'ROT_TWO' 3 | ROT_THREE = 'ROT_THREE' 4 | DUP_TOP = 'DUP_TOP' 5 | DUP_TOP_TWO = 'DUP_TOP_TWO' 6 | NOP = 'NOP' 7 | UNARY_POSITIVE = 'UNARY_POSITIVE' 8 | UNARY_NEGATIVE = 'UNARY_NEGATIVE' 9 | UNARY_NOT = 'UNARY_NOT' 10 | UNARY_INVERT = 'UNARY_INVERT' 11 | BINARY_MATRIX_MULTIPLY = 'BINARY_MATRIX_MULTIPLY' 12 | INPLACE_MATRIX_MULTIPLY = 'INPLACE_MATRIX_MULTIPLY' 13 | BINARY_POWER = 'BINARY_POWER' 14 | BINARY_MULTIPLY = 'BINARY_MULTIPLY' 15 | BINARY_MODULO = 'BINARY_MODULO' 16 | BINARY_ADD = 'BINARY_ADD' 17 | BINARY_SUBTRACT = 'BINARY_SUBTRACT' 18 | BINARY_SUBSCR = 'BINARY_SUBSCR' 19 | BINARY_FLOOR_DIVIDE = 'BINARY_FLOOR_DIVIDE' 20 | BINARY_TRUE_DIVIDE = 'BINARY_TRUE_DIVIDE' 21 | INPLACE_FLOOR_DIVIDE = 'INPLACE_FLOOR_DIVIDE' 22 | INPLACE_TRUE_DIVIDE = 'INPLACE_TRUE_DIVIDE' 23 | GET_AITER = 'GET_AITER' 24 | GET_ANEXT = 'GET_ANEXT' 25 | BEFORE_ASYNC_WITH = 'BEFORE_ASYNC_WITH' 26 | INPLACE_ADD = 'INPLACE_ADD' 27 | INPLACE_SUBTRACT = 'INPLACE_SUBTRACT' 28 | INPLACE_MULTIPLY = 'INPLACE_MULTIPLY' 29 | INPLACE_MODULO = 'INPLACE_MODULO' 30 | STORE_SUBSCR = 'STORE_SUBSCR' 31 | DELETE_SUBSCR = 'DELETE_SUBSCR' 32 | BINARY_LSHIFT = 'BINARY_LSHIFT' 33 | BINARY_RSHIFT = 'BINARY_RSHIFT' 34 | BINARY_AND = 'BINARY_AND' 35 | BINARY_XOR = 'BINARY_XOR' 36 | BINARY_OR = 'BINARY_OR' 37 | INPLACE_POWER = 'INPLACE_POWER' 38 | GET_ITER = 'GET_ITER' 39 | GET_YIELD_FROM_ITER = 'GET_YIELD_FROM_ITER' 40 | PRINT_EXPR = 'PRINT_EXPR' 41 | LOAD_BUILD_CLASS = 'LOAD_BUILD_CLASS' 42 | YIELD_FROM = 'YIELD_FROM' 43 | GET_AWAITABLE = 'GET_AWAITABLE' 44 | INPLACE_LSHIFT = 'INPLACE_LSHIFT' 45 | INPLACE_RSHIFT = 'INPLACE_RSHIFT' 46 | INPLACE_AND = 'INPLACE_AND' 47 | INPLACE_XOR = 'INPLACE_XOR' 48 | INPLACE_OR = 'INPLACE_OR' 49 | BREAK_LOOP = 'BREAK_LOOP' 50 | WITH_CLEANUP_START = 'WITH_CLEANUP_START' 51 | WITH_CLEANUP_FINISH = 'WITH_CLEANUP_FINISH' 52 | RETURN_VALUE = 'RETURN_VALUE' 53 | IMPORT_STAR = 'IMPORT_STAR' 54 | SETUP_ANNOTATIONS = 'SETUP_ANNOTATIONS' 55 | YIELD_VALUE = 'YIELD_VALUE' 56 | POP_BLOCK = 'POP_BLOCK' 57 | END_FINALLY = 'END_FINALLY' 58 | POP_EXCEPT = 'POP_EXCEPT' 59 | STORE_NAME = 'STORE_NAME' 60 | DELETE_NAME = 'DELETE_NAME' 61 | UNPACK_SEQUENCE = 'UNPACK_SEQUENCE' 62 | FOR_ITER = 'FOR_ITER' 63 | UNPACK_EX = 'UNPACK_EX' 64 | STORE_ATTR = 'STORE_ATTR' 65 | DELETE_ATTR = 'DELETE_ATTR' 66 | STORE_GLOBAL = 'STORE_GLOBAL' 67 | DELETE_GLOBAL = 'DELETE_GLOBAL' 68 | LOAD_CONST = 'LOAD_CONST' 69 | LOAD_NAME = 'LOAD_NAME' 70 | BUILD_TUPLE = 'BUILD_TUPLE' 71 | BUILD_LIST = 'BUILD_LIST' 72 | BUILD_SET = 'BUILD_SET' 73 | BUILD_MAP = 'BUILD_MAP' 74 | LOAD_ATTR = 'LOAD_ATTR' 75 | COMPARE_OP = 'COMPARE_OP' 76 | IMPORT_NAME = 'IMPORT_NAME' 77 | IMPORT_FROM = 'IMPORT_FROM' 78 | JUMP_FORWARD = 'JUMP_FORWARD' 79 | JUMP_IF_FALSE_OR_POP = 'JUMP_IF_FALSE_OR_POP' 80 | JUMP_IF_TRUE_OR_POP = 'JUMP_IF_TRUE_OR_POP' 81 | JUMP_ABSOLUTE = 'JUMP_ABSOLUTE' 82 | POP_JUMP_IF_FALSE = 'POP_JUMP_IF_FALSE' 83 | POP_JUMP_IF_TRUE = 'POP_JUMP_IF_TRUE' 84 | LOAD_GLOBAL = 'LOAD_GLOBAL' 85 | CONTINUE_LOOP = 'CONTINUE_LOOP' 86 | SETUP_LOOP = 'SETUP_LOOP' 87 | SETUP_EXCEPT = 'SETUP_EXCEPT' 88 | SETUP_FINALLY = 'SETUP_FINALLY' 89 | LOAD_FAST = 'LOAD_FAST' 90 | STORE_FAST = 'STORE_FAST' 91 | DELETE_FAST = 'DELETE_FAST' 92 | RAISE_VARARGS = 'RAISE_VARARGS' 93 | CALL_FUNCTION = 'CALL_FUNCTION' 94 | MAKE_FUNCTION = 'MAKE_FUNCTION' 95 | BUILD_SLICE = 'BUILD_SLICE' 96 | LOAD_CLOSURE = 'LOAD_CLOSURE' 97 | LOAD_DEREF = 'LOAD_DEREF' 98 | STORE_DEREF = 'STORE_DEREF' 99 | DELETE_DEREF = 'DELETE_DEREF' 100 | CALL_FUNCTION_KW = 'CALL_FUNCTION_KW' 101 | CALL_FUNCTION_EX = 'CALL_FUNCTION_EX' 102 | SETUP_WITH = 'SETUP_WITH' 103 | LIST_APPEND = 'LIST_APPEND' 104 | SET_ADD = 'SET_ADD' 105 | MAP_ADD = 'MAP_ADD' 106 | LOAD_CLASSDEREF = 'LOAD_CLASSDEREF' 107 | EXTENDED_ARG = 'EXTENDED_ARG' 108 | BUILD_LIST_UNPACK = 'BUILD_LIST_UNPACK' 109 | BUILD_MAP_UNPACK = 'BUILD_MAP_UNPACK' 110 | BUILD_MAP_UNPACK_WITH_CALL = 'BUILD_MAP_UNPACK_WITH_CALL' 111 | BUILD_TUPLE_UNPACK = 'BUILD_TUPLE_UNPACK' 112 | BUILD_SET_UNPACK = 'BUILD_SET_UNPACK' 113 | SETUP_ASYNC_WITH = 'SETUP_ASYNC_WITH' 114 | FORMAT_VALUE = 'FORMAT_VALUE' 115 | BUILD_CONST_KEY_MAP = 'BUILD_CONST_KEY_MAP' 116 | BUILD_STRING = 'BUILD_STRING' 117 | BUILD_TUPLE_UNPACK_WITH_CALL = 'BUILD_TUPLE_UNPACK_WITH_CALL' 118 | LOAD_METHOD = 'LOAD_METHOD' 119 | CALL_METHOD = 'CALL_METHOD' 120 | -------------------------------------------------------------------------------- /restrain_jit/abs_compiler/py_apis.py: -------------------------------------------------------------------------------- 1 | import types 2 | import bytecode as bc 3 | from restrain_jit.vm.am import * 4 | 5 | 6 | class NS: 7 | RestrainJIT = "RestrainJIT" 8 | 9 | 10 | def py_not(a: Repr): 11 | fn = yield from_lower(NS.RestrainJIT, py_not.__name__) 12 | a = yield app(fn, [a]) 13 | return a 14 | 15 | 16 | def py_pos(a: Repr): 17 | fn = yield from_lower(NS.RestrainJIT, py_pos.__name__) 18 | a = yield app(fn, [a]) 19 | return a 20 | 21 | 22 | def py_is_true(a: Repr): 23 | fn = yield from_lower(NS.RestrainJIT, py_is_true.__name__) 24 | a = yield app(fn, [a]) 25 | return a 26 | 27 | 28 | def py_inv(a: Repr): 29 | fn = yield from_lower(NS.RestrainJIT, py_inv.__name__) 30 | a = yield app(fn, [a]) 31 | return a 32 | 33 | 34 | def py_neg(a: Repr): 35 | fn = yield from_lower(NS.RestrainJIT, py_neg.__name__) 36 | a = yield app(fn, [a]) 37 | return a 38 | 39 | 40 | def py_iter(a: Repr): 41 | fn = yield from_lower(NS.RestrainJIT, py_iter.__name__) 42 | a = yield app(fn, [a]) 43 | return a 44 | 45 | 46 | def py_pow(a: Repr, b: Repr): 47 | fn = yield from_lower(NS.RestrainJIT, py_pow.__name__) 48 | a = yield app(fn, [a, b]) 49 | return a 50 | 51 | 52 | def py_mul(a: Repr, b: Repr): 53 | fn = yield from_lower(NS.RestrainJIT, py_mul.__name__) 54 | a = yield app(fn, [a, b]) 55 | return a 56 | 57 | 58 | def py_matmul(a: Repr, b: Repr): 59 | fn = yield from_lower(NS.RestrainJIT, py_matmul.__name__) 60 | a = yield app(fn, [a, b]) 61 | return a 62 | 63 | 64 | def py_floordiv(a: Repr, b: Repr): 65 | fn = yield from_lower(NS.RestrainJIT, py_floordiv.__name__) 66 | a = yield app(fn, [a, b]) 67 | return a 68 | 69 | 70 | def py_truediv(a: Repr, b: Repr): 71 | fn = yield from_lower(NS.RestrainJIT, py_truediv.__name__) 72 | a = yield app(fn, [a, b]) 73 | return a 74 | 75 | 76 | def py_mod(a: Repr, b: Repr): 77 | fn = yield from_lower(NS.RestrainJIT, py_mod.__name__) 78 | a = yield app(fn, [a, b]) 79 | return a 80 | 81 | 82 | def py_add(a: Repr, b: Repr): 83 | fn = yield from_lower(NS.RestrainJIT, py_add.__name__) 84 | a = yield app(fn, [a, b]) 85 | return a 86 | 87 | 88 | def py_sub(a: Repr, b: Repr): 89 | fn = yield from_lower(NS.RestrainJIT, py_sub.__name__) 90 | a = yield app(fn, [a, b]) 91 | return a 92 | 93 | 94 | def py_subscr(a: Repr, b: Repr): 95 | fn = yield from_lower(NS.RestrainJIT, py_subscr.__name__) 96 | a = yield app(fn, [a, b]) 97 | return a 98 | 99 | 100 | def py_lsh(a: Repr, b: Repr): 101 | fn = yield from_lower(NS.RestrainJIT, py_lsh.__name__) 102 | a = yield app(fn, [a, b]) 103 | return a 104 | 105 | 106 | def py_rsh(a: Repr, b: Repr): 107 | fn = yield from_lower(NS.RestrainJIT, py_rsh.__name__) 108 | a = yield app(fn, [a, b]) 109 | return a 110 | 111 | 112 | def py_and(a: Repr, b: Repr): 113 | fn = yield from_lower(NS.RestrainJIT, py_and.__name__) 114 | a = yield app(fn, [a, b]) 115 | return a 116 | 117 | 118 | def py_xor(a: Repr, b: Repr): 119 | fn = yield from_lower(NS.RestrainJIT, py_xor.__name__) 120 | a = yield app(fn, [a, b]) 121 | return a 122 | 123 | 124 | def py_or(a: Repr, b: Repr): 125 | fn = yield from_lower(NS.RestrainJIT, py_or.__name__) 126 | a = yield app(fn, [a, b]) 127 | return a 128 | 129 | 130 | def py_ipow(a: Repr, b: Repr): 131 | fn = yield from_lower(NS.RestrainJIT, py_ipow.__name__) 132 | a = yield app(fn, [a, b]) 133 | return a 134 | 135 | 136 | def py_imul(a: Repr, b: Repr): 137 | fn = yield from_lower(NS.RestrainJIT, py_imul.__name__) 138 | a = yield app(fn, [a, b]) 139 | return a 140 | 141 | 142 | def py_imatmul(a: Repr, b: Repr): 143 | fn = yield from_lower(NS.RestrainJIT, py_imatmul.__name__) 144 | a = yield app(fn, [a, b]) 145 | return a 146 | 147 | 148 | def py_ifloordiv(a: Repr, b: Repr): 149 | fn = yield from_lower(NS.RestrainJIT, py_ifloordiv.__name__) 150 | a = yield app(fn, [a, b]) 151 | return a 152 | 153 | 154 | def py_itruediv(a: Repr, b: Repr): 155 | fn = yield from_lower(NS.RestrainJIT, py_itruediv.__name__) 156 | a = yield app(fn, [a, b]) 157 | return a 158 | 159 | 160 | def py_imod(a: Repr, b: Repr): 161 | fn = yield from_lower(NS.RestrainJIT, py_imod.__name__) 162 | a = yield app(fn, [a, b]) 163 | return a 164 | 165 | 166 | def py_iadd(a: Repr, b: Repr): 167 | fn = yield from_lower(NS.RestrainJIT, py_iadd.__name__) 168 | a = yield app(fn, [a, b]) 169 | return a 170 | 171 | 172 | def py_isub(a: Repr, b: Repr): 173 | fn = yield from_lower(NS.RestrainJIT, py_isub.__name__) 174 | a = yield app(fn, [a, b]) 175 | return a 176 | 177 | 178 | def py_isubscr(a: Repr, b: Repr): 179 | fn = yield from_lower(NS.RestrainJIT, py_isubscr.__name__) 180 | a = yield app(fn, [a, b]) 181 | return a 182 | 183 | 184 | def py_ilsh(a: Repr, b: Repr): 185 | fn = yield from_lower(NS.RestrainJIT, py_ilsh.__name__) 186 | a = yield app(fn, [a, b]) 187 | return a 188 | 189 | 190 | def py_irsh(a: Repr, b: Repr): 191 | fn = yield from_lower(NS.RestrainJIT, py_irsh.__name__) 192 | a = yield app(fn, [a, b]) 193 | return a 194 | 195 | 196 | def py_iand(a: Repr, b: Repr): 197 | fn = yield from_lower(NS.RestrainJIT, py_iand.__name__) 198 | a = yield app(fn, [a, b]) 199 | return a 200 | 201 | 202 | def py_ixor(a: Repr, b: Repr): 203 | fn = yield from_lower(NS.RestrainJIT, py_ixor.__name__) 204 | a = yield app(fn, [a, b]) 205 | return a 206 | 207 | 208 | def py_ior(a: Repr, b: Repr): 209 | fn = yield from_lower(NS.RestrainJIT, py_ior.__name__) 210 | a = yield app(fn, [a, b]) 211 | return a 212 | 213 | 214 | def py_setitem(subj, key, val): 215 | fn = yield from_lower(NS.RestrainJIT, py_setitem.__name__) 216 | yield app(fn, [subj, key, val]) 217 | 218 | 219 | def py_delitem(subj, key): 220 | fn = yield from_lower(NS.RestrainJIT, py_delitem.__name__) 221 | yield app(fn, [subj, key]) 222 | 223 | 224 | def py_printexpr(tos): 225 | fn = yield from_lower(NS.RestrainJIT, py_printexpr.__name__) 226 | yield app(fn, [tos]) 227 | 228 | 229 | def py_set_add(subj, tos): 230 | fn = yield from_lower(NS.RestrainJIT, py_set_add.__name__) 231 | yield app(fn, [subj, tos]) 232 | 233 | 234 | def py_list_append(subj, tos): 235 | fn = yield from_lower(NS.RestrainJIT, py_list_append.__name__) 236 | yield app(fn, [subj, tos]) 237 | 238 | 239 | def py_map_add(subj, tos, tos1): 240 | fn = yield from_lower(NS.RestrainJIT, py_map_add.__name__) 241 | yield app(fn, [subj, tos, tos1]) 242 | 243 | 244 | def py_to_list(tos): 245 | fn = yield from_lower(NS.RestrainJIT, py_to_list.__name__) 246 | a = yield app(fn, [tos]) 247 | return a 248 | 249 | 250 | def py_to_tuple(tos): 251 | fn = yield from_lower(NS.RestrainJIT, py_to_tuple.__name__) 252 | a = yield app(fn, [tos]) 253 | return a 254 | 255 | 256 | def py_len(tos): 257 | fn = yield from_lower(NS.RestrainJIT, py_len.__name__) 258 | a = yield app(fn, [tos]) 259 | return a 260 | 261 | 262 | def py_build_slice(tos2, tos1, tos): 263 | fn = yield from_lower(NS.RestrainJIT, py_build_slice.__name__) 264 | a = yield app(fn, [tos2, tos1, tos]) 265 | return a 266 | 267 | 268 | def py_to_str(tos): 269 | fn = yield from_lower(NS.RestrainJIT, py_to_str.__name__) 270 | a = yield app(fn, [tos]) 271 | return a 272 | 273 | 274 | def py_to_repr(tos): 275 | fn = yield from_lower(NS.RestrainJIT, py_to_repr.__name__) 276 | a = yield app(fn, [tos]) 277 | return a 278 | 279 | 280 | def py_to_ascii(tos): 281 | fn = yield from_lower(NS.RestrainJIT, py_to_ascii.__name__) 282 | a = yield app(fn, [tos]) 283 | return a 284 | 285 | 286 | def py_format(tos): 287 | fn = yield from_lower(NS.RestrainJIT, py_format.__name__) 288 | a = yield app(fn, [tos]) 289 | return a 290 | 291 | 292 | def py_load_method_(tos, attr: str): 293 | attr = yield const(ValSymbol(attr)) 294 | fn = yield from_lower(NS.RestrainJIT, py_load_method_.__name__) 295 | a = yield app(fn, [tos, attr]) 296 | _0 = yield const(0) 297 | _1 = yield const(1) 298 | a1 = yield from py_subscr(a, _0) 299 | a2 = yield from py_subscr(a, _1) 300 | yield push(a1) 301 | yield push(a2) 302 | 303 | 304 | def py_get_attr(tos: Repr, attr: str): 305 | attr = yield const(ValSymbol(attr)) 306 | fn = yield from_lower(NS.RestrainJIT, py_get_attr.__name__) 307 | a = yield app(fn, [tos, attr]) 308 | return a 309 | 310 | 311 | def py_store_attr(tos: Repr, val: Repr, attr: str): 312 | attr = yield const(Symbol(attr)) 313 | fn = yield from_lower(NS.RestrainJIT, py_store_attr.__name__) 314 | a = yield app(fn, [tos, attr, val]) 315 | return a 316 | 317 | 318 | def py_del_attr(subj: Repr, attr: str): 319 | attr = yield const(Symbol(attr)) 320 | fn = yield from_lower(NS.RestrainJIT, py_del_attr.__name__) 321 | a = yield app(fn, [subj, attr]) 322 | return a 323 | 324 | 325 | def py_call_method(*params: Repr): 326 | fn = yield from_lower(NS.RestrainJIT, py_call_method.__name__) 327 | a = yield app(fn, list(params)) 328 | return a 329 | 330 | 331 | def py_mk_closure(closure_vars: Repr, native_fn_ptr: Repr): 332 | mk_closure = yield from_lower(NS.RestrainJIT, 333 | py_mk_closure.__name__) 334 | a = yield app(mk_closure, [closure_vars, native_fn_ptr]) 335 | return a 336 | 337 | 338 | def py_get_no_exception_iter(v: Repr): 339 | f = yield from_lower(NS.RestrainJIT, py_get_no_exception_iter.__name__) 340 | a = yield app(f, [v]) 341 | return a 342 | 343 | 344 | class Indirect: 345 | """ 346 | start a new VM object 347 | """ 348 | f: 't.Callable' 349 | 350 | 351 | def py_mk_func(name: str, code: types.CodeType): 352 | code = bc.Bytecode.from_code(code) 353 | a = yield code_info(code) 354 | for each in a.glob_deps: 355 | yield require_global(each) 356 | 357 | a = yield const(a) 358 | f = yield from_lower(NS.RestrainJIT, py_mk_func.__name__) 359 | n = yield const(name) 360 | a = yield app(f, [n, a]) 361 | return a 362 | 363 | 364 | def py_call_func_varargs(f, a): 365 | raise NotImplemented 366 | 367 | 368 | def py_call_func_varargs_kwargs(f, a, b): 369 | raise NotImplemented 370 | 371 | 372 | def py_call_func_kwargs(f, attrs: t.Tuple[str, ...], *args): 373 | raise NotImplemented 374 | 375 | 376 | def py_call_func(f: Repr, *args: Repr): 377 | a = yield app(f, args) 378 | return a 379 | 380 | 381 | def py_raise(*xs): 382 | """ 383 | len(xs) = 0, 1, 2 384 | """ 385 | raise NotImplemented 386 | 387 | 388 | def py_mk_tuple(xs: t.List[Repr]): 389 | fn = yield from_lower(NS.RestrainJIT, py_mk_tuple.__name__) 390 | a = yield app(fn, xs) 391 | return a 392 | 393 | 394 | def py_mk_list(xs: t.List[Repr]): 395 | fn = yield from_lower(NS.RestrainJIT, py_mk_list.__name__) 396 | a = yield app(fn, xs) 397 | return a 398 | 399 | 400 | def py_mk_set(xs: t.List[Repr]): 401 | fn = yield from_lower(NS.RestrainJIT, py_mk_set.__name__) 402 | a = yield app(fn, xs) 403 | return a 404 | 405 | 406 | def py_mk_map(keys: Repr, vals: Repr): 407 | fn = yield from_lower(NS.RestrainJIT, py_mk_map.__name__) 408 | a = yield app(fn, [keys, vals]) 409 | return a 410 | 411 | 412 | def py_cat_strs(vs: t.List[Repr]): 413 | fn = yield from_lower(NS.RestrainJIT, py_cat_strs.__name__) 414 | a = yield app(fn, vs) 415 | return a 416 | 417 | 418 | def py_is_none(v: Repr): 419 | fn = yield from_lower(NS.RestrainJIT, py_is_none.__name__) 420 | a = yield app(fn, [v]) 421 | return a 422 | 423 | def yield_from(a: Repr): 424 | raise NotImplemented 425 | 426 | 427 | def py_eq(a: Repr, b: Repr): 428 | fn = from_lower(NS.RestrainJIT, py_eq.__name__) 429 | a = yield from app(fn, [a, b]) 430 | return a 431 | 432 | 433 | def py_neq(a: Repr, b: Repr): 434 | fn = yield from_lower(NS.RestrainJIT, py_neq.__name__) 435 | a = yield app(fn, [a, b]) 436 | return a 437 | 438 | 439 | def py_is(a: Repr, b: Repr): 440 | fn = yield from_lower(NS.RestrainJIT, py_is.__name__) 441 | a = yield app(fn, [a, b]) 442 | return a 443 | 444 | 445 | def py_is_not(a: Repr, b: Repr): 446 | fn = yield from_lower(NS.RestrainJIT, py_is_not.__name__) 447 | a = yield app(fn, [a, b]) 448 | return a 449 | 450 | 451 | def py_lt(a: Repr, b: Repr): 452 | fn = yield from_lower(NS.RestrainJIT, py_lt.__name__) 453 | a = yield app(fn, [a, b]) 454 | return a 455 | 456 | 457 | def py_le(a: Repr, b: Repr): 458 | fn = yield from_lower(NS.RestrainJIT, py_le.__name__) 459 | a = yield app(fn, [a, b]) 460 | return a 461 | 462 | 463 | def py_gt(a: Repr, b: Repr): 464 | fn = yield from_lower(NS.RestrainJIT, py_gt.__name__) 465 | a = yield app(fn, [a, b]) 466 | return a 467 | 468 | 469 | def py_ge(a: Repr, b: Repr): 470 | fn = yield from_lower(NS.RestrainJIT, py_ge.__name__) 471 | a = yield app(fn, [a, b]) 472 | return a 473 | 474 | 475 | def py_in(a: Repr, b: Repr): 476 | fn = yield from_lower(NS.RestrainJIT, py_in.__name__) 477 | a = yield app(fn, [a, b]) 478 | return a 479 | 480 | 481 | def py_not_in(a: Repr, b: Repr): 482 | fn = yield from_lower(NS.RestrainJIT, py_not_in.__name__) 483 | a = yield app(fn, [a, b]) 484 | return a 485 | 486 | 487 | def py_enter(with_val: Repr): 488 | fn = yield from_lower(NS.RestrainJIT, py_enter.__name__) 489 | with_val = yield app(fn, [with_val]) 490 | return with_val 491 | 492 | 493 | def py_exc_match(exc: Repr, exc_ty): 494 | fn = yield from_lower(NS.RestrainJIT, py_exc_match.__name__) 495 | with_val = yield app(fn, [exc, exc_ty]) 496 | return with_val 497 | 498 | 499 | def py_throw(err: Repr): 500 | fn = yield from_lower(NS.RestrainJIT, py_throw.__name__) 501 | err = yield app(fn, [err]) 502 | return err 503 | 504 | 505 | def py_exit(with_val: Repr): 506 | fn = yield from_lower(NS.RestrainJIT, py_exit.__name__) 507 | yield app(fn, [with_val]) 508 | 509 | 510 | def py_none(): 511 | a = yield const(None) 512 | return a 513 | 514 | 515 | def jl_isa(a: Repr, b: Repr): 516 | fn = yield from_lower(NS.RestrainJIT, jl_isa.__name__) 517 | a = yield app(fn, [a, b]) 518 | return a 519 | -------------------------------------------------------------------------------- /restrain_jit/becython/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/restrain_jit/becython/__init__.py -------------------------------------------------------------------------------- /restrain_jit/becython/cy_jit.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_jit_hotspot_comp import Strategy, HitN, JITRecompilationDecision, extension_type_pxd_paths, is_jit_able_type 2 | from restrain_jit.becython.cy_method_codegen import CodeEmit, UndefGlobal 3 | from restrain_jit.becython.cy_jit_common import * 4 | from restrain_jit.becython.cython_vm import CyVM 5 | from restrain_jit.becython.cy_jit_ext_template import mk_module_code, mk_call_record_t, mk_hard_coded_method_getter_module 6 | from restrain_jit.becython.cy_loader import compile_module, JIT_FUNCTION_DIR, JIT_TYPE_DIR, ximport_module 7 | from restrain_jit.jit_info import PyFuncInfo, PyCodeInfo 8 | from restrain_jit.utils import CodeOut 9 | from restrain_jit.becython import cy_jit_common as cmc 10 | 11 | from types import FunctionType, ModuleType 12 | 13 | import typing as t 14 | 15 | 16 | class JITFunctionHoldPlace: 17 | memoize_partial_code: t.Optional[CodeOut] 18 | method_arguments: t.Optional[t.Tuple[str, ...]] 19 | function_module: ModuleType 20 | base_method_module: ModuleType 21 | methods: t.Dict[call_record_t, fptr_addr] 22 | 23 | def __init__(self, jit_sys: 'JITSystem', code_info: PyCodeInfo): 24 | self.sys = jit_sys 25 | self.globals = {} 26 | self.methods: t.Dict[call_record_t, fptr_addr] = {} 27 | # do not access code_info.glob_deps, 28 | # due to some global var like `type`, `len` are removed from here 29 | self.glob_deps = () 30 | self.name_that_makes_sense = code_info.name 31 | self.argc = code_info.argcount 32 | 33 | self.base_method_code = "" 34 | self.base_method_addr = None 35 | self._function_module = None 36 | self._fn_ptr_name = None 37 | self.call_recorder = None 38 | self.memoize_partial_code = None 39 | self.method_arguments = None 40 | 41 | @property 42 | def function_module(self): 43 | assert self._function_module 44 | return self._function_module 45 | 46 | @function_module.setter 47 | def function_module(self, v): 48 | self._function_module = v 49 | 50 | @property 51 | def fn_ptr_name(self): 52 | assert self._fn_ptr_name 53 | return self._fn_ptr_name 54 | 55 | @fn_ptr_name.setter 56 | def fn_ptr_name(self, v): 57 | self._fn_ptr_name = v 58 | 59 | @property 60 | def unique_index(self): 61 | return id(self) 62 | 63 | def counter_name(self) -> str: 64 | raise NotImplementedError 65 | 66 | def globals_from(self, globs: dict): 67 | self.globals = globs 68 | 69 | def add_methods(self, dispatcher: t.Dict[call_record_t, t.Tuple[type, ...]]): 70 | if not dispatcher: 71 | return 72 | for argtypeids, argtypes in dispatcher.items(): 73 | self.add_method(argtypeids, argtypes) 74 | 75 | self._rebuild_method_getter_and_set() 76 | 77 | def add_method(self, argtypeids: call_record_t, argtypes: t.Tuple[type, ...]): 78 | self.sys.add_jit_method(self, argtypeids, argtypes) 79 | 80 | def _rebuild_method_getter_and_set(self): 81 | code = mk_hard_coded_method_getter_module(self.methods, self.base_method_addr, 82 | self.argc) 83 | mod = compile_module(JIT_FUNCTION_DIR, "MethodGetter", code) 84 | self.function_module.f.mut_method_get(mod.method_get_addr) 85 | 86 | 87 | class JITSystem: 88 | 89 | def __init__(self, strategies: t.List[Strategy] = None): 90 | strategies = strategies or [HitN(200)] 91 | self.fn_count = 0 92 | self.jit_hotspot_analysis = JITRecompilationDecision(strategies) 93 | self.memoize_partial_code = {} 94 | self.jit_fptr_index = {} 95 | self.fn_place_index = {} # type: t.Dict[int, JITFunctionHoldPlace] 96 | self.store_base_method_log = False 97 | 98 | def jitdata(self, cls: type): 99 | undef = object() 100 | anns = {} 101 | code = [] 102 | imports = [] 103 | for i, (k, t) in enumerate(cls.__annotations__.items()): 104 | path = extension_type_pxd_paths.get(t, undef) 105 | if path is undef: 106 | anns[k] = 'object' 107 | elif path is None: 108 | anns[k] = t.__name__ 109 | else: 110 | typename = "{}{}".format(typed_head, i) 111 | anns[k] = typename 112 | imports.append('from {} cimport {} as {}'.format( 113 | path, t.__name__, typename)) 114 | code.append("cdef class {}:".format(cls.__name__)) 115 | for attr, typestr in anns.items(): 116 | code.append(" cdef {} x_{}".format(typestr, attr)) 117 | code.append(" def __init__(self, {}):".format(', '.join(anns))) 118 | for attr, _ in anns.items(): 119 | code.append(" self.x_{0} = {0}".format(attr)) 120 | 121 | for attr, typestr in anns.items(): 122 | code.append(" cpdef {1} get_{0}(self):".format(attr, typestr)) 123 | code.append(" return self.x_{0}".format(attr)) 124 | 125 | code.append(" @property") 126 | code.append(" def {0}(self):".format(attr)) 127 | code.append(" return self.x_{0}".format(attr)) 128 | 129 | code.append(" cpdef void set_{0}(self, {1} {0}):".format(attr, typestr)) 130 | code.append(" self.x_{0} = {0}".format(attr)) 131 | 132 | code.append(" @{}.setter".format(attr)) 133 | code.append(" def {0}(self, {1} {0}):".format(attr, typestr)) 134 | code.append(" self.x_{0} = {0}".format(attr)) 135 | 136 | pyx = '\n'.join(imports + code) 137 | 138 | code.clear() 139 | code.append("cdef class {}:".format(cls.__name__)) 140 | for attr, typestr in anns.items(): 141 | code.append(" cdef {} x_{}".format(typestr, attr)) 142 | for attr, typestr in anns.items(): 143 | code.append(" cpdef void set_{0}(self, {1})".format(attr, typestr)) 144 | code.append(" cpdef {1} get_{0}(self)".format(attr, typestr)) 145 | 146 | pxd = '\n'.join(imports + code) 147 | mod = ximport_module(JIT_TYPE_DIR, cls.__name__, pyx, pxd) 148 | return getattr(mod, cls.__name__) 149 | 150 | def jit(self, f: FunctionType): 151 | func_info = self.get_func_info(f) 152 | code_info = func_info.r_codeinfo 153 | fn_place = self.allocate_place_for_function(code_info) 154 | fn_place.globals_from(func_info.r_globals) 155 | CodeEmit(self, code_info, fn_place) 156 | f = fn_place.function_module.f 157 | self.fn_place_index[id(f)] = fn_place 158 | return f 159 | 160 | @staticmethod 161 | def get_func_info(f: FunctionType) -> PyFuncInfo: 162 | return CyVM.func_info(f) 163 | 164 | def generate_module_for_code(self, code_info: PyCodeInfo): 165 | code = mk_module_code(code_info) 166 | unique_module_name = "Functions_{}".format(self.fn_count) 167 | mod = compile_module(JIT_FUNCTION_DIR, unique_module_name, code) 168 | return mod 169 | 170 | def remember_partial_code(self, fn_place: JITFunctionHoldPlace, code_out: CodeOut): 171 | fn_place.memoize_partial_code = code_out 172 | self.memoize_partial_code[id(fn_place)] = code_out 173 | 174 | def setup(self): 175 | # FIXME: create the temporary directory, 176 | # as the place to hold all compiled extensions and, 177 | # the source codes required by re-JIT. 178 | # ONE RUNTIME USES ONE JIT DIRECTORY 179 | raise NotImplementedError 180 | 181 | def allocate_place_for_function(self, code_info: PyCodeInfo) -> JITFunctionHoldPlace: 182 | # use object id and function name and module name to 183 | # generate unique path as well as meaningful. 184 | fn_place = JITFunctionHoldPlace(self, code_info) 185 | fn_place.function_module = self.generate_module_for_code(code_info) 186 | self.jit_fptr_index[id(fn_place)] = fn_place.function_module.f 187 | return fn_place 188 | 189 | def generate_base_method(self, function_place: JITFunctionHoldPlace, code): 190 | argc = function_place.argc 191 | method_argtype_comma_lst = ', '.join('object' for _ in range(argc)) 192 | method_get_argument_comma_lst = ', '.join('int64_t _%d' % i for i in range(argc)) 193 | # TODO: use auto-evolution-able method lookup 194 | fn_ptr_name = function_place.fn_ptr_name 195 | code = """ 196 | {} 197 | ctypedef object (*method_t)({}) 198 | cdef method_t this_method_getter({}): 199 | return {} 200 | base_method_addr = reinterpret_cast[int64_t]({}) 201 | method_getter_addr = reinterpret_cast[int64_t](this_method_getter) 202 | """.format(code, method_argtype_comma_lst, method_get_argument_comma_lst, 203 | fn_ptr_name, fn_ptr_name) 204 | 205 | if self.store_base_method_log: 206 | function_place.base_method_code = code 207 | 208 | unique_module = "Methods_{}_{}_base_method".format( 209 | id(function_place), function_place.name_that_makes_sense.replace('.', '__')) 210 | mod = compile_module(JIT_FUNCTION_DIR, unique_module, code) 211 | 212 | method_init_fptrs = getattr(mod, cmc.method_init_fptrs) 213 | method_init_globals = getattr(mod, cmc.method_init_globals) 214 | init_notifier = getattr(mod, cmc.method_init_recorder_and_notifier) 215 | method_getter_addr = getattr(mod, 'method_getter_addr') 216 | call_recorder = mod.NonJITCallRecorder() 217 | function_place.call_recorder = call_recorder 218 | init_notifier( 219 | call_recorder, lambda: self.jit_hotspot_analysis.trigger_jit(function_place)) 220 | g = function_place.globals 221 | method_init_globals(**{k: g[k] for k in function_place.glob_deps}) 222 | method_init_fptrs(self.jit_fptr_index) 223 | function_place.base_method_module = mod 224 | function_place.function_module.f.mut_method_get(method_getter_addr) 225 | function_place.base_method_addr = mod.base_method_addr 226 | 227 | def add_jit_method(self, function_place: JITFunctionHoldPlace, 228 | argtypeids: call_record_t, argtypes: t.Tuple[type, ...]): 229 | method_arguments = function_place.method_arguments 230 | fn_ptr_name = function_place.fn_ptr_name 231 | actual_args = [cmc.typed_head + arg for arg in method_arguments] 232 | once_code_out = CodeOut() 233 | declaring: list = once_code_out[cmc.Import] 234 | typing: list = once_code_out[cmc.Customizable] 235 | 236 | typing.append("cdef {}({}):".format(fn_ptr_name, ', '.join(actual_args))) 237 | undef = object() 238 | for i, (actual_arg, arg, argtype) in enumerate( 239 | zip(actual_args, method_arguments, argtypes)): 240 | path = extension_type_pxd_paths.get(argtype, undef) 241 | if path is undef: 242 | # well, this type cannot JIT in fact, so still object and stop recording it. 243 | typing.append("{}{} = {}".format(cmc.IDENTATION_SECTION, arg, actual_arg)) 244 | elif path is None: 245 | # it's builtin extension types, like dict, list, etc. 246 | typing.append("{}cdef {} {} = {}".format(cmc.IDENTATION_SECTION, 247 | argtype.__name__, arg, actual_arg)) 248 | else: 249 | # Good, user-defined extension types! 250 | # You'll see how fast it'll be! 251 | 252 | # Firstly we import the required type 253 | typename = "{}{}".format(cmc.typed_head, i) 254 | declaring.append('from {} cimport {} as {}'.format( 255 | path, argtype.__name__, typename)) 256 | typing.append("{}cdef {} {} = {}".format(cmc.IDENTATION_SECTION, typename, 257 | arg, actual_arg)) 258 | 259 | once_code_out.merge_update(function_place.memoize_partial_code) 260 | code = '\n'.join(once_code_out.to_code_lines()) 261 | # TODO: use auto-evolution-able method lookup 262 | code = """ 263 | {} 264 | method_addr = reinterpret_cast[int64_t]({}) 265 | """.format(code, fn_ptr_name) 266 | 267 | unique_module = "Methods_{}_{}_JITed".format( 268 | id(function_place), function_place.name_that_makes_sense.replace('.', '__')) 269 | mod = compile_module(JIT_FUNCTION_DIR, unique_module, code) 270 | method_init_fptrs = getattr(mod, cmc.method_init_fptrs) 271 | method_init_globals = getattr(mod, cmc.method_init_globals) 272 | g = function_place.globals 273 | method_init_globals(**{k: g[k] for k in function_place.glob_deps}) 274 | method_init_fptrs(self.jit_fptr_index) 275 | function_place.methods[argtypeids] = mod.method_addr 276 | -------------------------------------------------------------------------------- /restrain_jit/becython/cy_jit_common.py: -------------------------------------------------------------------------------- 1 | import typing 2 | call_record_t = typing.Tuple[int, ...] 3 | call_records_t = typing.List[call_record_t] 4 | fptr_addr = int 5 | 6 | lineno_name = "restrain_lineno" 7 | cont_name = "res_cont" 8 | last_cont_name = "res_last_cont" 9 | mangle_head = 'res_mg_' 10 | gensym_head = "res_gen_" 11 | typed_head = "res_tobe_typed_" 12 | var_head = 'res_var_' 13 | sym_head = "res_symbol_" 14 | glob_head = "res_global_" 15 | fn_head = "res_localfn_" 16 | fn_addr = "res_address" 17 | method_init_globals = "method_init_globals" 18 | method_init_fptrs = "method_init_fptrs" 19 | 20 | recorder = "restrain_call_recorder" 21 | notifier = "restrain_recompilation_notifier" 22 | method_init_recorder_and_notifier = "method_init_recorder_and_notifier" 23 | 24 | Import = -10 25 | TypeDecl = -5 26 | Signature = -2 27 | Customizable = 0 28 | Normal = 5 29 | Finally = 10 30 | 31 | IDENTATION_SECTION = " " 32 | -------------------------------------------------------------------------------- /restrain_jit/becython/cy_jit_ext_template.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | 3 | from restrain_jit.jit_info import PyCodeInfo 4 | from restrain_jit.becython.cy_jit_common import * 5 | from string import Template 6 | import typing 7 | 8 | template = """ 9 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 10 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, inttoptr 11 | from libc.stdint cimport int64_t, int8_t 12 | from libcpp.vector cimport vector as std_vector 13 | from libcpp.cast cimport reinterpret_cast 14 | 15 | # method type 16 | ctypedef object (*method_t)($many_objects) 17 | 18 | # method lookup type 19 | ctypedef method_t (*method_get_t)($many_int64_t) 20 | 21 | # method look up ptr 22 | cdef method_get_t method_get 23 | 24 | # to avoid params have conflicts against 'type' 25 | cdef inline method_t $method_get_invoker($unnamed_args): 26 | func = method_get$typeids 27 | return func 28 | 29 | cdef class F: 30 | # python compatible method to change 'method look up ptr' 31 | cpdef mut_method_get(self, int64_t f): 32 | global method_get 33 | method_get = reinterpret_cast[method_get_t](inttoptr(f)) 34 | 35 | def __call__(self, $arguments): 36 | $method = $method_get_invoker($arguments) 37 | return $method($arguments) 38 | 39 | f = F() 40 | """ 41 | 42 | 43 | def mk_call_record_t(argc): 44 | if argc is 0: 45 | call_record_t = 'void*' 46 | elif argc is 1: 47 | call_record_t = '(int64_t, )' 48 | else: 49 | call_record_t = '(' + ', '.join(['int64_t'] * argc) + ')' 50 | return call_record_t 51 | 52 | 53 | def mk_call_record(args): 54 | if not args: 55 | return "()" 56 | elif len(args) is 1: 57 | call_record = '(pytoint(type(%s)), )' % args[0] 58 | else: 59 | call_record = '(' + ', '.join('pytoint(type(%s))' % arg for arg in args) + ')' 60 | return call_record 61 | 62 | 63 | def mk_module_code(code_info: PyCodeInfo): 64 | freevars = code_info.freevars 65 | argnames = code_info.argnames 66 | 67 | method_name = "method" 68 | method_getter_invoker_name = "invoke_method_get" 69 | 70 | while method_name in argnames: 71 | method_name += "_emm" 72 | 73 | while method_getter_invoker_name in argnames: 74 | method_name += "_emm" 75 | 76 | arguments = freevars + argnames 77 | argc = len(arguments) 78 | unnamed_args = ['a%d' % i for i in range(argc)] 79 | 80 | return Template(template).substitute( 81 | method=method_name, 82 | many_objects=', '.join(['object'] * argc), 83 | many_int64_t=', '.join(['int64_t'] * argc), 84 | unnamed_args=", ".join(unnamed_args), 85 | typeids=mk_call_record(unnamed_args), 86 | method_get_invoker=method_getter_invoker_name, 87 | arguments=', '.join(arguments)) 88 | 89 | 90 | build_method_getter_template = """ 91 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 92 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, inttoptr 93 | from libc.stdint cimport int64_t 94 | from libcpp.cast cimport reinterpret_cast 95 | 96 | ctypedef object (*method_t)($many_objects) 97 | cdef method_t base_method = reinterpret_cast[method_t](inttoptr($base_method_addr)) 98 | 99 | $declare_all_methods 100 | 101 | cdef method_t method_get($many_int64_args): 102 | $hard_coded_dynamic_dispatch 103 | 104 | method_get_addr = reinterpret_cast[int64_t](method_get) 105 | """ 106 | 107 | 108 | def group_by(seq, f): 109 | res = defaultdict(list) 110 | for each in seq: 111 | res[f(each)].append(each) 112 | return dict(res) 113 | 114 | 115 | def decision_recurse(m: typing.List[typing.Tuple[call_record_t, fptr_addr]], 116 | argc: int, 117 | depth=0): 118 | if depth == argc: 119 | assert len(m) is 1 120 | return m[0][1] 121 | groups = group_by(m, lambda x: x[0][depth]) 122 | return {k: decision_recurse(v, argc, depth + 1) for k, v in groups.items()} 123 | 124 | 125 | def mk_hard_coded_method_getter_module(methods: typing.Dict[call_record_t, fptr_addr], 126 | base_method_addr: fptr_addr, argc: int): 127 | 128 | declare_all_methods = [] 129 | hard_coded_dynamic_dispatch = [] 130 | many_objects = ', '.join(['object'] * argc) 131 | many_int64_args = ', '.join("int64_t arg%d" % i for i in range(argc)) 132 | 133 | def codegen_jit_method_reinterp(fptr_addrs: typing.Iterable[fptr_addr]): 134 | for each in fptr_addrs: 135 | each = 'cdef method_t jit_method_{} = reinterpret_cast[method_t](inttoptr({}LL))'.format( 136 | each, hex(each)) 137 | declare_all_methods.append(each) 138 | 139 | def codegen_dispatch_recurse(d: typing.Union[int, dict], indent=" ", argi=0): 140 | if isinstance(d, int): 141 | hard_coded_dynamic_dispatch.append("{}return jit_method_{}".format(indent, d)) 142 | return 143 | 144 | for k, v in d.items(): 145 | hard_coded_dynamic_dispatch.append("{}if {}LL == arg{}:".format(indent, hex(k), argi)) 146 | codegen_dispatch_recurse(v, indent + " ", argi + 1) 147 | 148 | m = list(methods.items()) 149 | codegen_jit_method_reinterp(x[1] for x in m) 150 | dispatch = decision_recurse(m, argc) 151 | codegen_dispatch_recurse(dispatch) 152 | hard_coded_dynamic_dispatch.append(' return base_method') 153 | 154 | code = Template(build_method_getter_template).substitute( 155 | many_objects=many_objects, 156 | many_int64_args=many_int64_args, 157 | base_method_addr=base_method_addr, 158 | declare_all_methods='\n'.join(declare_all_methods), 159 | hard_coded_dynamic_dispatch='\n'.join(hard_coded_dynamic_dispatch)) 160 | return code 161 | -------------------------------------------------------------------------------- /restrain_jit/becython/cy_jit_hotspot_comp.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | import abc 3 | import numpy as np 4 | from restrain_jit.becython.cy_jit_common import * 5 | from collections import Counter 6 | try: 7 | from .cy_jit import JITSystem, JITFunctionHoldPlace 8 | except ImportError: 9 | pass 10 | 11 | extension_type_pxd_paths = { 12 | int: None, 13 | float: None, 14 | str: None, 15 | bytes: None, 16 | list: None, 17 | tuple: None, 18 | dict: None, 19 | set: None, 20 | np.ndarray: 'numpy' 21 | } 22 | 23 | 24 | class Strategy(abc.ABC): 25 | 26 | @abc.abstractmethod 27 | def __call__(self, calls: call_records_t, freq: Counter, n: int): 28 | raise NotImplementedError 29 | 30 | 31 | def is_jit_able_type(t: type): 32 | return t in extension_type_pxd_paths 33 | 34 | 35 | class HitN(Strategy): 36 | 37 | def __init__(self, threshold: int): 38 | self.threshold = threshold 39 | 40 | def __call__(self, _, freq: Counter, n: int): 41 | threshold = self.threshold 42 | for k, v in freq.items(): 43 | if v > threshold: 44 | yield k 45 | 46 | 47 | class JITRecompilationDecision: 48 | 49 | def __init__(self, strategies: t.Iterable[Strategy]): 50 | self.strategies = list(strategies) 51 | 52 | def add_rules(self, strategies: t.Iterable[Strategy]): 53 | self.strategies.extend(strategies) 54 | 55 | def select(self, calls): 56 | freq = Counter(calls) 57 | n = len(calls) 58 | new_jt = set() 59 | for strategy in self.strategies: 60 | new_jt.update(strategy(calls, freq, n)) 61 | return new_jt 62 | 63 | def trigger_jit(self, fn_place: 'JITFunctionHoldPlace'): 64 | call_recorder = fn_place.call_recorder 65 | new_jt = self.select(call_recorder.get()) 66 | dispatcher = {} 67 | for argtypeids in new_jt: 68 | argtypes = tuple(map(call_recorder.load_type, argtypeids)) 69 | dispatcher[argtypeids] = argtypes 70 | fn_place.add_methods(dispatcher) 71 | -------------------------------------------------------------------------------- /restrain_jit/becython/cy_loader.py: -------------------------------------------------------------------------------- 1 | import tempfile 2 | import os 3 | import platform 4 | import sys 5 | import pyximport 6 | from importlib import util, import_module 7 | from restrain_jit.config import RESTRAIN_CONFIG 8 | from pathlib import Path 9 | from string import Template 10 | from restrain_jit.utils import exec_cc 11 | try: 12 | import Cython.Includes as includes 13 | except ImportError: 14 | raise RuntimeError("For using Cython backend, cython package is required.") 15 | include_paths = list(includes.__path__) 16 | 17 | JIT_DATA_DIR = Path(tempfile.TemporaryDirectory(prefix="restrain").name) 18 | JIT_DATA_DIR.mkdir() 19 | 20 | JIT_FUNCTION_DIR = JIT_DATA_DIR / "Restrain_paths_functions" 21 | JIT_TYPE_DIR = JIT_DATA_DIR / "Restrain_paths_types" 22 | sys.path.append(str(JIT_DATA_DIR)) 23 | 24 | JIT_FUNCTION_DIR.mkdir() 25 | JIT_TYPE_DIR.mkdir() 26 | with (JIT_TYPE_DIR / "__init__.py").open('w') as f: 27 | f.write('# So, this is a module, a cython module!') 28 | 29 | suffix = '.pyd' if platform.system() == 'Windows' else '.so' 30 | 31 | template = Template(r""" 32 | from distutils.core import setup 33 | from distutils.extension import Extension 34 | from Cython.Build import cythonize 35 | 36 | exts = [ 37 | Extension($module, 38 | [$module_path], 39 | include_dirs=$include_dirs, 40 | libraries=$libraries, 41 | library_dirs=$library_dirs, 42 | language="c++", 43 | extra_compile_args=["-std=c++11", '-O3'], 44 | extra_link_args=["-std=c++11", '-O3'] 45 | ) 46 | ] 47 | 48 | setup( 49 | ext_modules=cythonize( 50 | exts, 51 | compiler_directives=dict(language_level="3str", infer_types=True) 52 | )) 53 | """) 54 | 55 | 56 | def get_restrain_rts(): 57 | restrain_rts = Path(RESTRAIN_CONFIG.cython.rts).expanduser() 58 | restrain_include = Path(restrain_rts / "include") 59 | restrain_lib = Path(restrain_rts / "lib") 60 | return restrain_include, restrain_lib 61 | 62 | 63 | def get_includes_and_libs(extra_libs=()): 64 | restrain_include, restrain_lib = get_restrain_rts() 65 | return dict( 66 | include_dirs=[str(restrain_include), *include_paths], 67 | library_dirs=[str(restrain_lib)], 68 | libraries=[':typeint', *extra_libs]) 69 | 70 | 71 | def compile_module(under_dir: Path, mod_name: str, source_code: str, libs=()): 72 | # TODO: 73 | # tempfile.TemporaryDirectory will close unexpectedly before removing the generated module. 74 | # Since that we don't delete the temporary dir as a workaround. 75 | 76 | mod_name = mod_name 77 | dirname = tempfile.mkdtemp(dir=str(under_dir)) 78 | mod_path = mod_name + '.pyx' 79 | with open(os.path.join(dirname, mod_path), 'w') as pyx_file, open( 80 | os.path.join(dirname, 'setup.py'), 'w') as setup_file: 81 | pyx_file.write(source_code) 82 | 83 | setup_file.write( 84 | template.substitute( 85 | module=repr(mod_name), 86 | module_path=repr(mod_path), 87 | **get_includes_and_libs())) 88 | 89 | cwd = os.getcwd() 90 | try: 91 | os.chdir(dirname) 92 | 93 | args = ['setup.py', 'build_ext', '--inplace'] 94 | c = exec_cc(sys.executable, args) 95 | 96 | hd = next(c) 97 | if hd is not 0: 98 | sys.stderr.buffer.write(b''.join(c)) 99 | raise RuntimeError("Cython compiler failed.") 100 | 101 | # find the python extension module. 102 | # pyd_name = next(each for each in os.listdir(dirname) if each.endswith(suffix)) 103 | finally: 104 | os.chdir(cwd) 105 | 106 | mod = import_module("{}.{}.{}".format(under_dir.name, 107 | os.path.split(dirname)[1], mod_name)) 108 | return mod 109 | 110 | 111 | def ximport_module(under_dir: Path, mod_name: str, pyx_source: str, pxd_source): 112 | mod_name = mod_name 113 | directory = Path(tempfile.mkdtemp(dir=str(under_dir))) 114 | pyx_file = directory / (mod_name + '.pyx') 115 | pxd_file = directory / (mod_name + '.pxd') 116 | with open(pyx_file, 'w') as f: 117 | f.write(pyx_source) 118 | with open(pxd_file, 'w') as f: 119 | f.write(pxd_source) 120 | 121 | # find the python extension module. 122 | # pyd_name = next(each for each in os.listdir(dirname) if each.endswith(suffix)) 123 | mod = import_module("{}.{}.{}".format(under_dir.name, directory.name, mod_name)) 124 | return mod 125 | 126 | 127 | def setup_pyx_for_cpp(): 128 | """ 129 | This is to support --cplus for pyximport 130 | """ 131 | 132 | old_get_distutils_extension = pyximport.pyximport.get_distutils_extension 133 | 134 | def new_get_distutils_extension(modname, pyxfilename, language_level=None): 135 | extension_mod, setup_args = old_get_distutils_extension(modname, pyxfilename, 136 | language_level) 137 | extension_mod.language = 'c++' 138 | for k, v in get_includes_and_libs().items(): 139 | setattr(extension_mod, k, v) 140 | return extension_mod, setup_args 141 | 142 | pyximport.pyximport.get_distutils_extension = new_get_distutils_extension 143 | 144 | 145 | setup_pyx_for_cpp() 146 | pyximport.install() 147 | -------------------------------------------------------------------------------- /restrain_jit/becython/cy_method_codegen.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.phi_vm import * 2 | from restrain_jit.becython.cy_jit_common import * 3 | from restrain_jit.becython.cy_jit_ext_template import mk_call_record, mk_call_record_t 4 | from restrain_jit.abs_compiler.py_apis import NS 5 | from restrain_jit.jit_info import PyCodeInfo, PyFuncInfo 6 | from restrain_jit.vm.am import ValSymbol, Symbol 7 | from restrain_jit.utils import CodeOut, MissingDict 8 | from collections import OrderedDict 9 | from contextlib import contextmanager 10 | from functools import singledispatch 11 | from string import Template 12 | import typing as t 13 | 14 | try: 15 | from .cy_jit import JITSystem, JITFunctionHoldPlace 16 | except ImportError: 17 | pass 18 | 19 | 20 | class UndefGlobal: 21 | pass 22 | 23 | 24 | preludes = [ 25 | "cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT", 26 | "from restrain_jit.becython.cython_rts.RestrainJIT cimport Cell", 27 | "from libc.stdint cimport int64_t, int8_t", 28 | "from libcpp.cast cimport reinterpret_cast", 29 | "from restrain_jit.becython.cython_rts.hotspot cimport typeid, inttoptr, pytoint", 30 | ] 31 | 32 | base_call_prelude = """ 33 | from libcpp.vector cimport vector as std_vector 34 | ctypedef $call_record_t call_record_t 35 | ctypedef std_vector[call_record_t] call_records_t 36 | 37 | cdef class NonJITCallRecorder: 38 | cdef call_records_t records 39 | cdef int64_t bound 40 | 41 | def __init__(self): 42 | self.bound = 233 43 | self.records = call_records_t() 44 | 45 | cdef record(self, call_record_t m): 46 | self.records.push_back(m) 47 | 48 | cdef int8_t check_bounded(self): 49 | return self.records.size() % self.bound == 0 50 | 51 | def load_type(self, int64_t i): 52 | return (inttoptr(i)) 53 | 54 | def get(self): 55 | return self.records 56 | 57 | cdef NonJITCallRecorder {0} 58 | cdef object {1} 59 | 60 | def {2}(a, b): 61 | global {0}, {1} 62 | {0} = a 63 | {1} = b 64 | """.format(recorder, notifier, method_init_recorder_and_notifier) 65 | 66 | 67 | class CodeEmit: 68 | 69 | def __init__(self, jit_system: 'JITSystem', code_info: PyCodeInfo, 70 | function_place: 'JITFunctionHoldPlace'): 71 | self.priority = Import 72 | self.other_fptrs = {} 73 | self.prefixes = [""] 74 | self.gen_sym_id = 0 75 | self.function_place = function_place 76 | self.jit_system = jit_system 77 | 78 | def emit_line(line): 79 | code_out[self.priority].append(line) 80 | 81 | def persistent_code(): 82 | nonlocal code_out 83 | code_out = persistent_code_out 84 | 85 | self.persistent_code = persistent_code 86 | 87 | def base_method_code(): 88 | nonlocal code_out 89 | code_out = once_code_out 90 | 91 | self.base_method_code = base_method_code 92 | 93 | self.emit_line = emit_line 94 | 95 | # Name generator 96 | mangle_map = self.mangle_map = MissingDict(lambda: "{}{}".format( 97 | mangle_head, len(mangle_map))) 98 | 99 | self.symbol_map = symbom_map = MissingDict(lambda: "{}{}".format( 100 | sym_head, len(symbom_map))) 101 | 102 | self.fn_map = fn_map = MissingDict(lambda: "{}{}".format(fn_head, len(fn_map))) 103 | 104 | # End name generator 105 | code_out = persistent_code_out = CodeOut() 106 | for each in preludes: 107 | self += each 108 | 109 | glob_deps = code_info.glob_deps 110 | 111 | # TODO: some globals, like 'type' or others, don't need to initialize here. 112 | # Also, the initialization of those can badly hurt the performance, e.g., 'type'. 113 | 114 | def is_supported_builtin(s): 115 | if s in ('type', 'range', 'len', 'print'): 116 | return False 117 | raise NameError(s) 118 | 119 | actual_glob = function_place.globals 120 | self.glob = { 121 | each: self.gensym(each) 122 | for each in glob_deps 123 | if each in actual_glob or is_supported_builtin(each) 124 | } 125 | 126 | function_place.glob_deps = tuple(self.glob.keys()) 127 | fn_name = self.mangle(code_info.name) 128 | 129 | once_code_out = CodeOut() 130 | code_out = once_code_out 131 | self.priority = TypeDecl 132 | self += Template(base_call_prelude).substitute( 133 | call_record_t=mk_call_record_t(code_info.argcount)) 134 | 135 | # BASE_METHOD_CODE 136 | # different methods of a JIT function 137 | # varies from their behaviours/codegen rules of following block 138 | code_out = once_code_out 139 | self.priority = Customizable 140 | cellvars = code_info.cellvars 141 | arguments = OrderedDict((e, self.mangle(e) if e not in cellvars else self.gensym(e)) 142 | for e in code_info.freevars + code_info.argnames) 143 | function_place.method_arguments = tuple(arguments.values()) 144 | self += 'cdef {}({}):'.format(fn_name, ', '.join(arguments.values())) 145 | 146 | code_out = once_code_out 147 | self.priority = Customizable 148 | with indent(self): 149 | if code_info.argcount: 150 | typeids = mk_call_record(list(arguments.values())) 151 | self += "{}# BEGIN MONITOR".format(self.prefix) 152 | self += "{}cdef call_record_t res_call_record = {}".format( 153 | self.prefix, typeids) 154 | self += "{}{}.record(res_call_record)".format(self.prefix, recorder) 155 | self += "{}if {}.check_bounded(): {}()".format(self.prefix, recorder, 156 | notifier) 157 | self += "{}# END MONITOR".format(self.prefix) 158 | 159 | # PERSISTENT CODE 160 | code_out = persistent_code_out 161 | self.priority = Normal 162 | cellvars = code_info.cellvars 163 | for cell in cellvars: 164 | actual_cell = self.mangle(cell) 165 | if cell in code_info.argnames: 166 | cell_in_arg = self.gensym(cell) 167 | self += "{}cdef Cell {} = {}.Cell({})".format( 168 | self.prefix, actual_cell, NS.RestrainJIT, cell_in_arg) 169 | else: 170 | self += "{}cdef Cell {} = {}.Cell(None)".format( 171 | self.prefix, actual_cell, NS.RestrainJIT) 172 | 173 | code_out = persistent_code_out 174 | self.priority = Normal 175 | 176 | with indent(self): 177 | self.emit_instrs(code_info.instrs) 178 | 179 | self.declare_globals() 180 | jit_system.remember_partial_code(function_place, code_out=persistent_code_out) 181 | # dict.update not work here. 182 | once_code_out.merge_update(persistent_code_out) 183 | base_call_code = '\n'.join(once_code_out.to_code_lines()) 184 | function_place.fn_ptr_name = fn_name 185 | jit_system.generate_base_method(function_place, code=base_call_code) 186 | 187 | @property 188 | def prefix(self): 189 | return self.prefixes[-1] 190 | 191 | def gensym(self, s: str = ""): 192 | if not s: 193 | base = 'r3e' 194 | elif s.isidentifier(): 195 | base = s 196 | else: 197 | base = ''.join(e for e in s if e.isidentifier()) 198 | ret = "{}{}_{}".format(gensym_head, self.gen_sym_id, base) 199 | self.gen_sym_id += 1 200 | return ret 201 | 202 | def indent(self): 203 | self.prefixes.append(self.prefix + IDENTATION_SECTION) 204 | 205 | def dedent(self): 206 | self.prefixes.pop() 207 | 208 | def mangle(self, n: str): 209 | if n.isidentifier(): 210 | return var_head + n 211 | return self.mangle_map[n] 212 | 213 | def declare_globals(self): 214 | """ 215 | strongly-typed globals? 216 | """ 217 | p = self.priority 218 | self.priority = Finally 219 | 220 | prefix = " " 221 | 222 | arguments_comma_lst = ', '.join(self.glob.keys()) 223 | 224 | global_vars_comma_lst = ', '.join(self.glob.values()) 225 | fptrs_comma_lst = ', '.join(self.other_fptrs.keys()) 226 | 227 | # use as 'module_init_globals(**global_dict)' 228 | self += "cpdef method_init_globals({}):".format(arguments_comma_lst) 229 | self += '{}global {}'.format(prefix, global_vars_comma_lst) 230 | 231 | jit_fptr_index = self.gensym('ptr_index') 232 | # use as 'method_init_fptrs(jit_system.jit_fptr_index)' 233 | self += "cpdef method_init_fptrs({}):".format(jit_fptr_index) 234 | self += '{}global {}'.format(prefix, fptrs_comma_lst) 235 | 236 | for k, v in self.other_fptrs.items(): 237 | self += "{}{} = {}[{}]".format(prefix, k, jit_fptr_index, v) 238 | 239 | for k, v in self.glob.items(): 240 | self += '{}{} = {}'.format(prefix, k, v) 241 | 242 | self.priority = TypeDecl 243 | for var in self.glob.values(): 244 | self += "cdef object {}".format(var) 245 | 246 | for var in self.other_fptrs.keys(): 247 | self += "cdef object {}".format(var) 248 | 249 | self.priority = p 250 | 251 | def declare_symbol(self, s: str): 252 | if s in self.symbol_map: 253 | return self.symbol_map[s] 254 | n = self.symbol_map[s] 255 | p = self.priority 256 | self.priority = TypeDecl 257 | try: 258 | self += "cdef int64_t {} = {}.get_symbol({!r})".format(n, NS.RestrainJIT, s) 259 | finally: 260 | self.priority = p 261 | return n 262 | 263 | def __iadd__(self, other): 264 | self.emit_line(other) 265 | return self 266 | 267 | def set_cont(self, label_i): 268 | self += "{}{} = {}".format(self.prefix, cont_name, label_i) 269 | 270 | def set_last_cont(self, last=None): 271 | self += "{}{} = {}".format(self.prefix, last_cont_name, last or cont_name) 272 | 273 | def emit_instrs(self, instrs: t.List[Instr]): 274 | 275 | tmplt = "{}cdef int {} = 0" 276 | self += tmplt.format(self.prefix, last_cont_name) 277 | self += tmplt.format(self.prefix, cont_name) 278 | self += tmplt.format(self.prefix, lineno_name) 279 | self += "{}# Control flow graph splitted".format(self.prefix) 280 | self += "{}while True:".format(self.prefix) 281 | with indent(self): 282 | for each in instrs: 283 | emit(each, self) 284 | 285 | 286 | @singledispatch 287 | def emit(a: Instr, self: CodeEmit): 288 | raise NotImplementedError(a) 289 | 290 | 291 | def emit_const(self: CodeEmit, r: object): 292 | if isinstance(r, (ValSymbol, Symbol)): 293 | return r.s 294 | # var_name = self.declare_symbol(r.s) 295 | # return var_name 296 | elif isinstance(r, PyCodeInfo): 297 | fn_place = self.jit_system.allocate_place_for_function(r) 298 | fn_place.globals_from(self.function_place.globals) 299 | CodeEmit(self.jit_system, r, fn_place) 300 | n = self.gensym(fn_place.name_that_makes_sense) 301 | self.other_fptrs[n] = fn_place.unique_index 302 | return n 303 | return repr(r) 304 | 305 | 306 | def emit_repr(self: CodeEmit, r: Repr) -> str: 307 | if isinstance(r, Const): 308 | return emit_const(self, r.val) 309 | elif isinstance(r, Reg): 310 | return self.mangle(r.n) 311 | elif isinstance(r, Prim): 312 | if not r.qual: 313 | # python const 314 | name = r.n 315 | glob_name = self.glob.get(name, name) 316 | return glob_name 317 | return "{}.{}".format(r.qual, r.n) 318 | else: 319 | raise TypeError(r) 320 | 321 | 322 | @contextmanager 323 | def indent(self: CodeEmit): 324 | self.indent() 325 | try: 326 | yield 327 | finally: 328 | self.dedent() 329 | 330 | 331 | @emit.register 332 | def set_lineno(n: SetLineno, self: CodeEmit): 333 | self += "{}{} = {}".format(self.prefix, lineno_name, n.lineno) 334 | 335 | 336 | def phi_for_given_branch(self, var_dict): 337 | for k, r in var_dict.items(): 338 | v = emit_repr(self, r) 339 | k = self.mangle(k) 340 | self += "{}{} = {}".format(self.prefix, k, v) 341 | 342 | 343 | def phi(self: CodeEmit, phi_dict): 344 | 345 | def app(_self, _token, _var_dict): 346 | _self += "{}{}:".format(_self.prefix, _token) 347 | with indent(_self): 348 | phi_for_given_branch(_self, _var_dict) 349 | 350 | head, *init, end = phi_dict.items() 351 | 352 | token = "if {} == {}".format(head[0], last_cont_name) 353 | app(self, token, head[1]) 354 | 355 | for come_from, var_dict in init: 356 | token = "elif {} == {}".format(come_from, last_cont_name) 357 | app(self, token, var_dict) 358 | 359 | self += "# must from label {}".format(head[0]) 360 | app(self, "else", end[1]) 361 | 362 | 363 | @emit.register 364 | def label(n: BeginBlock, self: CodeEmit): 365 | self += "{}if {} == {}:".format(self.prefix, n.label, cont_name) 366 | 367 | self.indent() 368 | self.set_cont(n.label) 369 | phi_dict = n.phi 370 | if phi_dict: 371 | self += "{}# Phi".format(self.prefix) 372 | if len(phi_dict) is 1: 373 | l, var_dict = next(iter(phi_dict.items())) 374 | self += "# must from label {}".format(l) 375 | phi_for_given_branch(self, var_dict) 376 | else: 377 | phi(self, phi_dict) 378 | 379 | 380 | @emit.register 381 | def end(_: EndBlock, self: CodeEmit): 382 | self.set_last_cont(cont_name) 383 | self.set_cont('{} + 1'.format(cont_name)) 384 | self += "{}continue".format(self.prefix) 385 | self.dedent() 386 | 387 | 388 | @emit.register 389 | def ret(r: Return, self: CodeEmit): 390 | self += "{}return {}".format(self.prefix, emit_repr(self, r.val)) 391 | 392 | 393 | @emit.register 394 | def set_cont(c: Jmp, self: CodeEmit): 395 | prefix = self.prefix 396 | self.set_last_cont() 397 | self.set_cont(c.label) 398 | self += "{}continue".format(prefix) 399 | 400 | 401 | @emit.register 402 | def set_conf_if(c: JmpIf, self: CodeEmit): 403 | self += "{}if {}:".format(self.prefix, emit_repr(self, c.cond)) 404 | with indent(self): 405 | prefix = self.prefix 406 | self += "{}{} = {}".format(prefix, cont_name, c.label) 407 | self.set_last_cont() 408 | self += "{}continue".format(prefix) 409 | 410 | 411 | @emit.register 412 | def store(st: Store, self: CodeEmit): 413 | ptr_tag = self.mangle(st.target) 414 | val = emit_repr(self, st.val) 415 | self += "{}{}.setref_cell({}) = {}".format(self.prefix, NS.RestrainJIT, ptr_tag, val) 416 | 417 | 418 | @emit.register 419 | def load(l: Load, self: CodeEmit): 420 | tag = self.mangle(l.target) 421 | ptr = emit_repr(self, l.reg) 422 | self += "{}{} = {}.deref_cell({})".format(self.prefix, tag, NS.RestrainJIT, ptr) 423 | 424 | 425 | @emit.register 426 | def ass(a: Ass, self: CodeEmit): 427 | tag = self.mangle(a.target) 428 | val = emit_repr(self, a.val) 429 | self += "{}{} = {}".format(self.prefix, tag, val) 430 | 431 | 432 | @emit.register 433 | def app(a: App, self: CodeEmit): 434 | tag = "" 435 | if a.target: 436 | tag = "{} = ".format(self.mangle(a.target)) 437 | args = list(map(lambda x: emit_repr(self, x), a.args)) 438 | if isinstance(a.f, Prim) and a.f.qual == NS.RestrainJIT: 439 | if a.f.n == "py_subscr": 440 | assert len(args) is 2 441 | self += "{}{}({}[{}])".format(self.prefix, tag, args[0], args[1]) 442 | return 443 | if a.f.n == "py_get_attr": 444 | assert len(args) is 2 445 | self += "{}{}{}.{}".format(self.prefix, tag, args[0], args[1]) 446 | return 447 | if a.f.n == "py_store_attr": 448 | assert len(args) is 3 449 | self += "{}{}.{} = {}".format(self.prefix, args[0], args[1], args[2]) 450 | return 451 | op = { 452 | 'py_add': '+', 453 | 'py_sub': '-', 454 | 'py_mul': '*', 455 | 'py_truediv': '/', 456 | 'py_floordiv': '//', 457 | 'py_mod': '%', 458 | 'py_lsh': '<<', 459 | 'py_rsh': '>>', 460 | 'py_xor': '^', 461 | 'py_or': '|', 462 | 'py_and': '&', 463 | 'py_pow': '**', 464 | 'py_is': 'is', 465 | 'py_gt': '>', 466 | 'py_lt': '<', 467 | 'py_in': 'in', 468 | 'py_not_in': 'not in', 469 | 'py_ge': '>=', 470 | 'py_le': '<=', 471 | 'py_neq': '!=', 472 | 'py_eq': '==', 473 | }.get(a.f.n, None) 474 | 475 | if op is not None: 476 | assert len(args) is 2 477 | self += "{}{}({} {} {})".format(self.prefix, tag, args[0], op, args[1]) 478 | return 479 | op = { 480 | "py_is_none": '{} is None', 481 | 'py_is_true': '{} is True', 482 | 'py_not': 'not {}', 483 | 'py_inv': '~{}', 484 | 'py_neg': '-{}', 485 | }.get(a.f.n, None) 486 | if op is not None: 487 | assert len(args) is 1 488 | self += "{}{}({})".format(self.prefix, tag, op.format(args[0])) 489 | return 490 | f = emit_repr(self, a.f) 491 | self += "{}{}{}({})".format(self.prefix, tag, f, ', '.join(args)) 492 | -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/Design1.md: -------------------------------------------------------------------------------- 1 | RestrainJIT Generated Cython Spec: 2 | 3 | 4 | # Prelude 5 | 6 | ```cython 7 | from restrain_jit.becython.cython_rts cimport RestrainJIT 8 | from restrain_jit.becython.cython_rts.hotspot cimport inttopy, pytoint, JITCounter 9 | from libc.stdint cimport int64_t, int32_t, int16_t, int8_t 10 | from libcpp.map cimport map as std_map 11 | from libcpp.vector cimport vector as std_vector 12 | from cython cimport typeof 13 | ``` 14 | 15 | 16 | # Functions 17 | 18 | A JITed function is represented as a `cpdef` function, whose arguments are typed with cython's fused types, 19 | for instance, given 20 | 21 | ```python 22 | def f(x, y, z): 23 | return x + y + z 24 | ``` 25 | 26 | The JIT version(haven't compiled JITed methods) is 27 | ``` 28 | cdef fused Arg1: 29 | object 30 | 31 | cdef fused Arg2: 32 | object 33 | 34 | cdef fused Arg3: 35 | object 36 | 37 | 38 | cdef JITCounter counter 39 | cdef object recompile_handler 40 | cdef object global_abs 41 | 42 | 43 | cpdef f(Arg1 x, Arg2 y, Arg3 z): 44 | if typeof(x) == typeof(object) or typeof(x) == typeof(object) or typeof(z) == typeof(object): 45 | counter[(jittype(x), jittype(y), jittype(z))] += 1 46 | if counter.times % 100 == 0: 47 | recompile_handler() 48 | 49 | # the next statement means the fake code. 50 | # * the actual code will be retrieved from 51 | # Python function object(including the bytecode) 52 | # and processed by several JIT compiler passes 53 | return x + y + z 54 | 55 | cpdef init(dict globs, dict _counter, _handler): 56 | global global_abs, counter, recompile_handler 57 | global_abs = globs['abs'] 58 | counter = JITCounter(_counter) 59 | recompile_handler = _handler 60 | ``` 61 | 62 | # JIT 63 | 64 | When the `recompile_handler` is notified, it'll use the counter to 65 | check which combinations of arguments' types are the hottest used, 66 | and add them to the fused types, and recompile the file and reload them. 67 | 68 | 69 | e.g, if `f(int, float, int)` is found out to be used the most frequently, 70 | we'll change the definition of fused types, 71 | 72 | ``` 73 | cdef fused Arg1: 74 | object 75 | int 76 | 77 | cdef fused Arg2: 78 | object 79 | float 80 | 81 | cdef fused Arg3: 82 | object 83 | int 84 | ``` 85 | 86 | The use of fused types could lead to some problems in current stage, 87 | and there's a workaround: https://github.com/cython/cython/issues/3204 88 | 89 | 90 | # Dependencies of JIT 91 | 92 | Sometimes, the type is not the trivial primitive types, 93 | they might be user-defined types. 94 | 95 | To support JIT, a type should meet some conditions, and can be translated 96 | into Cython cdef class. 97 | 98 | *To be continue* 99 | -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/Design2.md: -------------------------------------------------------------------------------- 1 | 2 | # Deps 3 | 4 | ```cython 5 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 6 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, inttoptr 7 | from libc.stdint cimport uint64_t, int64_t, int32_t, int16_t, int8_t 8 | from libcpp.map cimport map as std_map 9 | from libcpp.vector cimport vector as std_vector 10 | from libcpp.cast cimport static_cast, reinterpret_cast 11 | from cython cimport typeof 12 | ``` 13 | 14 | # Functions 15 | 16 | A JITed function is a `cdef class`, each call has a overhead of method lookup. 17 | 18 | The methods can be registered dynamically, thus re-compilation don't need to renew all stuffs. 19 | 20 | 21 | Given this example 22 | 23 | ```python 24 | def f(x, y, z): 25 | return x + y + z 26 | ``` 27 | 28 | translates to 29 | 30 | ```cython 31 | # method type 32 | ctypedef object (*method_t)(object, object, object) 33 | 34 | # method lookup type 35 | ctypedef method_t (*method_get_t)(int64_t, int64_t, int64_t) 36 | 37 | # method look up ptr 38 | cdef method_get_t method_get = get_method 39 | 40 | # to avoid params have conflicts against 'type' 41 | cdef inline method_t __invoke_method_get(x, y, z): 42 | func = method_get(pytoint(type(x)), pytoint(type(y)), pytoint(type(x))) 43 | return func 44 | 45 | cdef class ty_f: 46 | # python compatible method to change 'method look up ptr' 47 | cpdef mut_method_get(self, int64_t f): 48 | global method_get 49 | method_get = reinterpret_cast[method_get_t](inttoptr(f)) 50 | def __call__(self, x, y, z): 51 | """ 52 | '_func__manglexxxx' and '__invoke_method_get' 53 | should not conflict against the prospective argument name. 54 | """ 55 | _func__manglexxxx = __invoke_method_get(x, y, z) 56 | return _func__manglexxxx(x, y, z) 57 | ``` 58 | -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/RestrainJIT.pxd: -------------------------------------------------------------------------------- 1 | from cython cimport fused_type 2 | from libc.stdint cimport int64_t, int32_t, int16_t, int8_t, uint64_t 3 | from libcpp.map cimport map as std_map 4 | from libcpp.string cimport string 5 | 6 | NumOrObj1 = fused_type(short, int, object) 7 | NumOrObj2 = fused_type(short, int, object) 8 | 9 | cdef std_map[string, int64_t] internedstrings 10 | cdef int64_t get_symbol(string) 11 | 12 | 13 | cdef class Cell: 14 | cdef object cell_contents 15 | 16 | cdef object deref_cell(Cell) 17 | cdef void setref_cell(Cell cell, object o) 18 | 19 | cdef inline py_add(NumOrObj1 o1, NumOrObj2 o2) -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/RestrainJIT.pyx: -------------------------------------------------------------------------------- 1 | #distutils: language=c++ 2 | #cython: language_level=3 3 | from libc.stdint cimport int64_t, int32_t, int16_t, int8_t, uint64_t 4 | from libcpp.map cimport map as std_map 5 | from libcpp.string cimport string 6 | from cython.operator cimport dereference 7 | 8 | cdef std_map[string, int64_t] internedstrings = std_map[string, int64_t]() 9 | 10 | cdef int64_t get_symbol(string s): 11 | k = internedstrings.find(s) 12 | if k != internedstrings.end(): 13 | return dereference(k).second 14 | cdef uint64_t i = internedstrings.size() 15 | internedstrings[s] = i 16 | return i 17 | 18 | cdef class Cell: 19 | def __init__(self, x): 20 | self.cell_contents = x 21 | 22 | def get(self): 23 | return self.cell_contents 24 | 25 | def set(self, value): 26 | self.cell_contents = value 27 | 28 | cdef object deref_cell(Cell cell): 29 | return cell.cell_contents 30 | 31 | cdef void setref_cell(Cell cell, object o): 32 | cell.cell_contents = o 33 | 34 | cdef inline py_add(NumOrObj1 o1, NumOrObj2 o2): 35 | return o1 + o2 36 | -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/restrain_jit/becython/cython_rts/__init__.py -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/hotspot.pxd: -------------------------------------------------------------------------------- 1 | from libcpp.map cimport map as std_map 2 | from libcpp.vector cimport vector as std_vector 3 | from libc.stdint cimport int64_t, int32_t, int16_t, int8_t 4 | from cpython.ref cimport PyObject 5 | 6 | cdef extern from "typeint.h": 7 | object inttopy "inttoptr"(int64_t) 8 | void* inttoptr(int64_t) 9 | int64_t pytoint "ptrtoint"(object) 10 | int64_t ptrtoint(void*) 11 | int check_ptr_eq(object, object) 12 | object unsafe_cast(void*) 13 | 14 | cdef int64_t typeid(object) -------------------------------------------------------------------------------- /restrain_jit/becython/cython_rts/hotspot.pyx: -------------------------------------------------------------------------------- 1 | #distutils: language=c++ 2 | #cython: language_level=3 3 | 4 | cdef int64_t typeid(object x): 5 | return pytoint(type(x)) 6 | -------------------------------------------------------------------------------- /restrain_jit/becython/cython_vm.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.stack_vm_instructions import * 2 | from restrain_jit.becython.phi_elim import main as phi_elim 3 | from restrain_jit.becython.phi_node_analysis import main as phi_keep 4 | from restrain_jit.becython.relabel import apply as relabel 5 | from restrain_jit.becython.tools import show_instrs 6 | from restrain_jit.jit_info import PyCodeInfo, PyFuncInfo 7 | from restrain_jit.abs_compiler import instrnames as InstrNames 8 | from restrain_jit.abs_compiler.from_bc import Interpreter 9 | from restrain_jit.vm.am import AM, run_machine 10 | from dataclasses import dataclass, field 11 | from bytecode import Bytecode, ControlFlowGraph, Instr as PyInstr, CellVar, CompilerFlags 12 | import typing as t 13 | import types 14 | import sys 15 | 16 | Options = { 17 | 'log-stack-vm': False, 18 | 'log-phi': False, 19 | 'phi-pass': "keep-phi", 20 | } 21 | 22 | 23 | def load_arg(x, cellvars, lineno): 24 | if x in cellvars: 25 | return PyInstr(InstrNames.LOAD_DEREF, CellVar(x), lineno=lineno) 26 | 27 | return PyInstr(InstrNames.LOAD_FAST, x, lineno=lineno) 28 | 29 | 30 | def copy_func(f: types.FunctionType): 31 | # noinspection PyArgumentList 32 | nf = types.FunctionType(f.__code__, f.__globals__, None, None, f.__closure__) 33 | nf.__defaults__ = f.__defaults__ 34 | nf.__name__ = f.__name__ 35 | nf.__qualname__ = f.__qualname__ 36 | nf.__module__ = f.__module__ 37 | nf.__kwdefaults__ = f.__kwdefaults__ 38 | nf.__annotations__ = f.__annotations__ 39 | nf.__dict__ = f.__dict__ 40 | return nf 41 | 42 | 43 | @dataclass 44 | class CyVM(AM[Instr, Repr]): 45 | _meta: dict 46 | 47 | # stack 48 | st: t.List[Repr] 49 | 50 | # instructions 51 | blocks: t.List[t.Tuple[t.Optional[str], t.List[A]]] 52 | 53 | # allocated temporary 54 | used: t.Set[str] 55 | unused: t.Set[str] 56 | globals: t.Set[str] 57 | module: types.ModuleType 58 | 59 | def set_lineno(self, lineno: int): 60 | self.add_instr(None, SetLineno(lineno)) 61 | 62 | def get_module(self) -> types.ModuleType: 63 | return self.module 64 | 65 | def require_global(self, s: str): 66 | self.globals.add(s) 67 | 68 | @classmethod 69 | def func_info(cls, func: types.FunctionType): 70 | names = func.__code__.co_names 71 | code = Bytecode.from_code(func.__code__) 72 | codeinfo = cls.code_info(code) 73 | return PyFuncInfo(func.__name__, func.__module__, func.__defaults__, 74 | func.__kwdefaults__, func.__closure__, func.__globals__, codeinfo, 75 | func, {}, names) 76 | 77 | @classmethod 78 | def code_info(cls, code: Bytecode, *, debug_passes=()) -> PyCodeInfo[Repr]: 79 | 80 | cfg = ControlFlowGraph.from_bytecode(code) 81 | current = cls.empty() 82 | run_machine(Interpreter(code.first_lineno).abs_i_cfg(cfg), current) 83 | glob_deps = tuple(current.globals) 84 | instrs = current.instrs 85 | instrs = cls.pass_push_pop_inline(instrs) 86 | instrs = list(relabel(instrs)) 87 | if Options.get('log-stack-vm'): 88 | print('DEBUG: stack-vm'.center(20, '=')) 89 | show_instrs(instrs) 90 | 91 | phi_pass_name = Options['phi-pass'] 92 | e = None 93 | try: 94 | phi_pass = { 95 | 'phi-elim-by-move': phi_elim, 96 | 'keep-phi': phi_keep 97 | }[Options['phi-pass']] 98 | except KeyError as ke: 99 | e = Exception("Unknown phi pass {!r}".format(phi_pass_name)) 100 | if e is not None: 101 | raise e 102 | instrs = list(phi_pass(instrs)) 103 | if Options.get('log-phi'): 104 | print('DEBUG: phi'.center(20, '=')) 105 | show_instrs(instrs) 106 | return PyCodeInfo(code.name, tuple(glob_deps), code.argnames, code.freevars, 107 | code.cellvars, code.filename, code.first_lineno, code.argcount, 108 | code.kwonlyargcount, bool(code.flags & CompilerFlags.GENERATOR), 109 | bool(code.flags & CompilerFlags.VARKEYWORDS), 110 | bool(code.flags & CompilerFlags.VARARGS), instrs) 111 | 112 | def pop_exception(self, must: bool) -> Repr: 113 | raise NotImplemented 114 | 115 | def meta(self) -> dict: 116 | return self._meta 117 | 118 | def last_block_end(self) -> str: 119 | return self.end_label 120 | 121 | def push_block(self, end_label: str) -> None: 122 | self.blocks.append((end_label, [])) 123 | 124 | def pop_block(self) -> Repr: 125 | # end_label, instrs = self.blocks.pop() 126 | # self.add_instr(None, PushUnwind()) 127 | # for instr in instrs: 128 | # self.add_instr(instr.lhs, instr.rhs) 129 | # self.add_instr(None, PopUnwind()) 130 | # return Const(None) 131 | raise NotImplemented 132 | 133 | def from_const(self, val: Repr) -> object: 134 | assert isinstance(val, Const) 135 | return val.val 136 | 137 | def ret(self, val: Repr): 138 | return self.add_instr(None, Return(val)) 139 | 140 | def const(self, val: object): 141 | return Const(val) 142 | 143 | @classmethod 144 | def reg_of(cls, n: str): 145 | return Reg(n) 146 | 147 | def from_higher(self, qualifier: str, name: str): 148 | assert not qualifier 149 | return Prim('', name) 150 | 151 | def from_lower(self, qualifier: str, name: str): 152 | return Prim(qualifier, name) 153 | 154 | def app(self, f: Repr, args: t.List[Repr]) -> Repr: 155 | name = self.alloc() 156 | reg = Reg(name) 157 | self.add_instr(name, App(f, args)) 158 | return reg 159 | 160 | def store(self, n: str, val: Repr): 161 | self.add_instr(None, Store(Reg(n), val)) 162 | 163 | def load(self, n: str) -> Repr: 164 | r = Reg(n) 165 | name = self.alloc() 166 | self.add_instr(name, Load(r)) 167 | return Reg(name) 168 | 169 | def assign(self, n: str, v: Repr): 170 | self.add_instr(None, Ass(Reg(n), v)) 171 | 172 | def peek(self, n: int): 173 | try: 174 | return self.st[-n - 1] 175 | except IndexError: 176 | name = self.alloc() 177 | self.add_instr(name, Peek(n)) 178 | return name 179 | 180 | def jump(self, n: str): 181 | self.add_instr(None, Jmp(n)) 182 | 183 | def jump_if_push(self, n: str, cond: Repr, leave: Repr): 184 | self.add_instr(None, JmpIfPush(n, cond, leave)) 185 | 186 | def jump_if(self, n: str, cond: Repr): 187 | self.add_instr(None, JmpIf(n, cond)) 188 | 189 | def label(self, n: str) -> None: 190 | self.st.clear() 191 | self.add_instr(None, Label(n)) 192 | 193 | def push(self, r: Repr) -> None: 194 | self.st.append(r) 195 | self.add_instr(None, Push(r)) 196 | 197 | def pop(self) -> Repr: 198 | try: 199 | 200 | a = self.st.pop() 201 | self.add_instr(None, Pop()) 202 | except IndexError: 203 | name = self.alloc() 204 | self.add_instr(name, Pop()) 205 | a = Reg(name) 206 | return a 207 | 208 | def release(self, name: Repr): 209 | """ 210 | release temporary variable 211 | """ 212 | if not isinstance(name, Reg): 213 | return 214 | name = name.n 215 | if name in self.used: 216 | self.used.remove(name) 217 | self.unused.add(name) 218 | 219 | def alloc(self): 220 | """ 221 | allocate a new temporary variable 222 | """ 223 | if self.unused: 224 | return self.unused.pop() 225 | tmp_name = f"tmp-{len(self.used)}" 226 | self.used.add(tmp_name) 227 | return tmp_name 228 | 229 | def add_instr(self, tag, instr: Instr): 230 | self.instrs.append(A(tag, instr)) 231 | return None 232 | 233 | @property 234 | def instrs(self): 235 | return self.blocks[-1][1] 236 | 237 | @property 238 | def end_label(self) -> t.Optional[str]: 239 | return self.blocks[-1][0] 240 | 241 | @classmethod 242 | def empty(cls, module=None): 243 | return cls({}, [], [(None, [])], set(), set(), set(), module 244 | or sys.modules[cls.__module__]) 245 | 246 | @classmethod 247 | def pass_push_pop_inline(cls, instrs): 248 | blacklist = set() 249 | i = 0 250 | while True: 251 | try: 252 | assign = instrs[i] 253 | k, v = assign.lhs, assign.rhs 254 | except IndexError: 255 | break 256 | 257 | if k is None and isinstance(v, Pop): 258 | j = i - 1 259 | while True: 260 | assign = instrs[j] 261 | k, v = assign.lhs, assign.rhs 262 | if k is None and isinstance(v, Push): 263 | try: 264 | assign = instrs[i] 265 | k, v = assign.lhs, assign.rhs 266 | except IndexError: 267 | break 268 | 269 | if k is None and isinstance(v, Pop): 270 | pass 271 | else: 272 | break 273 | 274 | blacklist.add(j) 275 | blacklist.add(i) 276 | i += 1 277 | j -= 1 278 | 279 | try: 280 | assign = instrs[j] 281 | k, v = assign.lhs, assign.rhs 282 | except IndexError: 283 | break 284 | if k is None and isinstance(v, Push): 285 | continue 286 | break 287 | 288 | else: 289 | i += 1 290 | break 291 | else: 292 | i = i + 1 293 | 294 | return [each for i, each in enumerate(instrs) if i not in blacklist] 295 | -------------------------------------------------------------------------------- /restrain_jit/becython/mono_vm.gen: -------------------------------------------------------------------------------- 1 | import restrain_jit.becython.representations; 2 | 3 | abc Instr; 4 | 5 | data SetLineno(Instr) lineno:int; 6 | data App(Instr) target:t.Optional[str] f:Repr args:t.List[Repr]; 7 | data Ass(Instr) target:t.Optional[str] val:Repr; 8 | data Load(Instr) target:t.Optional[str] reg:Reg; 9 | data Store(Instr) target:t.Optional[str] val:Repr; 10 | data SetContIf(Instr) label:object cond:Repr; 11 | data SetCont(Instr) label:object; 12 | data BeginBlock(Instr) label:object; 13 | data EndBlock(Instr); 14 | data Return(Instr) val:Repr; 15 | data PyGlob(Instr) target:t.Optional[str] qual:str name:str; 16 | data CyGlob(Instr) target:t.Optional[str] qual:str name:str; 17 | data MoveOrAss(Instr) target:str reg:str; -------------------------------------------------------------------------------- /restrain_jit/becython/mono_vm.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | from restrain_jit.becython.representations import * 8 | 9 | 10 | class Instr: 11 | pass 12 | 13 | 14 | @dataclass(frozen=True, order=True) 15 | class SetLineno(Instr): 16 | lineno:int 17 | pass 18 | 19 | 20 | @dataclass(frozen=True, order=True) 21 | class App(Instr): 22 | target:t.Optional[str] 23 | f:Repr 24 | args:t.List[Repr] 25 | pass 26 | 27 | 28 | @dataclass(frozen=True, order=True) 29 | class Ass(Instr): 30 | target:t.Optional[str] 31 | val:Repr 32 | pass 33 | 34 | 35 | @dataclass(frozen=True, order=True) 36 | class Load(Instr): 37 | target:t.Optional[str] 38 | reg:Reg 39 | pass 40 | 41 | 42 | @dataclass(frozen=True, order=True) 43 | class Store(Instr): 44 | target:t.Optional[str] 45 | val:Repr 46 | pass 47 | 48 | 49 | @dataclass(frozen=True, order=True) 50 | class SetContIf(Instr): 51 | label:object 52 | cond:Repr 53 | pass 54 | 55 | 56 | @dataclass(frozen=True, order=True) 57 | class SetCont(Instr): 58 | label:object 59 | pass 60 | 61 | 62 | @dataclass(frozen=True, order=True) 63 | class BeginBlock(Instr): 64 | label:object 65 | pass 66 | 67 | 68 | @dataclass(frozen=True, order=True) 69 | class EndBlock(Instr): 70 | pass 71 | 72 | 73 | @dataclass(frozen=True, order=True) 74 | class Return(Instr): 75 | val:Repr 76 | pass 77 | 78 | 79 | @dataclass(frozen=True, order=True) 80 | class PyGlob(Instr): 81 | target:t.Optional[str] 82 | qual:str 83 | name:str 84 | pass 85 | 86 | 87 | @dataclass(frozen=True, order=True) 88 | class CyGlob(Instr): 89 | target:t.Optional[str] 90 | qual:str 91 | name:str 92 | pass 93 | 94 | 95 | @dataclass(frozen=True, order=True) 96 | class MoveOrAss(Instr): 97 | target:str 98 | reg:str 99 | pass 100 | -------------------------------------------------------------------------------- /restrain_jit/becython/phi_elim.py: -------------------------------------------------------------------------------- 1 | """ 2 | Convert stack-vm instructions to reg vm with Phi-node constructs, 3 | then using move-like semantics to eliminate Phi-nodes. 4 | """ 5 | from restrain_jit.becython.representations import Repr, Reg 6 | import restrain_jit.becython.stack_vm_instructions as sv 7 | import restrain_jit.becython.mono_vm as mono 8 | 9 | import typing as t 10 | from collections import defaultdict, OrderedDict 11 | 12 | Label = object 13 | Target = Label 14 | 15 | 16 | class LeftStack: 17 | name: Label 18 | objs: t.List[Repr] 19 | requested: int 20 | 21 | def __init__(self, name, objs, requested=0): 22 | self.name = name 23 | self.objs = objs 24 | self.requested = requested 25 | 26 | def __repr__(self): 27 | return '{!r}{!r}'.format(self.name, self.objs) 28 | 29 | 30 | FromLabel = object 31 | 32 | From = t.Union[FromLabel, LeftStack] 33 | 34 | 35 | class HigherOrderStackUsage(Exception): 36 | pass 37 | 38 | 39 | def pop_or_peek(self, lhs, rhs, val): 40 | if isinstance(val, Repr): 41 | if lhs is not None: 42 | # yield phi.A(None, phi.Ass(Reg(lhs), val)) 43 | self.add_mono_instr(mono.Ass(lhs, val)) 44 | 45 | else: 46 | cur_name = self.current_left_stack.name 47 | drawback_name = self.drawback_name(cur_name, val) 48 | # FIXME: should be max(val, ...) 49 | self.pops[cur_name] = val 50 | if lhs is not None: 51 | # yield phi.A(None, phi.Ass(Reg(lhs), Reg(drawback_name))) 52 | self.add_mono_instr(mono.Ass(lhs, Reg(drawback_name))) 53 | 54 | 55 | class PhiElimViaMove: 56 | current_left_stack: t.Optional[LeftStack] 57 | 58 | # when a block exits, the object lefted on the stack 59 | left_stacks: t.Dict[Label, LeftStack] 60 | 61 | # indicates the blocks that current block comes from 62 | come_from: t.Dict[Target, t.Set[From]] 63 | 64 | # indicates the use of stack objects from current block's parent blocks 65 | pops: t.Dict[Label, int] 66 | 67 | # is ended in current block 68 | come_to_end: bool 69 | 70 | blocks: t.Dict[Label, t.List[mono.Instr]] 71 | 72 | @staticmethod 73 | def drawback_name(from_label, drawback_n: int): 74 | return 'drawback_{}{}'.format(from_label, drawback_n) 75 | 76 | def __init__(self, sv_instrs: t.List[sv.A], elim_phi=True): 77 | self.left_stacks = {} 78 | self.current_left_stack = LeftStack("", []) 79 | self.come_from = defaultdict(set) 80 | 81 | self.sv_instrs = sv_instrs 82 | self.pops = {} 83 | self.phi_dict = {} 84 | self.come_to_end = True 85 | self.elim_phi = True 86 | 87 | self.blocks = OrderedDict() 88 | self.current_block = [] 89 | 90 | def add_mono_instr(self, instr: mono.Instr): 91 | self.current_block.append(instr) 92 | 93 | def __getitem__(self, item) -> t.Union[Repr, int]: 94 | it = self.current_left_stack 95 | try: 96 | return it.objs[item] 97 | except IndexError: 98 | assert item < 0 99 | return it.requested - item - len(it.objs) 100 | 101 | def pop(self) -> t.Union[Repr, int]: 102 | it = self.current_left_stack 103 | try: 104 | return it.objs.pop() 105 | except IndexError: 106 | it.requested += 1 107 | return it.requested 108 | 109 | def push(self, r: Repr): 110 | self.current_left_stack.objs.append(r) 111 | 112 | def check_if_has_entry_label(self): 113 | # if the first instruction is a Label 114 | pass 115 | 116 | def end_block(self, other_label_addr: object = None): 117 | if self.come_to_end: 118 | return 119 | assert self.current_left_stack 120 | if other_label_addr: 121 | self.come_from[other_label_addr].add( 122 | self.current_left_stack) 123 | self.blocks[self.current_left_stack.name] = self.current_block 124 | self.current_block = [] 125 | self.come_to_end = True 126 | 127 | def start_block(self, new_label_addr): 128 | if self.come_to_end: 129 | self.come_from[new_label_addr].add(self.current_left_stack) 130 | else: 131 | self.end_block(new_label_addr) 132 | left_stack = LeftStack(new_label_addr, []) 133 | self.left_stacks[ 134 | new_label_addr] = self.current_left_stack = left_stack 135 | self.come_to_end = False 136 | 137 | def peek_n(self, n: int, at: From): 138 | if not isinstance(at, LeftStack): 139 | at = self.left_stacks[at] 140 | 141 | try: 142 | return at.objs[-n] 143 | except IndexError: 144 | come_from = self.come_from[at.name] 145 | if len(come_from) is 1: 146 | come_from = next(iter(come_from)) 147 | n = n - len(at.objs) + self.pops[at.name] 148 | return self.peek_n(n, come_from) 149 | raise HigherOrderStackUsage 150 | 151 | 152 | def main(sv_instrs): 153 | self = PhiElimViaMove(sv_instrs) 154 | 155 | for ass in sv_instrs: 156 | 157 | rhs = ass.rhs 158 | lhs = ass.lhs 159 | if not lhs: 160 | # can be label or terminate instruction 161 | if isinstance(rhs, sv.Label): 162 | self.start_block(rhs.label) 163 | self.phi_dict[rhs.label] = defaultdict(dict) 164 | # yield phi.A(None, phi.Label(rhs.label, full)) 165 | continue 166 | elif isinstance(rhs, sv.Jmp): 167 | self.add_mono_instr(mono.SetCont(rhs.label)) 168 | self.end_block(rhs.label) 169 | # yield phi.A(None, phi.Jmp(rhs.label)) 170 | continue 171 | elif isinstance(rhs, sv.JmpIf): 172 | self.add_mono_instr(mono.SetContIf(rhs.label, rhs.cond)) 173 | self.end_block(rhs.label) 174 | continue 175 | elif isinstance(rhs, sv.JmpIfPush): 176 | objs = [*self.current_left_stack.objs, rhs.leave] 177 | requested = self.current_left_stack.requested 178 | name = self.current_left_stack.name 179 | tmp, self.current_left_stack = self.current_left_stack, LeftStack( 180 | name, objs, requested) 181 | self.add_mono_instr(mono.SetContIf(rhs.label, rhs.cond)) 182 | 183 | self.end_block(rhs.label) 184 | self.current_left_stack = tmp 185 | # yield phi.A(None, phi.JmpIf(rhs.label, rhs.cond)) 186 | continue 187 | if isinstance(rhs, sv.SetLineno): 188 | # yield phi.A(None, phi.SetLineno(rhs.lineno)) 189 | self.add_mono_instr(mono.SetLineno(rhs.lineno)) 190 | elif isinstance(rhs, sv.Return): 191 | self.add_mono_instr(mono.Return(rhs.val)) 192 | self.end_block() 193 | # yield phi.A(None, phi.Return(rhs.val)) 194 | elif isinstance(rhs, sv.App): 195 | # yield phi.A(lhs, phi.App(rhs.f, rhs.args)) 196 | self.add_mono_instr(mono.App(lhs, rhs.f, rhs.args)) 197 | elif isinstance(rhs, sv.Ass): 198 | # yield phi.A(lhs, phi.Ass(rhs.reg, rhs.val)) 199 | self.add_mono_instr(mono.Ass(rhs.reg.n, rhs.val)) 200 | if lhs is not None: 201 | self.add_mono_instr(mono.Ass(lhs, rhs.reg)) 202 | elif isinstance(rhs, sv.Store): 203 | # yield phi.A(lhs, phi.Store(rhs.reg, rhs.val)) 204 | self.add_mono_instr(mono.Store(rhs.reg.n, rhs.val)) 205 | if lhs is not None: 206 | self.add_mono_instr(mono.Store(lhs, rhs.reg)) 207 | 208 | elif isinstance(rhs, sv.PyGlob): 209 | # yield phi.A(lhs, phi.PyGlob(rhs.qual, rhs.name)) 210 | self.add_mono_instr(mono.PyGlob(lhs, rhs.qual, rhs.name)) 211 | 212 | elif isinstance(rhs, sv.CyGlob): 213 | # yield phi.A(lhs, phi.CyGlob(rhs.qual, rhs.name)) 214 | self.add_mono_instr(mono.CyGlob(lhs, rhs.qual, rhs.name)) 215 | 216 | elif isinstance(rhs, sv.Load): 217 | # yield phi.A(lhs, phi.Load(rhs.reg)) 218 | self.add_mono_instr(mono.Load(lhs, rhs.reg)) 219 | 220 | elif isinstance(rhs, sv.Push): 221 | self.push(rhs.val) 222 | else: 223 | if isinstance(rhs, sv.Peek): 224 | val = self[-rhs.offset] 225 | 226 | elif isinstance(rhs, sv.Pop): 227 | val = self.pop() 228 | else: 229 | raise NotImplementedError(rhs) 230 | pop_or_peek(self, lhs, rhs, val) 231 | 232 | self.end_block() 233 | 234 | pops = self.pops 235 | come_from = self.come_from 236 | 237 | for label, max_required in pops.items(): 238 | for peek_n in range(1, max_required + 1): 239 | reg_name = self.drawback_name(label, peek_n) 240 | for can_from in come_from[label]: 241 | can_from_label = can_from.name if isinstance( 242 | can_from, LeftStack) else can_from 243 | assert isinstance(peek_n, int) 244 | r = self.peek_n(peek_n, at=can_from) 245 | # use move to remove phi nodes 246 | self.blocks[can_from_label].append( 247 | mono.MoveOrAss(reg_name, r)) 248 | 249 | for label, instrs in self.blocks.items(): 250 | yield mono.BeginBlock(label) 251 | yield from instrs 252 | yield mono.EndBlock() 253 | -------------------------------------------------------------------------------- /restrain_jit/becython/phi_node_analysis.py: -------------------------------------------------------------------------------- 1 | """ 2 | Convert stack-vm instructions to reg vm with Phi-node constructs. 3 | This file is preserved for prospective Cython-like backends with Phi-node constructs. 4 | Say, LLVM IR has Phi-nodes, as well as Julia IR 5 | """ 6 | from restrain_jit.becython.representations import Repr, Reg 7 | import restrain_jit.becython.stack_vm_instructions as sv 8 | import restrain_jit.becython.phi_vm as phi 9 | 10 | import typing as t 11 | from collections import defaultdict, OrderedDict 12 | 13 | Label = object 14 | Target = Label 15 | 16 | 17 | class LeftStack: 18 | name: Label 19 | objs: t.List[Repr] 20 | requested: int 21 | 22 | def __init__(self, name, objs, requested=0): 23 | self.name = name 24 | self.objs = objs 25 | self.requested = requested 26 | 27 | def __repr__(self): 28 | return '{!r}{!r}'.format(self.name, self.objs) 29 | 30 | 31 | FromLabel = object 32 | 33 | From = t.Union[FromLabel, LeftStack] 34 | 35 | 36 | class HigherOrderStackUsage(Exception): 37 | pass 38 | 39 | 40 | def pop_or_peek(self, lhs, val): 41 | if isinstance(val, Repr): 42 | if lhs is not None: 43 | self += phi.Ass(lhs, val) 44 | 45 | elif lhs is not None: 46 | cur_name = self.current_left_stack.name 47 | drawback_name = self.drawback_name(cur_name, val) 48 | self.pops[cur_name] = max(self.pops[cur_name], val) 49 | self += phi.Ass(lhs, Reg(drawback_name)) 50 | 51 | 52 | class Phi: 53 | current_left_stack: t.Optional[LeftStack] 54 | 55 | # when a block exits, the object lefted on the stack 56 | left_stacks: t.Dict[Label, LeftStack] 57 | 58 | # indicates the blocks that current block comes from 59 | come_from: t.Dict[Target, t.Set[From]] 60 | 61 | # indicates the use of stack objects from current block's parent blocks 62 | pops: t.Dict[Label, int] 63 | 64 | # stores the instruction operands of label. 65 | phi_dict: t.Dict[Label, t.Dict[Label, t.Dict[str, Repr]]] 66 | 67 | # is ended in current block 68 | come_to_end: bool 69 | 70 | def drawback_name(self, from_label, drawback_n: int): 71 | return 'drawback_{}{}'.format(from_label, drawback_n) 72 | 73 | def __init__(self, sv_instrs: t.List[sv.A]): 74 | self.left_stacks = {} 75 | self.current_left_stack = LeftStack("", []) 76 | self.come_from = defaultdict(set) 77 | self.sv_instrs = sv_instrs 78 | self.pops = defaultdict(lambda: 0) 79 | self.phi_dict = {} 80 | self.come_to_end = True 81 | self.elim_phi = True 82 | 83 | self.blocks = OrderedDict() 84 | self.block = [] 85 | 86 | def __iadd__(self, other): 87 | self.block.append(other) 88 | return self 89 | 90 | def __getitem__(self, item) -> t.Union[Repr, int]: 91 | it = self.current_left_stack 92 | try: 93 | return it.objs[item] 94 | except IndexError: 95 | assert item < 0 96 | return it.requested - item - len(it.objs) 97 | 98 | def pop(self) -> t.Union[Repr, int]: 99 | it = self.current_left_stack 100 | try: 101 | return it.objs.pop() 102 | except IndexError: 103 | it.requested += 1 104 | return it.requested 105 | 106 | def push(self, r: Repr): 107 | self.current_left_stack.objs.append(r) 108 | 109 | def check_if_has_entry_label(self): 110 | # if the first instruction is a Label 111 | pass 112 | 113 | def end_block(self, other_label_addr: object = None): 114 | if self.come_to_end: 115 | return 116 | assert self.current_left_stack 117 | if other_label_addr: 118 | self.come_from[other_label_addr].add(self.current_left_stack) 119 | self.block = [] 120 | self.come_to_end = True 121 | 122 | def start_block(self, new_label_addr): 123 | if self.come_to_end: 124 | self.come_from[new_label_addr].add(self.current_left_stack) 125 | else: 126 | self.end_block(new_label_addr) 127 | self.left_stacks[new_label_addr] = self.current_left_stack = LeftStack( 128 | new_label_addr, []) 129 | self.blocks[self.current_left_stack.name] = self.block 130 | self.come_to_end = False 131 | 132 | def peek_n(self, n: int, at: From): 133 | if not isinstance(at, LeftStack): 134 | at = self.left_stacks[at] 135 | 136 | try: 137 | return at.objs[-n] 138 | except IndexError: 139 | come_from = self.come_from[at.name] 140 | if len(come_from) is 1: 141 | come_from = next(iter(come_from)) 142 | n = n - len(at.objs) + self.pops[at.name] 143 | return self.peek_n(n, come_from) 144 | raise HigherOrderStackUsage 145 | 146 | 147 | def main(sv_instrs): 148 | self = Phi(sv_instrs) 149 | sv_instrs = self.sv_instrs 150 | 151 | for ass in sv_instrs: 152 | 153 | rhs = ass.rhs 154 | lhs = ass.lhs 155 | if not lhs: 156 | # can be label or terminate instruction 157 | if isinstance(rhs, sv.Label): 158 | self.start_block(rhs.label) 159 | full = self.phi_dict[rhs.label] = defaultdict(dict) 160 | self += phi.BeginBlock(rhs.label, full) 161 | continue 162 | elif isinstance(rhs, sv.Jmp): 163 | self += phi.Jmp(rhs.label) 164 | self.end_block(rhs.label) 165 | continue 166 | elif isinstance(rhs, sv.JmpIf): 167 | self += phi.JmpIf(rhs.label, rhs.cond) 168 | self.end_block(rhs.label) 169 | continue 170 | elif isinstance(rhs, sv.JmpIfPush): 171 | objs = [*self.current_left_stack.objs, rhs.leave] 172 | requested = self.current_left_stack.requested 173 | name = self.current_left_stack.name 174 | tmp, self.current_left_stack = self.current_left_stack, LeftStack( 175 | name, objs, requested) 176 | self.current_left_stack = tmp 177 | self += phi.JmpIf(rhs.label, rhs.cond) 178 | self.end_block(rhs.label) 179 | continue 180 | if isinstance(rhs, sv.SetLineno): 181 | self += phi.SetLineno(rhs.lineno) 182 | 183 | elif isinstance(rhs, sv.Return): 184 | self += phi.Return(rhs.val) 185 | self.end_block() 186 | 187 | elif isinstance(rhs, sv.App): 188 | self += phi.App(lhs, rhs.f, rhs.args) 189 | 190 | elif isinstance(rhs, sv.Ass): 191 | self += phi.Ass(rhs.reg.n, rhs.val) 192 | if lhs: 193 | self += phi.Ass(lhs, rhs.val) 194 | 195 | elif isinstance(rhs, sv.Store): 196 | assert not lhs 197 | self += phi.Store(rhs.reg.n, rhs.val) 198 | 199 | elif isinstance(rhs, sv.Load): 200 | if lhs: 201 | self += phi.Load(lhs, rhs.reg) 202 | 203 | elif isinstance(rhs, sv.Push): 204 | self.push(rhs.val) 205 | else: 206 | if isinstance(rhs, sv.Peek): 207 | val = self[-rhs.offset] 208 | 209 | elif isinstance(rhs, sv.Pop): 210 | val = self.pop() 211 | else: 212 | raise NotImplementedError(rhs) 213 | pop_or_peek(self, lhs, val) 214 | 215 | self.end_block() 216 | 217 | pops = self.pops 218 | come_from = self.come_from 219 | fulls = self.phi_dict 220 | 221 | for label, max_required in pops.items(): 222 | phi_dispatch_cases = fulls[label] 223 | for peek_n in range(1, max_required + 1): 224 | reg_name = self.drawback_name(label, peek_n) 225 | for can_from in come_from[label]: 226 | can_from_label = can_from.name if isinstance(can_from, 227 | LeftStack) else can_from 228 | assert isinstance(peek_n, int) 229 | r = self.peek_n(peek_n, at=can_from) 230 | reg_dispatch_cases = phi_dispatch_cases[can_from_label] 231 | reg_dispatch_cases[reg_name] = r 232 | 233 | for label, instrs in self.blocks.items(): 234 | yield from instrs 235 | yield phi.EndBlock() 236 | -------------------------------------------------------------------------------- /restrain_jit/becython/phi_vm.gen: -------------------------------------------------------------------------------- 1 | import restrain_jit.becython.representations; 2 | 3 | abc Instr; 4 | data SetLineno(Instr) lineno:int; 5 | data App(Instr) target:t.Optional[str] f:Repr args:t.List[Repr]; 6 | data Ass(Instr) target:t.Optional[str] val:Repr; 7 | data Load(Instr) target:t.Optional[str] reg:Reg; 8 | data Store(Instr) target:t.Optional[str] val:Repr; 9 | data JmpIf(Instr) label:object cond:Repr; 10 | data Jmp(Instr) label:object; 11 | data BeginBlock(Instr) label:object phi:t.Dict[object,t.Dict[str,Repr]]; 12 | data EndBlock(Instr); 13 | data Return(Instr) val:Repr; -------------------------------------------------------------------------------- /restrain_jit/becython/phi_vm.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | from restrain_jit.becython.representations import * 8 | 9 | 10 | class Instr: 11 | pass 12 | 13 | 14 | @dataclass(frozen=True, order=True) 15 | class SetLineno(Instr): 16 | lineno:int 17 | pass 18 | 19 | 20 | @dataclass(frozen=True, order=True) 21 | class App(Instr): 22 | target:t.Optional[str] 23 | f:Repr 24 | args:t.List[Repr] 25 | pass 26 | 27 | 28 | @dataclass(frozen=True, order=True) 29 | class Ass(Instr): 30 | target:t.Optional[str] 31 | val:Repr 32 | pass 33 | 34 | 35 | @dataclass(frozen=True, order=True) 36 | class Load(Instr): 37 | target:t.Optional[str] 38 | reg:Reg 39 | pass 40 | 41 | 42 | @dataclass(frozen=True, order=True) 43 | class Store(Instr): 44 | target:t.Optional[str] 45 | val:Repr 46 | pass 47 | 48 | 49 | @dataclass(frozen=True, order=True) 50 | class JmpIf(Instr): 51 | label:object 52 | cond:Repr 53 | pass 54 | 55 | 56 | @dataclass(frozen=True, order=True) 57 | class Jmp(Instr): 58 | label:object 59 | pass 60 | 61 | 62 | @dataclass(frozen=True, order=True) 63 | class BeginBlock(Instr): 64 | label:object 65 | phi:t.Dict[object,t.Dict[str,Repr]] 66 | pass 67 | 68 | 69 | @dataclass(frozen=True, order=True) 70 | class EndBlock(Instr): 71 | pass 72 | 73 | 74 | @dataclass(frozen=True, order=True) 75 | class Return(Instr): 76 | val:Repr 77 | pass 78 | -------------------------------------------------------------------------------- /restrain_jit/becython/relabel.py: -------------------------------------------------------------------------------- 1 | import restrain_jit.becython.stack_vm_instructions as sv 2 | import typing as t 3 | 4 | # I don't know why, but using 5 | # `from restrain_jit.becy.tools import sv_jumps` 6 | # will lead to the failure of type checking at line 24 7 | # of this file. Seemingly a bug of PyCharm. 8 | sv_jumps = (sv.Jmp, sv.JmpIf, sv.JmpIfPush) 9 | 10 | 11 | def apply(instrs: t.List[sv.A]): 12 | target_labels = {} 13 | # TODO: using "whether in used_labels" to decide 14 | # "whether to codegen as a subroutine". 15 | used_labels = set() 16 | for ass in instrs: 17 | if ass.lhs: 18 | # cannot be jump or label 19 | continue 20 | rhs = ass.rhs 21 | if isinstance(rhs, sv.Label): 22 | target_labels[rhs.label] = len(target_labels) 23 | elif isinstance(rhs, sv_jumps): 24 | used_labels.add(rhs.label) 25 | 26 | for ass in instrs: 27 | if not ass.lhs: 28 | rhs = ass.rhs 29 | if isinstance(rhs, sv.Label): 30 | if rhs.label in target_labels: 31 | yield sv.A(None, sv.Label(target_labels[rhs.label])) 32 | continue 33 | elif isinstance(rhs, sv.JmpIf): 34 | yield sv.A(None, 35 | sv.JmpIf(target_labels[rhs.label], rhs.cond)) 36 | continue 37 | elif isinstance(rhs, sv.JmpIfPush): 38 | yield sv.A( 39 | None, 40 | sv.JmpIfPush(target_labels[rhs.label], rhs.cond, 41 | rhs.leave)) 42 | continue 43 | elif isinstance(rhs, sv.Jmp): 44 | yield sv.A(None, sv.Jmp(target_labels[rhs.label])) 45 | continue 46 | yield ass 47 | -------------------------------------------------------------------------------- /restrain_jit/becython/representations.gen: -------------------------------------------------------------------------------- 1 | abc Repr; 2 | 3 | data Reg(Repr) n:str; 4 | data Const(Repr) val:object; 5 | data Prim(Repr) qual:str n:str; -------------------------------------------------------------------------------- /restrain_jit/becython/representations.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | class Repr: 8 | pass 9 | 10 | 11 | @dataclass(frozen=True, order=True) 12 | class Reg(Repr): 13 | n:str 14 | pass 15 | 16 | 17 | @dataclass(frozen=True, order=True) 18 | class Const(Repr): 19 | val:object 20 | pass 21 | 22 | 23 | @dataclass(frozen=True, order=True) 24 | class Prim(Repr): 25 | qual:str 26 | n:str 27 | pass 28 | -------------------------------------------------------------------------------- /restrain_jit/becython/stack_vm_instructions.gen: -------------------------------------------------------------------------------- 1 | import restrain_jit.becython.representations; 2 | abc Instr; 3 | 4 | data A lhs:t.Optional[str] rhs:Instr; 5 | 6 | data SetLineno(Instr) lineno:int; 7 | data App(Instr) f:Repr args:t.List[Repr]; 8 | data Ass(Instr) reg:Reg val:Repr; 9 | data Load(Instr) reg:Reg; 10 | data Store(Instr) reg:Reg val:Repr; 11 | data JmpIf(Instr) label:object cond:Repr; 12 | data JmpIfPush(Instr) label:object cond:Repr leave:Repr; 13 | data Jmp(Instr) label:object; 14 | data Label(Instr) label:object; 15 | data Peek(Instr) offset:int; 16 | data Return(Instr) val:Repr; 17 | data Push(Instr) val:Repr; 18 | data Pop(Instr) ; 19 | -------------------------------------------------------------------------------- /restrain_jit/becython/stack_vm_instructions.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | from restrain_jit.becython.representations import * 8 | 9 | 10 | class Instr: 11 | pass 12 | 13 | 14 | @dataclass(frozen=True, order=True) 15 | class A: 16 | lhs:t.Optional[str] 17 | rhs:Instr 18 | pass 19 | 20 | 21 | @dataclass(frozen=True, order=True) 22 | class SetLineno(Instr): 23 | lineno:int 24 | pass 25 | 26 | 27 | @dataclass(frozen=True, order=True) 28 | class App(Instr): 29 | f:Repr 30 | args:t.List[Repr] 31 | pass 32 | 33 | 34 | @dataclass(frozen=True, order=True) 35 | class Ass(Instr): 36 | reg:Reg 37 | val:Repr 38 | pass 39 | 40 | 41 | @dataclass(frozen=True, order=True) 42 | class Load(Instr): 43 | reg:Reg 44 | pass 45 | 46 | 47 | @dataclass(frozen=True, order=True) 48 | class Store(Instr): 49 | reg:Reg 50 | val:Repr 51 | pass 52 | 53 | 54 | @dataclass(frozen=True, order=True) 55 | class JmpIf(Instr): 56 | label:object 57 | cond:Repr 58 | pass 59 | 60 | 61 | @dataclass(frozen=True, order=True) 62 | class JmpIfPush(Instr): 63 | label:object 64 | cond:Repr 65 | leave:Repr 66 | pass 67 | 68 | 69 | @dataclass(frozen=True, order=True) 70 | class Jmp(Instr): 71 | label:object 72 | pass 73 | 74 | 75 | @dataclass(frozen=True, order=True) 76 | class Label(Instr): 77 | label:object 78 | pass 79 | 80 | 81 | @dataclass(frozen=True, order=True) 82 | class Peek(Instr): 83 | offset:int 84 | pass 85 | 86 | 87 | @dataclass(frozen=True, order=True) 88 | class Return(Instr): 89 | val:Repr 90 | pass 91 | 92 | 93 | @dataclass(frozen=True, order=True) 94 | class Push(Instr): 95 | val:Repr 96 | pass 97 | 98 | 99 | @dataclass(frozen=True, order=True) 100 | class Pop(Instr): 101 | pass 102 | -------------------------------------------------------------------------------- /restrain_jit/becython/tools.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython import stack_vm_instructions as sv 2 | from restrain_jit.becython import phi_vm as phi 3 | from restrain_jit.becython import mono_vm as mono 4 | from restrain_jit.jit_info import PyCodeInfo 5 | import typing as t 6 | sv_jumps = (sv.JmpIfPush, sv.JmpIf, sv.Jmp) 7 | 8 | 9 | def show_mono_instrs(instrs: t.List[mono.Instr], indent=''): 10 | for v in instrs: 11 | k = getattr(v, 'target', None) 12 | if k is not None: 13 | print(indent + k, '=', end=' ') 14 | else: 15 | print(indent, end='') 16 | next_indent = indent + ' ' 17 | if isinstance(v, mono.App): 18 | print('call', v.f) 19 | for each in v.args: 20 | if isinstance(each, mono.Const) and isinstance( 21 | each.val, PyCodeInfo): 22 | print(next_indent, "function", each.val.name) 23 | show_instrs(each.val.instrs, 24 | next_indent + " ") 25 | else: 26 | print(next_indent, each, sep='') 27 | elif isinstance(v, mono.BeginBlock): 28 | print('label {}'.format(v.label).center(20, '=')) 29 | elif isinstance(v, mono.EndBlock): 30 | print('end'.center(20, '=')) 31 | else: 32 | print(v) 33 | 34 | 35 | def show_stack_instrs(instrs: t.List[sv.A], indent=''): 36 | for a in instrs: 37 | k = a.lhs 38 | v = a.rhs 39 | if k is not None: 40 | print(indent + k, '=', end=' ') 41 | else: 42 | print(indent, end='') 43 | next_indent = indent + ' ' 44 | if isinstance(v, sv.App): 45 | print('call', v.f) 46 | for each in v.args: 47 | if isinstance(each, sv.Const) and isinstance( 48 | each.val, PyCodeInfo): 49 | print(next_indent, "function", each.val.name) 50 | show_instrs(each.val.instrs, 51 | next_indent + " ") 52 | else: 53 | print(next_indent, each, sep='') 54 | else: 55 | print(v) 56 | 57 | 58 | def show_phi_instrs(instrs: t.List[phi.Instr], indent=''): 59 | for a in instrs: 60 | k = getattr(a, 'target', None) 61 | if k is not None: 62 | print(indent + k, '=', end=' ') 63 | else: 64 | print(indent, end='') 65 | next_indent = indent + ' ' 66 | if isinstance(a, phi.App): 67 | print('call', a.f) 68 | for each in a.args: 69 | if isinstance(each, phi.Const) and isinstance( 70 | each.val, PyCodeInfo): 71 | print(next_indent, "function", each.val.name) 72 | show_instrs(each.val.instrs, 73 | next_indent + " ") 74 | else: 75 | print(next_indent, each, sep='') 76 | elif isinstance(a, phi.BeginBlock): 77 | print(indent, 'label ', a.label, ':', sep='') 78 | for label_name, values in a.phi.items(): 79 | print(indent + ' ', label_name, ': ', sep='', end='') 80 | for reg_name, a in values.items(): 81 | print('{} = {}'.format(reg_name, a), end=', ') 82 | if values: 83 | print() 84 | else: 85 | print(a) 86 | 87 | 88 | def show_instrs(instrs, t=None): 89 | if t is None: 90 | t = type(instrs[0]) 91 | if issubclass(t, mono.Instr): 92 | show_mono_instrs(instrs) 93 | return 94 | elif issubclass(t, sv.A): 95 | show_stack_instrs(instrs) 96 | return 97 | elif issubclass(t, phi.Instr): 98 | show_phi_instrs(instrs) 99 | else: 100 | raise TypeError("Instructions is a List of {}".format(t)) 101 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/restrain_jit/bejulia/__init__.py -------------------------------------------------------------------------------- /restrain_jit/bejulia/basics.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | from typing import overload, Sequence 3 | 4 | T = t.TypeVar("T") 5 | 6 | 7 | class JList(t.Generic[T]): 8 | 9 | def __init__(self, jl): 10 | self.__jit__ = jl 11 | 12 | @overload 13 | def __getitem__(self, i: int) -> T: 14 | ... 15 | 16 | @overload 17 | def __getitem__(self, s: slice) -> Sequence[T]: 18 | ... 19 | 20 | def __getitem__(self, i: int) -> T: 21 | return JList.__getitem__.__jit__(self, i) 22 | 23 | def __setitem__(self, i: int, v: T): 24 | return JList.__setitem__.__jit__(self, i, v) 25 | 26 | def append(self, e: T) -> None: 27 | JList.append.__jit__(self, e) 28 | 29 | def pop(self) -> T: 30 | return JList.pop.__jit__(self) 31 | 32 | __jit__: object 33 | 34 | 35 | class JDict: 36 | __jit__: object 37 | 38 | 39 | class JTuple: 40 | __jit__: object 41 | 42 | 43 | class JSet: 44 | __jit__: object 45 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/functional.py: -------------------------------------------------------------------------------- 1 | """ 2 | As a workaround to the lack of optimizations on Julia's stack machine emulation, 3 | we provide some precompile Julia functions to avoid using coontrol flows like 4 | `for` for `try-catch` for they're not sufficiently efficient. 5 | """ 6 | 7 | 8 | class JitMap: 9 | __jit__: callable 10 | 11 | def __call__(self, xs): 12 | return self.__jit__(xs) 13 | 14 | 15 | select = JitMap() 16 | 17 | 18 | class JitForeach: 19 | __jit__: callable 20 | 21 | def __call__(self, xs): 22 | return self.__jit__(xs) 23 | 24 | 25 | foreach = JitForeach() 26 | 27 | 28 | class AsJuliaObject: 29 | __jit__: callable 30 | 31 | def __matmul__(self, a): 32 | return self.__jit__(a) 33 | 34 | 35 | J = AsJuliaObject() 36 | 37 | 38 | class SIMDMap: 39 | __jit__: callable 40 | 41 | def __call__(self, a): 42 | return self.__jit__(a) 43 | 44 | 45 | simd_select = SIMDMap() 46 | 47 | 48 | class SIMDForeach: 49 | __jit__: callable 50 | 51 | def __call__(self, a): 52 | return self.__jit__(a) 53 | 54 | 55 | simd_foreach = SIMDForeach() 56 | 57 | 58 | class Out: 59 | __jit__: callable 60 | 61 | def __call__(self, a): 62 | return self.__jit__(a) 63 | 64 | 65 | out = Out() 66 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/instructions: -------------------------------------------------------------------------------- 1 | import restrain_jit.bejulia.representations; 2 | 3 | 4 | data A lhs:t.Optional[str] rhs:'Instr'; 5 | 6 | abc Instr; 7 | data App(Instr) f:Repr args:t.List[Repr]; 8 | data Ass(Instr) reg:Reg val:Repr; 9 | data Load(Instr) reg:Reg; 10 | data Store(Instr) reg:Reg val:Repr; 11 | data JmpIf(Instr) label:str cond:Repr; 12 | data JmpIfPush(Instr) label:str cond:Repr leave:Repr; 13 | data Jmp(Instr) label:str; 14 | data Label(Instr) label:str; 15 | data Peek(Instr) offset:int; 16 | data Return(Instr) val:Repr; 17 | data Push(Instr) val:Repr; 18 | data Pop(Instr) ; 19 | data PyGlob(Instr) qual:str name:str; 20 | data JlGlob(Instr) qual:str name:str; 21 | data UnwindBlock(Instr) instrs:t.List[A]; 22 | data PopException(Instr) must:bool; -------------------------------------------------------------------------------- /restrain_jit/bejulia/instructions.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | from restrain_jit.bejulia.representations import * 8 | 9 | 10 | @dataclass 11 | class A: 12 | lhs:t.Optional[str] 13 | rhs:'Instr' 14 | pass 15 | 16 | 17 | class Instr: 18 | pass 19 | 20 | 21 | @dataclass 22 | class App(Instr): 23 | f:Repr 24 | args:t.List[Repr] 25 | pass 26 | 27 | 28 | @dataclass 29 | class Ass(Instr): 30 | reg:Reg 31 | val:Repr 32 | pass 33 | 34 | 35 | @dataclass 36 | class Load(Instr): 37 | reg:Reg 38 | pass 39 | 40 | 41 | @dataclass 42 | class Store(Instr): 43 | reg:Reg 44 | val:Repr 45 | pass 46 | 47 | 48 | @dataclass 49 | class JmpIf(Instr): 50 | label:str 51 | cond:Repr 52 | pass 53 | 54 | 55 | @dataclass 56 | class JmpIfPush(Instr): 57 | label:str 58 | cond:Repr 59 | leave:Repr 60 | pass 61 | 62 | 63 | @dataclass 64 | class Jmp(Instr): 65 | label:str 66 | pass 67 | 68 | 69 | @dataclass 70 | class Label(Instr): 71 | label:str 72 | pass 73 | 74 | 75 | @dataclass 76 | class Peek(Instr): 77 | offset:int 78 | pass 79 | 80 | 81 | @dataclass 82 | class Return(Instr): 83 | val:Repr 84 | pass 85 | 86 | 87 | @dataclass 88 | class Push(Instr): 89 | val:Repr 90 | pass 91 | 92 | 93 | @dataclass 94 | class Pop(Instr): 95 | pass 96 | 97 | 98 | @dataclass 99 | class PyGlob(Instr): 100 | qual:str 101 | name:str 102 | pass 103 | 104 | 105 | @dataclass 106 | class JlGlob(Instr): 107 | qual:str 108 | name:str 109 | pass 110 | 111 | 112 | @dataclass 113 | class UnwindBlock(Instr): 114 | instrs:t.List[A] 115 | pass 116 | 117 | 118 | @dataclass 119 | class PopException(Instr): 120 | must:bool 121 | pass 122 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/jl_init.py: -------------------------------------------------------------------------------- 1 | from julia import core 2 | from restrain_jit.bejulia.simple_julia_py import get_julia 3 | from restrain_jit.bejulia.jl_protocol import bridge, Aware 4 | import ctypes 5 | 6 | libjulia = get_julia().lib 7 | 8 | jl_main = ctypes.c_void_p.in_dll(libjulia, "jl_main_module") 9 | 10 | # libjulia.jl_exception_occurred.restype = ctypes.c_void_p 11 | 12 | # jl_get_global = libjulia.jl_get_global 13 | # jl_get_global.argtypes = [ctypes.c_void_p, ctypes.c_void_p] 14 | # jl_get_global.restype = ctypes.c_void_p 15 | # 16 | # jl_call0 = libjulia.jl_call0 17 | # jl_call0.argtypes = [ctypes.c_void_p] 18 | # jl_call0.restype = ctypes.c_void_p 19 | # 20 | # jl_symbol = libjulia.jl_symbol 21 | # jl_symbol.argtypes = [ctypes.c_char_p] 22 | # jl_symbol.restype = ctypes.c_void_p 23 | 24 | libjulia.jl_eval_string(b"import RestrainJIT") 25 | libjulia.jl_eval_string(b"aware! = RestrainJIT.init!()") 26 | libjulia.jl_eval_string(b'print(1)') 27 | libjulia.jl_eval_string(b'println(PyCall)') 28 | libjulia.jl_eval_string(b'print(2)') 29 | 30 | 31 | def init(): 32 | 33 | # restrain_jl_side_aware_ = jl_get_global(jl_main, 34 | # jl_symbol(b"aware!")) 35 | 36 | def aware_(val): 37 | bridge.append(val) 38 | libjulia.jl_eval_string(b"aware!()") 39 | return bridge.pop() 40 | 41 | Aware.f = aware_ 42 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/jl_protocol.py: -------------------------------------------------------------------------------- 1 | class Aware: 2 | f = None 3 | 4 | 5 | bridge = [] 6 | 7 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/julia_vm.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.instructions import * 2 | from restrain_jit.bejulia.representations import * 3 | from restrain_jit.bejulia.jl_protocol import bridge, Aware 4 | from restrain_jit.jit_info import PyCodeInfo, PyFuncInfo 5 | from restrain_jit.abs_compiler import instrnames as InstrNames 6 | from restrain_jit.abs_compiler.from_bc import Interpreter 7 | from restrain_jit.vm.am import AM, run_machine 8 | from dataclasses import dataclass 9 | from bytecode import Bytecode, ControlFlowGraph, Instr as PyInstr, CellVar, CompilerFlags 10 | import typing as t 11 | import types 12 | import sys 13 | 14 | 15 | def load_arg(x, cellvars, lineno): 16 | if x in cellvars: 17 | return PyInstr(InstrNames.LOAD_DEREF, CellVar(x), lineno=lineno) 18 | 19 | return PyInstr(InstrNames.LOAD_FAST, x, lineno=lineno) 20 | 21 | 22 | def copy_func(f: types.FunctionType): 23 | # noinspection PyArgumentList 24 | nf = types.FunctionType(f.__code__, f.__globals__, None, None, 25 | f.__closure__) 26 | nf.__defaults__ = f.__defaults__ 27 | nf.__name__ = f.__name__ 28 | nf.__qualname__ = f.__qualname__ 29 | nf.__module__ = f.__module__ 30 | nf.__kwdefaults__ = f.__kwdefaults__ 31 | nf.__annotations__ = f.__annotations__ 32 | nf.__dict__ = f.__dict__ 33 | return nf 34 | 35 | 36 | @dataclass 37 | class JuVM(AM[Instr, Repr]): 38 | 39 | def set_lineno(self, lineno: int): 40 | # TODO 41 | pass 42 | 43 | def get_module(self) -> types.ModuleType: 44 | return self.module 45 | 46 | def require_global(self, s: str): 47 | self.globals.add(s) 48 | 49 | @classmethod 50 | def func_info(cls, func: types.FunctionType) -> types.FunctionType: 51 | names = func.__code__.co_names 52 | code = Bytecode.from_code(func.__code__) 53 | codeinfo = cls.code_info(code) 54 | 55 | def r_compile(): 56 | jit_func = Aware.f(self) 57 | print("jit_func", type(jit_func)) 58 | bc = Bytecode() 59 | 60 | bc.append(PyInstr(InstrNames.LOAD_CONST, jit_func)) 61 | bc.extend( 62 | [load_arg(each, cellvars, lineno) for each in argnames]) 63 | bc.extend([ 64 | PyInstr(InstrNames.CALL_FUNCTION, len(argnames)), 65 | PyInstr(InstrNames.RETURN_VALUE) 66 | ]) 67 | bc._copy_attr_from(code) 68 | start_func.__code__ = bc.to_code() 69 | start_func.__jit__ = jit_func 70 | return jit_func 71 | 72 | start_func = copy_func(func) 73 | start_func_code = Bytecode() 74 | lineno = code.first_lineno 75 | argnames = code.argnames 76 | start_func_code.argnames = argnames 77 | cellvars = code.cellvars 78 | start_func_code.extend([ 79 | PyInstr(InstrNames.LOAD_CONST, r_compile, lineno=lineno), 80 | PyInstr(InstrNames.CALL_FUNCTION, 0, lineno=lineno), 81 | *(load_arg(each, cellvars, lineno) for each in argnames), 82 | PyInstr(InstrNames.CALL_FUNCTION, 83 | len(argnames), 84 | lineno=lineno), 85 | PyInstr(InstrNames.RETURN_VALUE, lineno=lineno) 86 | ]) 87 | start_func_code._copy_attr_from(code) 88 | self = PyFuncInfo(func.__name__, func.__module__, 89 | func.__defaults__, func.__kwdefaults__, 90 | func.__closure__, func.__globals__, codeinfo, 91 | func, {}, names) 92 | start_func.__code__ = start_func_code.to_code() 93 | start_func.__func_info__ = self 94 | start_func.__compile__ = r_compile 95 | start_func.__jit__ = None 96 | return start_func 97 | 98 | @classmethod 99 | def code_info(cls, code: Bytecode) -> PyCodeInfo[Repr]: 100 | 101 | cfg = ControlFlowGraph.from_bytecode(code) 102 | current = cls.empty() 103 | run_machine( 104 | Interpreter(code.first_lineno).abs_i_cfg(cfg), current) 105 | glob_deps = tuple(current.globals) 106 | instrs = current.instrs 107 | instrs = current.pass_push_pop_inline(instrs) 108 | return PyCodeInfo(code.name, tuple(glob_deps), code.argnames, 109 | code.freevars, code.cellvars, code.filename, 110 | code.first_lineno, code.argcount, 111 | code.kwonlyargcount, 112 | bool(code.flags & CompilerFlags.GENERATOR), 113 | bool(code.flags & CompilerFlags.VARKEYWORDS), 114 | bool(code.flags & CompilerFlags.VARARGS), 115 | instrs) 116 | 117 | def pop_exception(self, must: bool) -> Repr: 118 | name = self.alloc() 119 | self.add_instr(name, PopException(must)) 120 | return Reg(name) 121 | 122 | def meta(self) -> dict: 123 | return self._meta 124 | 125 | def last_block_end(self) -> str: 126 | return self.end_label 127 | 128 | def push_block(self, end_label: str) -> None: 129 | self.blocks.append((end_label, [])) 130 | 131 | def pop_block(self) -> Repr: 132 | end_label, instrs = self.blocks.pop() 133 | regname = self.alloc() 134 | instr = UnwindBlock(instrs) 135 | self.add_instr(regname, instr) 136 | return Reg(regname) 137 | 138 | def from_const(self, val: Repr) -> object: 139 | assert isinstance(val, Const) 140 | return val.val 141 | 142 | def ret(self, val: Repr): 143 | return self.add_instr(None, Return(val)) 144 | 145 | def const(self, val: object): 146 | return Const(val) 147 | 148 | @classmethod 149 | def reg_of(cls, n: str): 150 | return Reg(n) 151 | 152 | def from_higher(self, qualifier: str, name: str): 153 | regname = self.alloc() 154 | self.add_instr(regname, PyGlob(qualifier, name)) 155 | return Reg(regname) 156 | 157 | def from_lower(self, qualifier: str, name: str): 158 | regname = self.alloc() 159 | self.add_instr(regname, JlGlob(qualifier, name)) 160 | return Reg(regname) 161 | 162 | def app(self, f: Repr, args: t.List[Repr]) -> Repr: 163 | name = self.alloc() 164 | reg = Reg(name) 165 | self.add_instr(name, App(f, args)) 166 | return reg 167 | 168 | def store(self, n: str, val: Repr): 169 | self.add_instr(None, Store(Reg(n), val)) 170 | 171 | def load(self, n: str) -> Repr: 172 | r = Reg(n) 173 | name = self.alloc() 174 | self.add_instr(name, Load(r)) 175 | return Reg(name) 176 | 177 | def assign(self, n: str, v: Repr): 178 | self.add_instr(None, Ass(Reg(n), v)) 179 | 180 | def peek(self, n: int): 181 | try: 182 | return self.st[-n - 1] 183 | except IndexError: 184 | name = self.alloc() 185 | self.add_instr(name, Peek(n)) 186 | return name 187 | 188 | def jump(self, n: str): 189 | self.add_instr(None, Jmp(n)) 190 | 191 | def jump_if_push(self, n: str, cond: Repr, leave: Repr): 192 | self.add_instr(None, JmpIfPush(n, cond, leave)) 193 | 194 | def jump_if(self, n: str, cond: Repr): 195 | self.add_instr(None, JmpIf(n, cond)) 196 | 197 | def label(self, n: str) -> None: 198 | self.st.clear() 199 | self.add_instr(None, Label(n)) 200 | 201 | def push(self, r: Repr) -> None: 202 | self.st.append(r) 203 | self.add_instr(None, Push(r)) 204 | 205 | def pop(self) -> Repr: 206 | try: 207 | 208 | a = self.st.pop() 209 | self.add_instr(None, Pop()) 210 | except IndexError: 211 | name = self.alloc() 212 | self.add_instr(name, Pop()) 213 | a = Reg(name) 214 | return a 215 | 216 | def release(self, name: Repr): 217 | """ 218 | release temporary variable 219 | """ 220 | if not isinstance(name, Reg): 221 | return 222 | name = name.n 223 | if name in self.used: 224 | self.used.remove(name) 225 | self.unused.add(name) 226 | 227 | def alloc(self): 228 | """ 229 | allocate a new temporary variable 230 | """ 231 | if self.unused: 232 | return self.unused.pop() 233 | tmp_name = f"tmp-{len(self.used)}" 234 | self.used.add(tmp_name) 235 | return tmp_name 236 | 237 | def add_instr(self, tag, instr: Instr): 238 | self.instrs.append(A(tag, instr)) 239 | return None 240 | 241 | _meta: dict 242 | 243 | # stack 244 | st: t.List[Repr] 245 | 246 | # instructions 247 | blocks: t.List[t.Tuple[t.Optional[str], t.List[A]]] 248 | 249 | # allocated temporary 250 | used: t.Set[str] 251 | unused: t.Set[str] 252 | globals: t.Set[str] 253 | module: types.ModuleType 254 | 255 | @property 256 | def instrs(self): 257 | return self.blocks[-1][1] 258 | 259 | @property 260 | def end_label(self) -> t.Optional[str]: 261 | return self.blocks[-1][0] 262 | 263 | @classmethod 264 | def pass_push_pop_inline(cls, instrs): 265 | blacklist = set() 266 | i = 0 267 | while True: 268 | try: 269 | assign = instrs[i] 270 | k, v = assign.lhs, assign.rhs 271 | except IndexError: 272 | break 273 | if isinstance(v, UnwindBlock): 274 | v.instrs = cls.pass_push_pop_inline(v.instrs) 275 | if k is None and isinstance(v, Pop): 276 | j = i - 1 277 | while True: 278 | assign = instrs[j] 279 | k, v = assign.lhs, assign.rhs 280 | if k is None and isinstance(v, Push): 281 | try: 282 | assign = instrs[i] 283 | k, v = assign.lhs, assign.rhs 284 | except IndexError: 285 | break 286 | 287 | if k is None and isinstance(v, Pop): 288 | pass 289 | else: 290 | break 291 | 292 | blacklist.add(j) 293 | blacklist.add(i) 294 | i += 1 295 | j -= 1 296 | 297 | try: 298 | assign = instrs[j] 299 | k, v = assign.lhs, assign.rhs 300 | except IndexError: 301 | break 302 | if k is None and isinstance(v, Push): 303 | continue 304 | break 305 | 306 | else: 307 | i += 1 308 | break 309 | else: 310 | i = i + 1 311 | 312 | return [ 313 | each for i, each in enumerate(instrs) if i not in blacklist 314 | ] 315 | 316 | @classmethod 317 | def empty(cls, module=None): 318 | return cls({}, [], [(None, [])], set(), set(), set(), module 319 | or sys.modules[cls.__module__]) 320 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/pragmas.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | T = t.TypeVar("T") 3 | 4 | 5 | class _Const(t.Generic[T]): 6 | 7 | def __getitem__(self, item: T) -> T: 8 | return self 9 | 10 | 11 | const = _Const() 12 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/representations: -------------------------------------------------------------------------------- 1 | abc Repr; 2 | data Reg(Repr) n:str; 3 | data Const(Repr) val:object; -------------------------------------------------------------------------------- /restrain_jit/bejulia/representations.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, auto as _auto 2 | import abc 3 | import typing as t 4 | from dataclasses import dataclass 5 | 6 | 7 | class Repr: 8 | pass 9 | 10 | 11 | @dataclass 12 | class Reg(Repr): 13 | n:str 14 | pass 15 | 16 | 17 | @dataclass 18 | class Const(Repr): 19 | val:object 20 | pass 21 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/simple_julia_py.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | import json 3 | from pathlib import Path 4 | from functools import lru_cache 5 | from typing import Union 6 | 7 | 8 | def check_config(conf: dict): 9 | # TODO 10 | pass 11 | 12 | 13 | @lru_cache() 14 | def get_conf(path: str): 15 | try: 16 | path = Path(path).expanduser() / "info.json" 17 | with path.open() as info: 18 | conf = json.load(info) 19 | check_config(conf) 20 | return conf 21 | except IOError: 22 | raise IOError( 23 | "Didn't find out correct configurations for Restrain at {}". 24 | format(path.absolute())) 25 | 26 | 27 | class JuliaPreLoad: 28 | 29 | def __init__(self, init, lib): 30 | self.init = init 31 | self.lib = lib 32 | 33 | 34 | def get_jl_lib(conf) -> JuliaPreLoad: 35 | jl = conf["julia"] 36 | lib_path = jl["lib"] 37 | sys_image = jl["image"] 38 | binary = jl["bin"] 39 | 40 | lib = ctypes.PyDLL(lib_path, ctypes.RTLD_GLOBAL) 41 | lib.jl_eval_string.argtypes = [ctypes.c_char_p] 42 | lib.jl_eval_string.restype = ctypes.c_void_p 43 | 44 | try: 45 | init = lib.jl_init_with_image 46 | except AttributeError: 47 | init = lib.jl_init_with_image__threading 48 | 49 | return JuliaPreLoad( 50 | lambda: init(binary.encode(), sys_image.encode()), lib) 51 | 52 | 53 | def get_julia(julia: Union[JuliaPreLoad, str] = '~/.restrain'): 54 | if isinstance(julia, str): 55 | conf = get_conf(julia) 56 | return get_julia(get_jl_lib(conf)) 57 | assert isinstance(julia, JuliaPreLoad) 58 | julia.init() 59 | julia.lib.jl_eval_string(b'try using PyCall catch e; println(e) end') 60 | julia.lib.jl_eval_string(b'println(100)') 61 | 62 | return julia 63 | -------------------------------------------------------------------------------- /restrain_jit/bejulia/tools.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.instructions import UnwindBlock, App, Const 2 | from restrain_jit.jit_info import PyCodeInfo 3 | 4 | 5 | def show_instrs(instrs, indent=''): 6 | 7 | for a in instrs: 8 | k = a.lhs 9 | v = a.rhs 10 | if k is not None: 11 | print(indent + k, '=', end=' ') 12 | else: 13 | print(indent, end='') 14 | next_indent = indent + ' ' 15 | if isinstance(v, UnwindBlock): 16 | print() 17 | print(next_indent, 'Unwind', sep='') 18 | show_instrs(v.instrs, next_indent) 19 | elif isinstance(v, App): 20 | print('call', v.f) 21 | for each in v.args: 22 | if isinstance(each, Const) and isinstance( 23 | each.val, PyCodeInfo): 24 | print(next_indent, "function", each.val.name) 25 | show_instrs(each.val.instrs, 26 | next_indent + " ") 27 | else: 28 | print(next_indent, each) 29 | 30 | else: 31 | print(v) 32 | -------------------------------------------------------------------------------- /restrain_jit/config.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import json 3 | 4 | 5 | class Configuration: 6 | # FIXME: better config obj 7 | config_dict: dict 8 | 9 | def __init__(self, config_dict): 10 | object.__setattr__(self, 'config_dict', config_dict) 11 | 12 | def __getattr__(self, item): 13 | val = self.config_dict[item] 14 | if isinstance(val, dict): 15 | return Configuration(val) 16 | return val 17 | 18 | def __setattr__(self, key, value): 19 | self.config_dict[key] = value 20 | 21 | def __delattr__(self, item): 22 | del self.config_dict[item] 23 | 24 | 25 | RESTRAIN_DIR_PATH = Path("~/.restrain").expanduser() 26 | 27 | RESTRAIN_CONFIG_FILE_PATH = RESTRAIN_DIR_PATH / "info.json" 28 | 29 | if not RESTRAIN_DIR_PATH.exists( 30 | ) or not RESTRAIN_CONFIG_FILE_PATH.exists(): 31 | # FIXME: auto create config 32 | raise RuntimeError("Python Restrain JIT not configured.") 33 | 34 | with RESTRAIN_CONFIG_FILE_PATH.open() as f: 35 | RESTRAIN_CONFIG = Configuration(json.load(f)) 36 | del f 37 | -------------------------------------------------------------------------------- /restrain_jit/cpy_compat.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import types 3 | from restrain_jit.abs_compiler import instrnames as InstrNames 4 | from bytecode import Bytecode, Instr 5 | from importlib import _bootstrap 6 | from importlib._bootstrap import ModuleSpec 7 | from importlib.abc import Loader 8 | from contextlib import contextmanager 9 | from importlib._bootstrap_external import PathFinder, FileLoader, ExtensionFileLoader 10 | 11 | 12 | class RePyLoader(Loader): 13 | 14 | def __init__(self, loader: FileLoader): 15 | self.loader = loader 16 | 17 | def create_module(self, spec: ModuleSpec): 18 | # mod = RestrainModule(spec.name) 19 | mod = types.ModuleType(spec.name) 20 | _bootstrap._init_module_attrs(spec, mod) 21 | return mod 22 | 23 | def exec_module(self, module): 24 | code = self.loader.get_code(module.__name__) 25 | if code is None: 26 | raise ImportError('cannot load module {!r} when get_code() ' 27 | 'returns None'.format(module.__name__)) 28 | __glob_refs__ = module.__glob_refs__ = { 29 | } # from a global symbol to jit functions it's referenced 30 | bc = Bytecode.from_code(code) 31 | 32 | def update_generations(name): 33 | functions = __glob_refs__.get(name, None) 34 | if functions is None: 35 | return 36 | for fn in functions: 37 | fn.__update_global_ref__(name) 38 | 39 | module.__dict__['__update_generations__'] = update_generations 40 | 41 | def update_bc(): 42 | for each in bc: 43 | yield each 44 | if isinstance( 45 | each, 46 | Instr) and each.name == InstrNames.STORE_NAME: 47 | 48 | yield Instr( 49 | InstrNames.LOAD_NAME, 50 | '__update_generations__', 51 | lineno=each.lineno) 52 | yield Instr( 53 | InstrNames.LOAD_CONST, 54 | each.arg, 55 | lineno=each.lineno) 56 | yield Instr( 57 | InstrNames.CALL_FUNCTION, 1, lineno=each.lineno) 58 | yield Instr(InstrNames.POP_TOP, lineno=each.lineno) 59 | 60 | lst = list(update_bc()) 61 | bc.clear() 62 | bc.extend(lst) 63 | code = bc.to_code() 64 | exec(code, module.__dict__) 65 | 66 | 67 | class RePyFinder(PathFinder): 68 | 69 | @classmethod 70 | def find_spec(cls, fullname, path=None, target=None): 71 | spec: ModuleSpec = PathFinder.find_spec(fullname, path, target) 72 | if spec and spec.loader and isinstance( 73 | spec.loader, FileLoader) and not isinstance( 74 | spec.loader, ExtensionFileLoader): 75 | spec.loader = RePyLoader(spec.loader) 76 | return spec 77 | 78 | 79 | def unregister(): 80 | sys.meta_path.remove(RePyFinder) 81 | 82 | 83 | def register(): 84 | sys.meta_path.insert(0, RePyFinder) 85 | 86 | 87 | @contextmanager 88 | def with_registered(): 89 | try: 90 | register() 91 | yield 92 | finally: 93 | unregister() 94 | -------------------------------------------------------------------------------- /restrain_jit/jit_info.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | import types 3 | from dataclasses import dataclass 4 | 5 | Instr = t.TypeVar("Instr") 6 | 7 | 8 | @dataclass 9 | class PyCodeInfo(t.Generic[Instr]): 10 | name: str 11 | 12 | glob_deps: t.Tuple[str] 13 | argnames: t.List[str] 14 | freevars: t.List[str] 15 | cellvars: t.List[str] 16 | 17 | filename: str 18 | lineno: int 19 | 20 | argcount: int 21 | kwonlyargcount: int 22 | 23 | is_gen: bool 24 | has_var_kw: bool 25 | has_var_arg: bool 26 | 27 | instrs: t.List 28 | 29 | 30 | @dataclass(unsafe_hash=True, order=True, repr=False) 31 | class PyFuncInfo(t.Generic[Instr]): 32 | r_name: str 33 | r_module: str 34 | r_defaults: t.Optional[t.Tuple[object, ...]] 35 | r_kw_defaults: t.Optional[dict] 36 | r_closure: t.Optional[t.Tuple[object, ...]] 37 | r_globals: dict 38 | r_codeinfo: PyCodeInfo 39 | r_func: types.FunctionType 40 | r_options: dict 41 | r_attrnames: t.Tuple[str, ...] 42 | -------------------------------------------------------------------------------- /restrain_jit/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | 5 | class MissingDict(dict): 6 | 7 | def __init__(self, fact): 8 | dict.__init__(self) 9 | self.fact = fact 10 | 11 | def __missing__(self, key): 12 | val = self[key] = self.fact() 13 | return val 14 | 15 | 16 | class CodeOut(dict): 17 | 18 | def __missing__(self, key): 19 | v = self[key] = [] 20 | return v 21 | 22 | def merge_update(self, another: 'CodeOut'): 23 | for k, v in another.items(): 24 | self[k] += v 25 | 26 | def to_code_lines(self): 27 | 28 | def key(x): 29 | return x[0] 30 | 31 | for _, v in sorted(self.items(), key=key): 32 | yield from v 33 | 34 | 35 | def exec_cc(cmd, args): 36 | """ 37 | Execute with current context. 38 | Yes, you're right -- I'm naming it after call/cc. 39 | 40 | Return a generator. 41 | The first yielded one is the status of 42 | the execution of subprocess command. 43 | 44 | The following ones are the the buffer batches of stderr, 45 | each of which is a Python 'bytes' object 46 | """ 47 | file = cmd 48 | err_in, err_out = os.pipe() 49 | out_in, out_out = os.pipe() 50 | if os.fork(): 51 | _, status = os.wait() 52 | os.close(err_out) 53 | os.close(out_out) 54 | yield status 55 | while True: 56 | load = os.read(err_in, 1024) 57 | if not load: 58 | break 59 | yield load 60 | else: 61 | # for child process 62 | os.close(err_in) 63 | os.close(out_in) 64 | os.dup2(err_out, sys.stderr.fileno()) 65 | os.dup2(out_out, sys.stdout.fileno()) 66 | 67 | os.execvpe(file, [cmd, *args], dict(os.environ)) 68 | # in case that os.execvp fails 69 | sys.exit(127) 70 | -------------------------------------------------------------------------------- /restrain_jit/vm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/restrain_jit/vm/__init__.py -------------------------------------------------------------------------------- /restrain_jit/vm/am.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import types 3 | import typing as t 4 | import bytecode 5 | from restrain_jit.jit_info import PyCodeInfo, PyFuncInfo 6 | from dataclasses import dataclass 7 | 8 | 9 | @dataclass 10 | class Symbol: 11 | s: str 12 | 13 | 14 | @dataclass 15 | class ValSymbol: 16 | s: str 17 | 18 | 19 | Instr = t.TypeVar("Instr") 20 | Repr = t.TypeVar("Repr") 21 | 22 | 23 | class AM(t.Generic[Instr, Repr]): 24 | 25 | @abc.abstractmethod 26 | def yield_return(self, val: Repr): 27 | raise NotImplemented 28 | 29 | @abc.abstractmethod 30 | def set_lineno(self, lineno: int): 31 | raise NotImplemented 32 | 33 | @abc.abstractmethod 34 | def require_global(self, s: str): 35 | raise NotImplemented 36 | 37 | @abc.abstractmethod 38 | def meta(self) -> dict: 39 | raise NotImplemented 40 | 41 | @abc.abstractmethod 42 | def pop_exception(self, must: bool) -> Repr: 43 | raise NotImplemented 44 | 45 | @abc.abstractmethod 46 | def push_block(self, end_label: str) -> None: 47 | raise NotImplemented 48 | 49 | @abc.abstractmethod 50 | def pop_block(self) -> Repr: 51 | raise NotImplemented 52 | 53 | @abc.abstractmethod 54 | def last_block_end(self) -> str: 55 | raise NotImplemented 56 | 57 | @classmethod 58 | @abc.abstractmethod 59 | def reg_of(cls, n: str): 60 | raise NotImplemented 61 | 62 | @abc.abstractmethod 63 | def from_higher(self, qualifier: str, name: str): 64 | raise NotImplemented 65 | 66 | @abc.abstractmethod 67 | def from_lower(self, qualifier: str, name: str): 68 | raise NotImplemented 69 | 70 | @abc.abstractmethod 71 | def release(self, name: Repr) -> None: 72 | raise NotImplemented 73 | 74 | @abc.abstractmethod 75 | def alloc(self) -> str: 76 | raise NotImplemented 77 | 78 | @abc.abstractmethod 79 | def add_instr(self, tag: t.Union[None, str], instr: Instr) -> Repr: 80 | raise NotImplemented 81 | 82 | @abc.abstractmethod 83 | def pop(self) -> Repr: 84 | raise NotImplemented 85 | 86 | @abc.abstractmethod 87 | def push(self, r: Repr) -> None: 88 | raise NotImplemented 89 | 90 | @abc.abstractmethod 91 | def label(self, n: str) -> None: 92 | raise NotImplemented 93 | 94 | @abc.abstractmethod 95 | def jump_if(self, n: str, cond: Repr) -> None: 96 | raise NotImplemented 97 | 98 | @abc.abstractmethod 99 | def jump_if_push(self, n: str, cond: Repr, leave: Repr) -> None: 100 | raise NotImplemented 101 | 102 | @abc.abstractmethod 103 | def jump(self, n: str) -> None: 104 | raise NotImplemented 105 | 106 | @abc.abstractmethod 107 | def peek(self, n: int) -> Repr: 108 | raise NotImplemented 109 | 110 | @abc.abstractmethod 111 | def assign(self, reg: str, v: Repr): 112 | raise NotImplemented 113 | 114 | @abc.abstractmethod 115 | def load(self, reg: str) -> Repr: 116 | raise NotImplemented 117 | 118 | @abc.abstractmethod 119 | def store(self, reg: str, val: Repr) -> None: 120 | raise NotImplemented 121 | 122 | @abc.abstractmethod 123 | def app(self, f: Repr, args: t.List[Repr]) -> Repr: 124 | raise NotImplemented 125 | 126 | @abc.abstractmethod 127 | def const(self, val: object) -> Repr: 128 | raise NotImplemented 129 | 130 | @abc.abstractmethod 131 | def from_const(self, val: Repr) -> object: 132 | raise NotImplemented 133 | 134 | @abc.abstractmethod 135 | def ret(self, val: Repr) -> None: 136 | raise NotImplemented 137 | 138 | @classmethod 139 | @abc.abstractmethod 140 | def code_info(cls, code: bytecode.Bytecode) -> PyCodeInfo[Instr]: 141 | raise NotImplemented 142 | 143 | @classmethod 144 | @abc.abstractmethod 145 | def func_info(cls, func: types.FunctionType): 146 | raise NotImplemented 147 | 148 | @abc.abstractmethod 149 | def get_module(self) -> types.ModuleType: 150 | raise NotImplemented 151 | 152 | 153 | def code_info(code: bytecode.Bytecode): 154 | return lambda vm: vm.code_info(code) 155 | 156 | 157 | def func_info(fn: types.FunctionType): 158 | return lambda vm: vm.code_info(fn) 159 | 160 | 161 | def pop_exception(must: bool = False) -> Repr: 162 | return lambda vm: vm.pop_exception(must) 163 | 164 | 165 | def require_global(a: str): 166 | return lambda vm: vm.require_global(a) 167 | 168 | 169 | def meta(): 170 | return lambda vm: vm.meta() 171 | 172 | 173 | def last_block_end(): 174 | return lambda vm: vm.last_block_end() 175 | 176 | 177 | def push_block(r: str): 178 | return lambda vm: vm.push_block(r) 179 | 180 | 181 | def pop_block(): 182 | return lambda vm: vm.pop_block() 183 | 184 | 185 | def from_const(r: Repr): 186 | return lambda vm: vm.from_const(r) 187 | 188 | 189 | def ret(val: Repr): 190 | return lambda vm: vm.ret(val) 191 | 192 | 193 | def const(val: object): 194 | return lambda vm: vm.const(val) 195 | 196 | 197 | def reg_of(name: str): 198 | return lambda vm: vm.reg_of(name) 199 | 200 | 201 | def release(name: Repr): 202 | return lambda vm: vm.release(name) 203 | 204 | 205 | def alloc(): 206 | return lambda vm: vm.alloc() 207 | 208 | 209 | def add_instr(instr: Instr): 210 | 211 | def apply(vm): 212 | a = vm.alloc() 213 | vm.add_instr(a, instr) 214 | return vm.reg_of(vm) 215 | 216 | return apply 217 | 218 | 219 | def from_higher(qualifier: str, name: str): 220 | return lambda vm: vm.from_higher(qualifier, name) 221 | 222 | 223 | def from_lower(qualifier: str, name: str): 224 | return lambda vm: vm.from_lower(qualifier, name) 225 | 226 | 227 | def pop() -> Repr: 228 | return lambda vm: vm.pop() 229 | 230 | 231 | def push(r: Repr): 232 | return lambda vm: vm.push(r) 233 | 234 | 235 | def label(n: str): 236 | return lambda vm: vm.label(n) 237 | 238 | 239 | def jump_if(n: str, cond: Repr): 240 | return lambda vm: vm.jump_if(n, cond) 241 | 242 | 243 | def jump_if_push(n: str, cond: Repr, leave: Repr): 244 | return lambda vm: vm.jump_if_push(n, cond, leave) 245 | 246 | 247 | def jump(n: str): 248 | return lambda vm: vm.jump(n) 249 | 250 | 251 | def peek(n: int): 252 | return lambda vm: vm.peek(n) 253 | 254 | 255 | def assign(reg: str, v: Repr): 256 | return lambda vm: vm.assign(reg, v) 257 | 258 | 259 | def load(reg: str) -> Repr: 260 | return lambda vm: vm.load(reg) 261 | 262 | 263 | def store(reg: str, val: Repr): 264 | return lambda vm: vm.store(reg, val) 265 | 266 | 267 | def app(f: Repr, args: t.List[Repr]): 268 | return lambda vm: vm.app(f, args) 269 | 270 | 271 | def get_module(): 272 | return lambda vm: vm.get_module() 273 | 274 | 275 | def set_lineno(i): 276 | return lambda vm: vm.set_lineno(i) 277 | 278 | 279 | def yield_return(val): 280 | return lambda vm: vm.yield_return(val) 281 | 282 | 283 | def run_machine(gen: t.Generator, vm: AM): 284 | """ 285 | top level of abstract interpretion 286 | """ 287 | v = None 288 | send = gen.send 289 | try: 290 | while True: 291 | binder = send(v) 292 | v = binder(vm) 293 | except StopIteration as e: 294 | return e.value 295 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | from pathlib import Path 3 | 4 | with Path('README.md').open() as readme: 5 | readme = readme.read() 6 | 7 | setup( 8 | name='restrain-jit', 9 | version="0.1", 10 | keywords="JIT, CPython, Optimizations", # keywords of your project that separated by comma "," 11 | description="CPython compatible Python JIT", # a conceise introduction of your project 12 | long_description=readme, 13 | long_description_content_type="text/markdown", 14 | license='mit', 15 | python_requires='>=3.6.0', 16 | url='https://github.com/thautwarm/restrain-jit', 17 | author='thautwarm', 18 | author_email='twshere@outlook.com', 19 | packages=find_packages() + ['restrain_jit.becython.cython_rts'], 20 | package_data={'': ['*.pyx', '*.pxd', '*.h', '*.c', '*.cpp']}, 21 | entry_points={"console_scripts": []}, 22 | # above option specifies commands to be installed, 23 | # e.g: entry_points={"console_scripts": ["yapypy=yapypy.cmd.compiler"]} 24 | install_requires=["bytecode"], 25 | platforms="any", 26 | classifiers=[ 27 | "Programming Language :: Python :: 3.6", 28 | "Programming Language :: Python :: 3.7", 29 | "Programming Language :: Python :: Implementation :: CPython", 30 | ], 31 | zip_safe=False, 32 | ) -------------------------------------------------------------------------------- /static/donate-weixin.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/static/donate-weixin.png -------------------------------------------------------------------------------- /static/p1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/static/p1.png -------------------------------------------------------------------------------- /static/p2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/static/p2.png -------------------------------------------------------------------------------- /static/p3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/static/p3.png -------------------------------------------------------------------------------- /static/procedure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/static/procedure.png -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | from tests.bejl.test_functional import * -------------------------------------------------------------------------------- /tests/becy/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /tests/becy/cfg_and_phi.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import compile_module, setup_pyx_for_cpp 2 | from timeit import timeit 3 | # from pyximport import pyximport 4 | # setup_pyx_for_cpp() 5 | # pyximport.install() 6 | # import restrain_jit.becython.cython_lib.hotspot 7 | -------------------------------------------------------------------------------- /tests/becy/design1_proof_of_concepts.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 2 | from pyximport import pyximport 3 | setup_pyx_for_cpp() 4 | pyximport.install() 5 | import restrain_jit.becython.cython_rts.hotspot 6 | from restrain_jit.becython.cy_loader import compile_module 7 | mod = """ 8 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 9 | from restrain_jit.becython.cython_rts.hotspot cimport inttopy, pytoint, JITCounter 10 | from libc.stdint cimport int64_t, int32_t, int16_t, int8_t 11 | from libcpp.map cimport map as std_map 12 | from libcpp.vector cimport vector as std_vector 13 | from cython cimport typeof, final 14 | 15 | cdef fused Arg1: 16 | object 17 | 18 | cdef fused Arg2: 19 | object 20 | 21 | cdef fused Arg3: 22 | object 23 | 24 | 25 | cdef JITCounter counter 26 | cdef object recompile_handler 27 | cdef object global_abs 28 | 29 | 30 | cpdef f(Arg1 x, Arg2 y, Arg3 z): 31 | if typeof(x) == typeof(object) or typeof(x) == typeof(object) or typeof(z) == typeof(object): 32 | counter[(type(x), type(y), type(z))] += 1 33 | if counter.times % 100 == 0: 34 | recompile_handler() 35 | return x + y + z 36 | 37 | cpdef init(dict globs, dict _counter, _handler): 38 | global global_abs, counter, recompile_handler 39 | global_abs = globs['abs'] 40 | counter = JITCounter(_counter) 41 | recompile_handler = _handler 42 | """ 43 | 44 | mod = compile_module('m', mod) 45 | mod.init(dict(abs=abs), {}, lambda : print("jit started!")) 46 | print(mod.f(14514, 2, 3)) 47 | -------------------------------------------------------------------------------- /tests/becy/design2_practice_1.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 2 | from pyximport import pyximport 3 | 4 | setup_pyx_for_cpp() 5 | pyximport.install() 6 | import restrain_jit.becython.cython_rts.hotspot 7 | from restrain_jit.becython.cy_loader import compile_module 8 | 9 | mod = """ 10 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 11 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, inttoptr 12 | from libc.stdint cimport int64_t 13 | from libcpp.cast cimport reinterpret_cast 14 | 15 | # method type 16 | ctypedef object (*method_t)(object, object, object) 17 | 18 | # method lookup type 19 | ctypedef method_t (*method_get_t)(int64_t, int64_t, int64_t) 20 | 21 | # method look up ptr 22 | cdef method_get_t method_get 23 | 24 | # to avoid params have conflicts against 'type' 25 | cdef inline method_t method_get_invoker(x, y, z): 26 | func = method_get(pytoint(type(x)), pytoint(type(y)), pytoint(type(x))) 27 | return func 28 | 29 | cdef class ty_f: 30 | # python compatible method to change 'method look up ptr' 31 | cpdef mut_method_get(self, int64_t f): 32 | global method_get 33 | method_get = reinterpret_cast[method_get_t](inttoptr(f)) 34 | 35 | def __call__(self, x, y, z): 36 | method = method_get_invoker(x, y, z) 37 | return method(x, y, z) 38 | """ 39 | 40 | mod = compile_module("a", mod) 41 | -------------------------------------------------------------------------------- /tests/becy/design2_proof_of_concepts.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 2 | from pyximport import pyximport 3 | setup_pyx_for_cpp() 4 | pyximport.install() 5 | import restrain_jit.becython.cython_rts.hotspot 6 | from restrain_jit.becython.cy_loader import compile_module 7 | mod = """ 8 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 9 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, inttoptr 10 | from libc.stdint cimport int64_t 11 | from libcpp.cast cimport reinterpret_cast 12 | 13 | cdef object call(_, x, y, z): 14 | return x + y + z 15 | 16 | ctypedef object (*method_t)(object, object, object, object) 17 | ctypedef method_t (*method_get_t)(int64_t, int64_t, int64_t) 18 | 19 | cdef method_t get_method(int64_t x, int64_t y, int64_t z): 20 | if x == pytoint(int): 21 | return call 22 | return call 23 | 24 | cdef method_get_t method_get = get_method 25 | 26 | 27 | cdef inline method_t __invoke_method_get(x, y, z): 28 | func = method_get(pytoint(type(x)), pytoint(type(y)), pytoint(type(x))) 29 | return func 30 | 31 | cdef class ty_f: 32 | def __call__(self, x, y, z): 33 | func = __invoke_method_get(x, y, z) 34 | return func(self, x, y, z) 35 | 36 | cpdef mut_method_get(self, int64_t f): 37 | global method_get 38 | method_get = reinterpret_cast[method_get_t](inttoptr(f)) 39 | 40 | cdef method_t getter2(x, y, z): 41 | return call 42 | 43 | 44 | f2_addr = reinterpret_cast[int64_t](getter2) 45 | f = ty_f() 46 | """ 47 | 48 | from timeit import timeit 49 | mod = compile_module('m', mod) 50 | 51 | print(mod.f(y=14514, x=2, z=3)) 52 | 53 | template = "f(1, 2, 3)" 54 | 55 | 56 | def f(x, y, z): 57 | return x + y + z 58 | 59 | 60 | def test(f): 61 | t = timeit(template, number=10_000_000, globals=dict(f=f)) 62 | print(f, 'costs', t) 63 | 64 | 65 | test(mod.f) 66 | test(f) 67 | 68 | mod.f.mut_method_get(mod.f2_addr) 69 | 70 | test(mod.f) 71 | test(f) 72 | 73 | -------------------------------------------------------------------------------- /tests/becy/dev.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 2 | from pyximport import pyximport 3 | 4 | setup_pyx_for_cpp() 5 | pyximport.install() 6 | import restrain_jit.becython.cython_rts.hotspot 7 | from restrain_jit.becython.cy_loader import compile_module 8 | 9 | mod = """ 10 | cimport restrain_jit.becython.cython_rts.RestrainJIT as RestrainJIT 11 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint 12 | print(RestrainJIT.get_symbol("attr")) 13 | print(RestrainJIT.get_symbol("attra")) 14 | 15 | cpdef int f(x): 16 | return pytoint(type(x)) 17 | """ 18 | 19 | mod = compile_module('m', mod) 20 | print(mod.f(10)) 21 | print(mod.f(10)) 22 | print(mod.f("")) 23 | print(mod.f("")) 24 | print(mod.f(1.0)) 25 | print(mod.f(1.0)) 26 | 27 | -------------------------------------------------------------------------------- /tests/becy/one_func_one_class.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import compile_module, setup_pyx_for_cpp 2 | from timeit import timeit 3 | # from pyximport import pyximport 4 | # setup_pyx_for_cpp() 5 | # pyximport.install() 6 | # import restrain_jit.becython.cython_lib.hotspot 7 | 8 | mod = """ 9 | cimport cython 10 | cdef class ty_S: 11 | def __call__(self, x, y, z): 12 | return x + y + z 13 | 14 | @cython.final 15 | cdef class ty_F: 16 | def __call__(self, x, y, z): 17 | return x + y + z 18 | 19 | s = ty_S() 20 | f = ty_F() 21 | 22 | cpdef F(x, y, z): 23 | return x + y + z 24 | 25 | cdef class ty_Si: 26 | def __call__(self, int x, int y, int z): 27 | return x + y + z 28 | 29 | @cython.final 30 | cdef class ty_Fi: 31 | def __call__(self, int x, int y, int z): 32 | return x + y + z 33 | 34 | si = ty_Si() 35 | fi = ty_Fi() 36 | 37 | cpdef Fi(int x, int y, int z): 38 | return x + y + z 39 | 40 | 41 | cdef fused ty_f_Arg1: 42 | object 43 | 44 | cdef fused ty_f_Arg2: 45 | object 46 | 47 | cdef fused ty_f_Arg3: 48 | object 49 | 50 | 51 | cpdef ff(ty_f_Arg1 x, ty_f_Arg2 y, ty_f_Arg3 z): 52 | return x + y + z 53 | 54 | """ 55 | 56 | mod = compile_module("a", mod) 57 | print(mod.s(1, 2, 3)) 58 | 59 | 60 | def f(x, y, z): 61 | return x + y + z 62 | 63 | 64 | template = "f(1, 2, 3)" 65 | 66 | 67 | def test(f): 68 | t = timeit(template, number=10_000_000, globals=dict(f=f)) 69 | print(f, 'costs', t) 70 | print(mod.ff(1, 2, 3)) 71 | 72 | test(mod.s) 73 | test(mod.f) 74 | test(mod.F) 75 | test(f) 76 | test(mod.si) 77 | test(mod.fi) 78 | test(mod.Fi) 79 | test(mod.ff) 80 | -------------------------------------------------------------------------------- /tests/becy/phi.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.phi_elim import main 2 | from restrain_jit.becython.stack_vm_instructions import * 3 | from restrain_jit.becython.tools import show_instrs 4 | instrs = [ 5 | A(None, Label(label=0)), 6 | A(None, Push(Reg("x0"))), 7 | A(None, JmpIfPush(2, Reg("x0"), Reg("x0"))), 8 | A(None, Label(label=1)), 9 | A(None, Push(Reg("x1"))), 10 | A(None, Label(label=2)), 11 | A("c", Pop()), 12 | A(None, Return(Reg("c"))) 13 | ] 14 | 15 | instrs = list(main(instrs)) 16 | show_instrs(instrs) 17 | -------------------------------------------------------------------------------- /tests/becy/test_datatype.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 3 | from pyximport import pyximport 4 | from timeit import timeit 5 | 6 | from restrain_jit.becython.cython_vm import Options 7 | 8 | setup_pyx_for_cpp() 9 | pyximport.install() 10 | import restrain_jit.becython.cython_rts.hotspot 11 | from restrain_jit.becython.cy_jit_ext_template import mk_module_code 12 | 13 | from restrain_jit.becython.cy_jit import JITSystem 14 | jit_sys = JITSystem() 15 | 16 | 17 | @jit_sys.jitdata 18 | class IntRef: 19 | a: int 20 | 21 | 22 | intref = IntRef(1) 23 | print(intref.a) 24 | 25 | 26 | @jit_sys.jit 27 | def f(x, v): 28 | x.a = v 29 | 30 | 31 | class IntRef2: 32 | a: int 33 | 34 | 35 | intref2 = IntRef2() 36 | 37 | intref2.a = 0 38 | 39 | 40 | def g(x, v): 41 | x.a = v 42 | 43 | 44 | template = "f(x, 10)" 45 | 46 | 47 | def test(f, data): 48 | t = timeit(template, number=10000000, globals=dict(f=f, x=data)) 49 | print(f, 'costs', t) 50 | 51 | 52 | for i in range(1000): 53 | f(intref, 10) 54 | 55 | for i in range(1000): 56 | f(intref2, 10) 57 | 58 | print('stats:') 59 | test(f, intref) 60 | test(g, intref) 61 | 62 | print('stats:') 63 | test(f, intref2) 64 | test(g, intref2) 65 | 66 | fn_place = jit_sys.fn_place_index[id(f)] 67 | method, *_ = fn_place.methods.items() 68 | 69 | load_ty = fn_place.call_recorder.load_type 70 | print('argtype:', tuple(map(load_ty, method[0])), 'addr:', hex(method[1])) 71 | -------------------------------------------------------------------------------- /tests/becy/test_if.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 3 | from pyximport import pyximport 4 | from timeit import timeit 5 | 6 | from restrain_jit.becython.cython_vm import Options 7 | 8 | setup_pyx_for_cpp() 9 | pyximport.install() 10 | import restrain_jit.becython.cython_rts.hotspot 11 | 12 | from restrain_jit.becython.cy_jit_ext_template import mk_module_code 13 | from restrain_jit.becython.cy_jit import JITSystem 14 | jit_sys = JITSystem() 15 | 16 | # DEBUG['stack-vm'] = True 17 | # Options['log-phi'] = True 18 | 19 | # show generated code for debug 20 | jit_sys.store_base_method_log = True 21 | 22 | 23 | @jit_sys.jit 24 | def f(x, y): 25 | if x < y: 26 | return x + y + y + y + x 27 | return 10 28 | 29 | 30 | def g(x, y): 31 | if x < y: 32 | return x + y + y + y + x 33 | return 10 34 | 35 | 36 | template = "f(2, 3)" 37 | 38 | 39 | def test(f): 40 | t = timeit(template, number=10_000_000, globals=dict(f=f)) 41 | print(f, 'costs', t) 42 | 43 | 44 | print(f(1, 2)) 45 | test(f) 46 | test(g) 47 | 48 | test(f) 49 | test(g) 50 | 51 | template = "f(2.0, 3.0)" 52 | 53 | 54 | test(f) 55 | test(g) 56 | 57 | test(f) 58 | test(g) 59 | -------------------------------------------------------------------------------- /tests/becy/test_loop.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from restrain_jit.becython.cy_loader import setup_pyx_for_cpp 3 | from pyximport import pyximport 4 | from timeit import timeit 5 | 6 | from restrain_jit.becython.cython_vm import Options 7 | 8 | setup_pyx_for_cpp() 9 | pyximport.install() 10 | import restrain_jit.becython.cython_rts.hotspot 11 | 12 | from restrain_jit.becython.cy_jit_ext_template import mk_module_code 13 | from restrain_jit.becython.cy_jit import JITSystem 14 | jit_sys = JITSystem() 15 | 16 | # DEBUG['stack-vm'] = True 17 | # Options['log-phi'] = True 18 | 19 | # show generated code for debug 20 | jit_sys.store_base_method_log = True 21 | 22 | 23 | @jit_sys.jit 24 | def f(seq, init): 25 | n = len(seq) 26 | i = 0 27 | while i < n: 28 | init = init + seq[i] 29 | i = i + 1 30 | return init 31 | 32 | 33 | def g(seq, init): 34 | n = len(seq) 35 | i = 0 36 | while i < n: 37 | init = init + seq[i] 38 | i = i + 1 39 | 40 | return init 41 | 42 | 43 | template = "f(seq, 12)" 44 | 45 | 46 | def test(f): 47 | t = timeit(template, number=30000, globals=dict(f=f, seq=[*range(1000)])) 48 | print(f, 'costs', t) 49 | 50 | 51 | test(f) 52 | test(g) 53 | test(f) 54 | test(g) 55 | 56 | 57 | @jit_sys.jit 58 | def f(low, high, step): 59 | s = 0 60 | i = low 61 | while i < high: 62 | s = s + i 63 | i = i + step 64 | return s 65 | 66 | 67 | def g(low, high, step): 68 | s = 0 69 | i = low 70 | while i < high: 71 | s = s + i 72 | i = i + step 73 | return s 74 | 75 | 76 | template = "f(0, 1000, 3)" 77 | 78 | 79 | def test(f): 80 | t = timeit(template, number=30000, globals=dict(f=f, seq=[*range(1000)])) 81 | print(f, 'costs', t) 82 | 83 | 84 | test(f) 85 | test(g) 86 | test(f) 87 | test(g) 88 | -------------------------------------------------------------------------------- /tests/becy/typeid.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.becython.cy_loader import compile_module, setup_pyx_for_cpp 2 | from pyximport import pyximport 3 | setup_pyx_for_cpp() 4 | pyximport.install() 5 | import restrain_jit.becython.cython_rts.hotspot 6 | 7 | mod = """ 8 | cimport cython 9 | from libc cimport stdint 10 | from cpython.ref cimport PyObject 11 | from restrain_jit.becython.cython_rts.hotspot cimport pytoint, check_ptr_eq 12 | cdef object o = None 13 | ot = pytoint(cython.typeof(o)) 14 | cdef fused C: 15 | int 16 | float 17 | object 18 | 19 | cdef class MyClass: 20 | pass 21 | 22 | cdef f(C x): 23 | cdef int a = pytoint(cython.typeof(x)) 24 | return a 25 | 26 | def g(C x): 27 | return f(x) 28 | 29 | def h(): 30 | z = cython.typeof(h) 31 | return check_ptr_eq(z, cython.typeof(object)) 32 | """ 33 | 34 | mod = compile_module('m', mod) 35 | print(mod.h()) 36 | 37 | a = mod.g(mod.MyClass()) 38 | assert isinstance(a, int) 39 | assert a == mod.g(mod.MyClass()) 40 | 41 | b = mod.g(1) 42 | assert isinstance(b, int) 43 | assert b == mod.g(1) 44 | 45 | assert a != b 46 | -------------------------------------------------------------------------------- /tests/bejl/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/tests/bejl/__init__.py -------------------------------------------------------------------------------- /tests/bejl/__main__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thautwarm/restrain-jit/f76b3e9ae8a34d2eef87a42cc87197153f14634c/tests/bejl/__main__.py -------------------------------------------------------------------------------- /tests/bejl/test_apis.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.jl_init import init 2 | from restrain_jit.bejulia.julia_vm import JuVM 3 | 4 | init() 5 | 6 | jit = JuVM.func_info 7 | 8 | 9 | @jit 10 | def test_bool(value: bool): 11 | return not value 12 | 13 | 14 | print(test_bool(True)) 15 | -------------------------------------------------------------------------------- /tests/bejl/test_array_op.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.functional import foreach, select, J 2 | from restrain_jit.bejulia.tools import show_instrs 3 | from restrain_jit.bejulia.pragmas import const 4 | from restrain_jit.bejulia.julia_vm import JuVM 5 | from restrain_jit.bejulia.jl_init import init 6 | # 7 | import numpy as np 8 | init() 9 | jit = JuVM.func_info 10 | 11 | a = J @ [1, 2, 3] 12 | 13 | 14 | @jit 15 | def test_append(x): 16 | x.append(1) 17 | return x 18 | 19 | 20 | # show_instrs(f.__func_info__.r_codeinfo.instrs) 21 | print(test_append(a) == np.array([1, 2, 3, 1])) 22 | 23 | 24 | @jit 25 | def test_repeat_append_jit(): 26 | x = [1] 27 | x.pop() 28 | 29 | for i in range(10000): 30 | x.append(i) 31 | return x 32 | 33 | 34 | @jit 35 | def test_repeat_append_jit_foreach(): 36 | x = [1] 37 | x.pop() 38 | 39 | @foreach(range(10000)) 40 | def each(e): 41 | x.append(e) 42 | 43 | return x 44 | 45 | 46 | def test_repeat_append_nojit(): 47 | x = [] 48 | for i in range(100000): 49 | x.append(i) 50 | return x 51 | 52 | 53 | # print(test_repeat_append_jit()) 54 | 55 | # print(test_repeat_append_jit_foreach()) 56 | # 57 | # for e in (test_repeat_append_jit_foreach.__func_info__.r_codeinfo.instrs): 58 | # print(e) 59 | # show_instrs(test_repeat_append_jit_foreach.__func_info__.r_codeinfo.instrs) 60 | # 61 | # %timeit test_repeat_append_jit() 62 | # %timeit test_repeat_append_nojit() 63 | # %timeit test_repeat_append_jit_foreach() 64 | 65 | from julia import Main 66 | 67 | tt = Main.eval(""" 68 | function tt() 69 | x = [] 70 | for i in 0:9999 71 | push!(x, i) 72 | end 73 | x 74 | end 75 | """) 76 | tt() 77 | 78 | import timeit 79 | 80 | print(timeit.timeit("tt()", globals=dict(tt=tt), number=1000)) 81 | print(timeit.timeit( 82 | "tt()", globals=dict(tt=test_repeat_append_nojit), number=1000)) 83 | -------------------------------------------------------------------------------- /tests/bejl/test_from_bc.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.julia_vm import JuVM, UnwindBlock 2 | from restrain_jit.vm.am import run_machine 3 | from restrain_jit.abs_compiler.from_bc import abs_i_cfg 4 | from prettyprinter import pprint 5 | import bytecode as bc 6 | import typing as t 7 | 8 | 9 | class S(dict): 10 | 11 | def __missing__(self, key): 12 | v = self[key] = len(self) 13 | return v 14 | 15 | 16 | def show_block(x): 17 | s = S() 18 | for ins in x: 19 | print('===== block ', s[id(ins)]) 20 | for j in ins: 21 | if isinstance(j.arg, bc.BasicBlock): 22 | print(j.name, ' -> block', s[id(j.arg)]) 23 | else: 24 | print(j) 25 | 26 | 27 | def f1(x): 28 | with x: 29 | for each in x: 30 | a = each + 1 31 | if a < 2: 32 | k(a).d() 33 | 34 | 35 | def show(instrs, indent=''): 36 | for a in instrs: 37 | k = a.lhs 38 | v = a.rhs 39 | if k is not None: 40 | print(indent + k, '=', end='') 41 | else: 42 | print(indent, end='') 43 | if isinstance(v, UnwindBlock): 44 | next_indent = indent + ' ' 45 | print() 46 | show(v.instrs, next_indent) 47 | else: 48 | print(v) 49 | 50 | 51 | def f2(): 52 | try: 53 | 1 / 0 54 | except ZeroDivisionError: 55 | raise Exception 56 | except Exception: 57 | print(2) 58 | finally: 59 | print(3) 60 | 61 | 62 | d = 2 63 | print(f2) 64 | 65 | v = [1, 2, 3, 4] 66 | 67 | a1 = {1, 2} 68 | a2 = {1: 2} 69 | a3 = (1, 2) 70 | 71 | 72 | def f3(x): 73 | a1 74 | a2 75 | a3 76 | for each in v: 77 | return each + x 78 | 79 | 80 | def func(x): 81 | for i in [1, 2, 3, 3]: 82 | x = x + i 83 | return x 84 | 85 | # print(c.__func_info__.r_codeinfo.glob_deps) 86 | # # 87 | # import dis 88 | # dis.dis(func) 89 | 90 | 91 | # 92 | from restrain_jit.bejulia.jl_init import init 93 | init() 94 | jit = JuVM.func_info 95 | 96 | 97 | @jit 98 | def g(a): 99 | for i in range(1000): 100 | a = a + i 101 | return a 102 | 103 | print(g(1)) 104 | # x = {1, 2} 105 | # @jit 106 | # def g(a): 107 | # for i in x: 108 | # a = a + i 109 | # return a 110 | # print(g(1)) 111 | # # 112 | # # # # print() 113 | # c.__compile__() 114 | # print(c.__jit__) 115 | -------------------------------------------------------------------------------- /tests/bejl/test_functional.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from types import ModuleType 3 | import _heapq as hq 4 | print(sys.modules['_heapq']) 5 | headq = ModuleType('_heapq', hq.__doc__) 6 | headq.__dict__.update(hq.__dict__) 7 | sys.modules['_heapq'] = headq 8 | import _heapq 9 | print(_heapq is headq, _heapq is hq) 10 | # from restrain_jit.bejulia.functional import foreach, select, simd_select, J, out 11 | # from restrain_jit.bejulia.julia_vm import JuVM 12 | from restrain_jit.bejulia.jl_init import init 13 | 14 | # import timeit 15 | # import numpy as np 16 | # 17 | # init() 18 | # 19 | # jit = JuVM.func_info 20 | # 21 | # 22 | # @jit 23 | # def all_add2(lst): 24 | # 25 | # @select(lst) 26 | # def ret(elt): 27 | # return elt + 2 28 | # 29 | # return ret 30 | # 31 | # 32 | # # xs = np.arange(20000) 33 | # # zs = all_add2(xs) 34 | # 35 | # 36 | # @jit 37 | # def const(_): 38 | # return 5 39 | # 40 | # 41 | # print(type(const)) 42 | # print(const(1)) 43 | # # 44 | # # def py_all_add2(lst): 45 | # # 46 | # # # np.fromiter is the most efficient way to 47 | # # # create numpy array from python iterator 48 | # # return np.fromiter( 49 | # # map(lambda x: x + 2, lst), dtype=np.int32, count=len(lst)) 50 | # # 51 | # # 52 | # # @jit 53 | # # def all_add2_simd(lst, out): 54 | # # 55 | # # @simd_select(lst, out) 56 | # # def ret(elt): 57 | # # return elt + 2 58 | # # 59 | # # return ret 60 | # # 61 | # # 62 | # # ret = out(J @ np.ones(len(xs))) 63 | # # print(all_add2_simd(xs, ret)) 64 | # # 65 | # # jit_time = timeit.timeit( 66 | # # """ 67 | # # all_add2(xs)""", 68 | # # globals=dict(all_add2=all_add2, xs=xs), 69 | # # number=1000) 70 | # # 71 | # # cpy_time = (timeit.timeit( 72 | # # """ 73 | # # all_add2(xs)""", 74 | # # globals=dict(all_add2=py_all_add2, xs=xs), 75 | # # number=1000)) 76 | # # 77 | # # print(ret) 78 | # # 79 | # # simd_time = (timeit.timeit( 80 | # # """ 81 | # # all_add2(xs, out)""", 82 | # # globals=dict(all_add2=all_add2_simd, xs=xs, out=ret), 83 | # # number=1000)) 84 | # # print() 85 | # # print(cpy_time) 86 | # # print(jit_time) 87 | # # print(simd_time) 88 | -------------------------------------------------------------------------------- /tests/bejl/test_instructions.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.julia_vm import JuVM, UnwindBlock, App, Repr, Reg, Const 2 | from restrain_jit.jit_info import PyCodeInfo 3 | from restrain_jit.bejulia.tools import show_instrs 4 | 5 | jit = JuVM.func_info 6 | 7 | 8 | @jit 9 | def f(x, y): 10 | pass 11 | 12 | 13 | show_instrs(f.__func_info__.r_codeinfo.instrs) 14 | 15 | # 16 | # @jit 17 | # def func1(x): 18 | # x[:2] = 1 19 | # 20 | # 21 | # show_instrs(func1.__func_info__.r_codeinfo.instrs) 22 | -------------------------------------------------------------------------------- /tests/bejl/test_load.py: -------------------------------------------------------------------------------- 1 | import restrain_jit.cpy_compat 2 | import test_load as a 3 | a.b = 1 4 | print(type(a)) 5 | print(a.__py_module__.b) 6 | a.b = 2 7 | 8 | print(a.__state_ages__) 9 | -------------------------------------------------------------------------------- /tests/bejl/test_load_pyjulia.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.julia_vm import JuVM 2 | from restrain_jit.bejulia.jl_init import init 3 | 4 | init() 5 | 6 | jit = JuVM.func_info 7 | 8 | 9 | @jit 10 | def func1(x): 11 | for i in range(1000): 12 | x = x + i 13 | return x + 1 14 | 15 | 16 | def func2(x): 17 | for i in range(1000): 18 | x += i 19 | return x + 1 20 | 21 | 22 | assert func1(100) == func2(100) 23 | 24 | print(func1.__jit__) -------------------------------------------------------------------------------- /tests/bejl/test_simple_case.py: -------------------------------------------------------------------------------- 1 | from restrain_jit.bejulia.julia_vm import JuVM as jit 2 | 3 | 4 | @jit 5 | def f(arr): 6 | n = len(arr) 7 | 8 | @parallel_for(arr) 9 | def apply(refi): 10 | refi[i] = i + n 11 | 12 | return arr 13 | -------------------------------------------------------------------------------- /tests/bejl/test_stage_count_import_star.py: -------------------------------------------------------------------------------- 1 | a = 1 2 | print(globals().keys()) 3 | -------------------------------------------------------------------------------- /tests/bejl/test_stage_count_main.py: -------------------------------------------------------------------------------- 1 | import sklearn 2 | import test_stage_count_tested 3 | -------------------------------------------------------------------------------- /tests/bejl/test_stage_count_tested.py: -------------------------------------------------------------------------------- 1 | """ 2 | hhhh 3 | """ 4 | 5 | from test_stage_count_import_star import * 6 | x = 1 7 | x = 2 8 | print(globals().keys()) 9 | assert a == 1 10 | 11 | # print(a, x) 12 | --------------------------------------------------------------------------------