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

gstaichi/lang/ast/ast_transformers/function_def_transformer.py

@@ -0,0 +1,320 @@
+# type: ignore
+
+import ast
+import dataclasses
+from typing import Any, Callable
+
+from gstaichi.lang import (
+    _ndarray,
+    any_array,
+    expr,
+    impl,
+    kernel_arguments,
+    matrix,
+)
+from gstaichi.lang import ops as ti_ops
+from gstaichi.lang.argpack import ArgPackType
+from gstaichi.lang.ast.ast_transformer_utils import (
+    ASTTransformerContext,
+)
+from gstaichi.lang.exception import (
+    GsTaichiSyntaxError,
+)
+from gstaichi.lang.matrix import MatrixType
+from gstaichi.lang.struct import StructType
+from gstaichi.lang.util import to_gstaichi_type
+from gstaichi.types import annotations, ndarray_type, primitive_types, texture_type
+
+
+class FunctionDefTransformer:
+    @staticmethod
+    def _decl_and_create_variable(
+        ctx: ASTTransformerContext, annotation, name, arg_features, invoke_later_dict, prefix_name, arg_depth
+    ) -> tuple[bool, Any]:
+        full_name = prefix_name + "_" + name
+        if not isinstance(annotation, primitive_types.RefType):
+            ctx.kernel_args.append(name)
+        if isinstance(annotation, ArgPackType):
+            kernel_arguments.push_argpack_arg(name)
+            d = {}
+            items_to_put_in_dict = []
+            for j, (_name, anno) in enumerate(annotation.members.items()):
+                result, obj = FunctionDefTransformer._decl_and_create_variable(
+                    ctx, anno, _name, arg_features[j], invoke_later_dict, full_name, arg_depth + 1
+                )
+                if not result:
+                    d[_name] = None
+                    items_to_put_in_dict.append((full_name + "_" + _name, _name, obj))
+                else:
+                    d[_name] = obj
+            argpack = kernel_arguments.decl_argpack_arg(annotation, d)
+            for item in items_to_put_in_dict:
+                invoke_later_dict[item[0]] = argpack, item[1], *item[2]
+            return True, argpack
+        if annotation == annotations.template or isinstance(annotation, annotations.template):
+            return True, ctx.global_vars[name]
+        if isinstance(annotation, annotations.sparse_matrix_builder):
+            return False, (
+                kernel_arguments.decl_sparse_matrix,
+                (
+                    to_gstaichi_type(arg_features),
+                    full_name,
+                ),
+            )
+        if isinstance(annotation, ndarray_type.NdarrayType):
+            return False, (
+                kernel_arguments.decl_ndarray_arg,
+                (
+                    to_gstaichi_type(arg_features[0]),
+                    arg_features[1],
+                    full_name,
+                    arg_features[2],
+                    arg_features[3],
+                ),
+            )
+        if isinstance(annotation, texture_type.TextureType):
+            return False, (kernel_arguments.decl_texture_arg, (arg_features[0], full_name))
+        if isinstance(annotation, texture_type.RWTextureType):
+            return False, (
+                kernel_arguments.decl_rw_texture_arg,
+                (arg_features[0], arg_features[1], arg_features[2], full_name),
+            )
+        if isinstance(annotation, MatrixType):
+            return True, kernel_arguments.decl_matrix_arg(annotation, name, arg_depth)
+        if isinstance(annotation, StructType):
+            return True, kernel_arguments.decl_struct_arg(annotation, name, arg_depth)
+        return True, kernel_arguments.decl_scalar_arg(annotation, name, arg_depth)
+
+    @staticmethod
+    def _transform_kernel_arg(
+        ctx: ASTTransformerContext,
+        invoke_later_dict: dict[str, tuple[Any, str, Callable, list[Any]]],
+        create_variable_later: dict[str, Any],
+        argument_name: str,
+        argument_type: Any,
+        this_arg_features: tuple[Any, ...],
+    ) -> None:
+        if isinstance(argument_type, ArgPackType):
+            kernel_arguments.push_argpack_arg(argument_name)
+            d = {}
+            items_to_put_in_dict: list[tuple[str, str, Any]] = []
+            for j, (name, anno) in enumerate(argument_type.members.items()):
+                result, obj = FunctionDefTransformer._decl_and_create_variable(
+                    ctx, anno, name, this_arg_features[j], invoke_later_dict, "__argpack_" + name, 1
+                )
+                if not result:
+                    d[name] = None
+                    items_to_put_in_dict.append(("__argpack_" + name, name, obj))
+                else:
+                    d[name] = obj
+            argpack = kernel_arguments.decl_argpack_arg(argument_type, d)
+            for item in items_to_put_in_dict:
+                invoke_later_dict[item[0]] = argpack, item[1], *item[2]
+            create_variable_later[argument_name] = argpack
+        elif dataclasses.is_dataclass(argument_type):
+            arg_features = this_arg_features
+            ctx.create_variable(argument_name, argument_type)
+            for field_idx, field in enumerate(dataclasses.fields(argument_type)):
+                flat_name = f"__ti_{argument_name}_{field.name}"
+                result, obj = FunctionDefTransformer._decl_and_create_variable(
+                    ctx,
+                    field.type,
+                    flat_name,
+                    arg_features[field_idx],
+                    invoke_later_dict,
+                    "",
+                    0,
+                )
+                if result:
+                    ctx.create_variable(flat_name, obj)
+                else:
+                    decl_type_func, type_args = obj
+                    obj = decl_type_func(*type_args)
+                    ctx.create_variable(flat_name, obj)
+        else:
+            result, obj = FunctionDefTransformer._decl_and_create_variable(
+                ctx,
+                argument_type,
+                argument_name,
+                this_arg_features if ctx.arg_features is not None else None,
+                invoke_later_dict,
+                "",
+                0,
+            )
+            if result:
+                ctx.create_variable(argument_name, obj)
+            else:
+                decl_type_func, type_args = obj
+                obj = decl_type_func(*type_args)
+                ctx.create_variable(argument_name, obj)
+
+    @staticmethod
+    def _transform_as_kernel(ctx: ASTTransformerContext, node: ast.FunctionDef, args: ast.arguments) -> None:
+        if node.returns is not None:
+            if not isinstance(node.returns, ast.Constant):
+                for return_type in ctx.func.return_type:
+                    kernel_arguments.decl_ret(return_type)
+        impl.get_runtime().compiling_callable.finalize_rets()
+
+        invoke_later_dict: dict[str, tuple[Any, str, Any]] = dict()
+        create_variable_later = dict()
+        for i, arg in enumerate(args.args):
+            argument = ctx.func.arguments[i]
+            FunctionDefTransformer._transform_kernel_arg(
+                ctx,
+                invoke_later_dict,
+                create_variable_later,
+                argument.name,
+                argument.annotation,
+                ctx.arg_features[i] if ctx.arg_features is not None else (),
+            )
+
+        for k, v in invoke_later_dict.items():
+            argpack, name, func, params = v
+            argpack[name] = func(*params)
+        for k, v in create_variable_later.items():
+            ctx.create_variable(k, v)
+
+        impl.get_runtime().compiling_callable.finalize_params()
+        # remove original args
+        node.args.args = []
+
+    @staticmethod
+    def _transform_func_arg(
+        ctx: ASTTransformerContext,
+        argument_name: str,
+        argument_type: Any,
+        data: Any,
+    ) -> None:
+        if isinstance(argument_type, annotations.template):
+            ctx.create_variable(argument_name, data)
+            return None
+
+        if dataclasses.is_dataclass(argument_type):
+            dataclass_type = argument_type
+            for field in dataclasses.fields(dataclass_type):
+                data_child = getattr(data, field.name)
+                if not isinstance(
+                    data_child,
+                    (
+                        _ndarray.ScalarNdarray,
+                        matrix.VectorNdarray,
+                        matrix.MatrixNdarray,
+                        any_array.AnyArray,
+                    ),
+                ):
+                    raise GsTaichiSyntaxError(
+                        f"Argument {argument_name} of type {dataclass_type} {field.type} is not recognized."
+                    )
+                field.type.check_matched(data_child.get_type(), field.name)
+                var_name = f"__ti_{argument_name}_{field.name}"
+                ctx.create_variable(var_name, data_child)
+            return None
+
+        # Ndarray arguments are passed by reference.
+        if isinstance(argument_type, (ndarray_type.NdarrayType)):
+            if not isinstance(
+                data,
+                (
+                    _ndarray.ScalarNdarray,
+                    matrix.VectorNdarray,
+                    matrix.MatrixNdarray,
+                    any_array.AnyArray,
+                ),
+            ):
+                raise GsTaichiSyntaxError(f"Argument {arg.arg} of type {argument_type} is not recognized.")
+            argument_type.check_matched(data.get_type(), argument_name)
+            ctx.create_variable(argument_name, data)
+            return None
+
+        # Matrix arguments are passed by value.
+        if isinstance(argument_type, (MatrixType)):
+            var_name = argument_name
+            # "data" is expected to be an Expr here,
+            # so we simply call "impl.expr_init_func(data)" to perform:
+            #
+            # TensorType* t = alloca()
+            # assign(t, data)
+            #
+            # We created local variable "t" - a copy of the passed-in argument "data"
+            if not isinstance(data, expr.Expr) or not data.ptr.is_tensor():
+                raise GsTaichiSyntaxError(
+                    f"Argument {var_name} of type {argument_type} is expected to be a Matrix, but got {type(data)}."
+                )
+
+            element_shape = data.ptr.get_rvalue_type().shape()
+            if len(element_shape) != argument_type.ndim:
+                raise GsTaichiSyntaxError(
+                    f"Argument {var_name} of type {argument_type} is expected to be a Matrix with ndim {argument_type.ndim}, but got {len(element_shape)}."
+                )
+
+            assert argument_type.ndim > 0
+            if element_shape[0] != argument_type.n:
+                raise GsTaichiSyntaxError(
+                    f"Argument {var_name} of type {argument_type} is expected to be a Matrix with n {argument_type.n}, but got {element_shape[0]}."
+                )
+
+            if argument_type.ndim == 2 and element_shape[1] != argument_type.m:
+                raise GsTaichiSyntaxError(
+                    f"Argument {var_name} of type {argument_type} is expected to be a Matrix with m {argument_type.m}, but got {element_shape[0]}."
+                )
+
+            ctx.create_variable(var_name, impl.expr_init_func(data))
+            return None
+
+        if id(argument_type) in primitive_types.type_ids:
+            var_name = argument_name
+            ctx.create_variable(var_name, impl.expr_init_func(ti_ops.cast(data, argument_type)))
+            return None
+        # Create a copy for non-template arguments,
+        # so that they are passed by value.
+        var_name = argument_name
+        ctx.create_variable(var_name, impl.expr_init_func(data))
+        return None
+
+    @staticmethod
+    def _transform_as_func(ctx: ASTTransformerContext, node: ast.FunctionDef, args: ast.arguments) -> None:
+        for data_i, data in enumerate(ctx.argument_data):
+            argument = ctx.func.arguments[data_i]
+            FunctionDefTransformer._transform_func_arg(
+                ctx,
+                argument.name,
+                argument.annotation,
+                data,
+            )
+
+        for v in ctx.func.orig_arguments:
+            if dataclasses.is_dataclass(v.annotation):
+                ctx.create_variable(v.name, v.annotation)
+
+    @staticmethod
+    def build_FunctionDef(
+        ctx: ASTTransformerContext,
+        node: ast.FunctionDef,
+        build_stmts: Callable[[ASTTransformerContext, list[ast.stmt]], None],
+    ) -> None:
+        if ctx.visited_funcdef:
+            raise GsTaichiSyntaxError(
+                f"Function definition is not allowed in 'ti.{'kernel' if ctx.is_kernel else 'func'}'."
+            )
+        ctx.visited_funcdef = True
+
+        args = node.args
+        assert args.vararg is None
+        assert args.kwonlyargs == []
+        assert args.kw_defaults == []
+        assert args.kwarg is None
+
+        if ctx.is_kernel:  # ti.kernel
+            FunctionDefTransformer._transform_as_kernel(ctx, node, args)
+
+        else:  # ti.func
+            if ctx.is_real_function:
+                FunctionDefTransformer._transform_as_kernel(ctx, node, args)
+            else:
+                FunctionDefTransformer._transform_as_func(ctx, node, args)
+
+        with ctx.variable_scope_guard():
+            build_stmts(ctx, node.body)
+
+        return None

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

