gstaichi 0.1.25.dev0__cp311-cp311-win_amd64.whl

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

gstaichi/_version.py

@@ -0,0 +1 @@
+__version__ = '0.1.25.dev0'

gstaichi/_version_check.py

@@ -0,0 +1,103 @@
+# type: ignore
+
+import datetime
+import json
+import os
+import platform
+import threading
+import uuid
+from urllib import request
+
+from gstaichi._lib import core as _ti_core
+
+
+def check_version(cur_uuid):
+    # Check GsTaichi version for the user.
+    major = _ti_core.get_version_major()
+    minor = _ti_core.get_version_minor()
+    patch = _ti_core.get_version_patch()
+    version = f"{major}.{minor}.{patch}"
+    payload = {"version": version, "platform": "", "python": ""}
+
+    system = platform.system()
+    u = platform.uname()
+    if (u.system, u.machine) == ("Linux", "x86_64"):
+        payload["platform"] = "manylinux_2_27_x86_64"
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        payload["platform"] = "manylinux_2_27_aarch64"
+    elif system == "Windows":
+        payload["platform"] = "win_amd64"
+    elif system == "Darwin":
+        if platform.release() < "19.0.0":
+            payload["platform"] = "macosx_10_14_x86_64"
+        elif platform.machine() == "x86_64":
+            payload["platform"] = "macosx_10_15_x86_64"
+        else:
+            payload["platform"] = "macosx_11_0_arm64"
+
+    python_version = platform.python_version().split(".")
+    payload["python"] = "cp" + python_version[0] + python_version[1]
+
+    payload["uuid"] = cur_uuid
+    if os.getenv("TI_CI") == "1":
+        payload["type"] = "CI"
+    # We do not want request exceptions break users' usage of GsTaichi.
+    try:
+        payload = json.dumps(payload)
+        payload = payload.encode()
+        req = request.Request("https://metadata.gstaichi.graphics/check_version", method="POST")
+        req.add_header("Content-Type", "application/json")
+        with request.urlopen(req, data=payload, timeout=5) as response:
+            response = json.loads(response.read().decode("utf-8"))
+            return response
+    except:
+        return None
+
+
+def write_version_info(response, cur_uuid, version_info_path, cur_date):
+    if response is None:
+        return
+    with open(version_info_path, "w") as f:
+        f.write((cur_date).strftime("%Y-%m-%d"))
+        f.write("\n")
+        if response["status"] == 1:
+            f.write(response["latest_version"])
+        else:
+            f.write("0.0.0")
+        f.write("\n")
+        f.write(cur_uuid)
+        f.write("\n")
+
+
+def try_check_version():
+    try:
+        os.makedirs(_ti_core.get_repo_dir(), exist_ok=True)
+        version_info_path = os.path.join(_ti_core.get_repo_dir(), "version_info")
+        cur_date = datetime.date.today()
+        if os.path.exists(version_info_path):
+            with open(version_info_path, "r") as f:
+                version_info_file = f.readlines()
+                last_time = version_info_file[0].rstrip()
+                cur_uuid = version_info_file[2].rstrip()
+            if cur_date.strftime("%Y-%m-%d") > last_time:
+                response = check_version(cur_uuid)
+                write_version_info(response, cur_uuid, version_info_path, cur_date)
+        else:
+            cur_uuid = str(uuid.uuid4())
+            write_version_info({"status": 0}, cur_uuid, version_info_path, cur_date)
+            response = check_version(cur_uuid)
+            write_version_info(response, cur_uuid, version_info_path, cur_date)
+    # Wildcard exception to catch potential file writing errors.
+    except:
+        pass
+
+
+def start_version_check_thread():
+    skip = os.environ.get("TI_SKIP_VERSION_CHECK")
+    if skip != "ON":
+        # We don't join this thread because we do not wish to block users.
+        check_version_thread = threading.Thread(target=try_check_version, daemon=True)
+        check_version_thread.start()
+
+
+__all__ = []

gstaichi/ad/__init__.py

@@ -0,0 +1,3 @@
+# type: ignore
+
+from gstaichi.ad._ad import *