gstaichi 2.1.1rc3__cp310-cp310-macosx_11_0_arm64.whl

gstaichi/types/compound_types.py

@@ -0,0 +1,71 @@
+# type: ignore
+
+import gstaichi
+from gstaichi._lib.utils import ti_python_core as _ti_python_core
+
+_type_factory = _ti_python_core.get_type_factory_instance()
+
+
+class CompoundType:
+    def from_kernel_struct_ret(self, launch_ctx, index: tuple):
+        raise NotImplementedError()
+
+
+# TODO: maybe move MatrixType, StructType here to avoid the circular import?
+def matrix(n=None, m=None, dtype=None):
+    """Creates a matrix type with given shape and data type.
+
+    Args:
+        n (int): number of rows of the matrix.
+        m (int): number of columns of the matrix.
+        dtype (:mod:`~gstaichi.types.primitive_types`): matrix data type.
+
+    Returns:
+        A matrix type.
+
+    Example::
+
+        >>> mat2x2 = ti.types.matrix(2, 2, ti.f32)  # 2x2 matrix type
+        >>> M = mat2x2([[1., 2.], [3., 4.]])  # an instance of this type
+    """
+    return gstaichi.lang.matrix.MatrixType(n, m, 2, dtype)
+
+
+def vector(n=None, dtype=None):
+    """Creates a vector type with given shape and data type.
+
+    Args:
+        n (int): dimension of the vector.
+        dtype (:mod:`~gstaichi.types.primitive_types`): vector data type.
+
+    Returns:
+        A vector type.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)  # 3d vector type
+        >>> v = vec3([1., 2., 3.])  # an instance of this type
+    """
+    return gstaichi.lang.matrix.VectorType(n, dtype)
+
+
+def struct(**kwargs):
+    """Creates a struct type with given members.
+
+    Args:
+        kwargs (dict): a dictionary contains the names and types of the
+            struct members.
+
+    Returns:
+        A struct type.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)
+        >>> sphere = ti.types.struct(center=vec3, radius=float)
+        >>> s = sphere(center=vec3([0., 0., 0.]), radius=1.0)
+    """
+    return gstaichi.lang.struct.StructType(**kwargs)
+
+
+__all__ = ["matrix", "vector", "struct"]

gstaichi/types/enums.py

@@ -0,0 +1,49 @@
+# type: ignore
+
+from gstaichi._lib import core as _ti_core
+
+Layout = _ti_core.Layout
+AutodiffMode = _ti_core.AutodiffMode
+SNodeGradType = _ti_core.SNodeGradType
+Format = _ti_core.Format
+BoundaryMode = _ti_core.BoundaryMode
+
+
+def to_boundary_enum(boundary):
+    if boundary == "clamp":
+        return BoundaryMode.CLAMP
+    if boundary == "unsafe":
+        return BoundaryMode.UNSAFE
+    raise ValueError(f"Invalid boundary argument: {boundary}")
+
+
+class DeviceCapability:
+    spirv_version_1_3 = "spirv_version=66304"
+    spirv_version_1_4 = "spirv_version=66560"
+    spirv_version_1_5 = "spirv_version=66816"
+    spirv_has_int8 = "spirv_has_int8"
+    spirv_has_int16 = "spirv_has_int16"
+    spirv_has_int64 = "spirv_has_int64"
+    spirv_has_float16 = "spirv_has_float16"
+    spirv_has_float64 = "spirv_has_float64"
+    spirv_has_atomic_int64 = "spirv_has_atomic_int64"
+    spirv_has_atomic_float16 = "spirv_has_atomic_float16"
+    spirv_has_atomic_float16_add = "spirv_has_atomic_float16_add"
+    spirv_has_atomic_float16_minmax = "spirv_has_atomic_float16_minmax"
+    spirv_has_atomic_float = "spirv_has_atomic_float"
+    spirv_has_atomic_float_add = "spirv_has_atomic_float_add"
+    spirv_has_atomic_float_minmax = "spirv_has_atomic_float_minmax"
+    spirv_has_atomic_float64 = "spirv_has_atomic_float64"
+    spirv_has_atomic_float64_add = "spirv_has_atomic_float64_add"
+    spirv_has_atomic_float64_minmax = "spirv_has_atomic_float64_minmax"
+    spirv_has_variable_ptr = "spirv_has_variable_ptr"
+    spirv_has_physical_storage_buffer = "spirv_has_physical_storage_buffer"
+    spirv_has_subgroup_basic = "spirv_has_subgroup_basic"
+    spirv_has_subgroup_vote = "spirv_has_subgroup_vote"
+    spirv_has_subgroup_arithmetic = "spirv_has_subgroup_arithmetic"
+    spirv_has_subgroup_ballot = "spirv_has_subgroup_ballot"
+    spirv_has_non_semantic_info = "spirv_has_non_semantic_info"
+    spirv_has_no_integer_wrap_decoration = "spirv_has_no_integer_wrap_decoration"
+
+
+__all__ = ["Layout", "AutodiffMode", "SNodeGradType", "Format", "DeviceCapability"]

