gstaichi 0.1.25.dev0__cp313-cp313-macosx_15_0_arm64.whl → 2.0.0__cp313-cp313-macosx_15_0_arm64.whl

gstaichi/lang/_fast_caching/src_hasher.py

@@ -0,0 +1,75 @@
+from typing import Any, Iterable, Sequence
+
+from pydantic import BaseModel
+
+from .._wrap_inspect import FunctionSourceInfo
+from . import args_hasher, config_hasher, function_hasher
+from .fast_caching_types import HashedFunctionSourceInfo
+from .hash_utils import hash_iterable_strings
+from .python_side_cache import PythonSideCache
+
+
+def create_cache_key(kernel_source_info: FunctionSourceInfo, args: Sequence[Any]) -> str | None:
+    """
+    cache key takes into account:
+    - arg types
+    - cache value arg values
+    - kernel function (but not sub functions)
+    - compilation config (which includes arch, and debug)
+    """
+    args_hash = args_hasher.hash_args(args)
+    if args_hash is None:
+        return None
+    kernel_hash = function_hasher.hash_kernel(kernel_source_info)
+    config_hash = config_hasher.hash_compile_config()
+    cache_key = hash_iterable_strings((kernel_hash, args_hash, config_hash))
+    return cache_key
+
+
+class CacheValue(BaseModel):
+    hashed_function_source_infos: list[HashedFunctionSourceInfo]
+
+
+def store(cache_key: str, function_source_infos: Iterable[FunctionSourceInfo]) -> None:
+    """
+    Note that unlike other caches, this cache is not going to store the actual value we want.
+    This cache is only used for verification that our cache key is valid. Big picture:
+    - we have a cache key, based on args and top level kernel function
+    - we want to use this to look up LLVM IR, in C++ side cache
+    - however, before doing that, we first want to validate that the source code didn't change
+        - i.e. is our cache key still valid?
+    - the python side cache contains information we will use to verify that our cache key is valid
+        - ie the list of function source infos
+    """
+    if not cache_key:
+        return
+    cache = PythonSideCache()
+    hashed_function_source_infos = function_hasher.hash_functions(function_source_infos)
+    cache_value_obj = CacheValue(hashed_function_source_infos=list(hashed_function_source_infos))
+    cache.store(cache_key, cache_value_obj.json())
+
+
+def _try_load(cache_key: str) -> Sequence[HashedFunctionSourceInfo] | None:
+    cache = PythonSideCache()
+    maybe_cache_value_json = cache.try_load(cache_key)
+    if maybe_cache_value_json is None:
+        return None
+    cache_value_obj = CacheValue.parse_raw(maybe_cache_value_json)
+    return cache_value_obj.hashed_function_source_infos
+
+
+def validate_cache_key(cache_key: str) -> bool:
+    """
+    loads function source infos from cache, if available
+    checks the hashes against the current source code
+    """
+    maybe_hashed_function_source_infos = _try_load(cache_key)
+    if not maybe_hashed_function_source_infos:
+        return False
+    return function_hasher.validate_hashed_function_infos(maybe_hashed_function_source_infos)
+
+
+def dump_stats() -> None:
+    print("dump stats")
+    args_hasher.dump_stats()
+    function_hasher.dump_stats()

gstaichi/lang/_kernel_impl_dataclass.py

