gstaichi 2.1.1rc3__cp310-cp310-macosx_11_0_arm64.whl

gstaichi/lang/util.py

@@ -0,0 +1,312 @@
+import functools
+import os
+import traceback
+import warnings
+from typing import Any
+
+import numpy as np
+from colorama import Fore, Style
+
+from gstaichi._lib import core as _ti_core
+from gstaichi._logging import is_logging_effective
+from gstaichi.lang import impl
+from gstaichi.types import Template
+from gstaichi.types.primitive_types import (
+    f16,
+    f32,
+    f64,
+    i8,
+    i16,
+    i32,
+    i64,
+    u1,
+    u8,
+    u16,
+    u32,
+    u64,
+)
+
+
+def has_pytorch():
+    """Whether has pytorch in the current Python environment.
+
+    Returns:
+        bool: True if has pytorch else False.
+
+    """
+    _has_pytorch = False
+    _env_torch = os.environ.get("TI_ENABLE_TORCH", "1")
+    if not _env_torch or int(_env_torch):
+        try:
+            import torch  # pylint: disable=C0415
+
+            _has_pytorch = True
+        except:
+            pass
+    return _has_pytorch
+
+
+def get_clangpp():
+    from distutils.spawn import find_executable  # pylint: disable=C0415
+
+    # GsTaichi itself uses llvm-10.0.0 to compile.
+    # There will be some issues compiling CUDA with other clang++ version.
+    _clangpp_candidates = ["clang++-10"]
+    for c in _clangpp_candidates:
+        if find_executable(c) is not None:
+            _clangpp_presence = find_executable(c)
+            return _clangpp_presence
+    return None
+
+
+def has_clangpp():
+    return get_clangpp() is not None
+
+
+def is_matrix_class(rhs):
+    matrix_class = False
+    try:
+        if rhs._is_matrix_class:
+            matrix_class = True
+    except:
+        pass
+    return matrix_class
+
+
+def is_gstaichi_class(rhs):
+    gstaichi_class = False
+    try:
+        if rhs._is_gstaichi_class:
+            gstaichi_class = True
+    except:
+        pass
+    return gstaichi_class
+
+
+def to_numpy_type(dt):
+    """Convert gstaichi data type to its counterpart in numpy.
+
+    Args:
+        dt (DataType): The desired data type to convert.
+
+    Returns:
+        DataType: The counterpart data type in numpy.
+
+    """
+    if dt == f32:
+        return np.float32
+    if dt == f64:
+        return np.float64
+    if dt == i32:
+        return np.int32
+    if dt == i64:
+        return np.int64
+    if dt == i8:
+        return np.int8
+    if dt == i16:
+        return np.int16
+    if dt == u1:
+        return np.bool_
+    if dt == u8:
+        return np.uint8
+    if dt == u16:
+        return np.uint16
+    if dt == u32:
+        return np.uint32
+    if dt == u64:
+        return np.uint64
+    if dt == f16:
+        return np.half
+    assert False
+
+
+def to_pytorch_type(dt):
+    """Convert gstaichi data type to its counterpart in torch.
+
+    Args:
+        dt (DataType): The desired data type to convert.
+
+    Returns:
+        DataType: The counterpart data type in torch.
+
+    """
+    import torch  # pylint: disable=C0415
+
+    # pylint: disable=E1101
+    if dt == f32:
+        return torch.float32
+    if dt == f64:
+        return torch.float64
+    if dt == i32:
+        return torch.int32
+    if dt == i64:
+        return torch.int64
+    if dt == i8:
+        return torch.int8
+    if dt == i16:
+        return torch.int16
+    if dt == u1:
+        return torch.bool
+    if dt == u8:
+        return torch.uint8
+    if dt == f16:
+        return torch.float16
+
+    if dt in (u16, u32, u64):
+        if hasattr(torch, "uint16"):
+            if dt == u16:
+                return torch.uint16
+            if dt == u32:
+                return torch.uint32
+            if dt == u64:
+                return torch.uint64
+        raise RuntimeError(f"PyTorch doesn't support {dt.to_string()} data type before version 2.3.0.")
+
+    raise RuntimeError(f"PyTorch doesn't support {dt.to_string()} data type.")
+
+
+def to_gstaichi_type(dt):
+    """Convert numpy or torch data type to its counterpart in gstaichi.
+
+    Args:
+        dt (DataType): The desired data type to convert.
+
+    Returns:
+        DataType: The counterpart data type in gstaichi.
+
+    """
+    if type(dt) == _ti_core.DataTypeCxx:
+        return dt
+
+    if dt == np.float32:
+        return f32
+    if dt == np.float64:
+        return f64
+    if dt == np.int32:
+        return i32
+    if dt == np.int64:
+        return i64
+    if dt == np.int8:
+        return i8
+    if dt == np.int16:
+        return i16
+    if dt == np.bool_:
+        return u1
+    if dt == np.uint8:
+        return u8
+    if dt == np.uint16:
+        return u16
+    if dt == np.uint32:
+        return u32
+    if dt == np.uint64:
+        return u64
+    if dt == np.half:
+        return f16
+
+    if has_pytorch():
+        import torch  # pylint: disable=C0415
+
+        # pylint: disable=E1101
+        if dt == torch.float32:
+            return f32
+        if dt == torch.float64:
+            return f64
+        if dt == torch.int32:
+            return i32
+        if dt == torch.int64:
+            return i64
+        if dt == torch.int8:
+            return i8
+        if dt == torch.int16:
+            return i16
+        if dt == torch.bool:
+            return u1
+        if dt == torch.uint8:
+            return u8
+        if dt == torch.float16:
+            return f16
+
+        if hasattr(torch, "uint16"):
+            if dt == torch.uint16:
+                return u16
+            if dt == torch.uint32:
+                return u32
+            if dt == torch.uint64:
+                return u64
+
+        raise RuntimeError(f"PyTorch doesn't support {dt.to_string()} data type before version 2.3.0.")
+
+    raise AssertionError(f"Unknown type {dt}")
+
+
+def cook_dtype(dtype):
+    if isinstance(dtype, _ti_core.DataTypeCxx):
+        return dtype
+    if isinstance(dtype, _ti_core.Type):
+        return _ti_core.DataTypeCxx(dtype)
+    if dtype is float:
+        return impl.get_runtime().default_fp
+    if dtype is int:
+        return impl.get_runtime().default_ip
+    if dtype is bool:
+        return u1
+    raise ValueError(f"Invalid data type {dtype}")
+
+
+def in_gstaichi_scope():
+    return impl.inside_kernel()
+
+
+def in_python_scope():
+    return not in_gstaichi_scope()
+
+
+def gstaichi_scope(func):
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        assert in_gstaichi_scope(), f"{func.__name__} cannot be called in Python-scope"
+        return func(*args, **kwargs)
+
+    return wrapped
+
+
+def python_scope(func):
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        assert in_python_scope(), f"{func.__name__} cannot be called in GsTaichi-scope"
+        return func(*args, **kwargs)
+
+    return wrapped
+
+
+def warning(msg, warning_type=UserWarning, stacklevel=1, print_stack=True):
+    """Print a warning message. Note that the builtin `warnings` module is
+    unreliable since it may be suppressed by other packages such as IPython.
+
+    Args:
+        msg (str): message to print.
+        warning_type (Type[Warning]): type of warning.
+        stacklevel (int): warning stack level from the caller.
+        print_stack (bool): whether to print the stack
+    """
+    if not is_logging_effective("warn"):
+        return
+    if print_stack:
+        msg += f"\n{get_traceback(stacklevel)}"
+    warnings.warn(Fore.YELLOW + Style.BRIGHT + msg + Style.RESET_ALL, warning_type)
+
+
+def get_traceback(stacklevel=1):
+    s = traceback.extract_stack()[: -1 - stacklevel]
+    return "".join(traceback.format_list(s))
+
+
+def is_data_oriented(obj: Any) -> bool:
+    return getattr(obj, "_data_oriented", False)
+
+
+def is_ti_template(annotation: Any) -> bool:
+    return annotation == Template or isinstance(annotation, Template)
+
+
+__all__ = []