@@ -2,8 +2,8 @@
 
 import ast
 
-from taichi.lang._wrap_inspect import getsourcefile, getsourcelines
-from taichi.lang.exception import TaichiSyntaxError
+from gstaichi.lang._wrap_inspect import getsourcefile, getsourcelines
+from gstaichi.lang.exception import GsTaichiSyntaxError
 
 
 class KernelSimplicityASTChecker(ast.NodeVisitor):
@@ -62,7 +62,7 @@ class KernelSimplicityASTChecker(ast.NodeVisitor):
         if not isinstance(node, ast.stmt):
             return False
         # TODO(#536): Frontend pass should help make sure |func| is a valid AST for
-        # Taichi.
+        # GsTaichi.
         ignored = [ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef]
         return not any(map(lambda t: isinstance(node, t), ignored))
 
@@ -72,7 +72,7 @@ class KernelSimplicityASTChecker(ast.NodeVisitor):
             return
 
         if not (self.top_level or self.current_scope.allows_more_stmt):
-            raise TaichiSyntaxError(f"No more statements allowed, at {self.get_error_location(node)}")
+            raise GsTaichiSyntaxError(f"No more statements allowed, at {self.get_error_location(node)}")
         old_top_level = self.top_level
         if old_top_level:
             self._scope_guards.append(self.new_scope())
