gstaichi 0.1.23.dev0__cp310-cp310-win_amd64.whl → 0.1.25.dev0__cp310-cp310-win_amd64.whl

{taichi → gstaichi}/lang/ast/ast_transformer_utils.py

@@ -4,33 +4,34 @@ import ast
 import builtins
 import traceback
 from enum import Enum
-from sys import version_info
 from textwrap import TextWrapper
 from typing import TYPE_CHECKING, Any, List
 
-from taichi._lib.core.taichi_python import ASTBuilder
-from taichi.lang import impl
-from taichi.lang.exception import (
-    TaichiCompilationError,
-    TaichiNameError,
-    TaichiSyntaxError,
+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 taichi.lang.kernel_impl import (
+    from gstaichi.lang.kernel_impl import (
         Func,
         Kernel,
     )
 
 
 class Builder:
-    def __call__(self, ctx, node):
+    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 TaichiSyntaxError(error_msg)
+                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)
@@ -41,15 +42,15 @@ class Builder:
                 raise e.with_traceback(None)
             ctx.raised = True
             e = handle_exception_from_cpp(e)
-            if not isinstance(e, TaichiCompilationError):
+            if not isinstance(e, GsTaichiCompilationError):
                 msg = ctx.get_pos_info(node) + traceback.format_exc()
-                raise TaichiCompilationError(msg) from None
+                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):
+    def __init__(self, scopes: list[dict[str, Any]]):
         self.scopes = scopes
 
     def __enter__(self):
@@ -65,7 +66,7 @@ class StaticScopeStatus:
 
 
 class StaticScopeGuard:
-    def __init__(self, status):
+    def __init__(self, status: StaticScopeStatus):
         self.status = status
 
     def __enter__(self):
@@ -107,7 +108,7 @@ class LoopScopeAttribute:
 
 
 class LoopScopeGuard:
-    def __init__(self, scopes, non_static_guard=None):
+    def __init__(self, scopes: list[dict[str, Any]], non_static_guard=None):
         self.scopes = scopes
         self.non_static_guard = non_static_guard
 
@@ -167,7 +168,7 @@ class ASTTransformerContext:
         is_real_function: bool = False,
     ):
         self.func = func
-        self.local_scopes = []
+        self.local_scopes: list[dict[str, Any]] = []
         self.loop_scopes: List[LoopScopeAttribute] = []
         self.excluded_parameters = excluded_parameters
         self.is_kernel = is_kernel
@@ -192,7 +193,7 @@ class ASTTransformerContext:
         self.ast_builder = ast_builder
         self.visited_funcdef = False
         self.is_real_function = is_real_function
-        self.kernel_args = []
+        self.kernel_args: list = []
 
     # e.g.: FunctionDef, Module, Global
     def variable_scope_guard(self):
@@ -211,61 +212,61 @@ class ASTTransformerContext:
             self.non_static_control_flow_status,
         )
 
-    def non_static_control_flow_guard(self):
+    def non_static_control_flow_guard(self) -> NonStaticControlFlowGuard:
         return NonStaticControlFlowGuard(self.non_static_control_flow_status)
 
-    def static_scope_guard(self):
+    def static_scope_guard(self) -> StaticScopeGuard:
         return StaticScopeGuard(self.static_scope_status)
 
-    def current_scope(self):
+    def current_scope(self) -> dict[str, Any]:
         return self.local_scopes[-1]
 
-    def current_loop_scope(self):
+    def current_loop_scope(self) -> dict[str, Any]:
         return self.loop_scopes[-1]
 
-    def loop_status(self):
+    def loop_status(self) -> LoopStatus:
         if self.loop_scopes:
             return self.loop_scopes[-1].status
         return LoopStatus.Normal
 
-    def set_loop_status(self, status):
+    def set_loop_status(self, status: LoopStatus) -> None:
         self.loop_scopes[-1].status = status
 
-    def is_in_static_for(self):
+    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):
+    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):
+    def is_in_static_scope(self) -> bool:
         return self.static_scope_status.is_in_static_scope
 
-    def is_var_declared(self, name):
+    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, var):
+    def create_variable(self, name: str, var: Any) -> None:
         if name in self.current_scope():
-            raise TaichiSyntaxError("Recreating variables is not allowed")
+            raise GsTaichiSyntaxError("Recreating variables is not allowed")
         self.current_scope()[name] = var
 
-    def check_loop_var(self, loop_var):
+    def check_loop_var(self, loop_var: str) -> None:
         if self.is_var_declared(loop_var):
-            raise TaichiSyntaxError(
+            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):
+    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 taichi.lang.matrix import (  # pylint: disable-msg=C0415
+            from gstaichi.lang.matrix import (  # pylint: disable-msg=C0415
                 Matrix,
                 make_matrix,
             )
@@ -276,19 +277,16 @@ class ASTTransformerContext:
         try:
             return getattr(builtins, name)
         except AttributeError:
-            raise TaichiNameError(f'Name "{name}" is not defined')
+            raise GsTaichiNameError(f'Name "{name}" is not defined')
 
-    def get_pos_info(self, node) -> str:
+    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'
-        if version_info < (3, 8):
-            msg += self.src[node.lineno - 1] + "\n"
-            return msg
         col_offset = self.indent + node.col_offset
         end_col_offset = self.indent + node.end_col_offset
 
         wrapper = TextWrapper(width=80)
 
-        def gen_line(code, hint):
+        def gen_line(code: str, hint: str) -> str:
             hint += " " * (len(code) - len(hint))
             code = wrapper.wrap(code)
             hint = wrapper.wrap(hint)
@@ -297,8 +295,9 @@ class ASTTransformerContext:
             return "".join([c + "\n" + h + "\n" for c, h in zip(code, hint)])
 
         if node.lineno == node.end_lineno:
-            hint = " " * col_offset + "^" * (end_col_offset - col_offset)
-            msg += gen_line(self.src[node.lineno - 1], hint)
+            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__
 
@@ -326,3 +325,17 @@ class ASTTransformerContext:
                     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