@@ -0,0 +1,212 @@
+import ast
+import dataclasses
+import inspect
+from typing import Any
+
+from gstaichi.lang import util
+from gstaichi.lang._dataclass_util import create_flat_name
+from gstaichi.lang.ast import (
+    ASTTransformerContext,
+)
+from gstaichi.lang.kernel_arguments import ArgMetadata
+
+
+def _populate_struct_locals_from_params_dict(basename: str, struct_locals, struct_type) -> None:
+    """
+    We are populating struct locals from a type included in function parameters, or one of their subtypes
+
+    struct_locals will be a list of all possible unpacked variable names we can form from the struct.
+    basename is used to take into account the parent struct's name. For example, lets say we have:
+
+    @dataclasses.dataclass
+    class StructAB:
+        a:
+        b:
+        struct_cd: StructCD
+
+    @dataclasses.dataclass
+    class StructCD:
+        c:
+        d:
+        struct_ef: StructEF
+
+    @dataclasses.dataclass
+    class StructEF:
+        e:
+        f:
+
+    ... and the function parameters look like: `def foo(struct_ab: StructAB)`
+
+    then all possible variables we could form from this are:
+    - struct_ab.a
+    - struct_ab.b
+    - struct_ab.struct_cd.c
+    - struct_ab.struct_cd.d
+    - struct_ab.struct_cd.strucdt_ef.e
+    - struct_ab.struct_cd.strucdt_ef.f
+
+    And the members of struct_locals should be:
+    - __ti_struct_ab__ti_a
+    - __ti_struct_ab__ti_b
+    - __ti_struct_ab__ti_struct_cd__ti_c
+    - __ti_struct_ab__ti_struct_cd__ti_d
+    - __ti_struct_ab__ti_struct_cd__ti_struct_ef__ti_e
+    - __ti_struct_ab__ti_struct_cd__ti_struct_ef__ti_f
+    """
+    for field in dataclasses.fields(struct_type):
+        child_name = create_flat_name(basename, field.name)
+        if dataclasses.is_dataclass(field.type):
+            _populate_struct_locals_from_params_dict(child_name, struct_locals, field.type)
+        else:
+            struct_locals.add(child_name)
+
+
+def extract_struct_locals_from_context(ctx: ASTTransformerContext) -> set[str]:
+    """
+    Provides meta information for later tarnsformation of nodes in AST
+
+    - Uses ctx.func.func to get the function signature.
+    - Searches this for any dataclasses:
+      - If it finds any dataclasses, then converts them into expanded names.
+      - E.g. my_struct: MyStruct, and MyStruct contains a, b, c would become:
+          {"__ti_my_struct_a", "__ti_my_struct_b, "__ti_my_struct_c"}
+    """
+    struct_locals = set()
+    assert ctx.func is not None
+    sig = inspect.signature(ctx.func.func)
+    parameters = sig.parameters
+    for param_name, parameter in parameters.items():
+        if dataclasses.is_dataclass(parameter.annotation):
+            for field in dataclasses.fields(parameter.annotation):
+                child_name = create_flat_name(param_name, field.name)
+                # child_name = f"__ti_{param_name}__ti_{field.name}"
+                if dataclasses.is_dataclass(field.type):
+                    _populate_struct_locals_from_params_dict(child_name, struct_locals, field.type)
+                    continue
+                struct_locals.add(child_name)
+    return struct_locals
+
+
+def expand_func_arguments(arguments: list[ArgMetadata]) -> list[ArgMetadata]:
+    """
+    Used to expand arguments for @ti.func
+    """
+    expanded_arguments = []
+    for i, argument in enumerate(arguments):
+        if dataclasses.is_dataclass(argument.annotation):
+            for field in dataclasses.fields(argument.annotation):
+                child_name = create_flat_name(argument.name, field.name)
+                if dataclasses.is_dataclass(field.type):
+                    new_arg = ArgMetadata(
+                        annotation=field.type,
+                        name=child_name,
+                        default=argument.default,
+                    )
+                    child_args = expand_func_arguments([new_arg])
+                    expanded_arguments += child_args
+                else:
+                    new_argument = ArgMetadata(
+                        annotation=field.type,
+                        name=child_name,
+                    )
+                    expanded_arguments.append(new_argument)
+        else:
+            expanded_arguments.append(argument)
+    return expanded_arguments
+
+
+class FlattenAttributeNameTransformer(ast.NodeTransformer):
+    def __init__(self, struct_locals: set[str]) -> None:
+        self.struct_locals = struct_locals
+
+    def visit_Attribute(self, node):
+        flat_name = FlattenAttributeNameTransformer._flatten_attribute_name(node)
+        if not flat_name or flat_name not in self.struct_locals:
+            return self.generic_visit(node)
+        return ast.copy_location(ast.Name(id=flat_name, ctx=node.ctx), node)
+
+    @staticmethod
+    def _flatten_attribute_name(node: ast.Attribute) -> str | None:
+        """
+        see unpack_ast_struct_expressions docstring for more explanation
+        """
+        if isinstance(node.value, ast.Name):
+            return create_flat_name(node.value.id, node.attr)
+        if isinstance(node.value, ast.Attribute):
+            child_flat_name = FlattenAttributeNameTransformer._flatten_attribute_name(node.value)
+            if not child_flat_name:
+                return None
+            return create_flat_name(child_flat_name, node.attr)
+        return None
+
+
+def unpack_ast_struct_expressions(tree: ast.Module, struct_locals: set[str]) -> ast.Module:
+    """
+    Transform nodes in AST, to flatten access to struct members
+
+    Examples of things we will transform/flatten:
+
+    # my_struct_ab.a
+    # Attribute(value=Name())
+    Attribute(
+        value=Name(id='my_struct_ab', ctx=Load()),
+        attr='a',
+        ctx=Load())
+    =>
+    # __ti_my_struct_ab__ti_a
+    Name(id='__ti_my_struct_ab__ti_a', ctx=Load()
+
+    # my_struct_ab.struct_cd.d
+    # Attribute(value=Attribute(value=Name()))
+    Attribute(
+        value=Attribute(
+            value=Name(id='my_struct_ab', ctx=Load()),
+            attr='struct_cd',
+            ctx=Load()),
+        attr='d',
+        ctx=Load())
+        visit_attribute
+    =>
+    # __ti_my_struct_ab__ti_struct_cd__ti_d
+    Name(id='__ti_my_struct_ab__ti_struct_cd__ti_d', ctx=Load()
+
+    # my_struct_ab.struct_cd.struct_ef.f
+    # Attribute(value=Attribute(value=Name()))
+    Attribute(
+        value=Attribute(
+            value=Attribute(
+            value=Name(id='my_struct_ab', ctx=Load()),
+            attr='struct_cd',
+            ctx=Load()),
+            attr='struct_ef',
+            ctx=Load()),
+        attr='f',
+        ctx=Load())
+    =>
+    # __ti_my_struct_ab__ti_struct_cd__ti_struct_ef__ti_f
+    Name(id='__ti_my_struct_ab__ti_struct_cd__ti_struct_ef__ti_f', ctx=Load()
+    """
+    transformer = FlattenAttributeNameTransformer(struct_locals=struct_locals)
+    new_tree = transformer.visit(tree)
+    ast.fix_missing_locations(new_tree)
+    return new_tree
+
+
+def populate_global_vars_from_dataclass(
+    param_name: str,
+    param_type: Any,
+    py_arg: Any,
+    global_vars: dict[str, Any],
+):
+    for field in dataclasses.fields(param_type):
+        child_value = getattr(py_arg, field.name)
+        flat_name = create_flat_name(param_name, field.name)
+        if dataclasses.is_dataclass(field.type):
+            populate_global_vars_from_dataclass(
+                param_name=flat_name,
+                param_type=field.type,
+                py_arg=child_value,
+                global_vars=global_vars,
+            )
+        elif util.is_ti_template(field.type):
+            global_vars[flat_name] = child_value

