gstaichi 2.1.1rc3__cp311-cp311-macosx_11_0_arm64.whl

gstaichi/lang/common_ops.py

@@ -0,0 +1,310 @@
+# type: ignore
+
+from typing import TYPE_CHECKING
+
+from gstaichi.lang import ops
+from gstaichi.lang.util import in_python_scope
+from gstaichi.types import primitive_types
+
+
+class GsTaichiOperations:
+    """The base class of gstaichi operations of expressions. Subclasses: :class:`~gstaichi.lang.expr.Expr`, :class:`~gstaichi.lang.matrix.Matrix`"""
+
+    if TYPE_CHECKING:
+        # Make pylint happy
+        def __getattr__(self, item):
+            pass
+
+    def __neg__(self):
+        return ops.neg(self)
+
+    def __abs__(self):
+        return ops.abs(self)
+
+    def __add__(self, other):
+        return ops.add(self, other)
+
+    def __radd__(self, other):
+        return ops.add(other, self)
+
+    def __sub__(self, other):
+        return ops.sub(self, other)
+
+    def __rsub__(self, other):
+        return ops.sub(other, self)
+
+    def __mul__(self, other):
+        return ops.mul(self, other)
+
+    def __rmul__(self, other):
+        return ops.mul(other, self)
+
+    def __truediv__(self, other):
+        return ops.truediv(self, other)
+
+    def __rtruediv__(self, other):
+        return ops.truediv(other, self)
+
+    def __floordiv__(self, other):
+        return ops.floordiv(self, other)
+
+    def __rfloordiv__(self, other):
+        return ops.floordiv(other, self)
+
+    def __mod__(self, other):
+        return ops.mod(self, other)
+
+    def __rmod__(self, other):
+        return ops.mod(other, self)
+
+    def __pow__(self, other, modulo=None):
+        return ops.pow(self, other)
+
+    def __rpow__(self, other, modulo=None):
+        return ops.pow(other, self)
+
+    def __le__(self, other):
+        return ops.cmp_le(self, other)
+
+    def __lt__(self, other):
+        return ops.cmp_lt(self, other)
+
+    def __ge__(self, other):
+        return ops.cmp_ge(self, other)
+
+    def __gt__(self, other):
+        return ops.cmp_gt(self, other)
+
+    def __eq__(self, other):
+        return ops.cmp_eq(self, other)
+
+    def __ne__(self, other):
+        return ops.cmp_ne(self, other)
+
+    def __and__(self, other):
+        return ops.bit_and(self, other)
+
+    def __rand__(self, other):
+        return ops.bit_and(other, self)
+
+    def __or__(self, other):
+        return ops.bit_or(self, other)
+
+    def __ror__(self, other):
+        return ops.bit_or(other, self)
+
+    def __xor__(self, other):
+        return ops.bit_xor(self, other)
+
+    def __rxor__(self, other):
+        return ops.bit_xor(other, self)
+
+    def __lshift__(self, other):
+        return ops.bit_shl(self, other)
+
+    def __rlshift__(self, other):
+        return ops.bit_shl(other, self)
+
+    def __rshift__(self, other):
+        return ops.bit_sar(self, other)
+
+    def __rrshift__(self, other):
+        return ops.bit_sar(other, self)
+
+    def __invert__(self):  # ~a => a.__invert__()
+        return ops.bit_not(self)
+
+    def __not__(self):  # not a => a.__not__()
+        return ops.logical_not(self)
+
+    def _atomic_add(self, other):
+        """Return the new expression of computing atomic add between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic add."""
+        return ops.atomic_add(self, other)
+
+    def _atomic_mul(self, other):
+        """Return the new expression of computing atomic mul between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic mul."""
+        return ops.atomic_mul(self, other)
+
+    def _atomic_sub(self, other):
+        """Return the new expression of computing atomic sub between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic sub."""
+        return ops.atomic_sub(self, other)
+
+    def _atomic_and(self, other):
+        """Return the new expression of computing atomic and between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic and."""
+        return ops.atomic_and(self, other)
+
+    def _atomic_xor(self, other):
+        """Return the new expression of computing atomic xor between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic xor."""
+        return ops.atomic_xor(self, other)
+
+    def _atomic_or(self, other):
+        """Return the new expression of computing atomic or between self and a given operand.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The computing expression of atomic or."""
+        return ops.atomic_or(self, other)
+
+    # In-place operators in python scope returns NotImplemented to fall back to normal operators
+    def __iadd__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_add(other)
+        return self
+
+    def __imul__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_mul(other)
+        return self
+
+    def __isub__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_sub(other)
+        return self
+
+    def __iand__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_and(other)
+        return self
+
+    def __ixor__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_xor(other)
+        return self
+
+    def __ior__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._atomic_or(other)
+        return self
+
+    # we don't support atomic_mul/truediv/floordiv/mod yet:
+    def __imul__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.mul(self, other))
+        return self
+
+    def __itruediv__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.truediv(self, other))
+        return self
+
+    def __ifloordiv__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.floordiv(self, other))
+        return self
+
+    def __imod__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.mod(self, other))
+        return self
+
+    def __ilshift__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.bit_shl(self, other))
+        return self
+
+    def __irshift__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.bit_sar(self, other))
+        return self
+
+    def __ipow__(self, other):
+        if in_python_scope():
+            return NotImplemented
+        self._assign(ops.pow(self, other))
+        return self
+
+    def _assign(self, other):
+        """Assign the expression of the given operand to self.
+
+        Args:
+            other (Any): Given operand.
+
+        Returns:
+            :class:`~gstaichi.lang.expr.Expr`: The expression after assigning."""
+        return ops.assign(self, other)
+
+    def _augassign(self, x, op):
+        """Generate the computing expression between self and the given operand of given operator and assigned to self.
+
+        Args:
+            x (Any): Given operand.
+            op (str): The name of operator."""
+        if op == "Add":
+            self += x
+        elif op == "Sub":
+            self -= x
+        elif op == "Mult":
+            self *= x
+        elif op == "Div":
+            self /= x
+        elif op == "FloorDiv":
+            self //= x
+        elif op == "Mod":
+            self %= x
+        elif op == "BitAnd":
+            self &= x
+        elif op == "BitOr":
+            self |= x
+        elif op == "BitXor":
+            self ^= x
+        elif op == "RShift":
+            self >>= x
+        elif op == "LShift":
+            self <<= x
+        elif op == "Pow":
+            self **= x
+        else:
+            assert False, op
+
+    def __ti_int__(self):
+        return ops.cast(self, int)
+
+    def __ti_bool__(self):
+        return ops.cast(self, primitive_types.u1)
+
+    def __ti_float__(self):
+        return ops.cast(self, float)