gstaichi/types/ndarray_type.py

@@ -0,0 +1,169 @@
+from typing import Any
+
+from gstaichi.types.compound_types import CompoundType, matrix, vector
+from gstaichi.types.enums import Layout, to_boundary_enum
+
+
+class NdarrayTypeMetadata:
+    def __init__(self, element_type, shape=None, needs_grad=False):
+        self.element_type = element_type
+        self.shape = shape
+        self.layout = Layout.AOS
+        self.needs_grad = needs_grad
+
+
+# TODO(Haidong): This is a helper function that creates a MatrixType
+#                with respect to element_dim and element_shape.
+#                Remove this function when the two args are totally deprecated.
+def _make_matrix_dtype_from_element_shape(element_dim, element_shape, primitive_dtype):
+    if isinstance(primitive_dtype, CompoundType):
+        raise TypeError(f'Cannot specifiy matrix dtype "{primitive_dtype}" and element shape or dim at the same time.')
+
+    # Scalars
+    if element_dim == 0 or (element_shape is not None and len(element_shape) == 0):
+        return primitive_dtype
+
+    # Cook element dim and shape into matrix type.
+    mat_dtype = None
+    if element_dim is not None:
+        # TODO: expand use case with arbitary tensor dims!
+        if element_dim < 0 or element_dim > 2:
+            raise ValueError("Only scalars, vectors, and matrices are allowed as elements of ti.types.ndarray()")
+        # Check dim consistency. The matrix dtype will be cooked later.
+        if element_shape is not None and len(element_shape) != element_dim:
+            raise ValueError(
+                f"Both element_shape and element_dim are specified, but shape doesn't match specified dim: "
+                f"{len(element_shape)}!={element_dim}"
+            )
+        mat_dtype = vector(None, primitive_dtype) if element_dim == 1 else matrix(None, None, primitive_dtype)
+    elif element_shape is not None:
+        if len(element_shape) > 2:
+            raise ValueError("Only scalars, vectors, and matrices are allowed as elements of ti.types.ndarray()")
+        mat_dtype = (
+            vector(element_shape[0], primitive_dtype)
+            if len(element_shape) == 1
+            else matrix(element_shape[0], element_shape[1], primitive_dtype)
+        )
+    return mat_dtype
+
+
+class NdarrayType:
+    """Type annotation for arbitrary arrays, including external arrays (numpy ndarrays and torch tensors) and GsTaichi ndarrays.
+
+    For external arrays, we treat it as a GsTaichi data container with Scalar, Vector or Matrix elements.
+    For GsTaichi vector/matrix ndarrays, we will automatically identify element dimension and their corresponding axis by the
+    dimension of datatype, say scalars, matrices or vectors.
+    For example, given type annotation `ti.types.ndarray(dtype=ti.math.vec3)`, a numpy array `np.zeros(10, 10, 3)` will be
+    recognized as a 10x10 matrix composed of vec3 elements.
+
+    Args:
+        dtype (Union[PrimitiveType, VectorType, MatrixType, NoneType], optional): None if not speicified.
+        ndim (Union[Int, NoneType]): None if not specified, number of field dimensions. This argument is ignored for externa
+        arrays for now.
+        element_dim (Union[Int, NoneType], optional):
+          None if not specified (will be treated as 0 for external arrays),
+          0 if scalar elements,
+          1 if vector elements, and
+          2 if matrix elements.
+        element_shape (Union[Tuple[Int], NoneType]):
+          None if not specified, shapes of each element.
+          For example, element_shape must be 1d for vector and 2d tuple for matrix.
+          This argument is ignored for external arrays for now.
+    """
+
+    def __init__(
+        self,
+        dtype=None,
+        ndim=None,
+        element_dim=None,
+        element_shape=None,
+        field_dim=None,
+        needs_grad=None,
+        boundary="unsafe",
+    ):
+        if field_dim is not None:
+            raise ValueError("The field_dim argument for ndarray type is already deprecated. Please use ndim instead.")
+        if element_dim is not None or element_shape is not None:
+            self.dtype = _make_matrix_dtype_from_element_shape(element_dim, element_shape, dtype)
+        else:
+            self.dtype = dtype
+
+        self.ndim = ndim
+        self.layout = Layout.AOS
+        self.needs_grad = needs_grad
+        self.boundary = to_boundary_enum(boundary)
+
+    @classmethod
+    def __class_getitem__(cls, args, **kwargs):
+        return cls(*args, **kwargs)
+
+    def check_matched(self, ndarray_type: NdarrayTypeMetadata, arg_name: str):
+        # FIXME(Haidong) Cannot use Vector/MatrixType due to circular import
+        # Use the CompuoundType instead to determine the specific typs.
+        # TODO Replace CompoundType with MatrixType and VectorType
+
+        # Check dtype match
+        if isinstance(self.dtype, CompoundType):
+            if not self.dtype.check_matched(ndarray_type.element_type):  # type: ignore
+                raise ValueError(
+                    f"Invalid value for argument {arg_name} - required element type: {self.dtype.to_string()}, "  # type: ignore
+                    f"but {ndarray_type.element_type.to_string()} is provided"
+                )
+        else:
+            if self.dtype is not None:
+                # Check dtype match for scalar.
+                from gstaichi.lang import util  # pylint: disable=C0415
+
+                if not util.cook_dtype(self.dtype) == ndarray_type.element_type:
+                    raise TypeError(
+                        f"Expect element type {self.dtype} for argument {arg_name}, but get {ndarray_type.element_type}"
+                    )
+
+        # Check ndim match
+        if self.ndim is not None and ndarray_type.shape is not None and self.ndim != len(ndarray_type.shape):
+            raise ValueError(
+                f"Invalid value for argument {arg_name} - required ndim={self.ndim}, but {len(ndarray_type.shape)}d "
+                f"ndarray with shape {ndarray_type.shape} is provided"
+            )
+
+        # Check needs_grad
+        if self.needs_grad is not None and self.needs_grad > ndarray_type.needs_grad:
+            # It's okay to pass a needs_grad=True ndarray at runtime to a need_grad=False arg but not vice versa.
+            raise ValueError(
+                f"Invalid value for argument {arg_name} - required needs_grad={self.needs_grad}, but "
+                f"{ndarray_type.needs_grad} is provided"
+            )
+
+    def __repr__(self):
+        return f"NdarrayType(dtype={self.dtype}, ndim={self.ndim}, layout={self.layout}, needs_grad={self.needs_grad})"
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __getitem__(self, i: Any) -> Any:
+        # needed for pyright
+        raise NotImplemented
+
+    def __setitem__(self, i: Any, v: Any) -> None:
+        # needed for pyright
+        raise NotImplemented
+
+
+ndarray = NdarrayType
+NDArray = NdarrayType
+"""Alias for :class:`~gstaichi.types.ndarray_type.NdarrayType`.
+
+Example::
+
+    >>> @ti.kernel
+    >>> def to_numpy(x: ti.types.ndarray(), y: ti.types.ndarray()):
+    >>>     for i in range(n):
+    >>>         x[i] = y[i]
+    >>>
+    >>> y = ti.ndarray(ti.f64, shape=n)
+    >>> ... # calculate y
+    >>> x = numpy.zeros(n)
+    >>> to_numpy(x, y)  # `x` will be filled with `y`'s data.
+"""
+
+__all__ = ["ndarray", "NDArray"]