@@ -96,7 +96,7 @@ class KernelSimplicityASTChecker(ast.NodeVisitor):
         #              and node.iter.func.attr == 'static')
         # if not (self.top_level or self.current_scope.allows_for_loop
         #         or is_static):
-        #     raise TaichiSyntaxError(
+        #     raise GsTaichiSyntaxError(
         #         f'No more for loops allowed, at {self.get_error_location(node)}'
         #     )
         # with self.new_scope():

gstaichi/lang/ast/transform.py

@@ -0,0 +1,9 @@
+# type: ignore
+
+from gstaichi.lang.ast.ast_transformer import ASTTransformer
+from gstaichi.lang.ast.ast_transformer_utils import ASTTransformerContext
+
+
+def transform_tree(tree, ctx: ASTTransformerContext):
+    ASTTransformer()(ctx, tree)
+    return ctx.return_data

{taichi → gstaichi}/lang/common_ops.py

@@ -2,13 +2,13 @@
 
 from typing import TYPE_CHECKING
 
-from taichi.lang import ops
-from taichi.lang.util import in_python_scope
-from taichi.types import primitive_types
+from gstaichi.lang import ops
+from gstaichi.lang.util import in_python_scope
+from gstaichi.types import primitive_types
 
 
-class TaichiOperations:
-    """The base class of taichi operations of expressions. Subclasses: :class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`"""
+class GsTaichiOperations:
+    """The base class of gstaichi operations of expressions. Subclasses: :class:`~gstaichi.lang.expr.Expr`, :class:`~gstaichi.lang.matrix.Matrix`"""
 
     if TYPE_CHECKING:
         # Make pylint happy