gstaichi/lang/exception.py

@@ -0,0 +1,80 @@
+# type: ignore
+
+from gstaichi._lib import core
+
+
+class GsTaichiCompilationError(Exception):
+    """Base class for all compilation exceptions."""
+
+    pass
+
+
+class GsTaichiSyntaxError(GsTaichiCompilationError, SyntaxError):
+    """Thrown when a syntax error is found during compilation."""
+
+    pass
+
+
+class GsTaichiNameError(GsTaichiCompilationError, NameError):
+    """Thrown when an undefine name is found during compilation."""
+
+    pass
+
+
+class GsTaichiIndexError(GsTaichiCompilationError, IndexError):
+    """Thrown when an index error is found during compilation."""
+
+    pass
+
+
+class GsTaichiTypeError(GsTaichiCompilationError, TypeError):
+    """Thrown when a type mismatch is found during compilation."""
+
+    pass
+
+
+class GsTaichiRuntimeError(RuntimeError):
+    """Thrown when the compiled program cannot be executed due to unspecified reasons."""
+
+    pass
+
+
+class GsTaichiAssertionError(GsTaichiRuntimeError, AssertionError):
+    """Thrown when assertion fails at runtime."""
+
+    pass
+
+
+class GsTaichiRuntimeTypeError(GsTaichiRuntimeError, TypeError):
+    @staticmethod
+    def get(pos, needed, provided):
+        return GsTaichiRuntimeTypeError(
+            f"Argument {pos} (type={provided}) cannot be converted into required type {needed}"
+        )
+
+    @staticmethod
+    def get_ret(needed, provided):
+        return GsTaichiRuntimeTypeError(f"Return (type={provided}) cannot be converted into required type {needed}")
+
+
+def handle_exception_from_cpp(exc):
+    if isinstance(exc, core.GsTaichiTypeError):
+        return GsTaichiTypeError(str(exc))
+    if isinstance(exc, core.GsTaichiSyntaxError):
+        return GsTaichiSyntaxError(str(exc))
+    if isinstance(exc, core.GsTaichiIndexError):
+        return GsTaichiIndexError(str(exc))
+    if isinstance(exc, core.GsTaichiAssertionError):
+        return GsTaichiAssertionError(str(exc))
+    return exc
+
+
+__all__ = [
+    "GsTaichiSyntaxError",
+    "GsTaichiTypeError",
+    "GsTaichiCompilationError",
+    "GsTaichiNameError",
+    "GsTaichiRuntimeError",
+    "GsTaichiRuntimeTypeError",
+    "GsTaichiAssertionError",
+]