gstaichi/linalg/__init__.py

@@ -0,0 +1,8 @@
+# type: ignore
+
+"""GsTaichi support module for sparse matrix operations."""
+
+from gstaichi.linalg.matrixfree_cg import *
+from gstaichi.linalg.sparse_cg import SparseCG
+from gstaichi.linalg.sparse_matrix import *
+from gstaichi.linalg.sparse_solver import SparseSolver

gstaichi/linalg/matrixfree_cg.py

@@ -0,0 +1,310 @@
+# type: ignore
+
+from math import sqrt
+
+from gstaichi.lang import misc
+from gstaichi.lang.exception import GsTaichiRuntimeError, GsTaichiTypeError
+from gstaichi.lang.impl import FieldsBuilder, field, grouped
+from gstaichi.lang.kernel_impl import data_oriented, kernel
+from gstaichi.types import primitive_types, template
+
+
+@data_oriented
+class LinearOperator:
+    def __init__(self, matvec_kernel):
+        self._matvec = matvec_kernel
+
+    def matvec(self, x, Ax):
+        if x.shape != Ax.shape:
+            raise GsTaichiRuntimeError(f"Dimension mismatch x.shape{x.shape} != Ax.shape{Ax.shape}.")
+        self._matvec(x, Ax)
+
+
+def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True):
+    """Matrix-free conjugate-gradient solver.
+
+    Use conjugate-gradient method to solve the linear system Ax = b, where A is implicitly
+    represented as a LinearOperator.
+
+    Args:
+        A (LinearOperator): The coefficient matrix A of the linear system.
+        b (Field): The right-hand side of the linear system.
+        x (Field): The initial guess for the solution.
+        maxiter (int): Maximum number of iterations.
+        atol: Tolerance(absolute) for convergence.
+        quiet (bool): Switch to turn on/off iteration log.
+    """
+
+    if b.dtype != x.dtype:
+        raise GsTaichiTypeError(f"Dtype mismatch b.dtype({b.dtype}) != x.dtype({x.dtype}).")
+    if str(b.dtype) == "f32":
+        solver_dtype = primitive_types.f32
+    elif str(b.dtype) == "f64":
+        solver_dtype = primitive_types.f64
+    else:
+        raise GsTaichiTypeError(f"Not supported dtype: {b.dtype}")
+    if b.shape != x.shape:
+        raise GsTaichiRuntimeError(f"Dimension mismatch b.shape{b.shape} != x.shape{x.shape}.")
+
+    size = b.shape
+    vector_fields_builder = FieldsBuilder()
+    p = field(dtype=solver_dtype)
+    r = field(dtype=solver_dtype)
+    Ap = field(dtype=solver_dtype)
+    Ax = field(dtype=solver_dtype)
+    if len(size) == 1:
+        axes = misc.i
+    elif len(size) == 2:
+        axes = misc.ij
+    elif len(size) == 3:
+        axes = misc.ijk
+    else:
+        raise GsTaichiRuntimeError(f"MatrixFreeCG only support 1D, 2D, 3D inputs; your inputs is {len(size)}-D.")
+    vector_fields_builder.dense(axes, size).place(p, r, Ap, Ax)
+    vector_fields_snode_tree = vector_fields_builder.finalize()
+
+    scalar_builder = FieldsBuilder()
+    alpha = field(dtype=solver_dtype)
+    beta = field(dtype=solver_dtype)
+    scalar_builder.place(alpha, beta)
+    scalar_snode_tree = scalar_builder.finalize()
+
+    @kernel
+    def init():
+        for I in grouped(x):
+            r[I] = b[I] - Ax[I]
+            p[I] = 0.0
+            Ap[I] = 0.0
+
+    @kernel
+    def reduce(p: template(), q: template()) -> solver_dtype:
+        result = solver_dtype(0.0)
+        for I in grouped(p):
+            result += p[I] * q[I]
+        return result
+
+    @kernel
+    def update_x():
+        for I in grouped(x):
+            x[I] += alpha[None] * p[I]
+
+    @kernel
+    def update_r():
+        for I in grouped(r):
+            r[I] -= alpha[None] * Ap[I]
+
+    @kernel
+    def update_p():
+        for I in grouped(p):
+            p[I] = r[I] + beta[None] * p[I]
+
+    def solve():
+        succeeded = True
+        A._matvec(x, Ax)
+        init()
+        initial_rTr = reduce(r, r)
+        if not quiet:
+            print(f">>> Initial residual = {initial_rTr:e}")
+        old_rTr = initial_rTr
+        new_rTr = initial_rTr
+        update_p()
+        if sqrt(initial_rTr) >= tol:  # Do nothing if the initial residual is small enough
+            # -- Main loop --
+            for i in range(maxiter):
+                A._matvec(p, Ap)  # compute Ap = A x p
+                pAp = reduce(p, Ap)
+                alpha[None] = old_rTr / pAp
+                update_x()
+                update_r()
+                new_rTr = reduce(r, r)
+                if sqrt(new_rTr) < tol:
+                    if not quiet:
+                        print(">>> Conjugate Gradient method converged.")
+                        print(f">>> #iterations {i}")
+                    break
+                beta[None] = new_rTr / old_rTr
+                update_p()
+                old_rTr = new_rTr
+                if not quiet:
+                    print(f">>> Iter = {i+1:4}, Residual = {sqrt(new_rTr):e}")
+        if new_rTr >= tol:
+            if not quiet:
+                print(
+                    f">>> Conjugate Gradient method failed to converge in {maxiter} iterations: Residual = {sqrt(new_rTr):e}"
+                )
+            succeeded = False
+        return succeeded
+
+    succeeded = solve()
+    vector_fields_snode_tree.destroy()
+    scalar_snode_tree.destroy()
+    return succeeded
+
+
+def MatrixFreeBICGSTAB(A, b, x, tol=1e-6, maxiter=5000, quiet=True):
+    """Matrix-free biconjugate-gradient stabilized solver (BiCGSTAB).
+
+    Use BiCGSTAB method to solve the linear system Ax = b, where A is implicitly
+    represented as a LinearOperator.
+
+    Args:
+        A (LinearOperator): The coefficient matrix A of the linear system.
+        b (Field): The right-hand side of the linear system.
+        x (Field): The initial guess for the solution.
+        maxiter (int): Maximum number of iterations.
+        atol: Tolerance(absolute) for convergence.
+        quiet (bool): Switch to turn on/off iteration log.
+    """
+
+    if b.dtype != x.dtype:
+        raise GsTaichiTypeError(f"Dtype mismatch b.dtype({b.dtype}) != x.dtype({x.dtype}).")
+    if str(b.dtype) == "f32":
+        solver_dtype = primitive_types.f32
+    elif str(b.dtype) == "f64":
+        solver_dtype = primitive_types.f64
+    else:
+        raise GsTaichiTypeError(f"Not supported dtype: {b.dtype}")
+    if b.shape != x.shape:
+        raise GsTaichiRuntimeError(f"Dimension mismatch b.shape{b.shape} != x.shape{x.shape}.")
+
+    size = b.shape
+    vector_fields_builder = FieldsBuilder()
+    p = field(dtype=solver_dtype)
+    p_hat = field(dtype=solver_dtype)
+    r = field(dtype=solver_dtype)
+    r_tld = field(dtype=solver_dtype)
+    s = field(dtype=solver_dtype)
+    s_hat = field(dtype=solver_dtype)
+    t = field(dtype=solver_dtype)
+    Ap = field(dtype=solver_dtype)
+    Ax = field(dtype=solver_dtype)
+    Ashat = field(dtype=solver_dtype)
+    if len(size) == 1:
+        axes = misc.i
+    elif len(size) == 2:
+        axes = misc.ij
+    elif len(size) == 3:
+        axes = misc.ijk
+    else:
+        raise GsTaichiRuntimeError(f"MatrixFreeBICGSTAB only support 1D, 2D, 3D inputs; your inputs is {len(size)}-D.")
+    vector_fields_builder.dense(axes, size).place(p, p_hat, r, r_tld, s, s_hat, t, Ap, Ax, Ashat)
+    vector_fields_snode_tree = vector_fields_builder.finalize()
+
+    scalar_builder = FieldsBuilder()
+    alpha = field(dtype=solver_dtype)
+    beta = field(dtype=solver_dtype)
+    omega = field(dtype=solver_dtype)
+    rho = field(dtype=solver_dtype)
+    rho_1 = field(dtype=solver_dtype)
+    scalar_builder.place(alpha, beta, omega, rho, rho_1)
+    scalar_snode_tree = scalar_builder.finalize()
+    succeeded = True
+
+    @kernel
+    def init():
+        for I in grouped(x):
+            r[I] = b[I] - Ax[I]
+            r_tld[I] = b[I]
+            p[I] = 0.0
+            Ap[I] = 0.0
+            Ashat[I] = 0.0
+        rho[None] = 0.0
+        rho_1[None] = 1.0
+        alpha[None] = 1.0
+        beta[None] = 1.0
+        omega[None] = 1.0
+
+    @kernel
+    def reduce(p: template(), q: template()) -> solver_dtype:
+        result = solver_dtype(0.0)
+        for I in grouped(p):
+            result += p[I] * q[I]
+        return result
+
+    @kernel
+    def copy(orig: template(), dest: template()):
+        for I in grouped(orig):
+            dest[I] = orig[I]
+
+    @kernel
+    def update_p():
+        for I in grouped(p):
+            p[I] = r[I] + beta[None] * (p[I] - omega[None] * Ap[I])
+
+    @kernel
+    def update_phat():
+        for I in grouped(p_hat):
+            p_hat[I] = p[I]
+
+    @kernel
+    def update_s():
+        for I in grouped(s):
+            s[I] = r[I] - alpha[None] * Ap[I]
+
+    @kernel
+    def update_shat():
+        for I in grouped(s_hat):
+            s_hat[I] = s[I]
+
+    @kernel
+    def update_x():
+        for I in grouped(x):
+            x[I] += alpha[None] * p_hat[I] + omega[None] * s_hat[I]
+
+    @kernel
+    def update_r():
+        for I in grouped(r):
+            r[I] = s[I] - omega[None] * t[I]
+
+    def solve():
+        succeeded = True
+        A._matvec(x, Ax)
+        init()
+        initial_rTr = reduce(r, r)
+        rTr = initial_rTr
+        if not quiet:
+            print(f">>> Initial residual = {initial_rTr:e}")
+        if sqrt(initial_rTr) >= tol:  # Do nothing if the initial residual is small enough
+            for i in range(maxiter):
+                rho[None] = reduce(r, r_tld)
+                if rho[None] == 0.0:
+                    if not quiet:
+                        print(">>> BICGSTAB failed because r@r_tld = 0.")
+                    succeeded = False
+                    break
+                if i == 0:
+                    copy(orig=r, dest=p)
+                else:
+                    beta[None] = (rho[None] / rho_1[None]) * (alpha[None] / omega[None])
+                    update_p()
+                update_phat()
+                A._matvec(p, Ap)
+                alpha_lower = reduce(r_tld, Ap)
+                alpha[None] = rho[None] / alpha_lower
+                update_s()
+                update_shat()
+                A._matvec(s_hat, Ashat)
+                copy(orig=Ashat, dest=t)
+                omega_upper = reduce(t, s)
+                omega_lower = reduce(t, t)
+                omega[None] = omega_upper / (omega_lower + 1e-16) if omega_lower == 0.0 else omega_upper / omega_lower
+                update_x()
+                update_r()
+                rTr = reduce(r, r)
+                if not quiet:
+                    print(f">>> Iter = {i+1:4}, Residual = {sqrt(rTr):e}")
+                if sqrt(rTr) < tol:
+                    if not quiet:
+                        print(f">>> BICGSTAB method converged at #iterations {i}")
+                    break
+                rho_1[None] = rho[None]
+        if rTr >= tol:
+            if not quiet:
+                print(f">>> BICGSTAB failed to converge in {maxiter} iterations: Residual = {sqrt(rTr):e}")
+            succeeded = False
+        return succeeded
+
+    succeeded = solve()
+    vector_fields_snode_tree.destroy()
+    scalar_snode_tree.destroy()
+    return succeeded