gstaichi/types/primitive_types.py

@@ -0,0 +1,206 @@
+from typing import Union
+
+from gstaichi._lib import core as ti_python_core
+
+# ========================================
+# real types
+
+# ----------------------------------------
+
+float16 = ti_python_core.DataType_f16
+"""16-bit precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f16 = float16
+"""Alias for :const:`~gstaichi.types.primitive_types.float16`
+"""
+
+# ----------------------------------------
+
+float32 = ti_python_core.DataType_f32
+"""32-bit single precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f32 = float32
+"""Alias for :const:`~gstaichi.types.primitive_types.float32`
+"""
+
+# ----------------------------------------
+
+float64 = ti_python_core.DataType_f64
+"""64-bit double precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f64 = float64
+"""Alias for :const:`~gstaichi.types.primitive_types.float64`
+"""
+# ----------------------------------------
+
+# ========================================
+# Integer types
+
+# ----------------------------------------
+
+int8 = ti_python_core.DataType_i8
+"""8-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i8 = int8
+"""Alias for :const:`~gstaichi.types.primitive_types.int8`
+"""
+
+# ----------------------------------------
+
+int16 = ti_python_core.DataType_i16
+"""16-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i16 = int16
+"""Alias for :const:`~gstaichi.types.primitive_types.int16`
+"""
+
+# ----------------------------------------
+
+int32 = ti_python_core.DataType_i32
+"""32-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i32 = int32
+"""Alias for :const:`~gstaichi.types.primitive_types.int32`
+"""
+
+# ----------------------------------------
+
+int64 = ti_python_core.DataType_i64
+"""64-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i64 = int64
+"""Alias for :const:`~gstaichi.types.primitive_types.int64`
+"""
+
+# ----------------------------------------
+
+uint8 = ti_python_core.DataType_u8
+"""8-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+uint1 = ti_python_core.DataType_u1
+"""1-bit unsigned integer data type. Same as booleans.
+"""
+
+# ----------------------------------------
+
+u1 = uint1
+"""Alias for :const:`~gstaichi.types.primitive_types.uint1`
+"""
+
+# ----------------------------------------
+
+u8 = uint8
+"""Alias for :const:`~gstaichi.types.primitive_types.uint8`
+"""
+
+# ----------------------------------------
+
+uint16 = ti_python_core.DataType_u16
+"""16-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u16 = uint16
+"""Alias for :const:`~gstaichi.types.primitive_types.uint16`
+"""
+
+# ----------------------------------------
+
+uint32 = ti_python_core.DataType_u32
+"""32-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u32 = uint32
+"""Alias for :const:`~gstaichi.types.primitive_types.uint32`
+"""
+
+# ----------------------------------------
+
+uint64 = ti_python_core.DataType_u64
+"""64-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u64 = uint64
+"""Alias for :const:`~gstaichi.types.primitive_types.uint64`
+"""
+
+# ----------------------------------------
+
+
+class RefType:
+    def __init__(self, tp):
+        self.tp = tp
+
+
+def ref(tp):
+    return RefType(tp)
+
+
+real_types = [f16, f32, f64, float]
+real_type_ids = [id(t) for t in real_types]
+
+integer_types = [i8, i16, i32, i64, u1, u8, u16, u32, u64, int, bool]
+integer_type_ids = [id(t) for t in integer_types]
+
+all_types = real_types + integer_types
+type_ids = [id(t) for t in all_types]
+
+_python_primitive_types = Union[int, float, bool, str, None]
+
+__all__ = [
+    "float32",
+    "f32",
+    "float64",
+    "f64",
+    "float16",
+    "f16",
+    "int8",
+    "i8",
+    "int16",
+    "i16",
+    "int32",
+    "i32",
+    "int64",
+    "i64",
+    "uint1",
+    "u1",
+    "uint8",
+    "u8",
+    "uint16",
+    "u16",
+    "uint32",
+    "u32",
+    "uint64",
+    "u64",
+    "ref",
+    "_python_primitive_types",
+]

gstaichi/types/quant.py

@@ -0,0 +1,88 @@
+# type: ignore
+
+"""
+This module defines generators of quantized types.
+For more details, read https://yuanming.gstaichi.graphics/publication/2021-quangstaichi/quangstaichi.pdf.
+"""
+
+from gstaichi._lib.utils import ti_python_core as _ti_python_core
+from gstaichi.types.primitive_types import i32
+
+_type_factory = _ti_python_core.get_type_factory_instance()
+
+
+def int(bits, signed=True, compute=None):  # pylint: disable=W0622
+    """Generates a quantized type for integers.
+
+    Args:
+        bits (int): Number of bits.
+        signed (bool): Signed or unsigned.
+        compute (DataType): Type for computation.
+
+    Returns:
+        DataType: The specified type.
+    """
+    if compute is None:
+        from gstaichi.lang import impl  # pylint: disable=C0415
+
+        compute = impl.get_runtime().default_ip if signed else impl.get_runtime().default_up
+    if isinstance(compute, _ti_python_core.DataTypeCxx):
+        compute = compute.get_ptr()
+    return _type_factory.get_quant_int_type(bits, signed, compute)
+
+
+def fixed(bits, signed=True, max_value=1.0, compute=None, scale=None):
+    """Generates a quantized type for fixed-point real numbers.
+
+    Args:
+        bits (int): Number of bits.
+        signed (bool): Signed or unsigned.
+        max_value (float): Maximum value of the number.
+        compute (DataType): Type for computation.
+        scale (float): Scaling factor. The argument is prioritized over range.
+
+    Returns:
+        DataType: The specified type.
+    """
+    if compute is None:
+        from gstaichi.lang import impl  # pylint: disable=C0415
+
+        compute = impl.get_runtime().default_fp
+    if isinstance(compute, _ti_python_core.DataTypeCxx):
+        compute = compute.get_ptr()
+    # TODO: handle cases with bits > 32
+    underlying_type = int(bits=bits, signed=signed, compute=i32)
+    if scale is None:
+        if signed:
+            scale = max_value / 2 ** (bits - 1)
+        else:
+            scale = max_value / 2**bits
+    return _type_factory.get_quant_fixed_type(underlying_type, compute, scale)
+
+
+def float(exp, frac, signed=True, compute=None):  # pylint: disable=W0622
+    """Generates a quantized type for floating-point real numbers.
+
+    Args:
+        exp (int): Number of exponent bits.
+        frac (int): Number of fraction bits.
+        signed (bool): Signed or unsigned.
+        compute (DataType): Type for computation.
+
+    Returns:
+        DataType: The specified type.
+    """
+    if compute is None:
+        from gstaichi.lang import impl  # pylint: disable=C0415
+
+        compute = impl.get_runtime().default_fp
+    if isinstance(compute, _ti_python_core.DataTypeCxx):
+        compute = compute.get_ptr()
+    # Exponent is always unsigned
+    exp_type = int(bits=exp, signed=False, compute=i32)
+    # TODO: handle cases with frac > 32
+    frac_type = int(bits=frac, signed=signed, compute=i32)
+    return _type_factory.get_quant_float_type(frac_type, exp_type, compute)
+
+
+__all__ = ["int", "fixed", "float"]

gstaichi/types/texture_type.py

@@ -0,0 +1,85 @@
+# type: ignore
+
+from gstaichi.lang.exception import GsTaichiCompilationError
+from gstaichi.types.enums import Format
+from gstaichi.types.primitive_types import f16, f32, i8, i16, i32, u8, u16, u32
+
+FORMAT2TY_CH = {
+    Format.r8: (u8, 1),
+    Format.r8u: (u8, 1),
+    Format.r8i: (i8, 1),
+    Format.rg8: (u8, 2),
+    Format.rg8u: (u8, 2),
+    Format.rg8i: (i8, 2),
+    Format.rgba8: (u8, 4),
+    Format.rgba8u: (u8, 4),
+    Format.rgba8i: (i8, 4),
+    Format.r16: (u16, 1),
+    Format.r16u: (u16, 1),
+    Format.r16i: (i16, 1),
+    Format.r16f: (f16, 1),
+    Format.rg16: (u16, 2),
+    Format.rg16u: (u16, 2),
+    Format.rg16i: (i16, 2),
+    Format.rg16f: (f16, 2),
+    Format.rgb16: (u16, 3),
+    Format.rgb16u: (u16, 3),
+    Format.rgb16i: (i16, 3),
+    Format.rgb16f: (f16, 3),
+    Format.rgba16: (u16, 4),
+    Format.rgba16u: (u16, 4),
+    Format.rgba16i: (i16, 4),
+    Format.rgba16f: (f16, 4),
+    Format.r32u: (u32, 1),
+    Format.r32i: (i32, 1),
+    Format.r32f: (f32, 1),
+    Format.rg32u: (u32, 2),
+    Format.rg32i: (i32, 2),
+    Format.rg32f: (f32, 2),
+    Format.rgb32u: (u32, 3),
+    Format.rgb32i: (i32, 3),
+    Format.rgb32f: (f32, 3),
+    Format.rgba32u: (u32, 4),
+    Format.rgba32i: (i32, 4),
+    Format.rgba32f: (f32, 4),
+}
+
+# Reverse lookup by (channel_format, num_channels)
+TY_CH2FORMAT = {v: k for k, v in FORMAT2TY_CH.items()}
+
+
+class TextureType:
+    """Type annotation for Textures.
+
+    Args:
+        num_dimensions (int): Number of dimensions. For examples for a 2D texture this should be `2`.
+    """
+
+    def __init__(self, num_dimensions):
+        self.num_dimensions = num_dimensions
+
+
+class RWTextureType:
+    """Type annotation for RW Textures (image load store).
+
+    Args:
+        num_dimensions (int): Number of dimensions. For examples for a 2D texture this should be `2`.
+        lod (float): Specifies the explicit level-of-detail.
+        fmt (ti.Format): Color format of texture
+    """
+
+    def __init__(self, num_dimensions, lod=0, fmt=None):
+        self.num_dimensions = num_dimensions
+        if fmt is None:
+            raise GsTaichiCompilationError("fmt is required for rw_texture type")
+        else:
+            self.fmt = fmt
+        self.lod = lod
+
+
+texture = TextureType
+rw_texture = RWTextureType
+"""Alias for :class:`~gstaichi.types.ndarray_type.TextureType`.
+"""
+
+__all__ = ["texture", "rw_texture"]

gstaichi/types/utils.py

@@ -0,0 +1,11 @@
+from gstaichi._lib import core as ti_python_core
+
+is_signed = ti_python_core.is_signed
+
+is_integral = ti_python_core.is_integral
+
+is_real = ti_python_core.is_real
+
+is_tensor = ti_python_core.is_tensor
+
+__all__ = ["is_signed", "is_integral", "is_real", "is_tensor"]