gstaichi 0.1.25.dev0__cp313-cp313-win_amd64.whl

gstaichi/lang/ast/ast_transformer_utils.py

@@ -0,0 +1,341 @@
+# type: ignore
+
+import ast
+import builtins
+import traceback
+from enum import Enum
+from textwrap import TextWrapper
+from typing import TYPE_CHECKING, Any, List
+
+from gstaichi._lib.core.gstaichi_python import ASTBuilder
+from gstaichi.lang import impl
+from gstaichi.lang._ndrange import ndrange
+from gstaichi.lang.ast.symbol_resolver import ASTResolver
+from gstaichi.lang.exception import (
+    GsTaichiCompilationError,
+    GsTaichiNameError,
+    GsTaichiSyntaxError,
+    handle_exception_from_cpp,
+)
+
+if TYPE_CHECKING:
+    from gstaichi.lang.kernel_impl import (
+        Func,
+        Kernel,
+    )
+
+
+class Builder:
+    def __call__(self, ctx: "ASTTransformerContext", node: ast.AST):
+        method = getattr(self, "build_" + node.__class__.__name__, None)
+        try:
+            if method is None:
+                error_msg = f'Unsupported node "{node.__class__.__name__}"'
+                raise GsTaichiSyntaxError(error_msg)
+            info = ctx.get_pos_info(node) if isinstance(node, (ast.stmt, ast.expr)) else ""
+            with impl.get_runtime().src_info_guard(info):
+                return method(ctx, node)
+        except Exception as e:
+            if impl.get_runtime().print_full_traceback:
+                raise e
+            if ctx.raised or not isinstance(node, (ast.stmt, ast.expr)):
+                raise e.with_traceback(None)
+            ctx.raised = True
+            e = handle_exception_from_cpp(e)
+            if not isinstance(e, GsTaichiCompilationError):
+                msg = ctx.get_pos_info(node) + traceback.format_exc()
+                raise GsTaichiCompilationError(msg) from None
+            msg = ctx.get_pos_info(node) + str(e)
+            raise type(e)(msg) from None
+
+
+class VariableScopeGuard:
+    def __init__(self, scopes: list[dict[str, Any]]):
+        self.scopes = scopes
+
+    def __enter__(self):
+        self.scopes.append({})
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.scopes.pop()
+
+
+class StaticScopeStatus:
+    def __init__(self):
+        self.is_in_static_scope = False
+
+
+class StaticScopeGuard:
+    def __init__(self, status: StaticScopeStatus):
+        self.status = status
+
+    def __enter__(self):
+        self.prev = self.status.is_in_static_scope
+        self.status.is_in_static_scope = True
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.status.is_in_static_scope = self.prev
+
+
+class NonStaticControlFlowStatus:
+    def __init__(self):
+        self.is_in_non_static_control_flow = False
+
+
+class NonStaticControlFlowGuard:
+    def __init__(self, status: NonStaticControlFlowStatus):
+        self.status = status
+
+    def __enter__(self):
+        self.prev = self.status.is_in_non_static_control_flow
+        self.status.is_in_non_static_control_flow = True
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.status.is_in_non_static_control_flow = self.prev
+
+
+class LoopStatus(Enum):
+    Normal = 0
+    Break = 1
+    Continue = 2
+
+
+class LoopScopeAttribute:
+    def __init__(self, is_static: bool):
+        self.is_static = is_static
+        self.status: LoopStatus = LoopStatus.Normal
+        self.nearest_non_static_if: ast.If | None = None
+
+
+class LoopScopeGuard:
+    def __init__(self, scopes: list[dict[str, Any]], non_static_guard=None):
+        self.scopes = scopes
+        self.non_static_guard = non_static_guard
+
+    def __enter__(self):
+        self.scopes.append(LoopScopeAttribute(self.non_static_guard is None))
+        if self.non_static_guard:
+            self.non_static_guard.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.scopes.pop()
+        if self.non_static_guard:
+            self.non_static_guard.__exit__(exc_type, exc_val, exc_tb)
+
+
+class NonStaticIfGuard:
+    def __init__(
+        self,
+        if_node: ast.If,
+        loop_attribute: LoopScopeAttribute,
+        non_static_status: NonStaticControlFlowStatus,
+    ):
+        self.loop_attribute = loop_attribute
+        self.if_node = if_node
+        self.non_static_guard = NonStaticControlFlowGuard(non_static_status)
+
+    def __enter__(self):
+        if self.loop_attribute:
+            self.old_non_static_if = self.loop_attribute.nearest_non_static_if
+            self.loop_attribute.nearest_non_static_if = self.if_node
+        self.non_static_guard.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.loop_attribute:
+            self.loop_attribute.nearest_non_static_if = self.old_non_static_if
+        self.non_static_guard.__exit__(exc_type, exc_val, exc_tb)
+
+
+class ReturnStatus(Enum):
+    NoReturn = 0
+    ReturnedVoid = 1
+    ReturnedValue = 2
+
+
+class ASTTransformerContext:
+    def __init__(
+        self,
+        excluded_parameters=(),
+        is_kernel: bool = True,
+        func: "Func | Kernel | None" = None,
+        arg_features=None,
+        global_vars: dict[str, Any] | None = None,
+        argument_data=None,
+        file: str | None = None,
+        src: list[str] | None = None,
+        start_lineno: int | None = None,
+        ast_builder: ASTBuilder | None = None,
+        is_real_function: bool = False,
+    ):
+        self.func = func
+        self.local_scopes: list[dict[str, Any]] = []
+        self.loop_scopes: List[LoopScopeAttribute] = []
+        self.excluded_parameters = excluded_parameters
+        self.is_kernel = is_kernel
+        self.arg_features = arg_features
+        self.returns = None
+        self.global_vars = global_vars
+        self.argument_data = argument_data
+        self.return_data = None
+        self.file = file
+        self.src = src
+        self.indent = 0
+        for c in self.src[0]:
+            if c == " ":
+                self.indent += 1
+            else:
+                break
+        self.lineno_offset = start_lineno - 1
+        self.raised = False
+        self.non_static_control_flow_status = NonStaticControlFlowStatus()
+        self.static_scope_status = StaticScopeStatus()
+        self.returned = ReturnStatus.NoReturn
+        self.ast_builder = ast_builder
+        self.visited_funcdef = False
+        self.is_real_function = is_real_function
+        self.kernel_args: list = []
+
+    # e.g.: FunctionDef, Module, Global
+    def variable_scope_guard(self):
+        return VariableScopeGuard(self.local_scopes)
+
+    # e.g.: For, While
+    def loop_scope_guard(self, is_static=False):
+        if is_static:
+            return LoopScopeGuard(self.loop_scopes)
+        return LoopScopeGuard(self.loop_scopes, self.non_static_control_flow_guard())
+
+    def non_static_if_guard(self, if_node: ast.If):
+        return NonStaticIfGuard(
+            if_node,
+            self.current_loop_scope() if self.loop_scopes else None,
+            self.non_static_control_flow_status,
+        )
+
+    def non_static_control_flow_guard(self) -> NonStaticControlFlowGuard:
+        return NonStaticControlFlowGuard(self.non_static_control_flow_status)
+
+    def static_scope_guard(self) -> StaticScopeGuard:
+        return StaticScopeGuard(self.static_scope_status)
+
+    def current_scope(self) -> dict[str, Any]:
+        return self.local_scopes[-1]
+
+    def current_loop_scope(self) -> dict[str, Any]:
+        return self.loop_scopes[-1]
+
+    def loop_status(self) -> LoopStatus:
+        if self.loop_scopes:
+            return self.loop_scopes[-1].status
+        return LoopStatus.Normal
+
+    def set_loop_status(self, status: LoopStatus) -> None:
+        self.loop_scopes[-1].status = status
+
+    def is_in_static_for(self) -> bool:
+        if self.loop_scopes:
+            return self.loop_scopes[-1].is_static
+        return False
+
+    def is_in_non_static_control_flow(self) -> bool:
+        return self.non_static_control_flow_status.is_in_non_static_control_flow
+
+    def is_in_static_scope(self) -> bool:
+        return self.static_scope_status.is_in_static_scope
+
+    def is_var_declared(self, name: str) -> bool:
+        for s in self.local_scopes:
+            if name in s:
+                return True
+        return False
+
+    def create_variable(self, name: str, var: Any) -> None:
+        if name in self.current_scope():
+            raise GsTaichiSyntaxError("Recreating variables is not allowed")
+        self.current_scope()[name] = var
+
+    def check_loop_var(self, loop_var: str) -> None:
+        if self.is_var_declared(loop_var):
+            raise GsTaichiSyntaxError(
+                f"Variable '{loop_var}' is already declared in the outer scope and cannot be used as loop variable"
+            )
+
+    def get_var_by_name(self, name: str) -> Any:
+        for s in reversed(self.local_scopes):
+            if name in s:
+                return s[name]
+        if name in self.global_vars:
+            var = self.global_vars[name]
+            from gstaichi.lang.matrix import (  # pylint: disable-msg=C0415
+                Matrix,
+                make_matrix,
+            )
+
+            if isinstance(var, Matrix):
+                return make_matrix(var.to_list())
+            return var
+        try:
+            return getattr(builtins, name)
+        except AttributeError:
+            raise GsTaichiNameError(f'Name "{name}" is not defined')
+
+    def get_pos_info(self, node: ast.AST) -> str:
+        msg = f'File "{self.file}", line {node.lineno + self.lineno_offset}, in {self.func.func.__name__}:\n'
+        col_offset = self.indent + node.col_offset
+        end_col_offset = self.indent + node.end_col_offset
+
+        wrapper = TextWrapper(width=80)
+
+        def gen_line(code: str, hint: str) -> str:
+            hint += " " * (len(code) - len(hint))
+            code = wrapper.wrap(code)
+            hint = wrapper.wrap(hint)
+            if not len(code):
+                return "\n\n"
+            return "".join([c + "\n" + h + "\n" for c, h in zip(code, hint)])
+
+        if node.lineno == node.end_lineno:
+            if node.lineno - 1 < len(self.src):
+                hint = " " * col_offset + "^" * (end_col_offset - col_offset)
+                msg += gen_line(self.src[node.lineno - 1], hint)
+        else:
+            node_type = node.__class__.__name__
+
+            if node_type in ["For", "While", "FunctionDef", "If"]:
+                end_lineno = max(node.body[0].lineno - 1, node.lineno)
+            else:
+                end_lineno = node.end_lineno
+
+            for i in range(node.lineno - 1, end_lineno):
+                last = len(self.src[i])
+                while last > 0 and (self.src[i][last - 1].isspace() or not self.src[i][last - 1].isprintable()):
+                    last -= 1
+                first = 0
+                while first < len(self.src[i]) and (
+                    self.src[i][first].isspace() or not self.src[i][first].isprintable()
+                ):
+                    first += 1
+                if i == node.lineno - 1:
+                    hint = " " * col_offset + "^" * (last - col_offset)
+                elif i == node.end_lineno - 1:
+                    hint = " " * first + "^" * (end_col_offset - first)
+                elif first < last:
+                    hint = " " * first + "^" * (last - first)
+                else:
+                    hint = ""
+                msg += gen_line(self.src[i], hint)
+        return msg
+
+
+def get_decorator(ctx: ASTTransformerContext, node) -> str:
+    if not isinstance(node, ast.Call):
+        return ""
+    for wanted, name in [
+        (impl.static, "static"),
+        (impl.static_assert, "static_assert"),
+        (impl.grouped, "grouped"),
+        (ndrange, "ndrange"),
+    ]:
+        if ASTResolver.resolve_to(node.func, wanted, ctx.global_vars):
+            return name
+    return ""