gstaichi/linalg/sparse_cg.py

@@ -0,0 +1,59 @@
+# type: ignore
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang._ndarray import Ndarray, ScalarNdarray
+from gstaichi.lang.exception import GsTaichiRuntimeError
+from gstaichi.lang.impl import get_runtime
+from gstaichi.types import f32, f64
+
+
+class SparseCG:
+    """Conjugate-gradient solver built for SparseMatrix.
+
+    Use conjugate-gradient method to solve the linear system Ax = b, where A is SparseMatrix.
+
+    Args:
+        A (SparseMatrix): The coefficient matrix A of the linear system.
+        b (numpy ndarray, gstaichi Ndarray): The right-hand side of the linear system.
+        x0 (numpy ndarray, gstaichi Ndarray): The initial guess for the solution.
+        max_iter (int): Maximum number of iterations.
+        atol: Tolerance(absolute) for convergence.
+    """
+
+    def __init__(self, A, b, x0=None, max_iter=50, atol=1e-6):
+        self.dtype = A.dtype
+        self.ti_arch = get_runtime().prog.config().arch
+        self.matrix = A
+        self.b = b
+        if self.ti_arch == _ti_core.Arch.cuda:
+            self.cg_solver = _ti_core.make_cucg_solver(A.matrix, max_iter, atol, True)
+        elif self.ti_arch == _ti_core.Arch.x64 or self.ti_arch == _ti_core.Arch.arm64:
+            if self.dtype == f32:
+                self.cg_solver = _ti_core.make_float_cg_solver(A.matrix, max_iter, atol, True)
+            elif self.dtype == f64:
+                self.cg_solver = _ti_core.make_double_cg_solver(A.matrix, max_iter, atol, True)
+            else:
+                raise GsTaichiRuntimeError(f"Unsupported CG dtype: {self.dtype}")
+            if isinstance(b, Ndarray):
+                self.cg_solver.set_b_ndarray(get_runtime().prog, b.arr)
+            elif isinstance(b, np.ndarray):
+                self.cg_solver.set_b(b)
+            if isinstance(x0, Ndarray):
+                self.cg_solver.set_x_ndarray(get_runtime().prog, x0.arr)
+            elif isinstance(x0, np.ndarray):
+                self.cg_solver.set_x(x0)
+        else:
+            raise GsTaichiRuntimeError(f"Unsupported CG arch: {self.ti_arch}")
+
+    def solve(self):
+        if self.ti_arch == _ti_core.Arch.cuda:
+            if isinstance(self.b, Ndarray):
+                x = ScalarNdarray(self.b.dtype, [self.matrix.m])
+                self.cg_solver.solve(get_runtime().prog, x.arr, self.b.arr)
+                return x, True
+            raise GsTaichiRuntimeError(f"Unsupported CG RHS type: {type(self.b)}")
+        else:
+            self.cg_solver.solve()
+            return self.cg_solver.get_x(), self.cg_solver.is_success()