gstaichi 2.1.1__cp310-cp310-macosx_11_0_arm64.whl

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/_ndarray.py

@@ -0,0 +1,352 @@
+# type: ignore
+
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.exception import GsTaichiIndexError
+from gstaichi.lang.util import cook_dtype, get_traceback, python_scope, to_numpy_type
+from gstaichi.types import primitive_types
+from gstaichi.types.enums import Layout
+from gstaichi.types.ndarray_type import NdarrayTypeMetadata
+from gstaichi.types.utils import is_real, is_signed
+
+if TYPE_CHECKING:
+    from gstaichi.lang.matrix import MatrixNdarray, VectorNdarray
+
+    TensorNdarray = Union["ScalarNdarray", VectorNdarray, MatrixNdarray]
+
+
+class Ndarray:
+    """GsTaichi ndarray class.
+
+    Args:
+        dtype (DataType): Data type of each value.
+        shape (Tuple[int]): Shape of the Ndarray.
+    """
+
+    def __init__(self):
+        self.host_accessor = None
+        self.shape = None
+        self.element_type = None
+        self.dtype = None
+        self.arr = None
+        self.layout = Layout.AOS
+        self.grad: "TensorNdarray | None" = None
+
+    def get_type(self):
+        return NdarrayTypeMetadata(self.element_type, self.shape, self.grad is not None)
+
+    @property
+    def element_shape(self):
+        """Gets ndarray element shape.
+
+        Returns:
+            Tuple[Int]: Ndarray element shape.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def __setitem__(self, key, value):
+        """Sets ndarray element in Python scope.
+
+        Args:
+            key (Union[List[int], int, None]): Coordinates of the ndarray element.
+            value (element type): Value to set.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def __getitem__(self, key):
+        """Gets ndarray element in Python scope.
+
+        Args:
+            key (Union[List[int], int, None]): Coordinates of the ndarray element.
+
+        Returns:
+            element type: Value retrieved.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def fill(self, val):
+        """Fills ndarray with a specific scalar value.
+
+        Args:
+            val (Union[int, float]): Value to fill.
+        """
+        if impl.current_cfg().arch != _ti_core.Arch.cuda and impl.current_cfg().arch != _ti_core.Arch.x64:
+            self._fill_by_kernel(val)
+        elif _ti_core.is_tensor(self.element_type):
+            self._fill_by_kernel(val)
+        elif self.dtype == primitive_types.f32:
+            impl.get_runtime().prog.fill_float(self.arr, val)
+        elif self.dtype == primitive_types.i32:
+            impl.get_runtime().prog.fill_int(self.arr, val)
+        elif self.dtype == primitive_types.u32:
+            impl.get_runtime().prog.fill_uint(self.arr, val)
+        else:
+            self._fill_by_kernel(val)
+
+    @python_scope
+    def _ndarray_to_numpy(self):
+        """Converts ndarray to a numpy array.
+
+        Returns:
+            numpy.ndarray: The result numpy array.
+        """
+        arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype))
+        from gstaichi._kernels import ndarray_to_ext_arr  # pylint: disable=C0415
+
+        ndarray_to_ext_arr(self, arr)
+        impl.get_runtime().sync()
+        return arr
+
+    @python_scope
+    def _ndarray_matrix_to_numpy(self, as_vector):
+        """Converts matrix ndarray to a numpy array.
+
+        Returns:
+            numpy.ndarray: The result numpy array.
+        """
+        arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype))
+        from gstaichi._kernels import ndarray_matrix_to_ext_arr  # pylint: disable=C0415
+
+        layout_is_aos = 1
+        ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector)
+        impl.get_runtime().sync()
+        return arr
+
+    @python_scope
+    def _ndarray_from_numpy(self, arr):
+        """Loads all values from a numpy array.
+
+        Args:
+            arr (numpy.ndarray): The source numpy array.
+        """
+        if not isinstance(arr, np.ndarray):
+            raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided")
+        if tuple(self.arr.total_shape()) != tuple(arr.shape):
+            raise ValueError(f"Mismatch shape: {tuple(self.arr.shape)} expected, but {tuple(arr.shape)} provided")
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+
+        from gstaichi._kernels import ext_arr_to_ndarray  # pylint: disable=C0415
+
+        ext_arr_to_ndarray(arr, self)
+        impl.get_runtime().sync()
+
+    @python_scope
+    def _ndarray_matrix_from_numpy(self, arr, as_vector):
+        """Loads all values from a numpy array.
+
+        Args:
+            arr (numpy.ndarray): The source numpy array.
+        """
+        if not isinstance(arr, np.ndarray):
+            raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided")
+        if tuple(self.arr.total_shape()) != tuple(arr.shape):
+            raise ValueError(
+                f"Mismatch shape: {tuple(self.arr.total_shape())} expected, but {tuple(arr.shape)} provided"
+            )
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+
+        from gstaichi._kernels import ext_arr_to_ndarray_matrix  # pylint: disable=C0415
+
+        layout_is_aos = 1
+        ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector)
+        impl.get_runtime().sync()
+
+    @python_scope
+    def _get_element_size(self):
+        """Returns the size of one element in bytes.
+
+        Returns:
+            Size in bytes.
+        """
+        return self.arr.element_size()
+
+    @python_scope
+    def _get_nelement(self):
+        """Returns the total number of elements.
+
+        Returns:
+            Total number of elements.
+        """
+        return self.arr.nelement()
+
+    @python_scope
+    def copy_from(self, other):
+        """Copies all elements from another ndarray.
+
+        The shape of the other ndarray needs to be the same as `self`.
+
+        Args:
+            other (Ndarray): The source ndarray.
+        """
+        assert isinstance(other, Ndarray)
+        assert tuple(self.arr.shape) == tuple(other.arr.shape)
+        from gstaichi._kernels import ndarray_to_ndarray  # pylint: disable=C0415
+
+        ndarray_to_ndarray(self, other)
+        impl.get_runtime().sync()
+
+    def _set_grad(self, grad: "TensorNdarray"):
+        """Sets the gradient ndarray.
+
+        Args:
+            grad (Ndarray): The gradient ndarray.
+        """
+        self.grad = grad
+
+    def __deepcopy__(self, memo=None):
+        """Copies all elements to a new ndarray.
+
+        Returns:
+            Ndarray: The result ndarray.
+        """
+        raise NotImplementedError()
+
+    def _fill_by_kernel(self, val):
+        """Fills ndarray with a specific scalar value using a ti.kernel.
+
+        Args:
+            val (Union[int, float]): Value to fill.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def _pad_key(self, key):
+        if key is None:
+            key = ()
+        if not isinstance(key, (tuple, list)):
+            key = (key,)
+        if len(key) != len(self.arr.total_shape()):
+            raise GsTaichiIndexError(f"{len(self.arr.total_shape())}d ndarray indexed with {len(key)}d indices: {key}")
+        return key
+
+    @python_scope
+    def _initialize_host_accessor(self):
+        if self.host_accessor:
+            return
+        impl.get_runtime().materialize()
+        self.host_accessor = NdarrayHostAccessor(self.arr)
+
+
+class ScalarNdarray(Ndarray):
+    """GsTaichi ndarray with scalar elements.
+
+    Args:
+        dtype (DataType): Data type of each value.
+        shape (Tuple[int]): Shape of the ndarray.
+    """
+
+    def __init__(self, dtype, arr_shape):
+        super().__init__()
+        self.dtype = cook_dtype(dtype)
+        self.arr = impl.get_runtime().prog.create_ndarray(
+            self.dtype, arr_shape, layout=Layout.NULL, zero_fill=True, dbg_info=_ti_core.DebugInfo(get_traceback())
+        )
+        self.shape = tuple(self.arr.shape)
+        self.element_type = dtype
+
+    def __del__(self):
+        if impl is not None and impl.get_runtime is not None and impl.get_runtime() is not None:
+            prog = impl.get_runtime()._prog
+            if prog is not None:
+                prog.delete_ndarray(self.arr)
+
+    @property
+    def element_shape(self):
+        return ()
+
+    @python_scope
+    def __setitem__(self, key, value):
+        self._initialize_host_accessor()
+        self.host_accessor.setter(value, *self._pad_key(key))
+
+    @python_scope
+    def __getitem__(self, key):
+        self._initialize_host_accessor()
+        return self.host_accessor.getter(*self._pad_key(key))
+
+    @python_scope
+    def to_numpy(self):
+        return self._ndarray_to_numpy()
+
+    @python_scope
+    def from_numpy(self, arr):
+        self._ndarray_from_numpy(arr)
+
+    def __deepcopy__(self, memo=None):
+        ret_arr = ScalarNdarray(self.dtype, self.shape)
+        ret_arr.copy_from(self)
+        return ret_arr
+
+    def _fill_by_kernel(self, val):
+        from gstaichi._kernels import fill_ndarray  # pylint: disable=C0415
+
+        fill_ndarray(self, val)
+
+    def __repr__(self):
+        return "<ti.ndarray>"
+
+
+class NdarrayHostAccessor:
+    def __init__(self, ndarray):
+        dtype = ndarray.element_data_type()
+        if is_real(dtype):
+
+            def getter(*key):
+                return ndarray.read_float(key)
+
+            def setter(value, *key):
+                ndarray.write_float(key, value)
+
+        else:
+            if is_signed(dtype):
+
+                def getter(*key):
+                    return ndarray.read_int(key)
+
+            else:
+
+                def getter(*key):
+                    return ndarray.read_uint(key)
+
+            def setter(value, *key):
+                ndarray.write_int(key, value)
+
+        self.getter = getter
+        self.setter = setter
+
+
+class NdarrayHostAccess:
+    """Class for accessing VectorNdarray/MatrixNdarray in Python scope.
+    Args:
+        arr (Union[VectorNdarray, MatrixNdarray]): See above.
+        indices_first (Tuple[Int]): Indices of first-level access (coordinates in the field).
+        indices_second (Tuple[Int]): Indices of second-level access (indices in the vector/matrix).
+    """
+
+    def __init__(self, arr, indices_first, indices_second):
+        self.ndarr = arr
+        self.arr = arr.arr
+        self.indices = indices_first + indices_second
+
+        def getter():
+            self.ndarr._initialize_host_accessor()
+            return self.ndarr.host_accessor.getter(*self.ndarr._pad_key(self.indices))
+
+        def setter(value):
+            self.ndarr._initialize_host_accessor()
+            self.ndarr.host_accessor.setter(value, *self.ndarr._pad_key(self.indices))
+
+        self.getter = getter
+        self.setter = setter
+
+
+__all__ = ["Ndarray", "ScalarNdarray"]