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

{taichi → gstaichi}/lang/kernel_impl.py

@@ -1,5 +1,3 @@
-# type: ignore
-
 import ast
 import dataclasses
 import functools
@@ -15,60 +13,181 @@ import time
 import types
 import typing
 import warnings
-import weakref
-from typing import Any, Callable, Type, Union
+from typing import Any, Callable, Type
 
 import numpy as np
 
-import taichi.lang
-import taichi.lang._ndarray
-import taichi.lang._texture
-import taichi.lang.expr
-import taichi.lang.snode
-import taichi.types.annotations
-from taichi import _logging
-from taichi._lib import core as _ti_core
-from taichi._lib.core.taichi_python import ASTBuilder
-from taichi.lang import impl, ops, runtime_ops
-from taichi.lang._wrap_inspect import getsourcefile, getsourcelines
-from taichi.lang.any_array import AnyArray
-from taichi.lang.argpack import ArgPack, ArgPackType
-from taichi.lang.ast import (
+import gstaichi.lang
+import gstaichi.lang._ndarray
+import gstaichi.lang._texture
+import gstaichi.types.annotations
+from gstaichi import _logging
+from gstaichi._lib import core as _ti_core
+from gstaichi._lib.core.gstaichi_python import (
+    ASTBuilder,
+    FunctionKey,
+    KernelCxx,
+    KernelLaunchContext,
+)
+from gstaichi.lang import impl, ops, runtime_ops
+from gstaichi.lang._template_mapper import GsTaichiCallableTemplateMapper
+from gstaichi.lang._wrap_inspect import getsourcefile, getsourcelines
+from gstaichi.lang.any_array import AnyArray
+from gstaichi.lang.argpack import ArgPack, ArgPackType
+from gstaichi.lang.ast import (
     ASTTransformerContext,
     KernelSimplicityASTChecker,
     transform_tree,
 )
-from taichi.lang.ast.ast_transformer_utils import ReturnStatus
-from taichi.lang.exception import (
-    TaichiCompilationError,
-    TaichiRuntimeError,
-    TaichiRuntimeTypeError,
-    TaichiSyntaxError,
-    TaichiTypeError,
+from gstaichi.lang.ast.ast_transformer_utils import ReturnStatus
+from gstaichi.lang.exception import (
+    GsTaichiCompilationError,
+    GsTaichiRuntimeError,
+    GsTaichiRuntimeTypeError,
+    GsTaichiSyntaxError,
+    GsTaichiTypeError,
     handle_exception_from_cpp,
 )
-from taichi.lang.expr import Expr
-from taichi.lang.kernel_arguments import KernelArgument
-from taichi.lang.matrix import MatrixType
-from taichi.lang.shell import _shell_pop_print
-from taichi.lang.struct import StructType
-from taichi.lang.util import cook_dtype, has_paddle, has_pytorch, to_taichi_type
-from taichi.types import (
+from gstaichi.lang.expr import Expr
+from gstaichi.lang.kernel_arguments import KernelArgument
+from gstaichi.lang.matrix import MatrixType
+from gstaichi.lang.shell import _shell_pop_print
+from gstaichi.lang.struct import StructType
+from gstaichi.lang.util import cook_dtype, has_paddle, has_pytorch
+from gstaichi.types import (
     ndarray_type,
     primitive_types,
     sparse_matrix_builder,
     template,
     texture_type,
 )
-from taichi.types.compound_types import CompoundType
-from taichi.types.enums import AutodiffMode, Layout
-from taichi.types.utils import is_signed
+from gstaichi.types.compound_types import CompoundType
+from gstaichi.types.enums import AutodiffMode, Layout
+from gstaichi.types.utils import is_signed
+
+CompiledKernelKeyType = tuple[Callable, int, AutodiffMode]
 
 
-def func(fn: Callable, is_real_function: bool = False):
-    """Marks a function as callable in Taichi-scope.
+class GsTaichiCallable:
+    """
+    BoundGsTaichiCallable is used to enable wrapping a bindable function with a class.
+
+    Design requirements for GsTaichiCallable:
+    - wrap/contain a reference to a class Func instance, and allow (the GsTaichiCallable) being passed around
+      like normal function pointer
+    - expose attributes of the wrapped class Func, such as `_if_real_function`, `_primal`, etc
+    - allow for (now limited) strong typing, and enable type checkers, such as pyright/mypy
+        - currently GsTaichiCallable is a shared type used for all functions marked with @ti.func, @ti.kernel,
+          python functions (?)
+        - note: current type-checking implementation does not distinguish between different type flavors of
+          GsTaichiCallable, with different values of `_if_real_function`, `_primal`, etc
+    - handle not only class-less functions, but also class-instance methods (where determining the `self`
+      reference is a challenge)
+
+    Let's take the following example:
+
+    def test_ptr_class_func():
+    @ti.data_oriented
+    class MyClass:
+        def __init__(self):
+            self.a = ti.field(dtype=ti.f32, shape=(3))
+
+        def add2numbers_py(self, x, y):
+            return x + y
+
+        @ti.func
+        def add2numbers_func(self, x, y):
+            return x + y
+
+        @ti.kernel
+        def func(self):
+            a, add_py, add_func = ti.static(self.a, self.add2numbers_py, self.add2numbers_func)
+            a[0] = add_py(2, 3)
+            a[1] = add_func(3, 7)
+
+    (taken from test_ptr_assign.py).
+
+    When the @ti.func decorator is parsed, the function `add2numbers_func` exists, but there is not yet any `self`
+    - it is not possible for the method to be bound, to a `self` instance
+    - however, the @ti.func annotation, runs the kernel_imp.py::func function --- it is at this point
+      that GsTaichi's original code creates a class Func instance (that wraps the add2numbers_func)
+      and immediately we create a GsTaichiCallable instance that wraps the Func instance.
+    - effectively, we have two layers of wrapping GsTaichiCallable->Func->function pointer
+      (actual function definition)
+    - later on, when we call self.add2numbers_py, here:
+
+            a, add_py, add_func = ti.static(self.a, self.add2numbers_py, self.add2numbers_func)
+
+      ... we want to call the bound method, `self.add2numbers_py`.
+    - an actual python function reference, created by doing somevar = MyClass.add2numbers, can automatically
+      binds to self, when called from self in this way (however, add2numbers_py is actually a class
+      Func instance, wrapping python function reference -- now also all wrapped by a GsTaichiCallable
+      instance -- returned by the kernel_impl.py::func function, run by @ti.func)
+    - however, in order to be able to add strongly typed attributes to the wrapped python function, we need
+      to wrap the wrapped python function in a class
+    - the wrapped python function, wrapped in a GsTaichiCallable class (which is callable, and will
+      execute the underlying double-wrapped python function), will NOT automatically bind
+    - when we invoke GsTaichiCallable, the wrapped function is invoked. The wrapped function is unbound, and
+      so `self` is not automatically passed in, as an argument, and things break
+
+    To address this we need to use the `__get__` method, in our function wrapper, ie GsTaichiCallable,
+    and have the `__get__` method return the `BoundGsTaichiCallable` object. The `__get__` method handles
+    running the binding for us, and effectively binds `BoundFunc` object to `self` object, by passing
+    in the instance, as an argument into `BoundGsTaichiCallable.__init__`.
+
+    `BoundFunc` can then be used as a normal bound func - even though it's just an object instance -
+    using its `__call__` method. Effectively, at the time of actually invoking the underlying python
+    function, we have 3 layers of wrapper instances:
+        BoundGsTaichiCallabe -> GsTaichiCallable -> Func -> python function reference/definition
+    """
+
+    def __init__(self, fn: Callable, wrapper: Callable) -> None:
+        self.fn: Callable = fn
+        self.wrapper: Callable = wrapper
+        self._is_real_function: bool = False
+        self._is_gstaichi_function: bool = False
+        self._is_wrapped_kernel: bool = False
+        self._is_classkernel: bool = False
+        self._primal: Kernel | None = None
+        self._adjoint: Kernel | None = None
+        self.grad: Kernel | None = None
+        self._is_staticmethod: bool = False
+        functools.update_wrapper(self, fn)
+
+    def __call__(self, *args, **kwargs):
+        return self.wrapper.__call__(*args, **kwargs)
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+        return BoundGsTaichiCallable(instance, self)
+
+
+class BoundGsTaichiCallable:
+    def __init__(self, instance: Any, gstaichi_callable: "GsTaichiCallable"):
+        self.wrapper = gstaichi_callable.wrapper
+        self.instance = instance
+        self.gstaichi_callable = gstaichi_callable
+
+    def __call__(self, *args, **kwargs):
+        return self.wrapper(self.instance, *args, **kwargs)
+
+    def __getattr__(self, k: str) -> Any:
+        res = getattr(self.gstaichi_callable, k)
+        return res
+
+    def __setattr__(self, k: str, v: Any) -> None:
+        # Note: these have to match the name of any attributes on this class.
+        if k in ("wrapper", "instance", "gstaichi_callable"):
+            object.__setattr__(self, k, v)
+        else:
+            setattr(self.gstaichi_callable, k, v)
 
-    This decorator transforms a Python function into a Taichi one. Taichi
+
+def func(fn: Callable, is_real_function: bool = False) -> GsTaichiCallable:
+    """Marks a function as callable in GsTaichi-scope.
+
+    This decorator transforms a Python function into a GsTaichi one. GsTaichi
     will JIT compile it into native instructions.
 
     Args:
@@ -91,29 +210,24 @@ def func(fn: Callable, is_real_function: bool = False):
     is_classfunc = _inside_class(level_of_class_stackframe=3 + is_real_function)
 
     fun = Func(fn, _classfunc=is_classfunc, is_real_function=is_real_function)
-
-    @functools.wraps(fn)
-    def decorated(*args, **kwargs):
-        return fun.__call__(*args, **kwargs)
-
-    decorated._is_taichi_function = True
-    decorated._is_real_function = is_real_function
-    decorated.func = fun
-    return decorated
+    gstaichi_callable = GsTaichiCallable(fn, fun)
+    gstaichi_callable._is_gstaichi_function = True
+    gstaichi_callable._is_real_function = is_real_function
+    return gstaichi_callable
 
 
-def real_func(fn: Callable):
+def real_func(fn: Callable) -> GsTaichiCallable:
     return func(fn, is_real_function=True)
 
 
-def pyfunc(fn: Callable):
-    """Marks a function as callable in both Taichi and Python scopes.
+def pyfunc(fn: Callable) -> GsTaichiCallable:
+    """Marks a function as callable in both GsTaichi and Python scopes.
 
-    When called inside the Taichi scope, Taichi will JIT compile it into
+    When called inside the GsTaichi scope, GsTaichi will JIT compile it into
     native instructions. Otherwise it will be invoked directly as a
     Python function.
 
-    See also :func:`~taichi.lang.kernel_impl.func`.
+    See also :func:`~gstaichi.lang.kernel_impl.func`.
 
     Args:
         fn (Callable): The Python function to be decorated
@@ -123,33 +237,28 @@ def pyfunc(fn: Callable):
     """
     is_classfunc = _inside_class(level_of_class_stackframe=3)
     fun = Func(fn, _classfunc=is_classfunc, _pyfunc=True)
-
-    @functools.wraps(fn)
-    def decorated(*args, **kwargs):
-        return fun.__call__(*args, **kwargs)
-
-    decorated._is_taichi_function = True
-    decorated._is_real_function = False
-    decorated.func = fun
-    return decorated
+    gstaichi_callable = GsTaichiCallable(fn, fun)
+    gstaichi_callable._is_gstaichi_function = True
+    gstaichi_callable._is_real_function = False
+    return gstaichi_callable
 
 
 def _get_tree_and_ctx(
     self: "Func | Kernel",
+    args: tuple[Any, ...],
     excluded_parameters=(),
     is_kernel: bool = True,
     arg_features=None,
-    args=None,
     ast_builder: ASTBuilder | None = None,
     is_real_function: bool = False,
-):
+) -> tuple[ast.Module, ASTTransformerContext]:
     file = getsourcefile(self.func)
     src, start_lineno = getsourcelines(self.func)
     src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
     tree = ast.parse(textwrap.dedent("\n".join(src)))
 
     func_body = tree.body[0]
-    func_body.decorator_list = []
+    func_body.decorator_list = []  # type: ignore , kick that can down the road...
 
     global_vars = _get_global_vars(self.func)
 
@@ -196,17 +305,18 @@ def expand_func_arguments(arguments: list[KernelArgument]) -> list[KernelArgumen
     return new_arguments
 
 
-def _process_args(self: "Func | Kernel", is_func: bool, args, kwargs):
+def _process_args(self: "Func | Kernel", is_func: bool, args: tuple[Any, ...], kwargs) -> tuple[Any, ...]:
     if is_func:
         self.arguments = expand_func_arguments(self.arguments)
     fused_args = [argument.default for argument in self.arguments]
+    ret: list[Any] = [argument.default for argument in self.arguments]
     len_args = len(args)
 
     if len_args > len(fused_args):
         arg_str = ", ".join([str(arg) for arg in args])
         expected_str = ", ".join([f"{arg.name} : {arg.annotation}" for arg in self.arguments])
         msg = f"Too many arguments. Expected ({expected_str}), got ({arg_str})."
-        raise TaichiSyntaxError(msg)
+        raise GsTaichiSyntaxError(msg)
 
     for i, arg in enumerate(args):
         fused_args[i] = arg
@@ -216,19 +326,19 @@ def _process_args(self: "Func | Kernel", is_func: bool, args, kwargs):
         for i, arg in enumerate(self.arguments):
             if key == arg.name:
                 if i < len_args:
-                    raise TaichiSyntaxError(f"Multiple values for argument '{key}'.")
+                    raise GsTaichiSyntaxError(f"Multiple values for argument '{key}'.")
                 fused_args[i] = value
                 found = True
                 break
         if not found:
-            raise TaichiSyntaxError(f"Unexpected argument '{key}'.")
+            raise GsTaichiSyntaxError(f"Unexpected argument '{key}'.")
 
     for i, arg in enumerate(fused_args):
         if arg is inspect.Parameter.empty:
             if self.arguments[i].annotation is inspect._empty:
-                raise TaichiSyntaxError(f"Parameter `{self.arguments[i].name}` missing.")
+                raise GsTaichiSyntaxError(f"Parameter `{self.arguments[i].name}` missing.")
             else:
-                raise TaichiSyntaxError(
+                raise GsTaichiSyntaxError(
                     f"Parameter `{self.arguments[i].name} : {self.arguments[i].annotation}` missing."
                 )
 
@@ -237,7 +347,7 @@ def _process_args(self: "Func | Kernel", is_func: bool, args, kwargs):
 
 def unpack_ndarray_struct(tree: ast.Module, struct_locals: set[str]) -> ast.Module:
     class AttributeToNameTransformer(ast.NodeTransformer):
-        def visit_Attribute(self, node: ast.AST):
+        def visit_Attribute(self, node: ast.Attribute):
             if isinstance(node.value, ast.Attribute):
                 return node
             if not isinstance(node.value, ast.Name):
@@ -278,7 +388,7 @@ def extract_struct_locals_from_context(ctx: ASTTransformerContext):
 class Func:
     function_counter = 0
 
-    def __init__(self, _func: Callable, _classfunc=False, _pyfunc=False, is_real_function=False):
+    def __init__(self, _func: Callable, _classfunc=False, _pyfunc=False, is_real_function=False) -> None:
         self.func = _func
         self.func_id = Func.function_counter
         Func.function_counter += 1
@@ -294,22 +404,22 @@ class Func:
         for i, arg in enumerate(self.arguments):
             if arg.annotation == template or isinstance(arg.annotation, template):
                 self.template_slot_locations.append(i)
-        self.mapper = TaichiCallableTemplateMapper(self.arguments, self.template_slot_locations)
-        self.taichi_functions = {}  # The |Function| class in C++
+        self.mapper = GsTaichiCallableTemplateMapper(self.arguments, self.template_slot_locations)
+        self.gstaichi_functions = {}  # The |Function| class in C++
         self.has_print = False
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs) -> Any:
         args = _process_args(self, is_func=True, args=args, kwargs=kwargs)
 
         if not impl.inside_kernel():
             if not self.pyfunc:
-                raise TaichiSyntaxError("Taichi functions cannot be called from Python-scope.")
+                raise GsTaichiSyntaxError("GsTaichi functions cannot be called from Python-scope.")
             return self.func(*args)
 
         current_kernel = impl.get_runtime().current_kernel
         if self.is_real_function:
             if current_kernel.autodiff_mode != AutodiffMode.NONE:
-                raise TaichiSyntaxError("Real function in gradient kernels unsupported.")
+                raise GsTaichiSyntaxError("Real function in gradient kernels unsupported.")
             instance_id, arg_features = self.mapper.lookup(args)
             key = _ti_core.FunctionKey(self.func.__name__, self.func_id, instance_id)
             if key.instance_id not in self.compiled:
@@ -328,10 +438,10 @@ class Func:
         ret = transform_tree(tree, ctx)
         if not self.is_real_function:
             if self.return_type and ctx.returned != ReturnStatus.ReturnedValue:
-                raise TaichiSyntaxError("Function has a return type but does not have a return statement")
+                raise GsTaichiSyntaxError("Function has a return type but does not have a return statement")
         return ret
 
-    def func_call_rvalue(self, key, args):
+    def func_call_rvalue(self, key: FunctionKey, args: tuple[Any, ...]) -> Any:
         # Skip the template args, e.g., |self|
         assert self.is_real_function
         non_template_args = []
@@ -345,7 +455,7 @@ class Func:
                     non_template_args.append(_ti_core.make_reference(args[i].ptr, dbg_info))
                 elif isinstance(anno, ndarray_type.NdarrayType):
                     if not isinstance(args[i], AnyArray):
-                        raise TaichiTypeError(
+                        raise GsTaichiTypeError(
                             f"Expected ndarray in the kernel argument for argument {kernel_arg.name}, got {args[i]}"
                         )
                     non_template_args += _ti_core.get_external_tensor_real_func_args(args[i].ptr, dbg_info)
@@ -355,7 +465,7 @@ class Func:
         compiling_callable = impl.get_runtime().compiling_callable
         assert compiling_callable is not None
         func_call = compiling_callable.ast_builder().insert_func_call(
-            self.taichi_functions[key.instance_id], non_template_args, dbg_info
+            self.gstaichi_functions[key.instance_id], non_template_args, dbg_info
         )
         if self.return_type is None:
             return None
@@ -372,14 +482,14 @@ class Func:
                     )
                 )
             elif isinstance(return_type, (StructType, MatrixType)):
-                ret.append(return_type.from_taichi_object(func_call, (i,)))
+                ret.append(return_type.from_gstaichi_object(func_call, (i,)))
             else:
-                raise TaichiTypeError(f"Unsupported return type for return value {i}: {return_type}")
+                raise GsTaichiTypeError(f"Unsupported return type for return value {i}: {return_type}")
         if len(ret) == 1:
             return ret[0]
         return tuple(ret)
 
-    def do_compile(self, key, args, arg_features):
+    def do_compile(self, key: FunctionKey, args: tuple[Any, ...], arg_features: tuple[Any, ...]) -> None:
         tree, ctx = _get_tree_and_ctx(
             self, is_kernel=False, args=args, arg_features=arg_features, is_real_function=self.is_real_function
         )
@@ -392,36 +502,42 @@ class Func:
             transform_tree(tree, ctx)
             impl.get_runtime().compiling_callable = old_callable
 
-        self.taichi_functions[key.instance_id] = fn
+        self.gstaichi_functions[key.instance_id] = fn
         self.compiled[key.instance_id] = func_body
-        self.taichi_functions[key.instance_id].set_function_body(func_body)
+        self.gstaichi_functions[key.instance_id].set_function_body(func_body)
 
     def extract_arguments(self) -> None:
         sig = inspect.signature(self.func)
         if sig.return_annotation not in (inspect.Signature.empty, None):
             self.return_type = sig.return_annotation
             if (
-                isinstance(self.return_type, (types.GenericAlias, typing._GenericAlias))
-                and self.return_type.__origin__ is tuple
+                isinstance(self.return_type, (types.GenericAlias, typing._GenericAlias))  # type: ignore
+                and self.return_type.__origin__ is tuple  # type: ignore
             ):
-                self.return_type = self.return_type.__args__
+                self.return_type = self.return_type.__args__  # type: ignore
+            if self.return_type is None:
+                return
             if not isinstance(self.return_type, (list, tuple)):
                 self.return_type = (self.return_type,)
             for i, return_type in enumerate(self.return_type):
                 if return_type is Ellipsis:
-                    raise TaichiSyntaxError("Ellipsis is not supported in return type annotations")
+                    raise GsTaichiSyntaxError("Ellipsis is not supported in return type annotations")
         params = sig.parameters
         arg_names = params.keys()
         for i, arg_name in enumerate(arg_names):
             param = params[arg_name]
             if param.kind == inspect.Parameter.VAR_KEYWORD:
-                raise TaichiSyntaxError("Taichi functions do not support variable keyword parameters (i.e., **kwargs)")
+                raise GsTaichiSyntaxError(
+                    "GsTaichi functions do not support variable keyword parameters (i.e., **kwargs)"
+                )
             if param.kind == inspect.Parameter.VAR_POSITIONAL:
-                raise TaichiSyntaxError("Taichi functions do not support variable positional parameters (i.e., *args)")
+                raise GsTaichiSyntaxError(
+                    "GsTaichi functions do not support variable positional parameters (i.e., *args)"
+                )
             if param.kind == inspect.Parameter.KEYWORD_ONLY:
-                raise TaichiSyntaxError("Taichi functions do not support keyword parameters")
+                raise GsTaichiSyntaxError("GsTaichi functions do not support keyword parameters")
             if param.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD:
-                raise TaichiSyntaxError('Taichi functions only support "positional or keyword" parameters')
+                raise GsTaichiSyntaxError('GsTaichi functions only support "positional or keyword" parameters')
             annotation = param.annotation
             if annotation is inspect.Parameter.empty:
                 if i == 0 and self.classfunc:
@@ -429,8 +545,8 @@ class Func:
                 # TODO: pyfunc also need type annotation check when real function is enabled,
                 #       but that has to happen at runtime when we know which scope it's called from.
                 elif not self.pyfunc and self.is_real_function:
-                    raise TaichiSyntaxError(
-                        f"Taichi function `{self.func.__name__}` parameter `{arg_name}` must be type annotated"
+                    raise GsTaichiSyntaxError(
+                        f"GsTaichi function `{self.func.__name__}` parameter `{arg_name}` must be type annotated"
                     )
             else:
                 if isinstance(annotation, ndarray_type.NdarrayType):
@@ -441,198 +557,24 @@ class Func:
                     pass
                 elif id(annotation) in primitive_types.type_ids:
                     pass
-                elif type(annotation) == taichi.types.annotations.Template:
+                elif type(annotation) == gstaichi.types.annotations.Template:
                     pass
-                elif isinstance(annotation, template) or annotation == taichi.types.annotations.Template:
+                elif isinstance(annotation, template) or annotation == gstaichi.types.annotations.Template:
                     pass
                 elif isinstance(annotation, primitive_types.RefType):
                     pass
                 elif isinstance(annotation, type) and dataclasses.is_dataclass(annotation):
                     pass
                 else:
-                    raise TaichiSyntaxError(f"Invalid type annotation (argument {i}) of Taichi function: {annotation}")
+                    raise GsTaichiSyntaxError(
+                        f"Invalid type annotation (argument {i}) of GsTaichi function: {annotation}"
+                    )
             self.arguments.append(KernelArgument(annotation, param.name, param.default))
             self.orig_arguments.append(KernelArgument(annotation, param.name, param.default))
 
 
-AnnotationType = Union[
-    template,
-    ArgPackType,
-    "texture_type.TextureType",
-    "texture_type.RWTextureType",
-    ndarray_type.NdarrayType,
-    sparse_matrix_builder,
-    Any,
-]
-
-
-class TaichiCallableTemplateMapper:
-    """
-    This should probably be renamed to sometihng like FeatureMapper, or
-    FeatureExtractor, since:
-    - it's not specific to templates
-    - it extracts what are later called 'features', for example for ndarray this includes:
-        - element type
-        - number dimensions
-        - needs grad (or not)
-    - these are returned as a heterogeneous tuple, whose contents depends on the type
-    """
-
-    def __init__(self, arguments: list[KernelArgument], template_slot_locations: list[int]) -> None:
-        self.arguments = arguments
-        self.num_args = len(arguments)
-        self.template_slot_locations = template_slot_locations
-        self.mapping = {}
-
-    @staticmethod
-    def extract_arg(arg, annotation: AnnotationType, arg_name: str):
-        if annotation == template or isinstance(annotation, template):
-            if isinstance(arg, taichi.lang.snode.SNode):
-                return arg.ptr
-            if isinstance(arg, taichi.lang.expr.Expr):
-                return arg.ptr.get_underlying_ptr_address()
-            if isinstance(arg, _ti_core.Expr):
-                return arg.get_underlying_ptr_address()
-            if isinstance(arg, tuple):
-                return tuple(TaichiCallableTemplateMapper.extract_arg(item, annotation, arg_name) for item in arg)
-            if isinstance(arg, taichi.lang._ndarray.Ndarray):
-                raise TaichiRuntimeTypeError(
-                    "Ndarray shouldn't be passed in via `ti.template()`, please annotate your kernel using `ti.types.ndarray(...)` instead"
-                )
-
-            if isinstance(arg, (list, tuple, dict, set)) or hasattr(arg, "_data_oriented"):
-                # [Composite arguments] Return weak reference to the object
-                # Taichi kernel will cache the extracted arguments, thus we can't simply return the original argument.
-                # Instead, a weak reference to the original value is returned to avoid memory leak.
-
-                # TODO(zhanlue): replacing "tuple(args)" with "hash of argument values"
-                # This can resolve the following issues:
-                # 1. Invalid weak-ref will leave a dead(dangling) entry in both caches: "self.mapping" and "self.compiled_functions"
-                # 2. Different argument instances with same type and same value, will get templatized into seperate kernels.
-                return weakref.ref(arg)
-
-            # [Primitive arguments] Return the value
-            return arg
-        if isinstance(annotation, ArgPackType):
-            if not isinstance(arg, ArgPack):
-                raise TaichiRuntimeTypeError(f"Argument {arg_name} must be a argument pack, got {type(arg)}")
-            return tuple(
-                TaichiCallableTemplateMapper.extract_arg(arg[name], dtype, arg_name)
-                for index, (name, dtype) in enumerate(annotation.members.items())
-            )
-        if dataclasses.is_dataclass(annotation):
-            _res_l = []
-            for field in dataclasses.fields(annotation):
-                field_value = getattr(arg, field.name)
-                arg_name = f"__ti_{arg_name}_{field.name}"
-                field_extracted = TaichiCallableTemplateMapper.extract_arg(field_value, field.type, arg_name)
-                _res_l.append(field_extracted)
-            return tuple(_res_l)
-        if isinstance(annotation, texture_type.TextureType):
-            if not isinstance(arg, taichi.lang._texture.Texture):
-                raise TaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
-            if arg.num_dims != annotation.num_dimensions:
-                raise TaichiRuntimeTypeError(
-                    f"TextureType dimension mismatch for argument {arg_name}: expected {annotation.num_dimensions}, got {arg.num_dims}"
-                )
-            return (arg.num_dims,)
-        if isinstance(annotation, texture_type.RWTextureType):
-            if not isinstance(arg, taichi.lang._texture.Texture):
-                raise TaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
-            if arg.num_dims != annotation.num_dimensions:
-                raise TaichiRuntimeTypeError(
-                    f"RWTextureType dimension mismatch for argument {arg_name}: expected {annotation.num_dimensions}, got {arg.num_dims}"
-                )
-            if arg.fmt != annotation.fmt:
-                raise TaichiRuntimeTypeError(
-                    f"RWTextureType format mismatch for argument {arg_name}: expected {annotation.fmt}, got {arg.fmt}"
-                )
-            # (penguinliong) '0' is the assumed LOD level. We currently don't
-            # support mip-mapping.
-            return arg.num_dims, arg.fmt, 0
-        if isinstance(annotation, ndarray_type.NdarrayType):
-            if isinstance(arg, taichi.lang._ndarray.Ndarray):
-                annotation.check_matched(arg.get_type(), arg_name)
-                needs_grad = (arg.grad is not None) if annotation.needs_grad is None else annotation.needs_grad
-                assert arg.shape is not None
-                return arg.element_type, len(arg.shape), needs_grad, annotation.boundary
-            if isinstance(arg, AnyArray):
-                ty = arg.get_type()
-                annotation.check_matched(arg.get_type(), arg_name)
-                return ty.element_type, len(arg.shape), ty.needs_grad, annotation.boundary
-            # external arrays
-            shape = getattr(arg, "shape", None)
-            if shape is None:
-                raise TaichiRuntimeTypeError(f"Invalid type for argument {arg_name}, got {arg}")
-            shape = tuple(shape)
-            element_shape: tuple[int, ...] = ()
-            dtype = to_taichi_type(arg.dtype)
-            if isinstance(annotation.dtype, MatrixType):
-                if annotation.ndim is not None:
-                    if len(shape) != annotation.dtype.ndim + annotation.ndim:
-                        raise ValueError(
-                            f"Invalid value for argument {arg_name} - required array has ndim={annotation.ndim} element_dim={annotation.dtype.ndim}, "
-                            f"array with {len(shape)} dimensions is provided"
-                        )
-                else:
-                    if len(shape) < annotation.dtype.ndim:
-                        raise ValueError(
-                            f"Invalid value for argument {arg_name} - required element_dim={annotation.dtype.ndim}, "
-                            f"array with {len(shape)} dimensions is provided"
-                        )
-                element_shape = shape[-annotation.dtype.ndim :]
-                anno_element_shape = annotation.dtype.get_shape()
-                if None not in anno_element_shape and element_shape != anno_element_shape:
-                    raise ValueError(
-                        f"Invalid value for argument {arg_name} - required element_shape={anno_element_shape}, "
-                        f"array with element shape of {element_shape} is provided"
-                    )
-            elif annotation.dtype is not None:
-                # User specified scalar dtype
-                if annotation.dtype != dtype:
-                    raise ValueError(
-                        f"Invalid value for argument {arg_name} - required array has dtype={annotation.dtype.to_string()}, "
-                        f"array with dtype={dtype.to_string()} is provided"
-                    )
-
-                if annotation.ndim is not None and len(shape) != annotation.ndim:
-                    raise ValueError(
-                        f"Invalid value for argument {arg_name} - required array has ndim={annotation.ndim}, "
-                        f"array with {len(shape)} dimensions is provided"
-                    )
-            needs_grad = (
-                getattr(arg, "requires_grad", False) if annotation.needs_grad is None else annotation.needs_grad
-            )
-            element_type = (
-                _ti_core.get_type_factory_instance().get_tensor_type(element_shape, dtype)
-                if len(element_shape) != 0
-                else arg.dtype
-            )
-            return element_type, len(shape) - len(element_shape), needs_grad, annotation.boundary
-        if isinstance(annotation, sparse_matrix_builder):
-            return arg.dtype
-        # Use '#' as a placeholder because other kinds of arguments are not involved in template instantiation
-        return "#"
-
-    def extract(self, args):
-        extracted = []
-        for arg, kernel_arg in zip(args, self.arguments):
-            extracted.append(self.extract_arg(arg, kernel_arg.annotation, kernel_arg.name))
-        return tuple(extracted)
-
-    def lookup(self, args):
-        if len(args) != self.num_args:
-            raise TypeError(f"{self.num_args} argument(s) needed but {len(args)} provided.")
-
-        key = self.extract(args)
-        if key not in self.mapping:
-            count = len(self.mapping)
-            self.mapping[key] = count
-        return self.mapping[key], key
-
-
-def _get_global_vars(_func):
-    # Discussions: https://github.com/taichi-dev/taichi/issues/282
+def _get_global_vars(_func: Callable) -> dict[str, Any]:
+    # Discussions: https://github.com/taichi-dev/gstaichi/issues/282
     global_vars = _func.__globals__.copy()
 
     freevar_names = _func.__code__.co_freevars
@@ -648,7 +590,7 @@ def _get_global_vars(_func):
 class Kernel:
     counter = 0
 
-    def __init__(self, _func: Callable, autodiff_mode, _classkernel=False):
+    def __init__(self, _func: Callable, autodiff_mode: AutodiffMode, _classkernel=False) -> None:
         self.func = _func
         self.kernel_counter = Kernel.counter
         Kernel.counter += 1
@@ -668,27 +610,27 @@ class Kernel:
         for i, arg in enumerate(self.arguments):
             if arg.annotation == template or isinstance(arg.annotation, template):
                 self.template_slot_locations.append(i)
-        self.mapper = TaichiCallableTemplateMapper(self.arguments, self.template_slot_locations)
+        self.mapper = GsTaichiCallableTemplateMapper(self.arguments, self.template_slot_locations)
         impl.get_runtime().kernels.append(self)
         self.reset()
         self.kernel_cpp = None
-        self.compiled_kernels = {}
+        self.compiled_kernels: dict[CompiledKernelKeyType, KernelCxx] = {}
         self.has_print = False
 
     def ast_builder(self) -> ASTBuilder:
         assert self.kernel_cpp is not None
         return self.kernel_cpp.ast_builder()
 
-    def reset(self):
+    def reset(self) -> None:
         self.runtime = impl.get_runtime()
         self.compiled_kernels = {}
 
-    def extract_arguments(self):
+    def extract_arguments(self) -> None:
         sig = inspect.signature(self.func)
         if sig.return_annotation not in (inspect._empty, None):
             self.return_type = sig.return_annotation
             if (
-                isinstance(self.return_type, (types.GenericAlias, typing._GenericAlias))
+                isinstance(self.return_type, (types.GenericAlias, typing._GenericAlias))  # type: ignore
                 and self.return_type.__origin__ is tuple
             ):
                 self.return_type = self.return_type.__args__
@@ -696,27 +638,31 @@ class Kernel:
                 self.return_type = (self.return_type,)
             for return_type in self.return_type:
                 if return_type is Ellipsis:
-                    raise TaichiSyntaxError("Ellipsis is not supported in return type annotations")
+                    raise GsTaichiSyntaxError("Ellipsis is not supported in return type annotations")
         params = sig.parameters
         arg_names = params.keys()
         for i, arg_name in enumerate(arg_names):
             param = params[arg_name]
             if param.kind == inspect.Parameter.VAR_KEYWORD:
-                raise TaichiSyntaxError("Taichi kernels do not support variable keyword parameters (i.e., **kwargs)")
+                raise GsTaichiSyntaxError(
+                    "GsTaichi kernels do not support variable keyword parameters (i.e., **kwargs)"
+                )
             if param.kind == inspect.Parameter.VAR_POSITIONAL:
-                raise TaichiSyntaxError("Taichi kernels do not support variable positional parameters (i.e., *args)")
+                raise GsTaichiSyntaxError(
+                    "GsTaichi kernels do not support variable positional parameters (i.e., *args)"
+                )
             if param.default is not inspect.Parameter.empty:
-                raise TaichiSyntaxError("Taichi kernels do not support default values for arguments")
+                raise GsTaichiSyntaxError("GsTaichi kernels do not support default values for arguments")
             if param.kind == inspect.Parameter.KEYWORD_ONLY:
-                raise TaichiSyntaxError("Taichi kernels do not support keyword parameters")
+                raise GsTaichiSyntaxError("GsTaichi kernels do not support keyword parameters")
             if param.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD:
-                raise TaichiSyntaxError('Taichi kernels only support "positional or keyword" parameters')
+                raise GsTaichiSyntaxError('GsTaichi kernels only support "positional or keyword" parameters')
             annotation = param.annotation
             if param.annotation is inspect.Parameter.empty:
                 if i == 0 and self.classkernel:  # The |self| parameter
                     annotation = template()
                 else:
-                    raise TaichiSyntaxError("Taichi kernels parameters must be type annotated")
+                    raise GsTaichiSyntaxError("GsTaichi kernels parameters must be type annotated")
             else:
                 if isinstance(
                     annotation,
@@ -743,10 +689,12 @@ class Kernel:
                 elif isinstance(annotation, type) and dataclasses.is_dataclass(annotation):
                     pass
                 else:
-                    raise TaichiSyntaxError(f"Invalid type annotation (argument {i}) of Taichi kernel: {annotation}")
+                    raise GsTaichiSyntaxError(
+                        f"Invalid type annotation (argument {i}) of GsTaichi kernel: {annotation}"
+                    )
             self.arguments.append(KernelArgument(annotation, param.name, param.default))
 
-    def materialize(self, key, args: list[Any], arg_features):
+    def materialize(self, key: CompiledKernelKeyType | None, args: tuple[Any, ...], arg_features):
         if key is None:
             key = (self.func, 0, self.autodiff_mode)
         self.runtime.materialize()
@@ -767,15 +715,15 @@ class Kernel:
         if self.autodiff_mode != AutodiffMode.NONE:
             KernelSimplicityASTChecker(self.func).visit(tree)
 
-        # Do not change the name of 'taichi_ast_generator'
+        # Do not change the name of 'gstaichi_ast_generator'
         # The warning system needs this identifier to remove unnecessary messages
-        def taichi_ast_generator(kernel_cxx: Kernel):  # not sure if this type is correct, seems doubtful
+        def gstaichi_ast_generator(kernel_cxx: Kernel):  # not sure if this type is correct, seems doubtful
             nonlocal tree
             if self.runtime.inside_kernel:
-                raise TaichiSyntaxError(
+                raise GsTaichiSyntaxError(
                     "Kernels cannot call other kernels. I.e., nested kernels are not allowed. "
                     "Please check if you have direct/indirect invocation of kernels within kernels. "
-                    "Note that some methods provided by the Taichi standard library may invoke kernels, "
+                    "Note that some methods provided by the GsTaichi standard library may invoke kernels, "
                     "and please move their invocations to Python-scope."
                 )
             self.kernel_cpp = kernel_cxx
@@ -786,7 +734,7 @@ class Kernel:
             try:
                 ctx.ast_builder = kernel_cxx.ast_builder()
 
-                def ast_to_dict(node):
+                def ast_to_dict(node: ast.AST | list | primitive_types._python_primitive_types):
                     if isinstance(node, ast.AST):
                         fields = {k: ast_to_dict(v) for k, v in ast.iter_fields(node)}
                         return {
@@ -824,17 +772,17 @@ class Kernel:
                 transform_tree(tree, ctx)
                 if not ctx.is_real_function:
                     if self.return_type and ctx.returned != ReturnStatus.ReturnedValue:
-                        raise TaichiSyntaxError("Kernel has a return type but does not have a return statement")
+                        raise GsTaichiSyntaxError("Kernel has a return type but does not have a return statement")
             finally:
                 self.runtime.inside_kernel = False
                 self.runtime._current_kernel = None
                 self.runtime.compiling_callable = None
 
-        taichi_kernel = impl.get_runtime().prog.create_kernel(taichi_ast_generator, kernel_name, self.autodiff_mode)
+        gstaichi_kernel = impl.get_runtime().prog.create_kernel(gstaichi_ast_generator, kernel_name, self.autodiff_mode)
         assert key not in self.compiled_kernels
-        self.compiled_kernels[key] = taichi_kernel
+        self.compiled_kernels[key] = gstaichi_kernel
 
-    def launch_kernel(self, t_kernel, *args):
+    def launch_kernel(self, t_kernel: KernelCxx, *args) -> Any:
         assert len(args) == len(self.arguments), f"{len(self.arguments)} arguments needed but {len(args)} provided"
 
         tmps = []
@@ -842,25 +790,28 @@ class Kernel:
 
         actual_argument_slot = 0
         launch_ctx = t_kernel.make_launch_context()
-        max_arg_num = 64
+        max_arg_num = 512
         exceed_max_arg_num = False
 
-        def set_arg_ndarray(indices, v):
+        def set_arg_ndarray(indices: tuple[int, ...], v: gstaichi.lang._ndarray.Ndarray) -> None:
             v_primal = v.arr
             v_grad = v.grad.arr if v.grad else None
             if v_grad is None:
-                launch_ctx.set_arg_ndarray(indices, v_primal)
+                launch_ctx.set_arg_ndarray(indices, v_primal)  # type: ignore , solvable probably, just not today
             else:
-                launch_ctx.set_arg_ndarray_with_grad(indices, v_primal, v_grad)
+                launch_ctx.set_arg_ndarray_with_grad(indices, v_primal, v_grad)  # type: ignore
 
-        def set_arg_texture(indices, v):
+        def set_arg_texture(indices: tuple[int, ...], v: gstaichi.lang._texture.Texture) -> None:
             launch_ctx.set_arg_texture(indices, v.tex)
 
-        def set_arg_rw_texture(indices, v):
+        def set_arg_rw_texture(indices: tuple[int, ...], v: gstaichi.lang._texture.Texture) -> None:
             launch_ctx.set_arg_rw_texture(indices, v.tex)
 
-        def set_arg_ext_array(indices, v, needed):
-            # Element shapes are already specialized in Taichi codegen.
+        def set_arg_ext_array(indices: tuple[int, ...], v: Any, needed: ndarray_type.NdarrayType) -> None:
+            # v is things like torch Tensor and numpy array
+            # Not adding type for this, since adds additional dependencies
+            #
+            # Element shapes are already specialized in GsTaichi codegen.
             # The shape information for element dims are no longer needed.
             # Therefore we strip the element shapes from the shape vector,
             # so that it only holds "real" array shapes.
@@ -893,7 +844,7 @@ class Kernel:
                 else:
                     raise ValueError(
                         "Non contiguous numpy arrays are not supported, please call np.ascontiguousarray(arr) "
-                        "before passing it into taichi kernel."
+                        "before passing it into gstaichi kernel."
                     )
             elif has_pytorch():
                 import torch  # pylint: disable=C0415
@@ -902,9 +853,9 @@ class Kernel:
                     if not v.is_contiguous():
                         raise ValueError(
                             "Non contiguous tensors are not supported, please call tensor.contiguous() before "
-                            "passing it into taichi kernel."
+                            "passing it into gstaichi kernel."
                         )
-                    taichi_arch = self.runtime.prog.config().arch
+                    gstaichi_arch = self.runtime.prog.config().arch
 
                     def get_call_back(u, v):
                         def call_back():
@@ -923,14 +874,14 @@ class Kernel:
                             )
                         if not v.grad.is_contiguous():
                             raise ValueError(
-                                "Non contiguous gradient tensors are not supported, please call tensor.grad.contiguous() before passing it into taichi kernel."
+                                "Non contiguous gradient tensors are not supported, please call tensor.grad.contiguous() before passing it into gstaichi kernel."
                             )
 
                     tmp = v
                     if (str(v.device) != "cpu") and not (
-                        str(v.device).startswith("cuda") and taichi_arch == _ti_core.Arch.cuda
+                        str(v.device).startswith("cuda") and gstaichi_arch == _ti_core.Arch.cuda
                     ):
-                        # Getting a torch CUDA tensor on Taichi non-cuda arch:
+                        # Getting a torch CUDA tensor on GsTaichi non-cuda arch:
                         # We just replace it with a CPU tensor and by the end of kernel execution we'll use the
                         # callback to copy the values back to the original CUDA tensor.
                         host_v = v.to(device="cpu", copy=True)
@@ -945,8 +896,12 @@ class Kernel:
                         int(v.grad.data_ptr()) if v.grad is not None else 0,
                     )
                 else:
-                    raise TaichiRuntimeTypeError(f"Argument {needed} cannot be converted into required type {type(v)}")
+                    raise GsTaichiRuntimeTypeError(
+                        f"Argument {needed} cannot be converted into required type {type(v)}"
+                    )
             elif has_paddle():
+                # Do we want to continue to support paddle? :thinking_face:
+                # #maybeprunable
                 import paddle  # pylint: disable=C0415  # type: ignore
 
                 if isinstance(v, paddle.Tensor):
@@ -958,41 +913,41 @@ class Kernel:
                         return call_back
 
                     tmp = v.value().get_tensor()
-                    taichi_arch = self.runtime.prog.config().arch
+                    gstaichi_arch = self.runtime.prog.config().arch
                     if v.place.is_gpu_place():
-                        if taichi_arch != _ti_core.Arch.cuda:
-                            # Paddle cuda tensor on Taichi non-cuda arch
+                        if gstaichi_arch != _ti_core.Arch.cuda:
+                            # Paddle cuda tensor on GsTaichi non-cuda arch
                             host_v = v.cpu()
                             tmp = host_v.value().get_tensor()
                             callbacks.append(get_call_back(v, host_v))
                     elif v.place.is_cpu_place():
-                        if taichi_arch == _ti_core.Arch.cuda:
-                            # Paddle cpu tensor on Taichi cuda arch
+                        if gstaichi_arch == _ti_core.Arch.cuda:
+                            # Paddle cpu tensor on GsTaichi cuda arch
                             gpu_v = v.cuda()
                             tmp = gpu_v.value().get_tensor()
                             callbacks.append(get_call_back(v, gpu_v))
                     else:
                         # Paddle do support many other backends like XPU, NPU, MLU, IPU
-                        raise TaichiRuntimeTypeError(f"Taichi do not support backend {v.place} that Paddle support")
+                        raise GsTaichiRuntimeTypeError(f"GsTaichi do not support backend {v.place} that Paddle support")
                     launch_ctx.set_arg_external_array_with_shape(
                         indices, int(tmp._ptr()), v.element_size() * v.size, array_shape, 0
                     )
                 else:
-                    raise TaichiRuntimeTypeError(f"Argument {needed} cannot be converted into required type {v}")
+                    raise GsTaichiRuntimeTypeError(f"Argument {needed} cannot be converted into required type {v}")
             else:
-                raise TaichiRuntimeTypeError(f"Argument {needed} cannot be converted into required type {v}")
+                raise GsTaichiRuntimeTypeError(f"Argument {needed} cannot be converted into required type {v}")
 
-        def set_arg_matrix(indices, v, needed):
-            def cast_float(x):
+        def set_arg_matrix(indices: tuple[int, ...], v, needed) -> None:
+            def cast_float(x: float | np.floating | np.integer | int) -> float:
                 if not isinstance(x, (int, float, np.integer, np.floating)):
-                    raise TaichiRuntimeTypeError(
+                    raise GsTaichiRuntimeTypeError(
                         f"Argument {needed.dtype} cannot be converted into required type {type(x)}"
                     )
                 return float(x)
 
-            def cast_int(x):
+            def cast_int(x: int | np.integer) -> int:
                 if not isinstance(x, (int, np.integer)):
-                    raise TaichiRuntimeTypeError(
+                    raise GsTaichiRuntimeTypeError(
                         f"Argument {needed.dtype} cannot be converted into required type {type(x)}"
                     )
                 return int(x)
@@ -1012,13 +967,13 @@ class Kernel:
             v = needed(*v)
             needed.set_kernel_struct_args(v, launch_ctx, indices)
 
-        def set_arg_sparse_matrix_builder(indices, v):
+        def set_arg_sparse_matrix_builder(indices: tuple[int, ...], v) -> None:
             # Pass only the base pointer of the ti.types.sparse_matrix_builder() argument
             launch_ctx.set_arg_uint(indices, v._get_ndarray_addr())
 
         set_later_list = []
 
-        def recursive_set_args(needed_arg_type, provided_arg_type, v, indices):
+        def recursive_set_args(needed_arg_type: Type, provided_arg_type: Type, v: Any, indices: tuple[int, ...]) -> int:
             """
             Returns the number of kernel args set
             e.g. templates don't set kernel args, so returns 0
@@ -1033,7 +988,7 @@ class Kernel:
             actual_argument_slot += 1
             if isinstance(needed_arg_type, ArgPackType):
                 if not isinstance(v, ArgPack):
-                    raise TaichiRuntimeTypeError.get(indices, str(needed_arg_type), str(provided_arg_type))
+                    raise GsTaichiRuntimeTypeError.get(indices, str(needed_arg_type), str(provided_arg_type))
                 idx_new = 0
                 for j, (name, anno) in enumerate(needed_arg_type.members.items()):
                     idx_new += recursive_set_args(anno, type(v[name]), v[name], indices + (idx_new,))
@@ -1042,14 +997,14 @@ class Kernel:
             # Note: do not use sth like "needed == f32". That would be slow.
             if id(needed_arg_type) in primitive_types.real_type_ids:
                 if not isinstance(v, (float, int, np.floating, np.integer)):
-                    raise TaichiRuntimeTypeError.get(indices, needed_arg_type.to_string(), provided_arg_type)
+                    raise GsTaichiRuntimeTypeError.get(indices, needed_arg_type.to_string(), provided_arg_type)
                 if in_argpack:
                     return 1
                 launch_ctx.set_arg_float(indices, float(v))
                 return 1
             if id(needed_arg_type) in primitive_types.integer_type_ids:
                 if not isinstance(v, (int, np.integer)):
-                    raise TaichiRuntimeTypeError.get(indices, needed_arg_type.to_string(), provided_arg_type)
+                    raise GsTaichiRuntimeTypeError.get(indices, needed_arg_type.to_string(), provided_arg_type)
                 if in_argpack:
                     return 1
                 if is_signed(cook_dtype(needed_arg_type)):
@@ -1071,19 +1026,21 @@ class Kernel:
                     field_value = getattr(v, field.name)
                     idx += recursive_set_args(field.type, field.type, field_value, (indices[0] + idx,))
                 return idx
-            if isinstance(needed_arg_type, ndarray_type.NdarrayType) and isinstance(v, taichi.lang._ndarray.Ndarray):
+            if isinstance(needed_arg_type, ndarray_type.NdarrayType) and isinstance(v, gstaichi.lang._ndarray.Ndarray):
                 if in_argpack:
                     set_later_list.append((set_arg_ndarray, (v,)))
                     return 0
                 set_arg_ndarray(indices, v)
                 return 1
-            if isinstance(needed_arg_type, texture_type.TextureType) and isinstance(v, taichi.lang._texture.Texture):
+            if isinstance(needed_arg_type, texture_type.TextureType) and isinstance(v, gstaichi.lang._texture.Texture):
                 if in_argpack:
                     set_later_list.append((set_arg_texture, (v,)))
                     return 0
                 set_arg_texture(indices, v)
                 return 1
-            if isinstance(needed_arg_type, texture_type.RWTextureType) and isinstance(v, taichi.lang._texture.Texture):
+            if isinstance(needed_arg_type, texture_type.RWTextureType) and isinstance(
+                v, gstaichi.lang._texture.Texture
+            ):
                 if in_argpack:
                     set_later_list.append((set_arg_rw_texture, (v,)))
                     return 0
@@ -1103,8 +1060,12 @@ class Kernel:
             if isinstance(needed_arg_type, StructType):
                 if in_argpack:
                     return 1
-                if not isinstance(v, needed_arg_type):
-                    raise TaichiRuntimeTypeError(
+                # Unclear how to make the following pass typing checks
+                # StructType implements __instancecheck__, which should be a classmethod, but
+                # is currently an instance method
+                # TODO: look into this more deeply at some point
+                if not isinstance(v, needed_arg_type):  # type: ignore
+                    raise GsTaichiRuntimeTypeError(
                         f"Argument {provided_arg_type} cannot be converted into required type {needed_arg_type}"
                     )
                 needed_arg_type.set_kernel_struct_args(v, launch_ctx, indices)
@@ -1127,7 +1088,7 @@ class Kernel:
             set_arg_func((len(args) - template_num + i,), *params)
 
         if exceed_max_arg_num:
-            raise TaichiRuntimeError(
+            raise GsTaichiRuntimeError(
                 f"The number of elements in kernel arguments is too big! Do not exceed {max_arg_num} on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
             )
 
@@ -1162,7 +1123,7 @@ class Kernel:
 
         return ret
 
-    def construct_kernel_ret(self, launch_ctx, ret_type, index=()):
+    def construct_kernel_ret(self, launch_ctx: KernelLaunchContext, ret_type: Any, index: tuple[int, ...] = ()):
         if isinstance(ret_type, CompoundType):
             return ret_type.from_kernel_struct_ret(launch_ctx, index)
         if ret_type in primitive_types.integer_types:
@@ -1171,9 +1132,9 @@ class Kernel:
             return launch_ctx.get_struct_ret_uint(index)
         if ret_type in primitive_types.real_types:
             return launch_ctx.get_struct_ret_float(index)
-        raise TaichiRuntimeTypeError(f"Invalid return type on index={index}")
+        raise GsTaichiRuntimeTypeError(f"Invalid return type on index={index}")
 
-    def ensure_compiled(self, *args):
+    def ensure_compiled(self, *args: tuple[Any, ...]) -> tuple[Callable, int, AutodiffMode]:
         instance_id, arg_features = self.mapper.lookup(args)
         key = (self.func, instance_id, self.autodiff_mode)
         self.materialize(key=key, args=args, arg_features=arg_features)
@@ -1182,7 +1143,7 @@ class Kernel:
     # For small kernels (< 3us), the performance can be pretty sensitive to overhead in __call__
     # Thus this part needs to be fast. (i.e. < 3us on a 4 GHz x64 CPU)
     @_shell_pop_print
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs) -> Any:
         args = _process_args(self, is_func=False, args=args, kwargs=kwargs)
 
         # Transform the primal kernel to forward mode grad kernel
@@ -1213,7 +1174,7 @@ class Kernel:
         return self.launch_kernel(kernel_cpp, *args)
 
 
-# For a Taichi class definition like below:
+# For a GsTaichi class definition like below:
 #
 # @ti.data_oriented
 # class X:
@@ -1232,7 +1193,7 @@ _KERNEL_CLASS_STACKFRAME_STMT_RES = [
 ]
 
 
-def _inside_class(level_of_class_stackframe):
+def _inside_class(level_of_class_stackframe: int) -> bool:
     try:
         maybe_class_frame = sys._getframe(level_of_class_stackframe)
         statement_list = inspect.getframeinfo(maybe_class_frame)[3]
@@ -1247,7 +1208,7 @@ def _inside_class(level_of_class_stackframe):
     return False
 
 
-def _kernel_impl(_func: Callable, level_of_class_stackframe: int, verbose: bool = False):
+def _kernel_impl(_func: Callable, level_of_class_stackframe: int, verbose: bool = False) -> GsTaichiCallable:
     # Can decorators determine if a function is being defined inside a class?
     # https://stackoverflow.com/a/8793684/12003165
     is_classkernel = _inside_class(level_of_class_stackframe + 1)
@@ -1259,6 +1220,7 @@ def _kernel_impl(_func: Callable, level_of_class_stackframe: int, verbose: bool
     # Having |primal| contains |grad| makes the tape work.
     primal.grad = adjoint
 
+    wrapped: GsTaichiCallable
     if is_classkernel:
         # For class kernels, their primal/adjoint callables are constructed
         # when the kernel is accessed via the instance inside
@@ -1268,24 +1230,26 @@ def _kernel_impl(_func: Callable, level_of_class_stackframe: int, verbose: bool
         #
         # See also: _BoundedDifferentiableMethod, data_oriented.
         @functools.wraps(_func)
-        def wrapped(*args, **kwargs):
+        def wrapped_classkernel(*args, **kwargs):
             # If we reach here (we should never), it means the class is not decorated
             # with @ti.data_oriented, otherwise getattr would have intercepted the call.
             clsobj = type(args[0])
             assert not hasattr(clsobj, "_data_oriented")
-            raise TaichiSyntaxError(f"Please decorate class {clsobj.__name__} with @ti.data_oriented")
+            raise GsTaichiSyntaxError(f"Please decorate class {clsobj.__name__} with @ti.data_oriented")
 
+        wrapped = GsTaichiCallable(_func, wrapped_classkernel)
     else:
 
         @functools.wraps(_func)
-        def wrapped(*args, **kwargs):
+        def wrapped_func(*args, **kwargs):
             try:
                 return primal(*args, **kwargs)
-            except (TaichiCompilationError, TaichiRuntimeError) as e:
+            except (GsTaichiCompilationError, GsTaichiRuntimeError) as e:
                 if impl.get_runtime().print_full_traceback:
                     raise e
                 raise type(e)("\n" + str(e)) from None
 
+        wrapped = GsTaichiCallable(_func, wrapped_func)
         wrapped.grad = adjoint
 
     wrapped._is_wrapped_kernel = True
@@ -1296,10 +1260,10 @@ def _kernel_impl(_func: Callable, level_of_class_stackframe: int, verbose: bool
 
 
 def kernel(fn: Callable):
-    """Marks a function as a Taichi kernel.
+    """Marks a function as a GsTaichi kernel.
 
-    A Taichi kernel is a function written in Python, and gets JIT compiled by
-    Taichi into native CPU/GPU instructions (e.g. a series of CUDA kernels).
+    A GsTaichi kernel is a function written in Python, and gets JIT compiled by
+    GsTaichi into native CPU/GPU instructions (e.g. a series of CUDA kernels).
     The top-level ``for`` loops are automatically parallelized, and distributed
     to either a CPU thread pool or massively parallel GPUs.
 
@@ -1327,10 +1291,10 @@ def kernel(fn: Callable):
 
 
 class _BoundedDifferentiableMethod:
-    def __init__(self, kernel_owner, wrapped_kernel_func):
+    def __init__(self, kernel_owner: Any, wrapped_kernel_func: GsTaichiCallable | BoundGsTaichiCallable):
         clsobj = type(kernel_owner)
         if not getattr(clsobj, "_data_oriented", False):
-            raise TaichiSyntaxError(f"Please decorate class {clsobj.__name__} with @ti.data_oriented")
+            raise GsTaichiSyntaxError(f"Please decorate class {clsobj.__name__} with @ti.data_oriented")
         self._kernel_owner = kernel_owner
         self._primal = wrapped_kernel_func._primal
         self._adjoint = wrapped_kernel_func._adjoint
@@ -1339,23 +1303,26 @@ class _BoundedDifferentiableMethod:
 
     def __call__(self, *args, **kwargs):
         try:
+            assert self._primal is not None
             if self._is_staticmethod:
                 return self._primal(*args, **kwargs)
             return self._primal(self._kernel_owner, *args, **kwargs)
-        except (TaichiCompilationError, TaichiRuntimeError) as e:
+
+        except (GsTaichiCompilationError, GsTaichiRuntimeError) as e:
             if impl.get_runtime().print_full_traceback:
                 raise e
             raise type(e)("\n" + str(e)) from None
 
-    def grad(self, *args, **kwargs):
+    def grad(self, *args, **kwargs) -> Kernel:
+        assert self._adjoint is not None
         return self._adjoint(self._kernel_owner, *args, **kwargs)
 
 
 def data_oriented(cls):
-    """Marks a class as Taichi compatible.
+    """Marks a class as GsTaichi compatible.
 
-    To allow for modularized code, Taichi provides this decorator so that
-    Taichi kernels can be defined inside a class.
+    To allow for modularized code, GsTaichi provides this decorator so that
+    GsTaichi kernels can be defined inside a class.
 
     See also https://docs.taichi-lang.org/docs/odop
 
@@ -1394,11 +1361,11 @@ def data_oriented(cls):
                 wrapped = x.__func__
             else:
                 wrapped = x
+            assert isinstance(wrapped, (BoundGsTaichiCallable, GsTaichiCallable))
             wrapped._is_staticmethod = is_staticmethod
-            assert inspect.isfunction(wrapped)
             if wrapped._is_classkernel:
                 ret = _BoundedDifferentiableMethod(self, wrapped)
-                ret.__name__ = wrapped.__name__
+                ret.__name__ = wrapped.__name__  # type: ignore
                 if is_property:
                     return ret()
                 return ret