gstaichi 2.1.1__cp311-cp311-macosx_11_0_arm64.whl

gstaichi/_lib/utils.py

@@ -0,0 +1,243 @@
+# type: ignore
+
+import os
+import platform
+import re
+import sys
+import warnings
+
+from colorama import Fore, Style
+
+if sys.version_info[0] < 3 or sys.version_info[1] <= 5:
+    raise RuntimeError(
+        "\nPlease restart with Python 3.6+\n" + "Current Python version:",
+        sys.version_info,
+    )
+
+
+def in_docker():
+    if os.environ.get("TI_IN_DOCKER", "") == "":
+        return False
+    return True
+
+
+def get_os_name():
+    name = platform.platform()
+    # in python 3.8, platform.platform() uses mac_ver() on macOS
+    # it will return 'macOS-XXXX' instead of 'Darwin-XXXX'
+    if name.lower().startswith("darwin") or name.lower().startswith("macos"):
+        return "osx"
+    if name.lower().startswith("windows"):
+        return "win"
+    if name.lower().startswith("linux"):
+        return "linux"
+    if "bsd" in name.lower():
+        return "unix"
+    assert False, f"Unknown platform name {name}"
+
+
+def import_ti_python_core():
+    if get_os_name() != "win":
+        # pylint: disable=E1101
+        old_flags = sys.getdlopenflags()
+        sys.setdlopenflags(2 | 8)  # RTLD_NOW | RTLD_DEEPBIND
+    else:
+        pyddir = os.path.dirname(os.path.realpath(__file__))
+        os.environ["PATH"] += os.pathsep + pyddir
+    try:
+        from gstaichi._lib.core import gstaichi_python as core  # pylint: disable=C0415
+    except Exception as e:
+        if isinstance(e, ImportError):
+            print(
+                Fore.YELLOW + "Share object gstaichi_python import failed, "
+                "check this page for possible solutions:\n"
+                "https://docs.taichi-lang.org/docs/install" + Fore.RESET
+            )
+            if get_os_name() == "win":
+                # pylint: disable=E1101
+                e.msg += "\nConsider installing Microsoft Visual C++ Redistributable: https://aka.ms/vs/16/release/vc_redist.x64.exe"
+        raise e from None
+
+    if get_os_name() != "win":
+        sys.setdlopenflags(old_flags)  # pylint: disable=E1101
+    lib_dir = os.path.join(package_root, "_lib", "runtime")
+    core.set_lib_dir(locale_encode(lib_dir))
+    return core
+
+
+def locale_encode(path):
+    try:
+        import locale  # pylint: disable=C0415
+
+        return path.encode(locale.getdefaultlocale()[1])
+    except (UnicodeEncodeError, TypeError):
+        try:
+            return path.encode(sys.getfilesystemencoding())
+        except UnicodeEncodeError:
+            try:
+                return path.encode()
+            except UnicodeEncodeError:
+                return path
+
+
+def is_ci():
+    return os.environ.get("TI_CI", "") == "1"
+
+
+package_root = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+
+
+def get_core_shared_object():
+    directory = os.path.join(package_root, "_lib")
+    return os.path.join(directory, "libgstaichi_python.so")
+
+
+def print_red_bold(*args, **kwargs):
+    print(Fore.RED + Style.BRIGHT, end="")
+    print(*args, **kwargs)
+    print(Style.RESET_ALL, end="")
+
+
+def print_yellow_bold(*args, **kwargs):
+    print(Fore.YELLOW + Style.BRIGHT, end="")
+    print(*args, **kwargs)
+    print(Style.RESET_ALL, end="")
+
+
+def check_exists(src):
+    if not os.path.exists(src):
+        raise FileNotFoundError(f'File "{src}" not exist. Installation corrupted or build incomplete?')
+
+
+ti_python_core = import_ti_python_core()
+
+ti_python_core.set_python_package_dir(package_root)
+
+log_level = os.environ.get("TI_LOG_LEVEL", "")
+if log_level:
+    ti_python_core.set_logging_level(log_level)
+
+
+def get_dll_name(name):
+    if get_os_name() == "linux":
+        return f"libgstaichi_{name}.so"
+    if get_os_name() == "osx":
+        return f"libgstaichi_{name}.dylib"
+    if get_os_name() == "win":
+        return f"gstaichi_{name}.dll"
+    raise Exception(f"Unknown OS: {get_os_name()}")
+
+
+def at_startup():
+    ti_python_core.set_core_state_python_imported(True)
+
+
+at_startup()
+
+
+def compare_version(latest, current):
+    latest_num = map(int, latest.split("."))
+    current_num = map(int, current.split("."))
+    return tuple(latest_num) > tuple(current_num)
+
+
+def _print_gstaichi_header():
+    header = "[GsTaichi] "
+    header += f"version {ti_python_core.get_version_string()}, "
+
+    try:
+        timestamp_path = os.path.join(ti_python_core.get_repo_dir(), "timestamp")
+        if os.path.exists(timestamp_path):
+            latest_version = ""
+            with open(timestamp_path, "r") as f:
+                latest_version = f.readlines()[1].rstrip()
+            if compare_version(latest_version, ti_python_core.get_version_string()):
+                header += f"latest version {latest_version}, "
+    except:
+        pass
+
+    llvm_target_support = ti_python_core.get_llvm_target_support()
+    header += f"llvm {llvm_target_support}, "
+
+    commit_hash = ti_python_core.get_commit_hash()
+    commit_hash = commit_hash[:8]
+    header += f"commit {commit_hash}, "
+
+    header += f"{get_os_name()}, "
+
+    py_ver = ".".join(str(x) for x in sys.version_info[:3])
+    header += f"python {py_ver}"
+
+    print(header)
+
+
+if os.getenv("ENABLE_GSTAICHI_HEADER_PRINT", "True").lower() not in ("false", "0", "f"):
+    _print_gstaichi_header()
+
+
+def try_get_wheel_tag(module):
+    try:
+        from email.parser import Parser  # pylint: disable=import-outside-toplevel
+
+        wheel_path = f'{module.__path__[0]}-{".".join(map(str, module.__version__))}.dist-info/WHEEL'
+        with open(wheel_path, "r") as f:
+            meta = Parser().parse(f)
+        return meta.get("Tag")
+    except Exception:
+        return None
+
+
+def try_get_loaded_libc_version():
+    assert platform.system() == "Linux"
+    with open("/proc/self/maps") as f:
+        content = f.read()
+
+    try:
+        libc_path = next(v for v in content.split() if "libc-" in v)
+        ver = re.findall(r"\d+\.\d+", libc_path)
+        if not ver:
+            return None
+        return tuple([int(v) for v in ver[0].split(".")])
+    except StopIteration:
+        return None
+
+
+def try_get_pip_version():
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            import pip  # pylint: disable=import-outside-toplevel
+        return tuple([int(v) for v in pip.__version__.split(".")])
+    except ImportError:
+        return None
+
+
+def warn_restricted_version():
+    if os.environ.get("TI_MANYLINUX2014_OK", ""):
+        return
+
+    if get_os_name() == "linux":
+        try:
+            import gstaichi as ti  # pylint: disable=import-outside-toplevel
+
+            wheel_tag = try_get_wheel_tag(ti)
+            if wheel_tag and "manylinux" in wheel_tag:
+                libc_ver = try_get_loaded_libc_version()
+                if libc_ver and libc_ver < (2, 27):
+                    print_yellow_bold(
+                        "!! GsTaichi requires glibc >= 2.27 to run, please try upgrading your OS to a recent one (e.g. Ubuntu 18.04 or later) if possible."
+                    )
+
+                pip_ver = try_get_pip_version()
+                if pip_ver and pip_ver < (20, 3, 0):
+                    print_yellow_bold(
+                        f"!! Your pip (version {'.'.join(map(str, pip_ver))}) is outdated (20.3.0 or later required), "
+                        "try upgrading pip and install gstaichi again."
+                    )
+                    print()
+                    print_yellow_bold("    $ python3 -m pip install --upgrade pip")
+                    print_yellow_bold("    $ python3 -m pip install --force-reinstall gstaichi")
+                    print()
+
+        except Exception:
+            pass