@@ -124,7 +124,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic add."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic add."""
         return ops.atomic_add(self, other)
 
     def _atomic_mul(self, other):
@@ -134,7 +134,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic mul."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic mul."""
         return ops.atomic_mul(self, other)
 
     def _atomic_sub(self, other):
@@ -144,7 +144,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic sub."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic sub."""
         return ops.atomic_sub(self, other)
 
     def _atomic_and(self, other):
@@ -154,7 +154,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic and."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic and."""
         return ops.atomic_and(self, other)
 
     def _atomic_xor(self, other):
@@ -164,7 +164,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic xor."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic xor."""
         return ops.atomic_xor(self, other)
 
     def _atomic_or(self, other):
@@ -174,7 +174,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The computing expression of atomic or."""
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic or."""
         return ops.atomic_or(self, other)
 
     # In-place operators in python scope returns NotImplemented to fall back to normal operators
@@ -264,7 +264,7 @@ class TaichiOperations:
             other (Any): Given operand.
 
         Returns:
-            :class:`~taichi.lang.expr.Expr`: The expression after assigning."""
+            :class:`~gstaichi.lang.expr.Expr`: The expression after assigning."""
         return ops.assign(self, other)
 
     def _augassign(self, x, op):

gstaichi/lang/exception.py