gstaichi/lang/_template_mapper.py

@@ -1,6 +1,6 @@
 import dataclasses
 import weakref
-from typing import Any, Union
+from typing import Any, Callable, Union
 
 import gstaichi.lang
 import gstaichi.lang._ndarray
@@ -8,24 +8,27 @@ import gstaichi.lang._texture
 import gstaichi.lang.expr
 import gstaichi.lang.snode
 from gstaichi._lib import core as _ti_core
+from gstaichi.lang import _dataclass_util
 from gstaichi.lang.any_array import AnyArray
-from gstaichi.lang.argpack import ArgPack, ArgPackType
 from gstaichi.lang.exception import (
     GsTaichiRuntimeTypeError,
 )
-from gstaichi.lang.kernel_arguments import KernelArgument
+from gstaichi.lang.kernel_arguments import ArgMetadata
 from gstaichi.lang.matrix import MatrixType
-from gstaichi.lang.util import to_gstaichi_type
+from gstaichi.lang.util import is_ti_template, to_gstaichi_type
 from gstaichi.types import (
     ndarray_type,
     sparse_matrix_builder,
     template,
     texture_type,
 )
+from gstaichi.types.enums import AutodiffMode
+
+CompiledKernelKeyType = tuple[Callable, int, AutodiffMode]
+
 
 AnnotationType = Union[
     template,
-    ArgPackType,
     "texture_type.TextureType",
     "texture_type.RWTextureType",
     ndarray_type.NdarrayType,
@@ -34,7 +37,7 @@ AnnotationType = Union[
 ]
 
 
-class GsTaichiCallableTemplateMapper:
+class TemplateMapper:
     """
     This should probably be renamed to sometihng like FeatureMapper, or
     FeatureExtractor, since:
@@ -46,15 +49,15 @@ class GsTaichiCallableTemplateMapper:
     - 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: list[KernelArgument] = arguments
+    def __init__(self, arguments: list[ArgMetadata], template_slot_locations: list[int]) -> None:
+        self.arguments: list[ArgMetadata] = arguments
         self.num_args: int = len(arguments)
         self.template_slot_locations: list[int] = template_slot_locations
         self.mapping: dict[tuple[Any, ...], int] = {}
 
     @staticmethod
-    def extract_arg(arg, annotation: AnnotationType, arg_name: str) -> Any:
-        if annotation == template or isinstance(annotation, template):
+    def extract_arg(arg: Any, annotation: AnnotationType, arg_name: str) -> Any:
+        if is_ti_template(annotation):
             if isinstance(arg, gstaichi.lang.snode.SNode):
                 return arg.ptr
             if isinstance(arg, gstaichi.lang.expr.Expr):
@@ -62,7 +65,7 @@ class GsTaichiCallableTemplateMapper:
             if isinstance(arg, _ti_core.ExprCxx):
                 return arg.get_underlying_ptr_address()
             if isinstance(arg, tuple):
-                return tuple(GsTaichiCallableTemplateMapper.extract_arg(item, annotation, arg_name) for item in arg)
+                return tuple(TemplateMapper.extract_arg(item, annotation, arg_name) for item in arg)
             if isinstance(arg, gstaichi.lang._ndarray.Ndarray):
                 raise GsTaichiRuntimeTypeError(
                     "Ndarray shouldn't be passed in via `ti.template()`, please annotate your kernel using `ti.types.ndarray(...)` instead"
@@ -81,19 +84,12 @@ class GsTaichiCallableTemplateMapper:
 
             # [Primitive arguments] Return the value
             return arg
-        if isinstance(annotation, ArgPackType):
-            if not isinstance(arg, ArgPack):
-                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a argument pack, got {type(arg)}")
-            return tuple(
-                GsTaichiCallableTemplateMapper.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 = GsTaichiCallableTemplateMapper.extract_arg(field_value, field.type, arg_name)
+                child_name = _dataclass_util.create_flat_name(arg_name, field.name)
+                field_extracted = TemplateMapper.extract_arg(field_value, field.type, child_name)
                 _res_l.append(field_extracted)
             return tuple(_res_l)
         if isinstance(annotation, texture_type.TextureType):

gstaichi/lang/_wrap_inspect.py

@@ -19,8 +19,10 @@ import atexit
 import inspect
 import os
 import tempfile
+from typing import Callable
 
 import dill
+from pydantic import BaseModel
 
 _builtin_getfile = inspect.getfile
 _builtin_findsource = inspect.findsource
@@ -186,4 +188,28 @@ def getsourcefile(obj):
         return ret
 
 
-__all__ = ["getsourcelines", "getsourcefile"]
+class FunctionSourceInfo(BaseModel):
+    function_name: str
+    filepath: str
+    start_lineno: int
+    end_lineno: int
+
+    class Config:
+        frozen = True
+
+
+def get_source_info_and_src(func: Callable) -> tuple[FunctionSourceInfo, list[str]]:
+    file = getsourcefile(func)
+    name = func.__name__
+    src, start_lineno = getsourcelines(func)
+    end_lineno = start_lineno + len(src) - 1
+    func_info = FunctionSourceInfo(
+        function_name=name,
+        filepath=file,
+        start_lineno=start_lineno,
+        end_lineno=end_lineno,
+    )
+    return (func_info, src)
+
+
+__all__ = ["getsourcelines", "getsourcefile", "get_source_info_and_src"]

gstaichi/lang/ast/ast_transformer.py

@@ -2,10 +2,11 @@
 
 import ast
 import collections.abc
+import dataclasses
 import itertools
 import warnings
 from ast import unparse
-from typing import Any, Iterable, Type
+from typing import Any, Sequence, Type
 
 import numpy as np
 
@@ -40,7 +41,7 @@ from gstaichi.types import primitive_types
 from gstaichi.types.utils import is_integral
 
 
-def reshape_list(flat_list: list[Any], target_shape: Iterable[int]) -> list[Any]:
+def reshape_list(flat_list: list[Any], target_shape: Sequence[int]) -> list[Any]:
     if len(target_shape) < 2:
         return flat_list
 
@@ -645,6 +646,8 @@ class ASTTransformer(Builder):
 
                 node.ptr = getattr(tensor_ops, node.attr)
                 setattr(node, "caller", node.value.ptr)
+        elif dataclasses.is_dataclass(node.value.ptr):
+            node.ptr = next(field.type for field in dataclasses.fields(node.value.ptr))
         else:
             node.ptr = getattr(node.value.ptr, node.attr)
         return node.ptr
@@ -1309,6 +1312,8 @@ build_stmt = ASTTransformer()
 
 
 def build_stmts(ctx: ASTTransformerContext, stmts: list[ast.stmt]):
+    # TODO: Should we just make this part of ASTTransformer? Then, easier to pass around (just
+    # pass the ASTTransformer object around)
     with ctx.variable_scope_guard():
         for stmt in stmts:
             if ctx.returned != ReturnStatus.NoReturn or ctx.loop_status() != LoopStatus.Normal:

gstaichi/lang/ast/ast_transformer_utils.py

@@ -27,7 +27,8 @@ if TYPE_CHECKING:
 
 class Builder:
     def __call__(self, ctx: "ASTTransformerContext", node: ast.AST):
-        method = getattr(self, "build_" + node.__class__.__name__, None)
+        method_name = "build_" + node.__class__.__name__
+        method = getattr(self, method_name, None)
         try:
             if method is None:
                 error_msg = f'Unsupported node "{node.__class__.__name__}"'
@@ -155,17 +156,18 @@ class ReturnStatus(Enum):
 class ASTTransformerContext:
     def __init__(
         self,
-        excluded_parameters=(),
-        is_kernel: bool = True,
-        func: "Func | Kernel | None" = None,
-        arg_features=None,
-        global_vars: dict[str, Any] | None = None,
-        argument_data=None,
-        file: str | None = None,
-        src: list[str] | None = None,
-        start_lineno: int | None = None,
-        ast_builder: ASTBuilder | None = None,
-        is_real_function: bool = False,
+        excluded_parameters,
+        end_lineno: int,
+        is_kernel: bool,
+        func: "Func | Kernel",
+        arg_features: list[tuple[Any, ...]] | None,
+        global_vars: dict[str, Any],
+        argument_data,
+        file: str,
+        src: list[str],
+        start_lineno: int,
+        ast_builder: ASTBuilder | None,
+        is_real_function: bool,
     ):
         self.func = func
         self.local_scopes: list[dict[str, Any]] = []
@@ -176,7 +178,7 @@ class ASTTransformerContext:
         self.returns = None
         self.global_vars = global_vars
         self.argument_data = argument_data
-        self.return_data = None
+        self.return_data: tuple[Any, ...] | Any | None = None
         self.file = file
         self.src = src
         self.indent = 0
@@ -186,6 +188,8 @@ class ASTTransformerContext:
             else:
                 break
         self.lineno_offset = start_lineno - 1
+        self.start_lineno = start_lineno
+        self.end_lineno = end_lineno
         self.raised = False
         self.non_static_control_flow_status = NonStaticControlFlowStatus()
         self.static_scope_status = StaticScopeStatus()
@@ -194,6 +198,7 @@ class ASTTransformerContext:
         self.visited_funcdef = False
         self.is_real_function = is_real_function
         self.kernel_args: list = []
+        self.only_parse_function_def: bool = False
 
     # e.g.: FunctionDef, Module, Global
     def variable_scope_guard(self):