gstaichi 0.1.25.dev0__cp313-cp313-win_amd64.whl

gstaichi/lang/argpack.py

@@ -0,0 +1,411 @@
+# type: ignore
+
+import numpy as np
+
+import gstaichi.lang
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl, ops
+from gstaichi.lang.exception import (
+    GsTaichiRuntimeTypeError,
+    GsTaichiSyntaxError,
+)
+from gstaichi.lang.matrix import Matrix, MatrixType
+from gstaichi.lang.struct import Struct, StructType
+from gstaichi.lang.util import cook_dtype, in_python_scope, python_scope
+from gstaichi.types import (
+    ndarray_type,
+    primitive_types,
+    sparse_matrix_builder,
+    texture_type,
+)
+from gstaichi.types.compound_types import CompoundType
+from gstaichi.types.utils import is_signed
+
+
+class ArgPack:
+    """ The `ArgPack` Type Class.
+
+    The `ArgPack` operates as a dictionary-like data pack, storing members as (key, value) pairs. Members stored can
+    range from scalars and matrices to other dictionary-like structures. Distinguished from structs, `ArgPack` can
+    accommodate buffer types such as `NdarrayType` and `TextureType` from GsTaichi. However, unlike `ti.Struct` which
+    serves as a data container, `ArgPack` functions as a reference container. It's important to note that `ArgPack`
+    cannot be nested within other types except for another `ArgPack`, and can only be utilized as kernel parameters.
+
+    Args:
+        annotations (Dict[str, Union[Dict, Matrix, Struct]]): \
+            The keys and types for `ArgPack` members.
+        dtype (ArgPackType): \
+            The ArgPackType class of this ArgPack object.
+        entries (Dict[str, Union[Dict, Matrix, Struct]]): \
+            The keys and corresponding values for `ArgPack` members.
+
+    Returns:
+        An instance of this `ArgPack`.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)
+        >>> pack_type = ti.ArgPackType(v=vec3, t=ti.f32)
+        >>> a = pack_type(v=vec3([0, 0, 0]), t=1.0)
+        >>> print(a.items)
+        dict_items([('v', [0. 0. 0.]), ('t', 1.0)])
+    """
+
+    _instance_count = 0
+
+    def __init__(self, annotations, dtype, *args, **kwargs):
+        # converts dicts to argument packs
+        if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
+            self.__entries = args[0]
+        elif len(args) == 0:
+            self.__entries = kwargs
+        else:
+            raise GsTaichiSyntaxError(
+                "Custom argument packs need to be initialized using either dictionary or keyword arguments"
+            )
+        if annotations.keys() != self.__entries.keys():
+            raise GsTaichiSyntaxError("ArgPack annotations keys not equals to entries keys.")
+        self.__annotations = annotations
+        for k, v in self.__entries.items():
+            self.__entries[k] = v if in_python_scope() else impl.expr_init(v)
+        self._register_members()
+        self.__dtype = dtype
+        self.__argpack = impl.get_runtime().prog.create_argpack(self.__dtype)
+        for i, (k, v) in enumerate(self.__entries.items()):
+            self._write_to_device(self.__annotations[k], type(v), v, self._calc_element_true_index(i))
+
+    def __del__(self):
+        if impl is not None and impl.get_runtime() is not None and impl.get_runtime().prog is not None:
+            impl.get_runtime().prog.delete_argpack(self.__argpack)
+
+    @property
+    def keys(self):
+        """Returns the list of member names in string format.
+
+        Example::
+
+           >>> vec3 = ti.types.vector(3, ti.f32)
+           >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
+           >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
+           >>> sphere.keys
+           ['center', 'radius']
+        """
+        return list(self.__entries.keys())
+
+    @property
+    def _members(self):
+        return list(self.__entries.values())
+
+    @property
+    def _annotations(self):
+        return list(self.__annotations.values())
+
+    @property
+    def items(self):
+        """Returns the items in this argument pack.
+
+        Example::
+
+           >>> vec3 = ti.types.vector(3, ti.f32)
+           >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
+           >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
+           >>> sphere.items
+            dict_items([('center', 2), ('radius', 1.0)])
+        """
+        return self.__entries.items()
+
+    def __getitem__(self, key):
+        ret = self.__entries[key]
+        return ret
+
+    def __setitem__(self, key, value):
+        self.__entries[key] = value
+        index = self._calc_element_true_index(list(self.__annotations).index(key))
+        self._write_to_device(self.__annotations[key], type(value), value, index)
+
+    def _set_entries(self, value):
+        if isinstance(value, dict):
+            value = ArgPack(self.__annotations, value)
+        for k in self.keys:
+            self[k] = value[k]
+
+    @staticmethod
+    def _make_getter(key):
+        def getter(self):
+            """Get an entry from custom argument pack by name."""
+            return self[key]
+
+        return getter
+
+    @staticmethod
+    def _make_setter(key):
+        @python_scope
+        def setter(self, value):
+            self[key] = value
+
+        return setter
+
+    def _register_members(self):
+        # https://stackoverflow.com/questions/48448074/adding-a-property-to-an-existing-object-instance
+        cls = self.__class__
+        new_cls_name = cls.__name__ + str(cls._instance_count)
+        cls._instance_count += 1
+        properties = {k: property(cls._make_getter(k), cls._make_setter(k)) for k in self.keys}
+        self.__class__ = type(new_cls_name, (cls,), properties)
+
+    def __len__(self):
+        """Get the number of entries in a custom argument pack."""
+        return len(self.__entries)
+
+    def __iter__(self):
+        return self.__entries.values()
+
+    def __str__(self):
+        """Python scope argument pack array print support."""
+        if impl.inside_kernel():
+            item_str = ", ".join([str(k) + "=" + str(v) for k, v in self.items])
+            return f"<ti.ArgPack {item_str}>"
+        return str(self.to_dict())
+
+    def __repr__(self):
+        return str(self.to_dict())
+
+    def to_dict(self):
+        """Converts the ArgPack to a dictionary.
+
+        Returns:
+            Dict: The result dictionary.
+        """
+        res_dict = {
+            k: v.to_dict() if isinstance(v, ArgPack) else v.to_list() if isinstance(v, Matrix) else v
+            for k, v in self.__entries.items()
+        }
+        return res_dict
+
+    def _calc_element_true_index(self, old_index):
+        for i in range(old_index):
+            anno = list(self.__annotations.values())[i]
+            if (
+                isinstance(anno, sparse_matrix_builder)
+                or isinstance(anno, ndarray_type.NdarrayType)
+                or isinstance(anno, texture_type.TextureType)
+                or isinstance(anno, texture_type.RWTextureType)
+                or isinstance(anno, ndarray_type.NdarrayType)
+            ):
+                old_index -= 1
+        return old_index
+
+    def _write_to_device(self, needed, provided, v, index):
+        if isinstance(needed, ArgPackType):
+            if not isinstance(v, ArgPack):
+                raise GsTaichiRuntimeTypeError.get(index, str(needed), str(provided))
+            self.__argpack.set_arg_nested_argpack(index, v.__argpack)
+        else:
+            # Note: do not use sth like "needed == f32". That would be slow.
+            if id(needed) in primitive_types.real_type_ids:
+                if not isinstance(v, (float, int, np.floating, np.integer)):
+                    raise GsTaichiRuntimeTypeError.get(index, needed.to_string(), provided)
+                self.__argpack.set_arg_float((index,), float(v))
+            elif id(needed) in primitive_types.integer_type_ids:
+                if not isinstance(v, (int, np.integer)):
+                    raise GsTaichiRuntimeTypeError.get(index, needed.to_string(), provided)
+                if is_signed(cook_dtype(needed)):
+                    self.__argpack.set_arg_int((index,), int(v))
+                else:
+                    self.__argpack.set_arg_uint((index,), int(v))
+            elif isinstance(needed, sparse_matrix_builder):
+                pass
+            elif isinstance(needed, ndarray_type.NdarrayType) and isinstance(v, gstaichi.lang._ndarray.Ndarray):
+                pass
+            elif isinstance(needed, texture_type.TextureType) and isinstance(v, gstaichi.lang._texture.Texture):
+                pass
+            elif isinstance(needed, texture_type.RWTextureType) and isinstance(v, gstaichi.lang._texture.Texture):
+                pass
+            elif isinstance(needed, ndarray_type.NdarrayType):
+                pass
+            elif isinstance(needed, MatrixType):
+                if needed.dtype in primitive_types.real_types:
+
+                    def cast_func(x):
+                        if not isinstance(x, (int, float, np.integer, np.floating)):
+                            raise GsTaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x))
+                        return float(x)
+
+                elif needed.dtype in primitive_types.integer_types:
+
+                    def cast_func(x):
+                        if not isinstance(x, (int, np.integer)):
+                            raise GsTaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x))
+                        return int(x)
+
+                else:
+                    raise ValueError(f"Matrix dtype {needed.dtype} is not integer type or real type.")
+
+                if needed.ndim == 2:
+                    v = [cast_func(v[i, j]) for i in range(needed.n) for j in range(needed.m)]
+                else:
+                    v = [cast_func(v[i]) for i in range(needed.n)]
+                v = needed(*v)
+                needed.set_argpack_struct_args(v, self.__argpack, (index,))
+            elif isinstance(needed, StructType):
+                if not isinstance(v, needed):
+                    raise GsTaichiRuntimeTypeError.get(index, str(needed), provided)
+                needed.set_argpack_struct_args(v, self.__argpack, (index,))
+            else:
+                raise ValueError(f"Argument type mismatch. Expecting {needed}, got {type(v)}.")
+
+
+class _IntermediateArgPack(ArgPack):
+    """Intermediate argument pack class for compiler internal use only.
+
+    Args:
+        annotations (Dict[str, Union[Expr, Matrix, Struct]]): keys and types for struct members.
+        entries (Dict[str, Union[Expr, Matrix, Struct]]): keys and values for struct members.
+    """
+
+    def __init__(self, annotations, dtype, *args, **kwargs):
+        # converts dicts to argument packs
+        if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
+            self._ArgPack__entries = args[0]
+        elif len(args) == 0:
+            self._ArgPack__entries = kwargs
+        else:
+            raise GsTaichiSyntaxError(
+                "Custom argument packs need to be initialized using either dictionary or keyword arguments"
+            )
+        if annotations.keys() != self._ArgPack__entries.keys():
+            raise GsTaichiSyntaxError("ArgPack annotations keys not equals to entries keys.")
+        self._ArgPack__annotations = annotations
+        self._register_members()
+        self._ArgPack__dtype = dtype
+        self._ArgPack__argpack = impl.get_runtime().prog.create_argpack(dtype)
+
+    def __del__(self):
+        pass
+
+
+class ArgPackType(CompoundType):
+    def __init__(self, **kwargs):
+        self.members = {}
+        elements = []
+        for k, dtype in kwargs.items():
+            if isinstance(dtype, StructType):
+                self.members[k] = dtype
+                elements.append([dtype.dtype, k])
+            elif isinstance(dtype, ArgPackType):
+                self.members[k] = dtype
+                elements.append(
+                    [
+                        _ti_core.DataTypeCxx(
+                            _ti_core.get_type_factory_instance().get_struct_type_for_argpack_ptr(dtype.dtype)
+                        ),
+                        k,
+                    ]
+                )
+            elif isinstance(dtype, MatrixType):
+                # Convert MatrixType to StructType
+                if dtype.ndim == 1:
+                    elements_ = [(dtype.dtype, f"{k}_{i}") for i in range(dtype.n)]
+                else:
+                    elements_ = [(dtype.dtype, f"{k}_{i}_{j}") for i in range(dtype.n) for j in range(dtype.m)]
+                self.members[k] = dtype
+                elements.append([_ti_core.get_type_factory_instance().get_struct_type(elements_), k])
+            elif isinstance(dtype, sparse_matrix_builder):
+                self.members[k] = dtype
+            elif isinstance(dtype, ndarray_type.NdarrayType):
+                self.members[k] = dtype
+            elif isinstance(dtype, texture_type.RWTextureType):
+                self.members[k] = dtype
+            elif isinstance(dtype, texture_type.TextureType):
+                self.members[k] = dtype
+            else:
+                dtype = cook_dtype(dtype)
+                self.members[k] = dtype
+                elements.append([dtype, k])
+        if len(elements) == 0:
+            # Use i32 as a placeholder for empty argpacks
+            elements.append([primitive_types.i32, k])
+        self.dtype = _ti_core.get_type_factory_instance().get_argpack_type(elements)
+
+    def __call__(self, *args, **kwargs):
+        """Create an instance of this argument pack type."""
+        d = {}
+        items = self.members.items()
+        # iterate over the members of this argument pack
+        for index, pair in enumerate(items):
+            name, dtype = pair  # (member name, member type))
+            if index < len(args):  # set from args
+                data = args[index]
+            else:  # set from kwargs
+                data = kwargs.get(name, None)
+
+            # If dtype is CompoundType and data is a scalar, it cannot be
+            # casted in the self.cast call later. We need an initialization here.
+            if isinstance(dtype, CompoundType) and not isinstance(data, (dict, ArgPack, Struct)):
+                data = dtype(data)
+
+            d[name] = data
+
+        entries = ArgPack(self.members, self.dtype, d)
+        pack = self.cast(entries)
+        return pack
+
+    def __instancecheck__(self, instance):
+        if not isinstance(instance, ArgPack):
+            return False
+        if list(self.members.keys()) != list(instance._ArgPack__entries.keys()):
+            return False
+        for k, v in self.members.items():
+            if isinstance(v, ArgPackType):
+                if not isinstance(instance._ArgPack__entries[k], v):
+                    return False
+            elif instance._ArgPack__annotations[k] != v:
+                return False
+        return True
+
+    def cast(self, pack):
+        # sanity check members
+        if self.members.keys() != pack._ArgPack__entries.keys():
+            raise GsTaichiSyntaxError("Incompatible arguments for custom argument pack members!")
+        entries = {}
+        for k, dtype in self.members.items():
+            if isinstance(dtype, MatrixType):
+                entries[k] = dtype(pack._ArgPack__entries[k])
+            elif isinstance(dtype, CompoundType):
+                entries[k] = dtype.cast(pack._ArgPack__entries[k])
+            elif isinstance(dtype, ArgPackType):
+                entries[k] = dtype.cast(pack._ArgPack__entries[k])
+            elif isinstance(dtype, ndarray_type.NdarrayType):
+                entries[k] = pack._ArgPack__entries[k]
+            elif isinstance(dtype, texture_type.RWTextureType):
+                entries[k] = pack._ArgPack__entries[k]
+            elif isinstance(dtype, texture_type.TextureType):
+                entries[k] = pack._ArgPack__entries[k]
+            elif isinstance(dtype, sparse_matrix_builder):
+                entries[k] = pack._ArgPack__entries[k]
+            else:
+                if in_python_scope():
+                    v = pack._ArgPack__entries[k]
+                    entries[k] = int(v) if dtype in primitive_types.integer_types else float(v)
+                else:
+                    entries[k] = ops.cast(pack._ArgPack__entries[k], dtype)
+        pack = ArgPack(self.members, self.dtype, entries)
+        return pack
+
+    def from_gstaichi_object(self, arg_load_dict: dict):
+        d = {}
+        items = self.members.items()
+        for index, pair in enumerate(items):
+            name, dtype = pair
+            d[name] = arg_load_dict[name]
+        pack = _IntermediateArgPack(self.members, self.dtype, d)
+        pack._ArgPack__dtype = self.dtype
+        return pack
+
+    def __str__(self):
+        """Python scope argpack type print support."""
+        item_str = ", ".join([str(k) + "=" + str(v) for k, v in self.members.items()])
+        return f"<ti.ArgPackType {item_str}>"
+
+
+__all__ = ["ArgPack"]

gstaichi/lang/ast/__init__.py

@@ -0,0 +1,5 @@
+# type: ignore
+
+from gstaichi.lang.ast.ast_transformer_utils import ASTTransformerContext
+from gstaichi.lang.ast.checkers import KernelSimplicityASTChecker
+from gstaichi.lang.ast.transform import transform_tree