fastpython-native 0.1.0__py3-none-any.whl

fastpython/codegen/transpiler.py

@@ -0,0 +1,545 @@
+"""
+AST Transpiler for FastPython.
+Translates Python numerical functions to ultra-high-performance C99/C++ code
+optimized for execution on the host CPU with AVX2 SIMD and OpenMP multi-threading.
+"""
+
+import ast
+import inspect
+import textwrap
+from typing import Dict, List, Tuple, Optional, Any, Set
+
+
+TYPE_MAP = {
+    int: "int64_t",
+    float: "double",
+    bool: "bool",
+    "int": "int64_t",
+    "float": "double",
+    "bool": "bool",
+    "int64_t": "int64_t",
+    "double": "double",
+    "list": "double* __restrict",
+    "buffer": "double* __restrict",
+    "array": "double* __restrict",
+    "list[float]": "double* __restrict",
+    "list[int]": "int64_t* __restrict",
+}
+
+MATH_FUNCS = {
+    "sqrt": "sqrt",
+    "sin": "sin",
+    "cos": "cos",
+    "tan": "tan",
+    "exp": "exp",
+    "log": "log",
+    "log2": "log2",
+    "log10": "log10",
+    "fabs": "fabs",
+    "floor": "floor",
+    "ceil": "ceil",
+    "pow": "pow",
+    "hypot": "hypot",
+    "atan": "atan",
+    "atan2": "atan2",
+    "sinh": "sinh",
+    "cosh": "cosh",
+    "tanh": "tanh",
+}
+
+
+class TranspileError(Exception):
+    pass
+
+
+class PythonToCTranspiler(ast.NodeVisitor):
+    def __init__(self, func_name: str, type_hints: Dict[str, str], parallel: bool = False):
+        self.func_name = func_name
+        self.type_hints = type_hints.copy()
+        self.parallel = parallel
+        self.indent_level = 1
+        self.declared_vars: Dict[str, str] = {}  # var_name -> c_type
+        self.loop_depth = 0
+        self.code_lines: List[str] = []
+        self.return_type = self.type_hints.get("return", "double")
+        self.used_math_funcs: Set[str] = set()
+
+    def indent(self) -> str:
+        return "    " * self.indent_level
+
+    def emit(self, line: str):
+        self.code_lines.append(self.indent() + line)
+
+    def transpile_function(self, node: ast.FunctionDef) -> Tuple[str, Dict[str, str], str]:
+        # Identify arguments and assign C types
+        args_c_signature = []
+        for arg in node.args.args:
+            arg_name = arg.arg
+            if arg_name in self.type_hints:
+                c_type = self.type_hints[arg_name]
+            elif arg.annotation:
+                c_type = self._resolve_type_annotation(arg.annotation)
+            else:
+                c_type = "double"
+
+            self.type_hints[arg_name] = c_type
+            self.declared_vars[arg_name] = c_type
+            args_c_signature.append(f"{c_type} {arg_name}")
+
+        # Check return annotation
+        if node.returns:
+            self.return_type = self._resolve_type_annotation(node.returns)
+
+        # Process function body
+        for stmt in node.body:
+            self.visit(stmt)
+
+        body_code = "\n".join(self.code_lines)
+        sig = f"{self.return_type} kernel_{self.func_name}({', '.join(args_c_signature)})"
+        c_func = f"{sig} {{\n{body_code}\n}}"
+        return c_func, self.type_hints, self.return_type
+
+    def _resolve_type_annotation(self, node: ast.AST) -> str:
+        if isinstance(node, ast.Name):
+            return TYPE_MAP.get(node.id, "double")
+        elif isinstance(node, ast.Constant):
+            return TYPE_MAP.get(node.value, "double")
+        elif isinstance(node, ast.Subscript):
+            if isinstance(node.value, ast.Name) and node.value.id in ("list", "List", "Sequence", "array", "buffer"):
+                sub_t = self._resolve_type_annotation(node.slice)
+                base = "int64_t" if "int" in sub_t else "double"
+                return f"{base}* __restrict"
+        return "double"
+
+    def visit_Assign(self, node: ast.Assign):
+        if len(node.targets) != 1:
+            raise TranspileError("Multiple assignment targets not supported in fast kernel")
+
+        target = node.targets[0]
+        val_expr = self.visit(node.value)
+
+        if isinstance(target, ast.Name):
+            var_name = target.id
+            if var_name not in self.declared_vars:
+                inferred_type = self._infer_expr_type(node.value)
+                self.declared_vars[var_name] = inferred_type
+                self.emit(f"{inferred_type} {var_name} = {val_expr};")
+            else:
+                self.emit(f"{var_name} = {val_expr};")
+        elif isinstance(target, ast.Subscript):
+            target_expr = self.visit(target)
+            self.emit(f"{target_expr} = {val_expr};")
+        else:
+            raise TranspileError(f"Unsupported assignment target: {ast.dump(target)}")
+
+    def visit_AugAssign(self, node: ast.AugAssign):
+        target = self.visit(node.target)
+        val = self.visit(node.value)
+        op_map = {
+            ast.Add: "+=",
+            ast.Sub: "-=",
+            ast.Mult: "*=",
+            ast.Div: "/=",
+            ast.BitAnd: "&=",
+            ast.BitOr: "|=",
+            ast.BitXor: "^=",
+            ast.LShift: "<<=",
+            ast.RShift: ">>=",
+        }
+        op = op_map.get(type(node.op))
+        if not op:
+            raise TranspileError(f"Unsupported augmented operator: {type(node.op)}")
+        self.emit(f"{target} {op} {val};")
+
+    def _find_reductions(self, node: ast.For) -> List[str]:
+        reductions = []
+        for stmt in ast.walk(node):
+            if isinstance(stmt, ast.AugAssign) and isinstance(stmt.target, ast.Name):
+                var = stmt.target.id
+                if var in self.declared_vars and var not in reductions:
+                    reductions.append(var)
+        return reductions
+
+    def visit_For(self, node: ast.For):
+        self.loop_depth += 1
+        is_outer = (self.loop_depth == 1)
+
+        if isinstance(node.iter, ast.Call) and isinstance(node.iter.func, ast.Name) and node.iter.func.id == "range":
+            args = node.iter.args
+            var_name = node.target.id
+            self.declared_vars[var_name] = "int64_t"
+
+            if len(args) == 1:
+                start = "0"
+                stop = self.visit(args[0])
+                step = "1"
+            elif len(args) == 2:
+                start = self.visit(args[0])
+                stop = self.visit(args[1])
+                step = "1"
+            elif len(args) == 3:
+                start = self.visit(args[0])
+                stop = self.visit(args[1])
+                step = self.visit(args[2])
+            else:
+                raise TranspileError("range() with invalid arguments")
+
+            if is_outer and self.parallel:
+                reductions = self._find_reductions(node)
+                red_clause = ""
+                if reductions:
+                    red_clause = f" reduction(+:{', '.join(reductions)})"
+
+                self.emit(f"int {var_name};")
+                self.emit(f"#pragma omp parallel for schedule(guided){red_clause}")
+                step_expr = f"{var_name}++" if step == "1" else f"{var_name} += {step}"
+                self.emit(f"for ({var_name} = (int)({start}); {var_name} < (int)({stop}); {step_expr}) {{")
+            else:
+                # Add SIMD auto-vectorization hint for inner loops
+                self.emit("#pragma loop(ivdep)")
+                step_expr = f"{var_name}++" if step == "1" else f"{var_name} += {step}"
+                self.emit(f"for (int64_t {var_name} = {start}; {var_name} < {stop}; {step_expr}) {{")
+
+            self.indent_level += 1
+            for stmt in node.body:
+                self.visit(stmt)
+            self.indent_level -= 1
+            self.emit("}")
+        else:
+            raise TranspileError("Only range() iterations are currently supported in accelerated for-loops")
+
+        self.loop_depth -= 1
+
+    def visit_While(self, node: ast.While):
+        cond = self.visit(node.test)
+        self.emit(f"while ({cond}) {{")
+        self.indent_level += 1
+        for stmt in node.body:
+            self.visit(stmt)
+        self.indent_level -= 1
+        self.emit("}")
+
+    def visit_If(self, node: ast.If):
+        cond = self.visit(node.test)
+        self.emit(f"if ({cond}) {{")
+        self.indent_level += 1
+        for stmt in node.body:
+            self.visit(stmt)
+        self.indent_level -= 1
+
+        if node.orelse:
+            self.emit("} else {")
+            self.indent_level += 1
+            for stmt in node.orelse:
+                self.visit(stmt)
+            self.indent_level -= 1
+        self.emit("}")
+
+    def visit_Return(self, node: ast.Return):
+        if node.value is None:
+            self.emit("return;")
+        else:
+            val = self.visit(node.value)
+            self.emit(f"return ({self.return_type})({val});")
+
+    def visit_Break(self, node: ast.Break):
+        self.emit("break;")
+
+    def visit_Continue(self, node: ast.Continue):
+        self.emit("continue;")
+
+    def visit_Pass(self, node: ast.Pass):
+        pass
+
+    def visit_Subscript(self, node: ast.Subscript) -> str:
+        val = self.visit(node.value)
+        slc = self.visit(node.slice)
+        return f"{val}[{slc}]"
+
+    def visit_BinOp(self, node: ast.BinOp) -> str:
+        left = self.visit(node.left)
+        right = self.visit(node.right)
+        if isinstance(node.op, ast.Pow):
+            return f"pow((double)({left}), (double)({right}))"
+        elif isinstance(node.op, ast.FloorDiv):
+            return f"((int64_t)({left}) / (int64_t)({right}))"
+
+        op_symbols = {
+            ast.Add: "+",
+            ast.Sub: "-",
+            ast.Mult: "*",
+            ast.Div: "/",
+            ast.Mod: "%",
+            ast.BitAnd: "&",
+            ast.BitOr: "|",
+            ast.BitXor: "^",
+            ast.LShift: "<<",
+            ast.RShift: ">>",
+        }
+        sym = op_symbols.get(type(node.op), "+")
+        return f"({left} {sym} {right})"
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> str:
+        operand = self.visit(node.operand)
+        if isinstance(node.op, ast.USub):
+            return f"(-{operand})"
+        elif isinstance(node.op, ast.UAdd):
+            return f"(+{operand})"
+        elif isinstance(node.op, ast.Not):
+            return f"(!{operand})"
+        elif isinstance(node.op, ast.Invert):
+            return f"(~{operand})"
+        return operand
+
+    def visit_Compare(self, node: ast.Compare) -> str:
+        left = self.visit(node.left)
+        parts = []
+        curr_left = left
+        comp_ops = {
+            ast.Eq: "==",
+            ast.NotEq: "!=",
+            ast.Lt: "<",
+            ast.LtE: "<=",
+            ast.Gt: ">",
+            ast.GtE: ">=",
+        }
+        for op, comparator in zip(node.ops, node.comparators):
+            sym = comp_ops.get(type(op))
+            if not sym:
+                raise TranspileError(f"Unsupported comparison operator: {type(op)}")
+            right = self.visit(comparator)
+            parts.append(f"({curr_left} {sym} {right})")
+            curr_left = right
+        return " && ".join(parts)
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> str:
+        op = "&&" if isinstance(node.op, ast.And) else "||"
+        values = [self.visit(v) for v in node.values]
+        return f"({f' {op} '.join(values)})"
+
+    def visit_Name(self, node: ast.Name) -> str:
+        if node.id == "True":
+            return "true"
+        elif node.id == "False":
+            return "false"
+        return node.id
+
+    def visit_Constant(self, node: ast.Constant) -> str:
+        if isinstance(node.value, bool):
+            return "true" if node.value else "false"
+        elif isinstance(node.value, (int, float)):
+            return str(node.value)
+        elif isinstance(node.value, str):
+            return f'"{node.value}"'
+        return str(node.value)
+
+    def visit_Call(self, node: ast.Call) -> str:
+        func_name = None
+        if isinstance(node.func, ast.Attribute) and isinstance(node.func.value, ast.Name) and node.func.value.id == "math":
+            func_name = node.func.attr
+        elif isinstance(node.func, ast.Name):
+            func_name = node.func.id
+
+        if func_name in MATH_FUNCS:
+            c_name = MATH_FUNCS[func_name]
+            args = [self.visit(a) for a in node.args]
+            return f"{c_name}({', '.join(args)})"
+        elif func_name == "abs":
+            arg = self.visit(node.args[0])
+            return f"fabs((double)({arg}))"
+        elif func_name == "min":
+            a = self.visit(node.args[0])
+            b = self.visit(node.args[1])
+            return f"(({a}) < ({b}) ? ({a}) : ({b}))"
+        elif func_name == "max":
+            a = self.visit(node.args[0])
+            b = self.visit(node.args[1])
+            return f"(({a}) > ({b}) ? ({a}) : ({b}))"
+        elif func_name in ("int", "int64_t"):
+            arg = self.visit(node.args[0])
+            return f"((int64_t)({arg}))"
+        elif func_name in ("float", "double"):
+            arg = self.visit(node.args[0])
+            return f"((double)({arg}))"
+
+        raise TranspileError(f"Unsupported call: {func_name}()")
+
+    def _infer_expr_type(self, node: ast.AST) -> str:
+        if isinstance(node, ast.Constant):
+            if isinstance(node.value, bool):
+                return "bool"
+            elif isinstance(node.value, int):
+                return "int64_t"
+            elif isinstance(node.value, float):
+                return "double"
+        elif isinstance(node, ast.Name):
+            return self.declared_vars.get(node.id, "double")
+        elif isinstance(node, ast.BinOp):
+            left_t = self._infer_expr_type(node.left)
+            right_t = self._infer_expr_type(node.right)
+            if left_t == "double" or right_t == "double" or isinstance(node.op, ast.Div):
+                return "double"
+            return "int64_t"
+        elif isinstance(node, ast.UnaryOp):
+            return self._infer_expr_type(node.operand)
+        elif isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name) and node.func.id in ("int", "len"):
+                return "int64_t"
+            return "double"
+        elif isinstance(node, ast.Subscript):
+            if isinstance(node.value, ast.Name):
+                arr_t = self.declared_vars.get(node.value.id, "")
+                if "int64_t" in arr_t:
+                    return "int64_t"
+            return "double"
+        return "double"
+
+
+def generate_c_extension_source(
+    func: Any,
+    module_name: Optional[str] = None,
+    type_hints: Optional[Dict[str, Any]] = None,
+    parallel: bool = False,
+    fastmath: bool = True,
+    source_code: Optional[str] = None,
+    ast_node: Optional[ast.FunctionDef] = None
+) -> str:
+    if ast_node is not None:
+        func_def = ast_node
+    else:
+        if source_code is None:
+            source_code = inspect.getsource(func)
+        dedented = textwrap.dedent(source_code)
+        tree = ast.parse(dedented)
+
+        func_def = None
+        for item in tree.body:
+            if isinstance(item, ast.FunctionDef):
+                func_def = item
+                break
+
+    if not func_def:
+        raise TranspileError("Could not find function definition in source code")
+
+    if module_name is None:
+        module_name = f"native_{func_def.name}"
+
+    hints = {}
+    if type_hints:
+        for k, v in type_hints.items():
+            if v in TYPE_MAP:
+                hints[k] = TYPE_MAP[v]
+            elif isinstance(v, str):
+                if v in TYPE_MAP:
+                    hints[k] = TYPE_MAP[v]
+                elif "list" in v or "array" in v:
+                    hints[k] = "int64_t* __restrict" if "int" in v else "double* __restrict"
+
+    transpiler = PythonToCTranspiler(func_def.name, hints, parallel=parallel)
+    kernel_code, resolved_hints, ret_type = transpiler.transpile_function(func_def)
+
+    format_chars = []
+    c_var_decls = []
+    c_var_names = []
+    buffer_pre = []
+    buffer_post = []
+
+    for arg in func_def.args.args:
+        name = arg.arg
+        t = resolved_hints.get(name, "double")
+        if t == "int64_t":
+            format_chars.append("L")
+            c_var_decls.append(f"    long long {name};")
+            c_var_names.append(f"(int64_t){name}")
+        elif t == "double":
+            format_chars.append("d")
+            c_var_decls.append(f"    double {name};")
+            c_var_names.append(name)
+        elif t == "bool":
+            format_chars.append("p")
+            c_var_decls.append(f"    int {name};")
+            c_var_names.append(f"(bool){name}")
+        elif "*" in t:
+            # Buffer / array pointer parameter
+            format_chars.append("O")
+            c_var_decls.append(f"    PyObject* py_{name};")
+            c_var_decls.append(f"    Py_buffer view_{name};")
+            base_type = "int64_t" if "int64_t" in t else "double"
+            buffer_pre.append(
+                f"    if (PyObject_GetBuffer(py_{name}, &view_{name}, PyBUF_SIMPLE | PyBUF_WRITABLE) != 0) {{\n"
+                f"        return NULL;\n"
+                f"    }}\n"
+                f"    {base_type}* {name} = ({base_type}*)view_{name}.buf;"
+            )
+            buffer_post.append(f"    PyBuffer_Release(&view_{name});")
+            c_var_names.append(name)
+
+    parse_fmt = "".join(format_chars)
+    parse_ptrs = ", ".join([f"&py_{arg.arg}" if "*" in resolved_hints.get(arg.arg, "") else f"&{arg.arg}" for arg in func_def.args.args])
+
+    if ret_type == "int64_t":
+        ret_conv = "PyLong_FromLongLong((long long)result)"
+    elif ret_type == "double":
+        ret_conv = "PyFloat_FromDouble(result)"
+    elif ret_type == "bool":
+        ret_conv = "PyBool_FromLong(result ? 1 : 0)"
+    else:
+        ret_conv = "Py_None; Py_INCREF(Py_None)"
+
+    decls_str = "\n".join(c_var_decls)
+    buf_pre_str = "\n".join(buffer_pre)
+    buf_post_str = "\n".join(buffer_post)
+    kernel_call_str = f"kernel_{func_def.name}({', '.join(c_var_names)})"
+
+    c_source = f"""/* Auto-generated by FastPython Hardware-AOT Compiler */
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <stdint.h>
+#include <stdbool.h>
+#include <math.h>
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+/* Native Hardware-Optimized Kernel */
+{kernel_code}
+
+/* Python C-API Wrapper */
+static PyObject* py_{func_def.name}(PyObject* self, PyObject* args) {{
+{decls_str}
+    if (!PyArg_ParseTuple(args, "{parse_fmt}", {parse_ptrs})) {{
+        return NULL;
+    }}
+
+{buf_pre_str}
+
+    /* Release GIL for pure bare-metal execution across all host cores */
+    {ret_type} result;
+    Py_BEGIN_ALLOW_THREADS
+    result = {kernel_call_str};
+    Py_END_ALLOW_THREADS
+
+{buf_post_str}
+
+    return {ret_conv};
+}}
+
+static PyMethodDef FastMethods[] = {{
+    {{"{func_def.name}", py_{func_def.name}, METH_VARARGS, "FastPython Hardware-Optimized Native Function"}},
+    {{NULL, NULL, 0, NULL}}
+}};
+
+static struct PyModuleDef fastmodule = {{
+    PyModuleDef_HEAD_INIT,
+    "{module_name}",
+    "FastPython compiled hardware-native module",
+    -1,
+    FastMethods
+}};
+
+PyMODINIT_FUNC PyInit_{module_name}(void) {{
+    return PyModule_Create(&fastmodule);
+}}
+"""
+    return c_source

