gstaichi 0.1.25.dev0__cp311-cp311-win_amd64.whl

gstaichi/lang/__init__.py

@@ -0,0 +1,50 @@
+# type: ignore
+
+from gstaichi.lang import impl, simt
+from gstaichi.lang._ndarray import *
+from gstaichi.lang._ndrange import ndrange
+from gstaichi.lang._texture import Texture
+from gstaichi.lang.argpack import *
+from gstaichi.lang.exception import *
+from gstaichi.lang.field import *
+from gstaichi.lang.impl import *
+from gstaichi.lang.kernel_impl import *
+from gstaichi.lang.matrix import *
+from gstaichi.lang.mesh import *
+from gstaichi.lang.misc import *  # pylint: disable=W0622
+from gstaichi.lang.ops import *  # pylint: disable=W0622
+from gstaichi.lang.runtime_ops import *
+from gstaichi.lang.snode import *
+from gstaichi.lang.source_builder import *
+from gstaichi.lang.struct import *
+from gstaichi.types.enums import DeviceCapability, Format, Layout
+
+__all__ = [
+    s
+    for s in dir()
+    if not s.startswith("_")
+    and s
+    not in [
+        "any_array",
+        "ast",
+        "common_ops",
+        "enums",
+        "exception",
+        "expr",
+        "impl",
+        "inspect",
+        "kernel_arguments",
+        "kernel_impl",
+        "matrix",
+        "mesh",
+        "misc",
+        "ops",
+        "platform",
+        "runtime_ops",
+        "shell",
+        "snode",
+        "source_builder",
+        "struct",
+        "util",
+    ]
+]

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"]

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"]