mirror of https://github.com/exaloop/codon.git
Add Cython fixes (decorator arguments, pyvars support) (#42)
* Add Cython fixes (decorator arguments, pyvars support) * Update CI and README * Fix problematic test * Don't track Cython-generated cpp file [ci skip] * Fix test * Various codon.jit improvements * clang-format [ci skip] * Add docs [ci skip] * Add docs link Co-authored-by: A. R. Shajii <ars@ars.me>pull/44/head
Ibrahim Numanagić
GitHub
parent
a956a4d2a0
commit
48a2bfa57a
.github
actions/build-manylinux
workflows
codon
parser/visitors/typecheck
docs
interop
intro
stdlib/internal
test
parser
python
|
|
@ -31,7 +31,7 @@ cmake --build build --config Release -- VERBOSE=1
|
|||
export PATH=$PATH:$(pwd)/llvm/bin
|
||||
export LD_LIBRARY_PATH=$(pwd)/build:$LD_LIBRARY_PATH
|
||||
export CODON_DIR=$(pwd)/build
|
||||
python3 -m pip install cython wheel
|
||||
python3 -m pip install cython wheel astunparse
|
||||
python3 -m pip debug --verbose
|
||||
(cd extra/python; python3 setup.py sdist bdist_wheel --plat-name=manylinux2014_x86_64)
|
||||
python3 -m pip install -v extra/python/dist/*.whl
|
||||
|
|
|
|||
|
|
@ -126,7 +126,7 @@ jobs:
|
|||
- name: Set up Python
|
||||
run: |
|
||||
python -m pip install --upgrade pip setuptools wheel
|
||||
python -m pip install numpy cython wheel
|
||||
python -m pip install numpy cython wheel astunparse
|
||||
which python
|
||||
which pip
|
||||
echo "CODON_PYTHON=$(python test/python/find-python-library.py)" >> $GITHUB_ENV
|
||||
|
|
|
|||
|
|
@ -82,9 +82,11 @@ llvm::Error JIT::compile(const ir::Func *input) {
|
|||
return llvm::Error::success();
|
||||
}
|
||||
|
||||
llvm::Expected<ir::Func *> JIT::compile(const std::string &code) {
|
||||
llvm::Expected<ir::Func *> JIT::compile(const std::string &code,
|
||||
const std::string &file, int line) {
|
||||
auto *cache = compiler->getCache();
|
||||
ast::StmtPtr node = ast::parseCode(cache, JIT_FILENAME, code, /*startLine=*/0);
|
||||
ast::StmtPtr node = ast::parseCode(cache, file.empty() ? JIT_FILENAME : file, code,
|
||||
/*startLine=*/line);
|
||||
|
||||
auto sctx = cache->imports[MAIN_IMPORT].ctx;
|
||||
auto preamble = std::make_shared<std::vector<ast::StmtPtr>>();
|
||||
|
|
@ -166,8 +168,11 @@ llvm::Expected<std::string> JIT::run(const ir::Func *input) {
|
|||
return getCapturedOutput();
|
||||
}
|
||||
|
||||
llvm::Expected<std::string> JIT::execute(const std::string &code) {
|
||||
auto result = compile(code);
|
||||
llvm::Expected<std::string>
|
||||
JIT::execute(const std::string &code, const std::string &file, int line, bool debug) {
|
||||
if (debug)
|
||||
fmt::print(stderr, "[codon::jit::execute] code:\n{}-----\n", code);
|
||||
auto result = compile(code, file, line);
|
||||
if (auto err = result.takeError())
|
||||
return std::move(err);
|
||||
if (auto err = compile(result.get()))
|
||||
|
|
@ -198,7 +203,9 @@ std::string buildKey(const std::string &name, const std::vector<std::string> &ty
|
|||
}
|
||||
|
||||
std::string buildPythonWrapper(const std::string &name, const std::string &wrapname,
|
||||
const std::vector<std::string> &types) {
|
||||
const std::vector<std::string> &types,
|
||||
const std::string &pyModule,
|
||||
const std::vector<std::string> &pyVars) {
|
||||
std::stringstream wrap;
|
||||
wrap << "@export\n";
|
||||
wrap << "def " << wrapname << "(args: cobj) -> cobj:\n";
|
||||
|
|
@ -207,12 +214,21 @@ std::string buildPythonWrapper(const std::string &name, const std::string &wrapn
|
|||
<< "a" << i << " = " << types[i] << ".__from_py__(PyTuple_GetItem(args, " << i
|
||||
<< "))\n";
|
||||
}
|
||||
for (unsigned i = 0; i < pyVars.size(); i++) {
|
||||
wrap << " "
|
||||
<< "py" << i << " = pyobj._get_module(\"" << pyModule << "\")._getattr(\""
|
||||
<< pyVars[i] << "\")\n";
|
||||
}
|
||||
wrap << " return " << name << "(";
|
||||
for (unsigned i = 0; i < types.size(); i++) {
|
||||
if (i > 0)
|
||||
wrap << ", ";
|
||||
wrap << "a" << i;
|
||||
}
|
||||
for (unsigned i = 0; i < pyVars.size(); i++) {
|
||||
wrap << ", "
|
||||
<< "py" << i;
|
||||
}
|
||||
wrap << ").__to_py__()\n";
|
||||
|
||||
return wrap.str();
|
||||
|
|
@ -228,8 +244,9 @@ ir::types::Type *JIT::PythonData::getCObjType(ir::Module *M) {
|
|||
return cobj;
|
||||
}
|
||||
|
||||
JITResult JIT::executeSafe(const std::string &code) {
|
||||
auto result = execute(code);
|
||||
JITResult JIT::executeSafe(const std::string &code, const std::string &file, int line,
|
||||
bool debug) {
|
||||
auto result = execute(code, file, line, debug);
|
||||
if (auto err = result.takeError()) {
|
||||
auto errorInfo = llvm::toString(std::move(err));
|
||||
return JITResult::error(errorInfo);
|
||||
|
|
@ -238,8 +255,10 @@ JITResult JIT::executeSafe(const std::string &code) {
|
|||
}
|
||||
|
||||
JITResult JIT::executePython(const std::string &name,
|
||||
const std::vector<std::string> &types, void *arg) {
|
||||
|
||||
const std::vector<std::string> &types,
|
||||
const std::string &pyModule,
|
||||
const std::vector<std::string> &pyVars, void *arg,
|
||||
bool debug) {
|
||||
auto key = buildKey(name, types);
|
||||
auto &cache = pydata->cache;
|
||||
auto it = cache.find(key);
|
||||
|
|
@ -253,7 +272,9 @@ JITResult JIT::executePython(const std::string &name,
|
|||
} else {
|
||||
static int idx = 0;
|
||||
auto wrapname = "__codon_wrapped__" + name + "_" + std::to_string(idx++);
|
||||
auto wrapper = buildPythonWrapper(name, wrapname, types);
|
||||
auto wrapper = buildPythonWrapper(name, wrapname, types, pyModule, pyVars);
|
||||
if (debug)
|
||||
fmt::print(stderr, "[codon::jit::executePython] wrapper:\n{}-----\n", wrapper);
|
||||
if (auto err = compile(wrapper).takeError()) {
|
||||
auto errorInfo = llvm::toString(std::move(err));
|
||||
return JITResult::error(errorInfo);
|
||||
|
|
|
|||
|
|
@ -51,18 +51,25 @@ public:
|
|||
// General
|
||||
llvm::Error init();
|
||||
llvm::Error compile(const ir::Func *input);
|
||||
llvm::Expected<ir::Func *> compile(const std::string &code);
|
||||
llvm::Expected<ir::Func *> compile(const std::string &code,
|
||||
const std::string &file = "", int line = 0);
|
||||
llvm::Expected<void *> address(const ir::Func *input);
|
||||
llvm::Expected<std::string> run(const ir::Func *input);
|
||||
llvm::Expected<std::string> execute(const std::string &code);
|
||||
llvm::Expected<std::string> execute(const std::string &code,
|
||||
const std::string &file = "", int line = 0,
|
||||
bool debug = false);
|
||||
|
||||
// Python
|
||||
llvm::Expected<void *> runPythonWrapper(const ir::Func *wrapper, void *arg);
|
||||
llvm::Expected<ir::Func *> getWrapperFunc(const std::string &name,
|
||||
const std::vector<std::string> &types);
|
||||
JITResult executePython(const std::string &name,
|
||||
const std::vector<std::string> &types, void *arg);
|
||||
JITResult executeSafe(const std::string &code);
|
||||
const std::vector<std::string> &types,
|
||||
const std::string &pyModule,
|
||||
const std::vector<std::string> &pyVars, void *arg,
|
||||
bool debug);
|
||||
JITResult executeSafe(const std::string &code, const std::string &file, int line,
|
||||
bool debug);
|
||||
|
||||
// Errors
|
||||
llvm::Error handleJITError(const JITError &e);
|
||||
|
|
|
|||
|
|
@ -613,6 +613,13 @@ ExprPtr TypecheckVisitor::transformBinaryMagic(BinaryExpr *expr) {
|
|||
auto rt = expr->rexpr->getType()->getClass();
|
||||
seqassert(lt && rt, "lhs and rhs types not known");
|
||||
|
||||
if (!lt->is("pyobj") && rt->is("pyobj")) {
|
||||
// Special case: `obj op pyobj` -> `rhs.__rmagic__(lhs)` on lhs
|
||||
// Assumes that pyobj implements all left and right magics
|
||||
return transform(N<CallExpr>(N<DotExpr>(expr->rexpr, format("__{}__", rightMagic)),
|
||||
expr->lexpr));
|
||||
}
|
||||
|
||||
// Normal operations: check if `lhs.__magic__(lhs, rhs)` exists
|
||||
auto method = findBestMethod(lt, format("__{}__", magic), {lt, rt});
|
||||
|
||||
|
|
@ -626,16 +633,6 @@ ExprPtr TypecheckVisitor::transformBinaryMagic(BinaryExpr *expr) {
|
|||
// Normal case: `__magic__(lhs, rhs)`
|
||||
return transform(
|
||||
N<CallExpr>(N<IdExpr>(method->ast->name), expr->lexpr, expr->rexpr));
|
||||
} else if (lt->is("pyobj")) {
|
||||
// Special case: call `pyobj._getattr(magic)` on lhs
|
||||
return transform(N<CallExpr>(N<CallExpr>(N<DotExpr>(expr->lexpr, "_getattr"),
|
||||
N<StringExpr>(format("__{}__", magic))),
|
||||
expr->rexpr));
|
||||
} else if (rt->is("pyobj")) {
|
||||
// Special case: call `pyobj._getattr(magic)` on rhs
|
||||
return transform(N<CallExpr>(N<CallExpr>(N<DotExpr>(expr->rexpr, "_getattr"),
|
||||
N<StringExpr>(format("__r{}__", magic))),
|
||||
expr->lexpr));
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -250,7 +250,8 @@ bool TypecheckVisitor::wrapExpr(ExprPtr &expr, const TypePtr &expectedType,
|
|||
if (callee && expr->isType())
|
||||
expr = transform(N<CallExpr>(expr, N<EllipsisExpr>()));
|
||||
|
||||
std::unordered_set<std::string> hints = {"Generator", "float", TYPE_OPTIONAL};
|
||||
std::unordered_set<std::string> hints = {"Generator", "float", TYPE_OPTIONAL,
|
||||
"pyobj"};
|
||||
if (!exprClass && expectedClass && in(hints, expectedClass->name)) {
|
||||
return false; // argument type not yet known.
|
||||
} else if (expectedClass && expectedClass->name == "Generator" &&
|
||||
|
|
@ -267,6 +268,16 @@ bool TypecheckVisitor::wrapExpr(ExprPtr &expr, const TypePtr &expectedType,
|
|||
exprClass->name == TYPE_OPTIONAL &&
|
||||
exprClass->name != expectedClass->name) { // unwrap optional
|
||||
expr = transform(N<CallExpr>(N<IdExpr>(FN_UNWRAP), expr));
|
||||
} else if (expectedClass && expectedClass->name == "pyobj" &&
|
||||
exprClass->name != expectedClass->name) { // wrap to pyobj
|
||||
expr = transform(
|
||||
N<CallExpr>(N<IdExpr>("pyobj"), N<CallExpr>(N<DotExpr>(expr, "__to_py__"))));
|
||||
} else if (allowUnwrap && expectedClass && exprClass && exprClass->name == "pyobj" &&
|
||||
exprClass->name != expectedClass->name) { // unwrap pyobj
|
||||
auto texpr = N<IdExpr>(expectedClass->name);
|
||||
texpr->setType(expectedType);
|
||||
expr =
|
||||
transform(N<CallExpr>(N<DotExpr>(texpr, "__from_py__"), N<DotExpr>(expr, "p")));
|
||||
} else if (callee && exprClass && expr->type->getFunc() &&
|
||||
!(expectedClass && expectedClass->name == "Function")) {
|
||||
// Case 7: wrap raw Seq functions into Partial(...) call for easy realization.
|
||||
|
|
|
|||
|
|
@ -108,13 +108,68 @@ class Foo:
|
|||
def total(self):
|
||||
return self.a + self.b + self.c
|
||||
|
||||
print(Foo(10).total())
|
||||
print(Foo(10).total()) # 1110
|
||||
```
|
||||
|
||||
`@codon.convert` requires the annotated class to specify `__slots__`, which
|
||||
it uses to construct a generic Codon class (specifically, a named tuple) to
|
||||
store the class's converted fields.
|
||||
|
||||
# Passing globals to Codon
|
||||
|
||||
Global variables, functions or modules can be passed to JIT'd functions through
|
||||
the `pyvars` argument to `@codon.jit`:
|
||||
|
||||
``` python
|
||||
import codon
|
||||
|
||||
def foo(n):
|
||||
print(f'n is {n}')
|
||||
|
||||
@codon.jit(pyvars=['foo'])
|
||||
def bar(n):
|
||||
foo(n) # calls the Python function 'foo'
|
||||
return n ** 2
|
||||
|
||||
print(bar(9)) # 'n is 9' then '81'
|
||||
```
|
||||
|
||||
This also allows imported Python modules to be accessed by Codon. All `pyvars`
|
||||
are passed as Python objects. Note that JIT'd functions can call each other
|
||||
by default.
|
||||
|
||||
# Debugging
|
||||
|
||||
`@codon.jit` takes an optional `debug` parameter that can be used to print debug
|
||||
information such as generated Codon functions and data types:
|
||||
|
||||
``` python
|
||||
import codon
|
||||
|
||||
@codon.jit(debug=True)
|
||||
def sum_of_squares(v):
|
||||
return sum(i**2 for i in v)
|
||||
|
||||
print(sum_of_squares([1.4, 2.9, 3.14]))
|
||||
```
|
||||
|
||||
outputs:
|
||||
|
||||
```
|
||||
[codon::jit::execute] code:
|
||||
def sum_of_squares(v):
|
||||
return sum(i**2 for i in v)
|
||||
-----
|
||||
[python] sum_of_squares(['List[float]'])
|
||||
[codon::jit::executePython] wrapper:
|
||||
@export
|
||||
def __codon_wrapped__sum_of_squares_0(args: cobj) -> cobj:
|
||||
a0 = List[float].__from_py__(PyTuple_GetItem(args, 0))
|
||||
return sum_of_squares(a0).__to_py__()
|
||||
-----
|
||||
20.229599999999998
|
||||
```
|
||||
|
||||
# Internals and performance tips
|
||||
|
||||
Under the hood, the `codon` module maintains an instance of the Codon JIT,
|
||||
|
|
@ -127,4 +182,6 @@ but instead reuses the cached function pointer.
|
|||
|
||||
Although object conversions from Python to Codon are generally cheap, they do
|
||||
impose a small overhead, meaning **`@codon.jit` will work best on expensive and/or
|
||||
long-running operations** rather than short-lived operations.
|
||||
long-running operations** rather than short-lived operations. By the same token,
|
||||
**the more work can be done in Codon, the better,** as opposed to repeatedly
|
||||
transferring back and forth.
|
||||
|
|
|
|||
|
|
@ -96,7 +96,8 @@ for a in some_arbitrary_generator():
|
|||
|
||||
Included revamped `codon` module for Python, with `@codon.jit` decorator
|
||||
for compiling Python code in existing codebases. Further improved and
|
||||
optimized the Python bridge.
|
||||
optimized the Python bridge. Please see the [docs](../interop/decorator.md)
|
||||
for more information.
|
||||
|
||||
## Codon IR
|
||||
|
||||
|
|
|
|||
|
|
@ -1,2 +1,3 @@
|
|||
dist
|
||||
src/stdlib
|
||||
codon/stdlib
|
||||
codon/jit.cpp
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
recursive-include src/stdlib *.codon
|
||||
recursive-include src *.so
|
||||
recursive-include src *.dylib
|
||||
recursive-include codon/stdlib *.codon
|
||||
recursive-include codon *.so
|
||||
recursive-include codon *.dylib
|
||||
|
|
|
|||
|
|
@ -7,8 +7,8 @@ $ pip install extra/python
|
|||
To use:
|
||||
|
||||
```python
|
||||
from codon import codon, JitError
|
||||
import codon
|
||||
|
||||
@codon
|
||||
@codon.jit
|
||||
def ...
|
||||
```
|
||||
|
|
|
|||
|
|
@ -0,0 +1,176 @@
|
|||
from argparse import ArgumentError
|
||||
import ctypes
|
||||
import inspect
|
||||
import sys
|
||||
import os
|
||||
import functools
|
||||
import itertools
|
||||
import ast
|
||||
import astunparse
|
||||
|
||||
sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)
|
||||
|
||||
if "CODON_PATH" not in os.environ:
|
||||
os.environ["CODON_PATH"] = os.path.dirname(
|
||||
os.path.abspath(inspect.getfile(inspect.currentframe()))
|
||||
)
|
||||
os.environ["CODON_PATH"] += "/stdlib"
|
||||
|
||||
from .codon_jit import JITWrapper, JITError
|
||||
|
||||
pod_conversions = {type(None): "pyobj",
|
||||
int: "int",
|
||||
float: "float",
|
||||
bool: "bool",
|
||||
str: "str",
|
||||
complex: "complex",
|
||||
slice: "slice"}
|
||||
|
||||
custom_conversions = {}
|
||||
_error_msgs = set()
|
||||
|
||||
def _common_type(t, debug, sample_size):
|
||||
sub, is_optional = None, False
|
||||
for i in itertools.islice(t, sample_size):
|
||||
if i is None:
|
||||
is_optional = True
|
||||
else:
|
||||
s = _codon_type(i, debug=debug, sample_size=sample_size)
|
||||
if sub and sub != s:
|
||||
return "pyobj"
|
||||
sub = s
|
||||
if is_optional and sub and sub != "pyobj":
|
||||
sub = f"Optional[{sub}]"
|
||||
return sub if sub else "pyobj"
|
||||
|
||||
def _codon_type(arg, **kwargs):
|
||||
t = type(arg)
|
||||
|
||||
s = pod_conversions.get(t, "")
|
||||
if s:
|
||||
return s
|
||||
if issubclass(t, list):
|
||||
return f"List[{_common_type(arg, **kwargs)}]"
|
||||
if issubclass(t, set):
|
||||
return f"Set[{_common_type(arg, **kwargs)}]"
|
||||
if issubclass(t, dict):
|
||||
return f"Dict[{_common_type(arg.keys(), **kwargs)},{_common_type(arg.values(), **kwargs)}]"
|
||||
if issubclass(t, tuple):
|
||||
return f"Tuple[{','.join(_codon_type(a, **kwargs) for a in arg)}]"
|
||||
s = custom_conversions.get(t, "")
|
||||
if s:
|
||||
j = ','.join(_codon_type(getattr(arg, slot), **kwargs) for slot in t.__slots__)
|
||||
return f"{s}[{j}]"
|
||||
|
||||
debug = kwargs.get('debug', None)
|
||||
if debug:
|
||||
msg = f"cannot convert {t.__name__}"
|
||||
if msg not in _error_msgs:
|
||||
print(f"[python] {msg}", file=sys.stderr)
|
||||
_error_msgs.add(msg)
|
||||
return "pyobj"
|
||||
|
||||
def _codon_types(args, **kwargs):
|
||||
return tuple(_codon_type(arg, **kwargs) for arg in args)
|
||||
|
||||
def _reset_jit():
|
||||
global _jit
|
||||
_jit = JITWrapper()
|
||||
init_code = ("from internal.python import "
|
||||
"setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n"
|
||||
"setup_decorator()\n")
|
||||
_jit.execute(init_code, "", 0, False)
|
||||
return _jit
|
||||
|
||||
_jit = _reset_jit()
|
||||
|
||||
class RewriteFunctionArgs(ast.NodeTransformer):
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
for a in self.args:
|
||||
node.args.args.append(ast.arg(arg=a, annotation=None))
|
||||
return node
|
||||
|
||||
def _obj_to_str(obj, **kwargs) -> str:
|
||||
if inspect.isclass(obj):
|
||||
lines = inspect.getsourcelines(obj)[0]
|
||||
extra_spaces = lines[0].find("class")
|
||||
obj_str = "".join(l[extra_spaces:] for l in lines)
|
||||
elif callable(obj):
|
||||
lines = inspect.getsourcelines(obj)[0]
|
||||
extra_spaces = lines[0].find("@")
|
||||
obj_str = "".join(l[extra_spaces:] for l in lines[1:])
|
||||
if kwargs.get('pyvars', None):
|
||||
node = ast.fix_missing_locations(
|
||||
RewriteFunctionArgs(kwargs['pyvars']).visit(ast.parse(obj_str))
|
||||
)
|
||||
obj_str = astunparse.unparse(node)
|
||||
else:
|
||||
raise TypeError(f"Function or class expected, got {type(obj).__name__}.")
|
||||
return obj_str.replace("_@par", "@par")
|
||||
|
||||
def _obj_name(obj) -> str:
|
||||
if inspect.isclass(obj) or callable(obj):
|
||||
return obj.__name__
|
||||
else:
|
||||
raise TypeError(f"Function or class expected, got {type(obj).__name__}.")
|
||||
|
||||
def _parse_decorated(obj, **kwargs):
|
||||
return _obj_name(obj), _obj_to_str(obj, **kwargs)
|
||||
|
||||
def convert(t):
|
||||
if not hasattr(t, "__slots__"):
|
||||
raise JITError(f"class '{str(t)}' does not have '__slots__' attribute")
|
||||
|
||||
name = t.__name__
|
||||
slots = t.__slots__
|
||||
code = ("@tuple\n"
|
||||
"class " + name + "[" + ",".join(f"T{i}" for i in range(len(slots))) + "]:\n")
|
||||
for i, slot in enumerate(slots):
|
||||
code += f" {slot}: T{i}\n"
|
||||
|
||||
# PyObject_GetAttrString
|
||||
code += " def __from_py__(p: cobj):\n"
|
||||
for i, slot in enumerate(slots):
|
||||
code += f" a{i} = T{i}.__from_py__(PyObject_GetAttrString(p, '{slot}'.ptr))\n"
|
||||
code += f" return {name}({', '.join(f'a{i}' for i in range(len(slots)))})\n"
|
||||
|
||||
_jit.execute(code, "", 0, False)
|
||||
custom_conversions[t] = name
|
||||
return t
|
||||
|
||||
def jit(fn=None, debug=None, sample_size=5, pyvars=None):
|
||||
if not pyvars:
|
||||
pyvars = []
|
||||
if not isinstance(pyvars, list):
|
||||
raise ArgumentError("pyvars must be a list")
|
||||
def _decorate(f):
|
||||
try:
|
||||
obj_name, obj_str = _parse_decorated(f, pyvars=pyvars)
|
||||
_jit.execute(
|
||||
obj_str,
|
||||
f.__code__.co_filename,
|
||||
f.__code__.co_firstlineno,
|
||||
1 if debug else 0
|
||||
)
|
||||
except JITError:
|
||||
_reset_jit()
|
||||
raise
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrapped(*args, **kwargs):
|
||||
try:
|
||||
args = (*args, *kwargs.values())
|
||||
types = _codon_types(args, debug=debug, sample_size=sample_size)
|
||||
if debug:
|
||||
print(f"[python] {f.__name__}({list(types)})", file=sys.stderr)
|
||||
return _jit.run_wrapper(obj_name, types, f.__module__, pyvars, args, 1 if debug else 0)
|
||||
except JITError:
|
||||
_reset_jit()
|
||||
raise
|
||||
return wrapped
|
||||
if fn:
|
||||
return _decorate(fn)
|
||||
return _decorate
|
||||
|
|
@ -14,5 +14,5 @@ cdef extern from "codon/compiler/jit.h" namespace "codon::jit":
|
|||
cdef cppclass JIT:
|
||||
JIT(string)
|
||||
Error init()
|
||||
JITResult executeSafe(string)
|
||||
JITResult executePython(string, vector[string], object)
|
||||
JITResult executeSafe(string, string, int, char)
|
||||
JITResult executePython(string, vector[string], string, vector[string], object, char)
|
||||
|
|
@ -7,7 +7,7 @@ from cython.operator import dereference as dref
|
|||
from libcpp.string cimport string
|
||||
from libcpp.vector cimport vector
|
||||
|
||||
from src.jit cimport JIT, JITResult
|
||||
from codon.jit cimport JIT, JITResult
|
||||
|
||||
|
||||
class JITError(Exception):
|
||||
|
|
@ -24,16 +24,17 @@ cdef class JITWrapper:
|
|||
def __dealloc__(self):
|
||||
del self.jit
|
||||
|
||||
def execute(self, code: str) -> str:
|
||||
result = dref(self.jit).executeSafe(code)
|
||||
def execute(self, code: str, filename: str, fileno: int, debug: char) -> str:
|
||||
result = dref(self.jit).executeSafe(code, filename, fileno, <char>debug)
|
||||
if <bint>result:
|
||||
return None
|
||||
else:
|
||||
raise JITError(result.message)
|
||||
|
||||
def run_wrapper(self, name: str, types: list[str], args) -> object:
|
||||
def run_wrapper(self, name: str, types: list[str], module: str, pyvars: list[str], args, debug: char) -> object:
|
||||
cdef vector[string] types_vec = types
|
||||
result = dref(self.jit).executePython(name, types_vec, <object>args)
|
||||
cdef vector[string] pyvars_vec = pyvars
|
||||
result = dref(self.jit).executePython(name, types_vec, module, pyvars_vec, <object>args, <char>debug)
|
||||
if <bint>result:
|
||||
return <object>result.result
|
||||
else:
|
||||
|
|
@ -49,11 +49,27 @@ codon_dir = Path(os.environ.get("CODON_DIR", from_root("build")))
|
|||
codon_include_dir = os.environ.get("CODON_INCLUDE_DIR", codon_dir / "include")
|
||||
ext = "dylib" if sys.platform == "darwin" else "so"
|
||||
|
||||
def symlink(target, dest):
|
||||
tmp = "_tmp"
|
||||
os.symlink(str(target.resolve()), tmp)
|
||||
os.rename(tmp, str(dest))
|
||||
root = Path(os.path.dirname(os.path.realpath(__file__)))
|
||||
distutils.dir_util.copy_tree(str(codon_dir / ".." / "stdlib"), str(root / "src" / "stdlib"))
|
||||
shutil.copy(codon_dir / "lib" / "codon" / ("libcodonc." + ext), root / "src")
|
||||
shutil.copy(codon_dir / "lib" / "codon" / ("libcodonrt." + ext), root / "src")
|
||||
shutil.copy(codon_dir / "lib" / "codon" / ("libomp." + ext), root / "src")
|
||||
symlink(
|
||||
codon_dir / ".." / "stdlib",
|
||||
root / "codon" / "stdlib"
|
||||
)
|
||||
symlink(
|
||||
codon_dir / "lib" / "codon" / ("libcodonc." + ext),
|
||||
root / "codon" / ("libcodonc." + ext)
|
||||
)
|
||||
symlink(
|
||||
codon_dir / "lib" / "codon" / ("libcodonrt." + ext),
|
||||
root / "codon" / ("libcodonrt." + ext)
|
||||
)
|
||||
symlink(
|
||||
codon_dir / "lib" / "codon" / ("libomp." + ext),
|
||||
root / "codon" / ("libomp." + ext)
|
||||
)
|
||||
|
||||
print(f"<llvm> {llvm_include_dir}, {llvm_lib_dir}")
|
||||
print(f"<codon> {codon_include_dir}")
|
||||
|
|
@ -65,18 +81,19 @@ else:
|
|||
|
||||
jit_extension = Extension(
|
||||
"codon.codon_jit",
|
||||
sources=["src/jit.pyx"],
|
||||
sources=["codon/jit.pyx"],
|
||||
libraries=["codonc", "codonrt"],
|
||||
language="c++",
|
||||
extra_compile_args=["-w", "-std=c++17"],
|
||||
extra_link_args=[linker_args],
|
||||
include_dirs=[llvm_include_dir, str(codon_include_dir)],
|
||||
library_dirs=[llvm_lib_dir, str(root / "src")],
|
||||
library_dirs=[llvm_lib_dir, str(root / "codon")],
|
||||
)
|
||||
|
||||
setup(
|
||||
name="codon",
|
||||
version=CODON_VERSION,
|
||||
install_requires=["astunparse"],
|
||||
python_requires='>=3.6',
|
||||
description="Codon JIT decorator",
|
||||
url="https://exaloop.io",
|
||||
|
|
@ -87,6 +104,5 @@ setup(
|
|||
cmdclass={"build_ext": build_ext},
|
||||
ext_modules=[jit_extension],
|
||||
packages=["codon"],
|
||||
package_dir={"codon": "src"},
|
||||
include_package_data=True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,141 +0,0 @@
|
|||
import ctypes
|
||||
import inspect
|
||||
import importlib
|
||||
import importlib.util
|
||||
import sys
|
||||
import os
|
||||
|
||||
sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)
|
||||
|
||||
if "CODON_PATH" not in os.environ:
|
||||
os.environ["CODON_PATH"] = os.path.dirname(
|
||||
os.path.abspath(inspect.getfile(inspect.currentframe()))
|
||||
)
|
||||
os.environ["CODON_PATH"] += "/stdlib"
|
||||
|
||||
from .codon_jit import JITWrapper, JITError
|
||||
|
||||
pod_conversions = {type(None): "NoneType",
|
||||
int: "int",
|
||||
float: "float",
|
||||
bool: "bool",
|
||||
str: "str",
|
||||
complex: "complex",
|
||||
slice: "slice"}
|
||||
|
||||
custom_conversions = {}
|
||||
|
||||
def _codon_type(arg):
|
||||
t = type(arg)
|
||||
|
||||
s = pod_conversions.get(t, "")
|
||||
if s:
|
||||
return s
|
||||
|
||||
if issubclass(t, list):
|
||||
sub = "NoneType"
|
||||
x = next(iter(arg), None)
|
||||
if x is not None:
|
||||
sub = _codon_type(x)
|
||||
return f"List[{sub}]"
|
||||
|
||||
if issubclass(t, set):
|
||||
sub = "NoneType"
|
||||
x = next(iter(arg), None)
|
||||
if x is not None:
|
||||
sub = _codon_type(x)
|
||||
return f"Set[{sub}]"
|
||||
|
||||
if issubclass(t, dict):
|
||||
sub1 = "NoneType"
|
||||
sub2 = "NoneType"
|
||||
x = next(iter(arg.items()), None)
|
||||
if x is not None:
|
||||
sub1 = _codon_type(x[0])
|
||||
sub2 = _codon_type(x[1])
|
||||
return f"Dict[{sub1},{sub2}]"
|
||||
|
||||
if issubclass(t, tuple):
|
||||
return f"Tuple[{','.join(_codon_type(a) for a in arg)}]"
|
||||
|
||||
s = custom_conversions.get(t, "")
|
||||
if s:
|
||||
return f"{s}[{','.join(_codon_type(getattr(arg, slot)) for slot in t.__slots__)}]"
|
||||
|
||||
return "pyobj"
|
||||
|
||||
def _codon_types(args):
|
||||
return tuple(_codon_type(arg) for arg in args)
|
||||
|
||||
def _reset_jit():
|
||||
global _jit
|
||||
_jit = JITWrapper()
|
||||
init_code = ("from internal.python import "
|
||||
"setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n"
|
||||
"setup_decorator()\n")
|
||||
_jit.execute(init_code)
|
||||
return _jit
|
||||
|
||||
_jit = _reset_jit()
|
||||
|
||||
def _obj_to_str(obj) -> str:
|
||||
if inspect.isclass(obj):
|
||||
lines = inspect.getsourcelines(obj)[0]
|
||||
extra_spaces = lines[0].find("class")
|
||||
obj_str = "".join(l[extra_spaces:] for l in lines)
|
||||
elif callable(obj):
|
||||
lines = inspect.getsourcelines(obj)[0]
|
||||
extra_spaces = lines[0].find("@")
|
||||
obj_str = "".join(l[extra_spaces:] for l in lines[1:])
|
||||
else:
|
||||
raise TypeError(f"Function or class expected, got {type(obj).__name__}.")
|
||||
return obj_str.replace("_@par", "@par")
|
||||
|
||||
def _obj_name(obj) -> str:
|
||||
if inspect.isclass(obj) or callable(obj):
|
||||
return obj.__name__
|
||||
else:
|
||||
raise TypeError(f"Function or class expected, got {type(obj).__name__}.")
|
||||
|
||||
def _parse_decorated(obj):
|
||||
return _obj_name(obj), _obj_to_str(obj)
|
||||
|
||||
def convert(t):
|
||||
if not hasattr(t, "__slots__"):
|
||||
raise JITError(f"class '{str(t)}' does not have '__slots__' attribute")
|
||||
|
||||
name = t.__name__
|
||||
slots = t.__slots__
|
||||
code = ("@tuple\n"
|
||||
"class " + name + "[" + ",".join(f"T{i}" for i in range(len(slots))) + "]:\n")
|
||||
for i, slot in enumerate(slots):
|
||||
code += f" {slot}: T{i}\n"
|
||||
|
||||
# PyObject_GetAttrString
|
||||
code += " def __from_py__(p: cobj):\n"
|
||||
for i, slot in enumerate(slots):
|
||||
code += f" a{i} = T{i}.__from_py__(PyObject_GetAttrString(p, '{slot}'.ptr))\n"
|
||||
code += f" return {name}({', '.join(f'a{i}' for i in range(len(slots)))})\n"
|
||||
|
||||
_jit.execute(code)
|
||||
custom_conversions[t] = name
|
||||
return t
|
||||
|
||||
def jit(obj):
|
||||
try:
|
||||
obj_name, obj_str = _parse_decorated(obj)
|
||||
_jit.execute(obj_str)
|
||||
except JITError as e:
|
||||
_reset_jit()
|
||||
raise
|
||||
|
||||
def wrapped(*args, **kwargs):
|
||||
try:
|
||||
args = (*args, *kwargs.values())
|
||||
types = _codon_types(args)
|
||||
return _jit.run_wrapper(obj_name, types, args)
|
||||
except JITError as e:
|
||||
_reset_jit()
|
||||
raise
|
||||
|
||||
return wrapped
|
||||
|
|
@ -9,9 +9,11 @@ Py_DecRef = Function[[cobj], NoneType](cobj())
|
|||
Py_IncRef = Function[[cobj], NoneType](cobj())
|
||||
Py_Initialize = Function[[], NoneType](cobj())
|
||||
PyImport_AddModule = Function[[cobj], cobj](cobj())
|
||||
PyImport_AddModuleObject = Function[[cobj], cobj](cobj())
|
||||
PyImport_ImportModule = Function[[cobj], cobj](cobj())
|
||||
PyErr_Fetch = Function[[Ptr[cobj], Ptr[cobj], Ptr[cobj]], NoneType](cobj())
|
||||
PyRun_SimpleString = Function[[cobj], NoneType](cobj())
|
||||
PyEval_GetGlobals = Function[[], cobj](cobj())
|
||||
|
||||
# conversions
|
||||
PyLong_AsLong = Function[[cobj], int](cobj())
|
||||
|
|
@ -25,6 +27,7 @@ PyList_GetItem = Function[[cobj, int], cobj](cobj())
|
|||
PyList_SetItem = Function[[cobj, int, cobj], cobj](cobj())
|
||||
PyDict_New = Function[[], cobj](cobj())
|
||||
PyDict_Next = Function[[cobj, Ptr[int], Ptr[cobj], Ptr[cobj]], int](cobj())
|
||||
PyDict_GetItem = Function[[cobj, cobj], cobj](cobj())
|
||||
PyDict_SetItem = Function[[cobj, cobj, cobj], cobj](cobj())
|
||||
PySet_Add = Function[[cobj, cobj], cobj](cobj())
|
||||
PySet_New = Function[[cobj], cobj](cobj())
|
||||
|
|
@ -148,9 +151,11 @@ def init_dl_handles(py_handle: cobj):
|
|||
global Py_IncRef
|
||||
global Py_Initialize
|
||||
global PyImport_AddModule
|
||||
global PyImport_AddModuleObject
|
||||
global PyImport_ImportModule
|
||||
global PyErr_Fetch
|
||||
global PyRun_SimpleString
|
||||
global PyEval_GetGlobals
|
||||
global PyLong_AsLong
|
||||
global PyLong_FromLong
|
||||
global PyFloat_AsDouble
|
||||
|
|
@ -162,6 +167,7 @@ def init_dl_handles(py_handle: cobj):
|
|||
global PyList_SetItem
|
||||
global PyDict_New
|
||||
global PyDict_Next
|
||||
global PyDict_GetItem
|
||||
global PyDict_SetItem
|
||||
global PySet_Add
|
||||
global PySet_New
|
||||
|
|
@ -235,9 +241,11 @@ def init_dl_handles(py_handle: cobj):
|
|||
Py_IncRef = dlsym(py_handle, "Py_IncRef")
|
||||
Py_Initialize = dlsym(py_handle, "Py_Initialize")
|
||||
PyImport_AddModule = dlsym(py_handle, "PyImport_AddModule")
|
||||
PyImport_AddModuleObject = dlsym(py_handle, "PyImport_AddModuleObject")
|
||||
PyImport_ImportModule = dlsym(py_handle, "PyImport_ImportModule")
|
||||
PyErr_Fetch = dlsym(py_handle, "PyErr_Fetch")
|
||||
PyRun_SimpleString = dlsym(py_handle, "PyRun_SimpleString")
|
||||
PyEval_GetGlobals = dlsym(py_handle, "PyEval_GetGlobals")
|
||||
PyLong_AsLong = dlsym(py_handle, "PyLong_AsLong")
|
||||
PyLong_FromLong = dlsym(py_handle, "PyLong_FromLong")
|
||||
PyFloat_AsDouble = dlsym(py_handle, "PyFloat_AsDouble")
|
||||
|
|
@ -249,6 +257,7 @@ def init_dl_handles(py_handle: cobj):
|
|||
PyList_SetItem = dlsym(py_handle, "PyList_SetItem")
|
||||
PyDict_New = dlsym(py_handle, "PyDict_New")
|
||||
PyDict_Next = dlsym(py_handle, "PyDict_Next")
|
||||
PyDict_GetItem = dlsym(py_handle, "PyDict_GetItem")
|
||||
PyDict_SetItem = dlsym(py_handle, "PyDict_SetItem")
|
||||
PySet_Add = dlsym(py_handle, "PySet_Add")
|
||||
PySet_New = dlsym(py_handle, "PySet_New")
|
||||
|
|
@ -579,6 +588,9 @@ class pyobj:
|
|||
pyobj.decref(it)
|
||||
pyobj.exc_check()
|
||||
|
||||
def to_str(self, errors: str, empty: str = "") -> str:
|
||||
return pyobj.to_str(self.p, errors, empty)
|
||||
|
||||
def to_str(p: cobj, errors: str, empty: str = "") -> str:
|
||||
obj = PyUnicode_AsEncodedString(p, "utf-8".c_str(), errors.c_str())
|
||||
if obj == cobj():
|
||||
|
|
@ -588,9 +600,6 @@ class pyobj:
|
|||
pyobj.decref(obj)
|
||||
return res
|
||||
|
||||
def to_str(self, errors: str, empty: str = "") -> str:
|
||||
return pyobj.to_str(self.p, errors, empty)
|
||||
|
||||
def exc_check():
|
||||
ptype, pvalue, ptraceback = cobj(), cobj(), cobj()
|
||||
PyErr_Fetch(__ptr__(ptype), __ptr__(pvalue), __ptr__(ptraceback))
|
||||
|
|
@ -657,8 +666,15 @@ class pyobj:
|
|||
ensure_initialized()
|
||||
PyRun_SimpleString(code.c_str())
|
||||
|
||||
def _globals() -> pyobj:
|
||||
return pyobj(PyEval_GetGlobals())
|
||||
|
||||
def _get_module(name: str) -> pyobj:
|
||||
p = pyobj(pyobj.exc_wrap(PyImport_AddModule(name.c_str())))
|
||||
return p
|
||||
|
||||
def _main_module() -> pyobj:
|
||||
return pyobj(pyobj.exc_wrap(PyImport_AddModule("__main__".c_str())))
|
||||
return pyobj._get_module("__main__")
|
||||
|
||||
def _repr_mimebundle_(self, bundle=Set[str]()) -> Dict[str, str]:
|
||||
fn = pyobj._main_module()._getattr("__codon_repr__")
|
||||
|
|
|
|||
|
|
@ -16,10 +16,7 @@ class float:
|
|||
def __new__() -> float:
|
||||
return 0.0
|
||||
|
||||
def __new__(what) -> float:
|
||||
return what.__float__()
|
||||
|
||||
def __new__(s: str) -> float:
|
||||
def _from_str(s: str) -> float:
|
||||
from C import strtod(cobj, Ptr[cobj]) -> float
|
||||
|
||||
buf = __array__[byte](32)
|
||||
|
|
@ -41,6 +38,13 @@ class float:
|
|||
|
||||
return result
|
||||
|
||||
def __new__(what) -> float:
|
||||
# do not overload! (needed to avoid pyobj conversion)
|
||||
if isinstance(what, str):
|
||||
return float._from_str(what)
|
||||
else:
|
||||
return what.__float__()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
s = seq_str_float(self)
|
||||
return s if s != "-nan" else "nan"
|
||||
|
|
|
|||
|
|
@ -23,15 +23,16 @@ class int:
|
|||
def __new__() -> int:
|
||||
ret i64 0
|
||||
|
||||
def __new__(what) -> int:
|
||||
return what.__int__()
|
||||
|
||||
def __new__(s: str) -> int:
|
||||
return int._from_str(s, 10)
|
||||
|
||||
def __new__(s: str, base: int) -> int:
|
||||
return int._from_str(s, base)
|
||||
|
||||
def __new__(what) -> int:
|
||||
# do not overload! (needed to avoid pyobj conversion)
|
||||
if isinstance(what, str):
|
||||
return int._from_str(what, 10)
|
||||
else:
|
||||
return what.__int__()
|
||||
|
||||
def __int__(self) -> int:
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -917,6 +917,23 @@ def foo():
|
|||
return str(pyfoo3())
|
||||
print foo() #: 3
|
||||
|
||||
#%% python_pyobj
|
||||
@python
|
||||
def foofn() -> Dict[pyobj, pyobj]:
|
||||
return {"str": "hai", "int": 1}
|
||||
|
||||
foo = foofn()
|
||||
print(foo, foo.__class__.__name__)
|
||||
#: {'str': 'hai', 'int': 1} Dict[pyobj,pyobj]
|
||||
foo["codon"] = 5.15
|
||||
print(foo["codon"], foo["codon"].__class__.__name__, foo.__class__.__name__)
|
||||
#: 5.15 pyobj Dict[pyobj,pyobj]
|
||||
|
||||
a = {1: "s", 2: "t"}
|
||||
a[3] = foo["str"]
|
||||
print(a)
|
||||
#: {1: 's', 2: 't', 3: 'hai'}
|
||||
|
||||
#%% typeof_definition_error,barebones
|
||||
a = 1
|
||||
class X:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,3 @@
|
|||
import sys
|
||||
from io import StringIO
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import codon
|
||||
|
|
@ -45,7 +43,7 @@ def test_roundtrip():
|
|||
def roundtrip(x):
|
||||
return x
|
||||
|
||||
for i in range(5):
|
||||
for _ in range(5):
|
||||
assert roundtrip(None) == None
|
||||
assert roundtrip(42) == 42
|
||||
assert roundtrip(3.14) == 3.14
|
||||
|
|
@ -113,3 +111,43 @@ test_roundtrip()
|
|||
test_return_type()
|
||||
test_param_types()
|
||||
test_error_handling()
|
||||
|
||||
|
||||
@codon.jit
|
||||
def foo(y):
|
||||
return f"{y.__class__.__name__}; {y}"
|
||||
|
||||
|
||||
@codon.jit(debug=True)
|
||||
def foo2(y):
|
||||
return f"{y.__class__.__name__}; {y}"
|
||||
|
||||
class Foo:
|
||||
def __init__(self):
|
||||
self.x = 1
|
||||
|
||||
@codon.jit
|
||||
def a(x):
|
||||
return x+1
|
||||
|
||||
def b(x, z):
|
||||
y = a(x)
|
||||
return y * z
|
||||
|
||||
@codon.jit(pyvars=['b'])
|
||||
def c(x, y):
|
||||
n = b(x,y) ** a(1)
|
||||
return n
|
||||
|
||||
def test_cross_calls():
|
||||
assert foo([None, 1]) == "List[Optional[int]]; [None, 1]"
|
||||
assert foo([1, None, 1]) == "List[Optional[int]]; [1, None, 1]"
|
||||
assert foo([1, None, 1.2]) == "List[pyobj]; [1, None, 1.2]"
|
||||
assert foo({None: 1}) == "Dict[pyobj,int]; {None: 1}"
|
||||
assert foo2([None, Foo()]).startswith("List[pyobj]; [None, <__main__.Foo object at")
|
||||
|
||||
assert a(3) == 4
|
||||
assert b(3, 4) == 16
|
||||
assert round(c(5, 6.1), 2) == 1339.56
|
||||
|
||||
test_cross_calls()
|
||||
|
|
|
|||
Loading…
Reference in New Issue