fastpython/compiler/__init__.py

@@ -0,0 +1,4 @@
+from .toolchain import CompilerToolchain, CompilationError
+from .cache import ModuleCache
+
+__all__ = ["CompilerToolchain", "CompilationError", "ModuleCache"]

fastpython/compiler/cache.py

@@ -0,0 +1,57 @@
+"""
+Cache manager for FastPython.
+Manages disk caching of compiled native extensions using SHA-256 hashes
+of code AST, CPU instruction sets, and compilation options.
+"""
+
+import os
+import sys
+import hashlib
+import importlib.util
+from pathlib import Path
+from typing import Any, Callable, Optional
+from ..hardware.cpu import get_cpu_info
+
+
+class ModuleCache:
+    def __init__(self, cache_dir: Optional[str] = None):
+        if cache_dir is None:
+            self.cache_dir = Path.cwd() / ".fastpython_cache"
+        else:
+            self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def compute_hash(self, source_code: str, extra_meta: str = "") -> str:
+        cpu = get_cpu_info()
+        hasher = hashlib.sha256()
+        hasher.update(source_code.encode("utf-8"))
+        hasher.update(cpu.summary().encode("utf-8"))
+        hasher.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode("utf-8"))
+        hasher.update(extra_meta.encode("utf-8"))
+        return hasher.hexdigest()[:16]
+
+    def get_module_paths(self, module_name: str, code_hash: str):
+        suffix = ".pyd" if sys.platform == "win32" else ".so"
+        c_file = self.cache_dir / f"{module_name}.c"
+        pyd_file = self.cache_dir / f"{module_name}{suffix}"
+        return str(c_file), str(pyd_file)
+
+    def is_cached(self, pyd_path: str) -> bool:
+        return os.path.exists(pyd_path) and os.path.getsize(pyd_path) > 0
+
+    def load_extension(self, module_name: str, pyd_path: str, func_name: str) -> Callable[..., Any]:
+        """Dynamically load the compiled native shared library and return the function."""
+        # Use unique module spec name to avoid collisions
+        spec_name = Path(pyd_path).stem
+        spec = importlib.util.spec_from_file_location(spec_name, pyd_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not create module spec from {pyd_path}")
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[spec_name] = module
+        spec.loader.exec_module(module)
+
+        if not hasattr(module, func_name):
+            raise AttributeError(f"Module {spec_name} does not export function '{func_name}'")
+
+        return getattr(module, func_name)