gstaichi/_logging.py

@@ -0,0 +1,131 @@
+# type: ignore
+
+import inspect
+import os
+
+from gstaichi._lib import core as ti_python_core
+
+
+def _get_logging(name):
+    """Generates a decorator to decorate a specific logger function.
+
+    Args:
+        name (str): The string represents logging level.
+            Effective levels include: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+    Returns:
+        Callabe: The decorated function.
+    """
+
+    def logger(msg, *args, **kwargs):
+        # Python inspection takes time (~0.1ms) so avoid it as much as possible
+        if ti_python_core.logging_effective(name):
+            msg_formatted = msg.format(*args, **kwargs)
+            func = getattr(ti_python_core, name)
+            frame = inspect.currentframe().f_back
+            file_name, lineno, func_name, _, _ = inspect.getframeinfo(frame)
+            file_name = os.path.basename(file_name)
+            msg = f"[{file_name}:{func_name}@{lineno}] {msg_formatted}"
+            func(msg)
+
+    return logger
+
+
+def set_logging_level(level):
+    """Setting the logging level to a specified value.
+    Available levels are: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+    Note that after calling this function, logging levels below the specified one will
+    also be effective. For example if `level` is set to 'warn', then the levels below
+    it, which are 'error' and 'critical' in this case, will also be effective.
+
+    See also https://docs.taichi-lang.org/docs/developer_utilities#logging.
+
+    Args:
+        level (str): Logging level.
+
+    Example::
+
+        >>> set_logging_level('debug')
+    """
+    ti_python_core.set_logging_level(level)
+
+
+def is_logging_effective(level):
+    """Check if the specified logging level is effective.
+    All levels below current level will be effective.
+    The default level is 'info'.
+
+    See also https://docs.taichi-lang.org/docs/developer_utilities#logging.
+
+    Args:
+        level (str): The string represents logging level. \
+            Effective levels include: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+    Returns:
+        Bool: Indicate whether the logging level is effective.
+
+    Example::
+
+        >>> # assume current level is 'info'
+        >>> print(ti.is_logging_effective("trace"))     # False
+        >>> print(ti.is_logging_effective("debug"))     # False
+        >>> print(ti.is_logging_effective("info"))      # True
+        >>> print(ti.is_logging_effective("warn"))      # True
+        >>> print(ti.is_logging_effective("error"))     # True
+        >>> print(ti.is_logging_effective("critical"))  # True
+    """
+    return ti_python_core.logging_effective(level)
+
+
+# ------------------------
+
+DEBUG = "debug"
+"""The `str` 'debug', used for the `debug` logging level.
+"""
+# ------------------------
+
+TRACE = "trace"
+"""The `str` 'trace', used for the `debug` logging level.
+"""
+# ------------------------
+
+INFO = "info"
+"""The `str` 'info', used for the `info` logging level.
+"""
+# ------------------------
+
+WARN = "warn"
+"""The `str` 'warn', used for the `warn` logging level.
+"""
+# ------------------------
+
+ERROR = "error"
+"""The `str` 'error', used for the `error` logging level.
+"""
+# ------------------------
+
+CRITICAL = "critical"
+"""The `str` 'critical', used for the `critical` logging level.
+"""
+# ------------------------
+
+supported_log_levels = [DEBUG, TRACE, INFO, WARN, ERROR, CRITICAL]
+
+debug = _get_logging(DEBUG)
+trace = _get_logging(TRACE)
+info = _get_logging(INFO)
+warn = _get_logging(WARN)
+error = _get_logging(ERROR)
+critical = _get_logging(CRITICAL)
+
+__all__ = [
+    "DEBUG",
+    "TRACE",
+    "INFO",
+    "WARN",
+    "ERROR",
+    "CRITICAL",
+    "set_logging_level",
+    "is_logging_effective",
+]

