gstaichi 2.1.1__cp313-cp313-macosx_11_0_arm64.whl

gstaichi/lang/_ndrange.py

@@ -0,0 +1,152 @@
+# type: ignore
+
+import collections.abc
+from typing import Iterable
+
+import numpy as np
+
+from gstaichi.lang import ops
+from gstaichi.lang.exception import GsTaichiSyntaxError, GsTaichiTypeError
+from gstaichi.lang.expr import Expr
+from gstaichi.lang.matrix import Matrix
+from gstaichi.types.utils import is_integral
+
+
+class _Ndrange:
+    def __init__(self, *args):
+        args = list(args)
+        for i, arg in enumerate(args):
+            if not isinstance(arg, collections.abc.Sequence):
+                args[i] = (0, arg)
+            if len(args[i]) != 2:
+                raise GsTaichiSyntaxError(
+                    "Every argument of ndrange should be a scalar or a tuple/list like (begin, end)"
+                )
+            args[i] = (args[i][0], ops.max(args[i][0], args[i][1]))
+        for arg in args:
+            for bound in arg:
+                if not isinstance(bound, (int, np.integer)) and not (
+                    isinstance(bound, Expr) and is_integral(bound.ptr.get_rvalue_type())
+                ):
+                    raise GsTaichiTypeError(
+                        "Every argument of ndrange should be an integer scalar or a tuple/list of (int, int)"
+                    )
+        self.bounds = args
+
+        self.dimensions = [None] * len(args)
+        for i, bound in enumerate(self.bounds):
+            self.dimensions[i] = bound[1] - bound[0]
+
+        self.acc_dimensions = self.dimensions.copy()
+        for i in reversed(range(len(self.bounds) - 1)):
+            self.acc_dimensions[i] = self.acc_dimensions[i] * self.acc_dimensions[i + 1]
+        if len(self.acc_dimensions) == 0:  # for the empty case, e.g. ti.ndrange()
+            self.acc_dimensions = [1]
+
+    def __iter__(self):
+        def gen(d, prefix):
+            if d == len(self.bounds):
+                yield prefix
+            else:
+                for t in range(self.bounds[d][0], self.bounds[d][1]):
+                    yield from gen(d + 1, prefix + (t,))
+
+        yield from gen(0, ())
+
+    def grouped(self):
+        return GroupedNDRange(self)
+
+
+def ndrange(*args) -> Iterable:
+    """Return an immutable iterator object for looping over multi-dimensional indices.
+
+    This returned set of multi-dimensional indices is the direct product (in the set-theory sense)
+    of n groups of integers, where n equals the number of arguments in the input list, and looks like
+
+    range(x1, y1) x range(x2, y2) x ... x range(xn, yn)
+
+    The k-th argument corresponds to the k-th `range()` factor in the above product, and each
+    argument must be an integer or a pair of two integers. An integer argument n will be interpreted
+    as `range(0, n)`, and a pair of two integers (start, end) will be interpreted as `range(start, end)`.
+
+    You can loop over these multi-dimensonal indices in different ways, see the examples below.
+
+    Args:
+        entries: (int, tuple): Must be either an integer, or a tuple/list of two integers.
+
+    Returns:
+        An immutable iterator object.
+
+    Example::
+
+        You can loop over 1-D integers in range [start, end), as in native Python
+
+            >>> @ti.kernel
+            >>> def loop_1d():
+            >>>     start = 2
+            >>>     end = 5
+            >>>     for i in ti.ndrange((start, end)):
+            >>>         print(i)  # will print 2 3 4
+
+        Note the braces around `(start, end)` in the above code. If without them,
+        the parameter `2` will be interpreted as `range(0, 2)`, `5` will be
+        interpreted as `range(0, 5)`, and you will get a set of 2-D indices which
+        contains 2x5=10 elements, and need two indices i, j to loop over them:
+
+            >>> @ti.kernel
+            >>> def loop_2d():
+            >>>     for i, j in ti.ndrange(2, 5):
+            >>>         print(i, j)
+            0 0
+            ...
+            0 4
+            ...
+            1 4
+
+        But you do can use a single index i to loop over these 2-D indices, in this case
+        the indices are returned as a 1-D array `(0, 1, ..., 9)`:
+
+            >>> @ti.kernel
+            >>> def loop_2d_as_1d():
+            >>>     for i in ti.ndrange(2, 5):
+            >>>         print(i)
+            will print 0 1 2 3 4 5 6 7 8 9
+
+        In general, you can use any `1 <= k <= n` iterators to loop over a set of n-D
+        indices. For `k=n` all the indices are n-dimensional, and they are returned in
+        lexical order, but for `k<n` iterators the last n-k+1 dimensions will be collapsed into
+        a 1-D array of consecutive integers `(0, 1, 2, ...)` whose length equals the
+        total number of indices in the last n-k+1 dimensions:
+
+            >>> @ti.kernel
+            >>> def loop_3d_as_2d():
+            >>>     # use two iterators to loop over a set of 3-D indices
+            >>>     # the last two dimensions for 4, 5 will collapse into
+            >>>     # the array [0, 1, 2, ..., 19]
+            >>>     for i, j in ti.ndrange(3, 4, 5):
+            >>>         print(i, j)
+            will print 0 0, 0 1, ..., 0 19, ..., 2 19.
+
+        A typical usage of `ndrange` is when you want to loop over a tensor and process
+        its entries in parallel. You should avoid writing nested `for` loops here since
+        only top level `for` loops are paralleled in gstaichi, instead you can use `ndrange`
+        to hold all entries in one top level loop:
+
+            >>> @ti.kernel
+            >>> def loop_tensor():
+            >>>     for row, col, channel in ti.ndrange(image_height, image_width, channels):
+            >>>         image[row, col, channel] = ...
+    """
+    return _Ndrange(*args)
+
+
+class GroupedNDRange:
+    def __init__(self, r):
+        self.r = r
+
+    def __iter__(self):
+        for ind in self.r:
+            yield Matrix(list(ind))
+
+
+__all__ = ["ndrange"]