@@ -0,0 +1,80 @@
+# type: ignore
+
+from gstaichi._lib import core
+
+
+class GsTaichiCompilationError(Exception):
+    """Base class for all compilation exceptions."""
+
+    pass
+
+
+class GsTaichiSyntaxError(GsTaichiCompilationError, SyntaxError):
+    """Thrown when a syntax error is found during compilation."""
+
+    pass
+
+
+class GsTaichiNameError(GsTaichiCompilationError, NameError):
+    """Thrown when an undefine name is found during compilation."""
+
+    pass
+
+
+class GsTaichiIndexError(GsTaichiCompilationError, IndexError):
+    """Thrown when an index error is found during compilation."""
+
+    pass
+
+
+class GsTaichiTypeError(GsTaichiCompilationError, TypeError):
+    """Thrown when a type mismatch is found during compilation."""
+
+    pass
+
+
+class GsTaichiRuntimeError(RuntimeError):
+    """Thrown when the compiled program cannot be executed due to unspecified reasons."""
+
+    pass
+
+
+class GsTaichiAssertionError(GsTaichiRuntimeError, AssertionError):
+    """Thrown when assertion fails at runtime."""
+
+    pass
+
+
+class GsTaichiRuntimeTypeError(GsTaichiRuntimeError, TypeError):
+    @staticmethod
+    def get(pos, needed, provided):
+        return GsTaichiRuntimeTypeError(
+            f"Argument {pos} (type={provided}) cannot be converted into required type {needed}"
+        )
+
+    @staticmethod
+    def get_ret(needed, provided):
+        return GsTaichiRuntimeTypeError(f"Return (type={provided}) cannot be converted into required type {needed}")
+
+
+def handle_exception_from_cpp(exc):
+    if isinstance(exc, core.GsTaichiTypeError):
+        return GsTaichiTypeError(str(exc))
+    if isinstance(exc, core.GsTaichiSyntaxError):
+        return GsTaichiSyntaxError(str(exc))
+    if isinstance(exc, core.GsTaichiIndexError):
+        return GsTaichiIndexError(str(exc))
+    if isinstance(exc, core.GsTaichiAssertionError):
+        return GsTaichiAssertionError(str(exc))
+    return exc
+
+
+__all__ = [
+    "GsTaichiSyntaxError",
+    "GsTaichiTypeError",
+    "GsTaichiCompilationError",
+    "GsTaichiNameError",
+    "GsTaichiRuntimeError",
+    "GsTaichiRuntimeTypeError",
+    "GsTaichiAssertionError",
+]

{taichi → gstaichi}/lang/expr.py

@@ -2,21 +2,21 @@ from typing import TYPE_CHECKING
 
 import numpy as np
 
-from taichi._lib import core as _ti_core
-from taichi.lang import impl
-from taichi.lang.common_ops import TaichiOperations
-from taichi.lang.exception import TaichiCompilationError, TaichiTypeError
-from taichi.lang.matrix import make_matrix
-from taichi.lang.util import is_matrix_class, is_taichi_class, to_numpy_type
-from taichi.types import primitive_types
-from taichi.types.primitive_types import integer_types, real_types
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.common_ops import GsTaichiOperations
+from gstaichi.lang.exception import GsTaichiCompilationError, GsTaichiTypeError
+from gstaichi.lang.matrix import make_matrix
+from gstaichi.lang.util import is_gstaichi_class, is_matrix_class, to_numpy_type
+from gstaichi.types import primitive_types
+from gstaichi.types.primitive_types import integer_types, real_types
 
 if TYPE_CHECKING:
-    from taichi.lang.ast.ast_transformer_utils import ASTBuilder
+    from gstaichi.lang.ast.ast_transformer_utils import ASTBuilder
 
 
 # Scalar, basic data type
-class Expr(TaichiOperations):
+class Expr(GsTaichiOperations):
     """A Python-side Expr wrapper, whose member variable `ptr` is an instance of C++ Expr class. A C++ Expr object contains member variable `expr` which holds an instance of C++ Expression class."""
 
     def __init__(self, *args, dbg_info=None, dtype=None):
@@ -24,7 +24,7 @@ class Expr(TaichiOperations):
         self.ptr_type_checked = False
         self.declaration_tb: str = ""
         if len(args) == 1:
-            if isinstance(args[0], _ti_core.Expr):
+            if isinstance(args[0], _ti_core.ExprCxx):
                 self.ptr = args[0]
             elif isinstance(args[0], Expr):
                 self.ptr = args[0].ptr