gstaichi/_snode/__init__.py

@@ -0,0 +1,5 @@
+# type: ignore
+
+from gstaichi._snode.fields_builder import FieldsBuilder
+
+__all__ = ["FieldsBuilder"]

gstaichi/_snode/fields_builder.py

@@ -0,0 +1,187 @@
+# type: ignore
+
+from typing import Any, Optional, Sequence, Union
+
+from gstaichi._lib import core as _ti_core
+from gstaichi._lib.core.gstaichi_python import SNodeCxx
+from gstaichi._snode.snode_tree import SNodeTree
+from gstaichi.lang import impl, snode
+from gstaichi.lang.exception import GsTaichiRuntimeError
+from gstaichi.lang.util import warning
+
+_snode_registry = _ti_core.SNodeRegistry()
+
+_Axis = _ti_core.Axis
+
+
+class FieldsBuilder:
+    """A builder that constructs a SNodeTree instance.
+
+    Example::
+
+        x = ti.field(ti.i32)
+        y = ti.field(ti.f32)
+        fb = ti.FieldsBuilder()
+        fb.dense(ti.ij, 8).place(x)
+        fb.pointer(ti.ij, 8).dense(ti.ij, 4).place(y)
+
+        # After this line, `x` and `y` are placed. No more fields can be placed
+        # into `fb`.
+        #
+        # The tree looks like the following:
+        # (implicit root)
+        #  |
+        #  +-- dense +-- place(x)
+        #  |
+        #  +-- pointer +-- dense +-- place(y)
+        fb.finalize()
+    """
+
+    def __init__(self):
+        self.ptr: SNodeCxx = _snode_registry.create_root(impl.get_runtime().prog)
+        self.root = snode.SNode(self.ptr)
+        self.finalized = False
+        self.empty = True
+        impl.get_runtime().initialize_fields_builder(self)
+
+    # TODO: move this into SNodeTree
+    @classmethod
+    def _finalized_roots(cls):
+        """Gets all the roots of the finalized SNodeTree.
+
+        Returns:
+            A list of the roots of the finalized SNodeTree.
+        """
+        roots_ptr = []
+        size = impl.get_runtime().prog.get_snode_tree_size()
+        for i in range(size):
+            res = impl.get_runtime().prog.get_snode_root(i)
+            roots_ptr.append(snode.SNode(res))
+        return roots_ptr
+
+    # TODO: move this to SNodeTree class.
+    def deactivate_all(self):
+        """Same as :func:`gstaichi.lang.snode.SNode.deactivate_all`"""
+        if self.finalized:
+            self.root.deactivate_all()
+        else:
+            warning("""'deactivate_all()' would do nothing if FieldsBuilder is not finalized""")
+
+    def dense(
+        self,
+        indices: Union[Sequence[_Axis], _Axis],
+        dimensions: Union[Sequence[int], int],
+    ):
+        """Same as :func:`gstaichi.lang.snode.SNode.dense`"""
+        self._check_not_finalized()
+        self.empty = False
+        return self.root.dense(indices, dimensions)
+
+    def pointer(
+        self,
+        indices: Union[Sequence[_Axis], _Axis],
+        dimensions: Union[Sequence[int], int],
+    ):
+        """Same as :func:`gstaichi.lang.snode.SNode.pointer`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise GsTaichiRuntimeError("Pointer SNode is not supported on this backend.")
+        self._check_not_finalized()
+        self.empty = False
+        return self.root.pointer(indices, dimensions)
+
+    def _hash(self, indices, dimensions):
+        """Same as :func:`gstaichi.lang.snode.SNode.hash`"""
+        raise NotImplementedError()
+
+    def dynamic(
+        self,
+        index: Union[Sequence[_Axis], _Axis],
+        dimension: Union[Sequence[int], int],
+        chunk_size: Optional[int] = None,
+    ):
+        """Same as :func:`gstaichi.lang.snode.SNode.dynamic`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise GsTaichiRuntimeError("Dynamic SNode is not supported on this backend.")
+
+        if dimension >= 2**31:
+            raise GsTaichiRuntimeError(
+                f"The maximum dimension of a dynamic SNode cannot exceed the maximum value of a 32-bit signed integer: Got {dimension} > 2**31-1"
+            )
+        if chunk_size is not None and chunk_size >= 2**31:
+            raise GsTaichiRuntimeError(
+                f"Chunk size cannot exceed the maximum value of a 32-bit signed integer: Got {chunk_size} > 2**31-1"
+            )
+
+        self._check_not_finalized()
+        self.empty = False
+        return self.root.dynamic(index, dimension, chunk_size)
+
+    def bitmasked(
+        self,
+        indices: Union[Sequence[_Axis], _Axis],
+        dimensions: Union[Sequence[int], int],
+    ):
+        """Same as :func:`gstaichi.lang.snode.SNode.bitmasked`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise GsTaichiRuntimeError("Bitmasked SNode is not supported on this backend.")
+        self._check_not_finalized()
+        self.empty = False
+        return self.root.bitmasked(indices, dimensions)
+
+    def quant_array(
+        self,
+        indices: Union[Sequence[_Axis], _Axis],
+        dimensions: Union[Sequence[int], int],
+        max_num_bits: int,
+    ):
+        """Same as :func:`gstaichi.lang.snode.SNode.quant_array`"""
+        self._check_not_finalized()
+        self.empty = False
+        return self.root.quant_array(indices, dimensions, max_num_bits)
+
+    def place(self, *args: Any, offset: Optional[Union[Sequence[int], int]] = None):
+        """Same as :func:`gstaichi.lang.snode.SNode.place`"""
+        self._check_not_finalized()
+        self.empty = False
+        self.root.place(*args, offset=offset)
+
+    def lazy_grad(self):
+        """Same as :func:`gstaichi.lang.snode.SNode.lazy_grad`"""
+        # TODO: This complicates the implementation. Figure out why we need this
+        self._check_not_finalized()
+        self.empty = False
+        self.root.lazy_grad()
+
+    def _allocate_adjoint_checkbit(self):
+        """Same as :func:`gstaichi.lang.snode.SNode._allocate_adjoint_checkbit`"""
+        self._check_not_finalized()
+        self.empty = False
+        self.root._allocate_adjoint_checkbit()
+
+    def lazy_dual(self):
+        """Same as :func:`gstaichi.lang.snode.SNode.lazy_dual`"""
+        # TODO: This complicates the implementation. Figure out why we need this
+        self._check_not_finalized()
+        self.empty = False
+        self.root.lazy_dual()
+
+    def finalize(self, raise_warning=True):
+        """Constructs the SNodeTree and finalizes this builder.
+
+        Args:
+            raise_warning (bool): Raise warning or not."""
+        return self._finalize(raise_warning, compile_only=False)
+
+    def _finalize(self, raise_warning, compile_only) -> SNodeTree:
+        self._check_not_finalized()
+        if self.empty and raise_warning:
+            warning("Finalizing an empty FieldsBuilder!")
+        self.finalized = True
+        impl.get_runtime().finalize_fields_builder(self)
+        return SNodeTree(
+            _ti_core.finalize_snode_tree(_snode_registry, self.ptr, impl.get_runtime()._prog, compile_only)
+        )
+
+    def _check_not_finalized(self):
+        if self.finalized:
+            raise GsTaichiRuntimeError("FieldsBuilder finalized")

gstaichi/_snode/snode_tree.py

@@ -0,0 +1,34 @@
+# type: ignore
+
+# The reason we import just the gstaichi.core.util module, instead of the ti_python_core
+# object within it, is that ti_python_core is stateful. While in practice ti_python_core is
+# loaded during the import procedure, it's probably still good to delay the
+# access to it.
+
+from gstaichi.lang import impl
+from gstaichi.lang.exception import GsTaichiRuntimeError
+
+
+class SNodeTree:
+    def __init__(self, ptr):
+        self.prog = impl.get_runtime().prog
+        self.ptr = ptr
+        self.destroyed = False
+
+    def destroy(self):
+        if self.destroyed:
+            raise GsTaichiRuntimeError("SNode tree has been destroyed")
+        if self.prog != impl.get_runtime().prog:
+            return
+        self.ptr.destroy_snode_tree(impl.get_runtime().prog)
+
+        # FieldExpression holds a SNode* to the place-SNode associated with a SNodeTree
+        # Therefore, we have to recompile all the kernels after destroying a SNodeTree
+        impl.get_runtime().clear_compiled_functions()
+        self.destroyed = True
+
+    @property
+    def id(self):
+        if self.destroyed:
+            raise GsTaichiRuntimeError("SNode tree has been destroyed")
+        return self.ptr.id()

gstaichi/_test_tools/__init__.py

@@ -0,0 +1,18 @@
+import gstaichi as ti
+
+from . import textwrap2
+
+
+def ti_init_same_arch(**options) -> None:
+    """
+    Used in tests to call ti.init, passing in the same arch as currently
+    configured. Since it's fairly fiddly to do that, extracting this out
+    to this helper function.
+    """
+    assert ti.cfg is not None
+    options = dict(options)
+    options["arch"] = getattr(ti, ti.cfg.arch.name)
+    ti.init(**options)
+
+
+__all__ = ["textwrap2"]

gstaichi/_test_tools/dataclass_test_tools.py

@@ -0,0 +1,36 @@
+import dataclasses
+from typing import Any, cast
+
+import gstaichi as ti
+
+
+def _make_child_obj(obj_type: Any) -> Any:
+    if isinstance(obj_type, ti.types.NDArray):
+        ndarray_type = cast(ti.types.ndarray, obj_type)
+        assert ndarray_type.ndim is not None
+        shape = tuple([10] * ndarray_type.ndim)
+        child_obj = ti.ndarray(ndarray_type.dtype, shape=shape)
+    elif dataclasses.is_dataclass(obj_type):
+        child_obj = build_struct(obj_type)
+    elif isinstance(obj_type, ti.Template) or obj_type == ti.Template:
+        child_obj = ti.field(ti.i32, (10,))
+    else:
+        raise Exception("unknown type ", obj_type)
+    return child_obj
+
+
+def build_struct(struct_type: Any) -> Any:
+    member_objects = {}
+    for field in dataclasses.fields(struct_type):
+        child_obj = _make_child_obj(field.type)
+        member_objects[field.name] = child_obj
+    dataclass_object = struct_type(**member_objects)
+    return dataclass_object
+
+
+def build_obj_tuple_from_type_dict(name_to_type: dict[str, Any]) -> tuple[Any, ...]:
+    obj_l = []
+    for _name, param_type in name_to_type.items():
+        child_obj = _make_child_obj(param_type)
+        obj_l.append(child_obj)
+    return tuple(obj_l)

gstaichi/_test_tools/load_kernel_string.py

@@ -0,0 +1,30 @@
+import importlib.util
+import sys
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+
+
+def import_kernel_from_file(kernel_filepath: Path, kernel_name: str):
+    spec = importlib.util.spec_from_file_location(kernel_name, kernel_filepath)
+    assert spec is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[kernel_name] = module
+    loader = spec.loader
+    assert loader is not None
+    loader.exec_module(module)
+    return getattr(module, kernel_name)
+
+
+@contextmanager
+def load_kernel_from_string(kernel_str: str, kernel_name: str):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        filepath = Path(temp_dir) / f"{kernel_name}.py"
+        with open(filepath, "w") as f:
+            f.write(kernel_str)
+        try:
+            kernel = import_kernel_from_file(kernel_filepath=filepath, kernel_name=kernel_name)
+            yield kernel
+        finally:
+            if kernel_name in sys.modules:
+                del sys.modules[kernel_name]