gstaichi/lang/_template_mapper.py

@@ -0,0 +1,195 @@
+import dataclasses
+import weakref
+from typing import Any, Callable, Union
+
+import gstaichi.lang
+import gstaichi.lang._ndarray
+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.exception import (
+    GsTaichiRuntimeTypeError,
+)
+from gstaichi.lang.kernel_arguments import ArgMetadata
+from gstaichi.lang.matrix import MatrixType
+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,
+    "texture_type.TextureType",
+    "texture_type.RWTextureType",
+    ndarray_type.NdarrayType,
+    sparse_matrix_builder,
+    Any,
+]
+
+
+class TemplateMapper:
+    """
+    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[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: 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):
+                return arg.ptr.get_underlying_ptr_address()
+            if isinstance(arg, _ti_core.ExprCxx):
+                return arg.get_underlying_ptr_address()
+            if isinstance(arg, tuple):
+                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"
+                )
+
+            if isinstance(arg, (list, tuple, dict, set)) or hasattr(arg, "_data_oriented"):
+                # [Composite arguments] Return weak reference to the object
+                # GsTaichi 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 dataclasses.is_dataclass(annotation):
+            _res_l = []
+            for field in dataclasses.fields(annotation):
+                field_value = getattr(arg, field.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):
+            if not isinstance(arg, gstaichi.lang._texture.Texture):
+                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
+            if arg.num_dims != annotation.num_dimensions:
+                raise GsTaichiRuntimeTypeError(
+                    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, gstaichi.lang._texture.Texture):
+                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
+            if arg.num_dims != annotation.num_dimensions:
+                raise GsTaichiRuntimeTypeError(
+                    f"RWTextureType dimension mismatch for argument {arg_name}: expected {annotation.num_dimensions}, got {arg.num_dims}"
+                )
+            if arg.fmt != annotation.fmt:
+                raise GsTaichiRuntimeTypeError(
+                    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, gstaichi.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 GsTaichiRuntimeTypeError(f"Invalid type for argument {arg_name}, got {arg}")
+            shape = tuple(shape)
+            element_shape: tuple[int, ...] = ()
+            dtype = to_gstaichi_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: tuple[Any, ...]) -> tuple[Any, ...]:
+        extracted: list[Any] = []
+        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: tuple[Any, ...]) -> tuple[int, tuple[Any, ...]]:
+        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

gstaichi/lang/_texture.py

@@ -0,0 +1,172 @@
+# type: ignore
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.expr import Expr, make_expr_group
+from gstaichi.lang.matrix import Matrix
+from gstaichi.lang.util import gstaichi_scope
+from gstaichi.types import vector
+from gstaichi.types.primitive_types import f32
+
+
+def _get_entries(mat):
+    if isinstance(mat, Matrix):
+        return mat.entries
+    return [mat]
+
+
+class TextureSampler:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @gstaichi_scope
+    def sample_lod(self, uv, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(uv), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kSampleLod, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @gstaichi_scope
+    def fetch(self, index, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kFetchTexel, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+
+class RWTextureAccessor:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        # gstaichi_python.TexturePtrExpression.
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @gstaichi_scope
+    def load(self, index):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index))
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kLoad, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @gstaichi_scope
+    def store(self, index, value):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), *_get_entries(value))
+        impl.expr_init(
+            ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kStore, self.ptr_expr, args_group, dbg_info)
+        )
+
+    @property
+    @gstaichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr_expr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        ret = [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr_expr, i, dbg_info)) for i in range(dim)]
+        return ret
+
+    @gstaichi_scope
+    def _loop_range(self):
+        """Gets the corresponding gstaichi_python.Expr to serve as loop range.
+
+        Returns:
+            gstaichi_python.Expr: See above.
+        """
+        return self.ptr_expr
+
+
+class Texture:
+    """GsTaichi Texture class.
+
+    Args:
+        fmt (ti.Format): Color format of the texture.
+        shape (Tuple[int]): Shape of the Texture.
+    """
+
+    def __init__(self, fmt, arr_shape):
+        self.tex = impl.get_runtime().prog.create_texture(fmt, arr_shape)
+        self.fmt = fmt
+        self.num_dims = len(arr_shape)
+        self.shape = arr_shape
+
+    def from_ndarray(self, ndarray):
+        """Loads an ndarray to texture.
+
+        Args:
+            ndarray (ti.Ndarray): Source ndarray to load from.
+        """
+        self.tex.from_ndarray(ndarray.arr)
+
+    def from_field(self, field):
+        """Loads a field to texture.
+
+        Args:
+            field (ti.Field): Source field to load from.
+        """
+        self.tex.from_snode(field.snode.ptr)
+
+    def _device_allocation_ptr(self):
+        return self.tex.device_allocation_ptr()
+
+    def from_image(self, image):
+        """Loads a PIL image to texture. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Args:
+            image (PIL.Image.Image): Source PIL image to load from.
+
+        """
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        assert isinstance(image, Image.Image)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        assert image.size == tuple(self.shape)
+
+        assert self.num_dims == 2
+        # Don't use transpose method since its enums are too new
+        image = image.rotate(90, expand=True)
+        arr = np.asarray(image)
+        from gstaichi._kernels import (  # pylint: disable=import-outside-toplevel
+            load_texture_from_numpy,
+        )
+
+        load_texture_from_numpy(self, arr)
+
+    def to_image(self):
+        """Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Returns:
+            img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture.
+        """
+        assert self.num_dims == 2
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        res = np.zeros(self.shape + (3,), np.uint8)
+        from gstaichi._kernels import (  # pylint: disable=import-outside-toplevel
+            save_texture_to_numpy,
+        )
+
+        save_texture_to_numpy(self, res)
+        return Image.fromarray(res).rotate(270, expand=True)