gstaichi/lang/ast/ast_transformers/call_transformer.py

@@ -0,0 +1,267 @@
+# type: ignore
+
+import ast
+import dataclasses
+import inspect
+import math
+import operator
+import re
+import warnings
+from ast import unparse
+from collections import ChainMap
+
+import numpy as np
+
+from gstaichi.lang import (
+    expr,
+    impl,
+    matrix,
+)
+from gstaichi.lang import ops as ti_ops
+from gstaichi.lang.ast.ast_transformer_utils import (
+    ASTTransformerContext,
+    get_decorator,
+)
+from gstaichi.lang.exception import (
+    GsTaichiSyntaxError,
+    GsTaichiTypeError,
+)
+from gstaichi.lang.expr import Expr
+from gstaichi.lang.matrix import Matrix, Vector
+from gstaichi.lang.util import is_gstaichi_class
+from gstaichi.types import primitive_types
+
+
+class CallTransformer:
+    @staticmethod
+    def build_call_if_is_builtin(ctx: ASTTransformerContext, node, args, keywords):
+        from gstaichi.lang import matrix_ops  # pylint: disable=C0415
+
+        func = node.func.ptr
+        replace_func = {
+            id(print): impl.ti_print,
+            id(min): ti_ops.min,
+            id(max): ti_ops.max,
+            id(int): impl.ti_int,
+            id(bool): impl.ti_bool,
+            id(float): impl.ti_float,
+            id(any): matrix_ops.any,
+            id(all): matrix_ops.all,
+            id(abs): abs,
+            id(pow): pow,
+            id(operator.matmul): matrix_ops.matmul,
+        }
+
+        # Builtin 'len' function on Matrix Expr
+        if id(func) == id(len) and len(args) == 1:
+            if isinstance(args[0], Expr) and args[0].ptr.is_tensor():
+                node.ptr = args[0].get_shape()[0]
+                return True
+
+        if id(func) in replace_func:
+            node.ptr = replace_func[id(func)](*args, **keywords)
+            return True
+        return False
+
+    @staticmethod
+    def build_call_if_is_type(ctx: ASTTransformerContext, node, args, keywords):
+        func = node.func.ptr
+        if id(func) in primitive_types.type_ids:
+            if len(args) != 1 or keywords:
+                raise GsTaichiSyntaxError("A primitive type can only decorate a single expression.")
+            if is_gstaichi_class(args[0]):
+                raise GsTaichiSyntaxError("A primitive type cannot decorate an expression with a compound type.")
+
+            if isinstance(args[0], expr.Expr):
+                if args[0].ptr.is_tensor():
+                    raise GsTaichiSyntaxError("A primitive type cannot decorate an expression with a compound type.")
+                node.ptr = ti_ops.cast(args[0], func)
+            else:
+                node.ptr = expr.Expr(args[0], dtype=func)
+            return True
+        return False
+
+    @staticmethod
+    def is_external_func(ctx: ASTTransformerContext, func) -> bool:
+        if ctx.is_in_static_scope():  # allow external function in static scope
+            return False
+        if hasattr(func, "_is_gstaichi_function") or hasattr(func, "_is_wrapped_kernel"):  # gstaichi func/kernel
+            return False
+        if hasattr(func, "__module__") and func.__module__ and func.__module__.startswith("gstaichi."):
+            return False
+        return True
+
+    @staticmethod
+    def warn_if_is_external_func(ctx: ASTTransformerContext, node):
+        func = node.func.ptr
+        if not CallTransformer.is_external_func(ctx, func):
+            return
+        name = unparse(node.func).strip()
+        warnings.warn_explicit(
+            f"\x1b[38;5;226m"  # Yellow
+            f'Calling non-gstaichi function "{name}". '
+            f"Scope inside the function is not processed by the GsTaichi AST transformer. "
+            f"The function may not work as expected. Proceed with caution! "
+            f"Maybe you can consider turning it into a @ti.func?"
+            f"\x1b[0m",  # Reset
+            SyntaxWarning,
+            ctx.file,
+            node.lineno + ctx.lineno_offset,
+            module="gstaichi",
+        )
+
+    @staticmethod
+    # Parses a formatted string and extracts format specifiers from it, along with positional and keyword arguments.
+    # This function produces a canonicalized formatted string that includes solely empty replacement fields, e.g. 'qwerty {} {} {} {} {}'.
+    # Note that the arguments can be used multiple times in the string.
+    # e.g.:
+    # origin input: 'qwerty {1} {} {1:.3f} {k:.4f} {k:}'.format(1.0, 2.0, k=k)
+    # raw_string: 'qwerty {1} {} {1:.3f} {k:.4f} {k:}'
+    # raw_args: [1.0, 2.0]
+    # raw_keywords: {'k': <ti.Expr>}
+    # return value: ['qwerty {} {} {} {} {}', 2.0, 1.0, ['__ti_fmt_value__', 2.0, '.3f'], ['__ti_fmt_value__', <ti.Expr>, '.4f'], <ti.Expr>]
+    def canonicalize_formatted_string(raw_string: str, *raw_args: list, **raw_keywords: dict):
+        raw_brackets = re.findall(r"{(.*?)}", raw_string)
+        brackets = []
+        unnamed = 0
+        for bracket in raw_brackets:
+            item, spec = bracket.split(":") if ":" in bracket else (bracket, None)
+            if item.isdigit():
+                item = int(item)
+            # handle unnamed positional args
+            if item == "":
+                item = unnamed
+                unnamed += 1
+            # handle empty spec
+            if spec == "":
+                spec = None
+            brackets.append((item, spec))
+
+        # check for errors in the arguments
+        max_args_index = max([t[0] for t in brackets if isinstance(t[0], int)], default=-1)
+        if max_args_index + 1 != len(raw_args):
+            raise GsTaichiSyntaxError(
+                f"Expected {max_args_index + 1} positional argument(s), but received {len(raw_args)} instead."
+            )
+        brackets_keywords = [t[0] for t in brackets if isinstance(t[0], str)]
+        for item in brackets_keywords:
+            if item not in raw_keywords:
+                raise GsTaichiSyntaxError(f"Keyword '{item}' not found.")
+        for item in raw_keywords:
+            if item not in brackets_keywords:
+                raise GsTaichiSyntaxError(f"Keyword '{item}' not used.")
+
+        # reorganize the arguments based on their positions, keywords, and format specifiers
+        args = []
+        for item, spec in brackets:
+            new_arg = raw_args[item] if isinstance(item, int) else raw_keywords[item]
+            if spec is not None:
+                args.append(["__ti_fmt_value__", new_arg, spec])
+            else:
+                args.append(new_arg)
+        # put the formatted string as the first argument to make ti.format() happy
+        args.insert(0, re.sub(r"{.*?}", "{}", raw_string))
+        return args
+
+    @staticmethod
+    def expand_node_args_dataclasses(args: tuple[ast.AST, ...]) -> tuple[ast.AST, ...]:
+        args_new = []
+        for arg in args:
+            val = arg.ptr
+            if dataclasses.is_dataclass(val):
+                dataclass_type = val
+                for field in dataclasses.fields(dataclass_type):
+                    child_name = f"__ti_{arg.id}_{field.name}"
+                    load_ctx = ast.Load()
+                    arg_node = ast.Name(
+                        id=child_name,
+                        ctx=load_ctx,
+                        lineno=arg.lineno,
+                        end_lineno=arg.end_lineno,
+                        col_offset=arg.col_offset,
+                        end_col_offset=arg.end_col_offset,
+                    )
+                    args_new.append(arg_node)
+            else:
+                args_new.append(arg)
+        return tuple(args_new)
+
+    @staticmethod
+    def build_Call(ctx: ASTTransformerContext, node: ast.Call, build_stmt, build_stmts):
+        if get_decorator(ctx, node) in ["static", "static_assert"]:
+            with ctx.static_scope_guard():
+                build_stmt(ctx, node.func)
+                build_stmts(ctx, node.args)
+                build_stmts(ctx, node.keywords)
+        else:
+            build_stmt(ctx, node.func)
+            # creates variable for the dataclass itself (as well as other variables,
+            # not related to dataclasses). Necessary for calling further child functions
+            build_stmts(ctx, node.args)
+            node.args = CallTransformer.expand_node_args_dataclasses(node.args)
+            # create variables for the now-expanded dataclass members
+            build_stmts(ctx, node.args)
+            build_stmts(ctx, node.keywords)
+
+        args = []
+        for arg in node.args:
+            if isinstance(arg, ast.Starred):
+                arg_list = arg.ptr
+                if isinstance(arg_list, Expr) and arg_list.is_tensor():
+                    # Expand Expr with Matrix-type return into list of Exprs
+                    arg_list = [Expr(x) for x in ctx.ast_builder.expand_exprs([arg_list.ptr])]
+
+                for i in arg_list:
+                    args.append(i)
+            else:
+                args.append(arg.ptr)
+        keywords = dict(ChainMap(*[keyword.ptr for keyword in node.keywords]))
+        func = node.func.ptr
+
+        if id(func) in [id(print), id(impl.ti_print)]:
+            ctx.func.has_print = True
+
+        if isinstance(node.func, ast.Attribute) and isinstance(node.func.value.ptr, str) and node.func.attr == "format":
+            raw_string = node.func.value.ptr
+            args = CallTransformer.canonicalize_formatted_string(raw_string, *args, **keywords)
+            node.ptr = impl.ti_format(*args)
+            return node.ptr
+
+        if id(func) == id(Matrix) or id(func) == id(Vector):
+            node.ptr = matrix.make_matrix(*args, **keywords)
+            return node.ptr
+
+        if CallTransformer.build_call_if_is_builtin(ctx, node, args, keywords):
+            return node.ptr
+
+        if CallTransformer.build_call_if_is_type(ctx, node, args, keywords):
+            return node.ptr
+
+        if hasattr(node.func, "caller"):
+            node.ptr = func(node.func.caller, *args, **keywords)
+            return node.ptr
+
+        CallTransformer.warn_if_is_external_func(ctx, node)
+        try:
+            node.ptr = func(*args, **keywords)
+        except TypeError as e:
+            module = inspect.getmodule(func)
+            error_msg = re.sub(r"\bExpr\b", "GsTaichi Expression", str(e))
+            func_name = getattr(func, "__name__", func.__class__.__name__)
+            msg = f"TypeError when calling `{func_name}`: {error_msg}."
+            if CallTransformer.is_external_func(ctx, node.func.ptr):
+                args_has_expr = any([isinstance(arg, Expr) for arg in args])
+                if args_has_expr and (module == math or module == np):
+                    exec_str = f"from gstaichi import {func.__name__}"
+                    try:
+                        exec(exec_str, {})
+                    except:
+                        pass
+                    else:
+                        msg += f"\nDid you mean to use `ti.{func.__name__}` instead of `{module.__name__}.{func.__name__}`?"
+            raise GsTaichiTypeError(msg)
+
+        if getattr(func, "_is_gstaichi_function", False):
+            ctx.func.has_print |= func.wrapper.has_print
+
+        return node.ptr