gstaichi/lang/expr.py

@@ -0,0 +1,180 @@
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.common_ops import GsTaichiOperations
+from gstaichi.lang.exception import GsTaichiCompilationError, GsTaichiTypeError
+from gstaichi.lang.matrix import make_matrix
+from gstaichi.lang.util import is_gstaichi_class, is_matrix_class, to_numpy_type
+from gstaichi.types import primitive_types
+from gstaichi.types.primitive_types import integer_types, real_types
+
+if TYPE_CHECKING:
+    from gstaichi.lang.ast.ast_transformer_utils import ASTBuilder
+
+
+# Scalar, basic data type
+class Expr(GsTaichiOperations):
+    """A Python-side Expr wrapper, whose member variable `ptr` is an instance of C++ Expr class. A C++ Expr object contains member variable `expr` which holds an instance of C++ Expression class."""
+
+    def __init__(self, *args, dbg_info=None, dtype=None):
+        self.dbg_info = dbg_info
+        self.ptr_type_checked = False
+        self.declaration_tb: str = ""
+        if len(args) == 1:
+            if isinstance(args[0], _ti_core.ExprCxx):
+                self.ptr = args[0]
+            elif isinstance(args[0], Expr):
+                self.ptr = args[0].ptr
+                self.ptr_type_checked = args[0].ptr_type_checked
+                self.dbg_info = args[0].dbg_info
+            elif is_matrix_class(args[0]):
+                self.ptr = make_matrix(args[0].to_list()).ptr
+            elif isinstance(args[0], (list, tuple)):
+                self.ptr = make_matrix(args[0]).ptr
+            else:
+                # assume to be constant
+                arg = args[0]
+                if isinstance(arg, np.ndarray):
+                    if arg.shape:
+                        raise GsTaichiTypeError(
+                            "Only 0-dimensional numpy array can be used to initialize a scalar expression"
+                        )
+                    arg = arg.dtype.type(arg)
+                self.ptr = make_constant_expr(arg, dtype).ptr
+        else:
+            assert False
+        if self.dbg_info:
+            self.ptr.set_dbg_info(self.dbg_info)
+        if not self.ptr_type_checked:
+            self.ptr.type_check(impl.get_runtime().prog.config())
+            self.ptr_type_checked = True
+
+    def is_tensor(self):
+        return self.ptr.is_tensor()
+
+    def is_struct(self):
+        return self.ptr.is_struct()
+
+    def element_type(self):
+        return self.ptr.get_rvalue_type().element_type()
+
+    def get_shape(self):
+        if not self.is_tensor():
+            raise GsTaichiCompilationError(f"Getting shape of non-tensor type: {self.ptr.get_rvalue_type()}")
+        shape = self.ptr.get_shape()
+        assert shape is not None
+        return tuple(shape)
+
+    @property
+    def n(self):
+        shape = self.get_shape()
+        if len(shape) < 1:
+            raise GsTaichiCompilationError(f"Getting n of tensor type < 1D: {self.ptr.get_rvalue_type()}")
+        return shape[0]
+
+    @property
+    def m(self):
+        shape = self.get_shape()
+        if len(shape) < 2:
+            raise GsTaichiCompilationError(f"Getting m of tensor type < 2D: {self.ptr.get_rvalue_type()}")
+        return shape[1]
+
+    def __hash__(self):
+        return self.ptr.get_raw_address()
+
+    def __str__(self):
+        return "<ti.Expr>"
+
+    def __repr__(self):
+        return "<ti.Expr>"
+
+
+def _check_in_range(npty, val):
+    iif = np.iinfo(npty)
+    return iif.min <= val <= iif.max
+
+
+def _clamp_unsigned_to_range(npty, val: np.integer | int) -> np.integer | int:
+    # npty: np.int32 or np.int64
+    iif = np.iinfo(npty)
+    if iif.min <= val <= iif.max:
+        return val
+    cap = 1 << iif.bits
+    assert 0 <= val < cap
+    new_val = val - cap
+    return new_val
+
+
+def make_constant_expr(val, dtype):
+    if isinstance(val, (bool, np.bool_)):
+        constant_dtype = primitive_types.u1
+        return Expr(_ti_core.make_const_expr_bool(constant_dtype, val))
+
+    if isinstance(val, (float, np.floating)):
+        constant_dtype = impl.get_runtime().default_fp if dtype is None else dtype
+        if constant_dtype not in real_types:
+            raise GsTaichiTypeError(
+                "Floating-point literals must be annotated with a floating-point type. For type casting, use `ti.cast`."
+            )
+        return Expr(_ti_core.make_const_expr_fp(constant_dtype, val))
+
+    if isinstance(val, (int, np.integer)):
+        constant_dtype = impl.get_runtime().default_ip if dtype is None else dtype
+        if constant_dtype not in integer_types:
+            raise GsTaichiTypeError(
+                "Integer literals must be annotated with a integer type. For type casting, use `ti.cast`."
+            )
+        if _check_in_range(to_numpy_type(constant_dtype), val):
+            return Expr(_ti_core.make_const_expr_int(constant_dtype, _clamp_unsigned_to_range(np.int64, val)))
+        if dtype is None:
+            raise GsTaichiTypeError(
+                f"Integer literal {val} exceeded the range of default_ip: {impl.get_runtime().default_ip}, please specify the dtype via e.g. `ti.u64({val})` or set a different `default_ip` in `ti.init()`"
+            )
+        else:
+            raise GsTaichiTypeError(f"Integer literal {val} exceeded the range of specified dtype: {dtype}")
+
+    raise GsTaichiTypeError(f"Invalid constant scalar data type: {type(val)}")
+
+
+def make_var_list(size: int, ast_builder: "ASTBuilder | None" = None):
+    exprs = []
+    prog = impl.get_runtime().prog
+    for _ in range(size):
+        if ast_builder is None:
+            exprs.append(prog.make_id_expr(""))
+        else:
+            exprs.append(ast_builder.make_id_expr(""))
+    return exprs
+
+
+def make_expr_group(*exprs):
+    from gstaichi.lang.matrix import Matrix  # pylint: disable=C0415
+
+    if len(exprs) == 1:
+        if isinstance(exprs[0], (list, tuple)):
+            exprs = exprs[0]
+        elif isinstance(exprs[0], Matrix):
+            mat = exprs[0]
+            assert mat.m == 1
+            exprs = mat.entries
+    expr_group = _ti_core.ExprGroup()
+    for i in exprs:
+        flattened = _get_flattened_ptrs(i)
+        for item in flattened:
+            expr_group.push_back(item)
+    return expr_group
+
+
+def _get_flattened_ptrs(val):
+    if is_gstaichi_class(val):
+        ptrs = []
+        for item in val._members:
+            ptrs.extend(_get_flattened_ptrs(item))
+        return ptrs
+    return [Expr(val).ptr]
+
+
+__all__ = []