@@ -39,7 +39,7 @@ class Expr(TaichiOperations):
                 arg = args[0]
                 if isinstance(arg, np.ndarray):
                     if arg.shape:
-                        raise TaichiTypeError(
+                        raise GsTaichiTypeError(
                             "Only 0-dimensional numpy array can be used to initialize a scalar expression"
                         )
                     arg = arg.dtype.type(arg)
@@ -63,7 +63,7 @@ class Expr(TaichiOperations):
 
     def get_shape(self):
         if not self.is_tensor():
-            raise TaichiCompilationError(f"Getting shape of non-tensor type: {self.ptr.get_rvalue_type()}")
+            raise GsTaichiCompilationError(f"Getting shape of non-tensor type: {self.ptr.get_rvalue_type()}")
         shape = self.ptr.get_shape()
         assert shape is not None
         return tuple(shape)
@@ -72,14 +72,14 @@ class Expr(TaichiOperations):
     def n(self):
         shape = self.get_shape()
         if len(shape) < 1:
-            raise TaichiCompilationError(f"Getting n of tensor type < 1D: {self.ptr.get_rvalue_type()}")
+            raise GsTaichiCompilationError(f"Getting n of tensor type < 1D: {self.ptr.get_rvalue_type()}")
         return shape[0]
 
     @property
     def m(self):
         shape = self.get_shape()
         if len(shape) < 2:
-            raise TaichiCompilationError(f"Getting m of tensor type < 2D: {self.ptr.get_rvalue_type()}")
+            raise GsTaichiCompilationError(f"Getting m of tensor type < 2D: {self.ptr.get_rvalue_type()}")
         return shape[1]
 
     def __hash__(self):
@@ -116,7 +116,7 @@ def make_constant_expr(val, dtype):
     if isinstance(val, (float, np.floating)):
         constant_dtype = impl.get_runtime().default_fp if dtype is None else dtype
         if constant_dtype not in real_types:
-            raise TaichiTypeError(
+            raise GsTaichiTypeError(
                 "Floating-point literals must be annotated with a floating-point type. For type casting, use `ti.cast`."
             )
         return Expr(_ti_core.make_const_expr_fp(constant_dtype, val))
@@ -124,19 +124,19 @@ def make_constant_expr(val, dtype):
     if isinstance(val, (int, np.integer)):
         constant_dtype = impl.get_runtime().default_ip if dtype is None else dtype
         if constant_dtype not in integer_types:
-            raise TaichiTypeError(
+            raise GsTaichiTypeError(
                 "Integer literals must be annotated with a integer type. For type casting, use `ti.cast`."
             )
         if _check_in_range(to_numpy_type(constant_dtype), val):
             return Expr(_ti_core.make_const_expr_int(constant_dtype, _clamp_unsigned_to_range(np.int64, val)))
         if dtype is None:
-            raise TaichiTypeError(
+            raise GsTaichiTypeError(
                 f"Integer literal {val} exceeded the range of default_ip: {impl.get_runtime().default_ip}, please specify the dtype via e.g. `ti.u64({val})` or set a different `default_ip` in `ti.init()`"
             )
         else:
-            raise TaichiTypeError(f"Integer literal {val} exceeded the range of specified dtype: {dtype}")
+            raise GsTaichiTypeError(f"Integer literal {val} exceeded the range of specified dtype: {dtype}")
 
-    raise TaichiTypeError(f"Invalid constant scalar data type: {type(val)}")
+    raise GsTaichiTypeError(f"Invalid constant scalar data type: {type(val)}")
 
 
 def make_var_list(size: int, ast_builder: "ASTBuilder | None" = None):
@@ -151,7 +151,7 @@ def make_var_list(size: int, ast_builder: "ASTBuilder | None" = None):
 
 
 def make_expr_group(*exprs):
-    from taichi.lang.matrix import Matrix  # pylint: disable=C0415
+    from gstaichi.lang.matrix import Matrix  # pylint: disable=C0415
 
     if len(exprs) == 1:
         if isinstance(exprs[0], (list, tuple)):
@@ -169,7 +169,7 @@ def make_expr_group(*exprs):
 
 
 def _get_flattened_ptrs(val):
-    if is_taichi_class(val):
+    if is_gstaichi_class(val):
         ptrs = []
         for item in val._members:
             ptrs.extend(_get_flattened_ptrs(item))