gstaichi 2.1.1__cp313-cp313-macosx_11_0_arm64.whl

gstaichi/algorithms/_algorithms.py

@@ -0,0 +1,117 @@
+# type: ignore
+
+from gstaichi._kernels import (
+    blit_from_field_to_field,
+    scan_add_inclusive,
+    sort_stage,
+    uniform_add,
+    warp_shfl_up_i32,
+)
+from gstaichi.lang.impl import current_cfg, field
+from gstaichi.lang.kernel_impl import data_oriented
+from gstaichi.lang.misc import cuda, vulkan
+from gstaichi.lang.runtime_ops import sync
+from gstaichi.lang.simt import subgroup
+from gstaichi.types.primitive_types import i32
+
+
+def parallel_sort(keys, values=None):
+    """Odd-even merge sort
+
+    References:
+        https://developer.nvidia.com/gpugems/gpugems2/part-vi-simulation-and-numerical-algorithms/chapter-46-improved-gpu-sorting
+        https://en.wikipedia.org/wiki/Batcher_odd%E2%80%93even_mergesort
+    """
+    N = keys.shape[0]
+
+    num_stages = 0
+    p = 1
+    while p < N:
+        k = p
+        while k >= 1:
+            invocations = int((N - k - k % p) / (2 * k)) + 1
+            if values is None:
+                sort_stage(keys, 0, keys, N, p, k, invocations)
+            else:
+                sort_stage(keys, 1, values, N, p, k, invocations)
+            num_stages += 1
+            sync()
+            k = int(k / 2)
+        p = int(p * 2)
+
+
+@data_oriented
+class PrefixSumExecutor:
+    """Parallel Prefix Sum (Scan) Helper
+
+    Use this helper to perform an inclusive in-place's parallel prefix sum.
+
+    References:
+        https://developer.download.nvidia.com/compute/cuda/1.1-Beta/x86_website/projects/scan/doc/scan.pdf
+        https://github.com/NVIDIA/cuda-samples/blob/master/Samples/2_Concepts_and_Techniques/shfl_scan/shfl_scan.cu
+    """
+
+    def __init__(self, length):
+        self.sorting_length = length
+
+        BLOCK_SZ = 64
+        GRID_SZ = int((length + BLOCK_SZ - 1) / BLOCK_SZ)
+
+        # Buffer position and length
+        # This is a single buffer implementation for ease of aot usage
+        ele_num = length
+        self.ele_nums = [ele_num]
+        start_pos = 0
+        self.ele_nums_pos = [start_pos]
+
+        while ele_num > 1:
+            ele_num = int((ele_num + BLOCK_SZ - 1) / BLOCK_SZ)
+            self.ele_nums.append(ele_num)
+            start_pos += BLOCK_SZ * ele_num
+            self.ele_nums_pos.append(start_pos)
+
+        self.large_arr = field(i32, shape=start_pos)
+
+    def run(self, input_arr):
+        length = self.sorting_length
+        ele_nums = self.ele_nums
+        ele_nums_pos = self.ele_nums_pos
+
+        if input_arr.dtype != i32:
+            raise RuntimeError("Only ti.i32 type is supported for prefix sum.")
+
+        if current_cfg().arch == cuda:
+            inclusive_add = warp_shfl_up_i32
+        elif current_cfg().arch == vulkan:
+            inclusive_add = subgroup.inclusive_add
+        else:
+            raise RuntimeError(f"{str(current_cfg().arch)} is not supported for prefix sum.")
+
+        blit_from_field_to_field(self.large_arr, input_arr, 0, length)
+
+        # Kogge-Stone construction
+        for i in range(len(ele_nums) - 1):
+            if i == len(ele_nums) - 2:
+                scan_add_inclusive(
+                    self.large_arr,
+                    ele_nums_pos[i],
+                    ele_nums_pos[i + 1],
+                    True,
+                    inclusive_add,
+                )
+            else:
+                scan_add_inclusive(
+                    self.large_arr,
+                    ele_nums_pos[i],
+                    ele_nums_pos[i + 1],
+                    False,
+                    inclusive_add,
+                )
+
+        for i in range(len(ele_nums) - 3, -1, -1):
+            uniform_add(self.large_arr, ele_nums_pos[i], ele_nums_pos[i + 1])
+
+        blit_from_field_to_field(input_arr, self.large_arr, 0, length)
+
+
+__all__ = ["parallel_sort", "PrefixSumExecutor"]

gstaichi/assets/.git

@@ -0,0 +1 @@
+gitdir: ../../.git/modules/external/assets

gstaichi/examples/lcg_python.py

@@ -0,0 +1,26 @@
+import time
+
+import numpy as np
+import numpy.typing as npt
+
+
+def lcg_np(B: int, lcg_its: int, a: npt.NDArray) -> None:
+    for i in range(B):
+        x = a[i]
+        for j in range(lcg_its):
+            x = (1664525 * x + 1013904223) % 2147483647
+        a[i] = x
+
+
+def main() -> None:
+    B = 16000
+    a = np.ndarray((B,), np.int32)
+
+    start = time.time()
+    lcg_np(B, 1000, a)
+    end = time.time()
+    print("elapsed", end - start)
+    # elapsed 5.552601099014282 on macbook air m4
+
+
+main()

gstaichi/examples/lcg_taichi.py

@@ -0,0 +1,34 @@
+import time
+
+import gstaichi as ti
+
+
+@ti.kernel
+def lcg_ti(B: int, lcg_its: int, a: ti.types.NDArray[ti.i32, 1]) -> None:
+    for i in range(B):
+        x = a[i]
+        for j in range(lcg_its):
+            x = (1664525 * x + 1013904223) % 2147483647
+        a[i] = x
+
+
+def main() -> None:
+    ti.init(arch=ti.gpu)
+
+    B = 16000
+    a = ti.ndarray(ti.int32, (B,))
+
+    ti.sync()
+    start = time.time()
+    lcg_ti(B, 1000, a)
+    ti.sync()
+    end = time.time()
+    print("elapsed", end - start)
+
+    # [GsTaichi] version 1.8.0, llvm 15.0.7, commit 5afed1c9, osx, python 3.10.16
+    # [GsTaichi] Starting on arch=metal
+    # elapsed 0.04660296440124512
+    # (on mac air m4)
+
+
+main()

gstaichi/examples/minimal.py

@@ -0,0 +1,28 @@
+import gstaichi as ti
+
+
+@ti.kernel
+def lcg_ti(B: int, lcg_its: int, a: ti.types.NDArray[ti.i32, 1]) -> None:
+    """
+    Linear congruential generator https://en.wikipedia.org/wiki/Linear_congruential_generator
+    """
+    for i in range(B):
+        x = a[i]
+        for j in range(lcg_its):
+            x = (1664525 * x + 1013904223) % 2147483647
+        a[i] = x
+
+
+def main() -> None:
+    ti.init(arch=ti.cpu)
+
+    B = 10
+    lcg_its = 10
+
+    a = ti.ndarray(ti.int32, (B,))
+
+    lcg_ti(B, lcg_its, a)
+    print(f"LCG for B={B}, lcg_its={lcg_its}: ", a.to_numpy())  # pylint: disable=no-member
+
+
+main()

gstaichi/experimental.py

@@ -0,0 +1,16 @@
+# type: ignore
+
+import warnings
+
+from gstaichi.lang.kernel_impl import real_func as _real_func
+
+
+def real_func(func):
+    warnings.warn(
+        "ti.experimental.real_func is deprecated because it is no longer experimental. " "Use ti.real_func instead.",
+        DeprecationWarning,
+    )
+    return _real_func(func)
+
+
+__all__ = ["real_func"]

gstaichi/lang/__init__.py

@@ -0,0 +1,50 @@
+# type: ignore
+
+from gstaichi.lang import impl, simt
+from gstaichi.lang._fast_caching.function_hasher import pure
+from gstaichi.lang._ndarray import *
+from gstaichi.lang._ndrange import ndrange
+from gstaichi.lang._texture import Texture
+from gstaichi.lang.exception import *
+from gstaichi.lang.field import *
+from gstaichi.lang.impl import *
+from gstaichi.lang.kernel_impl import *
+from gstaichi.lang.matrix import *
+from gstaichi.lang.mesh import *
+from gstaichi.lang.misc import *  # pylint: disable=W0622
+from gstaichi.lang.ops import *  # pylint: disable=W0622
+from gstaichi.lang.runtime_ops import *
+from gstaichi.lang.snode import *
+from gstaichi.lang.source_builder import *
+from gstaichi.lang.struct import *
+from gstaichi.types.enums import DeviceCapability, Format, Layout
+
+__all__ = [
+    s
+    for s in dir()
+    if not s.startswith("_")
+    and s
+    not in [
+        "any_array",
+        "ast",
+        "common_ops",
+        "enums",
+        "exception",
+        "expr",
+        "impl",
+        "inspect",
+        "kernel_arguments",
+        "kernel_impl",
+        "matrix",
+        "mesh",
+        "misc",
+        "ops",
+        "platform",
+        "runtime_ops",
+        "shell",
+        "snode",
+        "source_builder",
+        "struct",
+        "util",
+    ]
+]

gstaichi/lang/_dataclass_util.py

@@ -0,0 +1,31 @@
+def create_flat_name(basename: str, child_name: str) -> str:
+    """
+    Appends child_name to basename, separated by __ti_.
+    If basename does not start with __ti_ then prefix the resulting string
+    with __ti_.
+
+    Note that we want to avoid adding prefix __ti_ if already included in `basename`,
+    to avoid duplicating said delimiter.
+
+    We'll use this when expanding py dataclass members, e.g.
+
+    @dataclasses.dataclass
+    def Foo:
+        a: int
+        b: int
+
+    foo = Foo(a=5, b=3)
+
+    When we expand out foo, we'll replace foo with the following names instead:
+    - __ti_foo__ti_a
+    - __ti_foo__ti_b
+
+    We use the __ti_ to ensure that it's easy to ensure no collision with existing user-defined
+    names. We require the user to not create any fields or variables which themselves are prefixed
+    with __ti_, and given this constraint, the names we create will not conflict with user-generated
+    names.
+    """
+    full_name = f"{basename}__ti_{child_name}"
+    if not full_name.startswith("__ti_"):
+        full_name = f"__ti_{full_name}"
+    return full_name

gstaichi/lang/_fast_caching/__init__.py

@@ -0,0 +1,3 @@
+from .args_hasher import FIELD_METADATA_CACHE_VALUE
+
+__all__ = ["FIELD_METADATA_CACHE_VALUE"]

gstaichi/lang/_fast_caching/args_hasher.py

@@ -0,0 +1,110 @@
+import dataclasses
+import enum
+import numbers
+import time
+from typing import Any, Sequence
+
+import numpy as np
+
+from .._ndarray import ScalarNdarray
+from ..field import ScalarField
+from ..matrix import MatrixField, MatrixNdarray, VectorNdarray
+from ..util import is_data_oriented
+from .hash_utils import hash_iterable_strings
+
+g_num_calls = 0
+g_num_args = 0
+g_hashing_time = 0
+g_repr_time = 0
+g_num_ignored_calls = 0
+
+
+FIELD_METADATA_CACHE_VALUE = "add_value_to_cache_key"
+
+
+def dataclass_to_repr(path: tuple[str, ...], arg: Any) -> str:
+    repr_l = []
+    for field in dataclasses.fields(arg):
+        child_value = getattr(arg, field.name)
+        _repr = stringify_obj_type(path + (field.name,), child_value)
+        full_repr = f"{field.name}: ({_repr})"
+        if field.metadata.get(FIELD_METADATA_CACHE_VALUE, False):
+            full_repr += f" = {child_value}"
+        repr_l.append(full_repr)
+    return "[" + ",".join(repr_l) + "]"
+
+
+def stringify_obj_type(path: tuple[str, ...], obj: Any) -> str | None:
+    """
+    Convert an object into a string representation that only depends on its type.
+
+    String should somehow represent the type of obj. Doesnt have to be hashed, nor does it have
+    to be the actual python type string, just a string that is representative of the type, and won't collide
+    with different (allowed) types.
+
+    `path` is used during debugging.
+    """
+    # TODO: We should have a way of printing this without having to hack the code really. Using logger perhaps?
+    # (I have another PR that addreses this https://github.com/Genesis-Embodied-AI/gstaichi/pull/144/files)
+    arg_type = type(obj)
+    if isinstance(obj, ScalarNdarray):
+        return f"[nd-{obj.dtype}-{len(obj.shape)}]"
+    if isinstance(obj, VectorNdarray):
+        return f"[ndv-{obj.n}-{obj.dtype}-{len(obj.shape)}]"
+    if isinstance(obj, ScalarField):
+        return f"[f-{obj.snode._id}-{obj.dtype}-{obj.shape}]"
+    if isinstance(obj, MatrixNdarray):
+        return f"[ndm-{obj.m}-{obj.n}-{obj.dtype}-{len(obj.shape)}]"
+    if "torch.Tensor" in str(arg_type):
+        return f"[pt-{obj.dtype}-{obj.ndim}]"
+    if isinstance(obj, np.ndarray):
+        return f"[np-{obj.dtype}-{obj.ndim}]"
+    if isinstance(obj, MatrixField):
+        return f"[fm-{obj.m}-{obj.n}-{obj.snode._id}-{obj.dtype}-{obj.shape}]"
+    if dataclasses.is_dataclass(obj):
+        return dataclass_to_repr(path, obj)
+    if is_data_oriented(obj):
+        child_repr_l = []
+        for k, v in obj.__dict__.items():
+            _child_repr = stringify_obj_type((*path, k), v)
+            if _child_repr is None:
+                print("not representable child", k, type(v), "path", path)
+                return None
+            child_repr_l.append(f"{k}: {_child_repr}")
+        return ", ".join(child_repr_l)
+    if issubclass(arg_type, (numbers.Number, np.number)):
+        return str(arg_type)
+    if arg_type is np.bool_:
+        # np is deprecating bool. Treat specially/carefully
+        return "np.bool_"
+    if isinstance(obj, enum.Enum):
+        return f"enum-{obj.name}-{obj.value}"
+    return None
+
+
+def hash_args(args: Sequence[Any]) -> str | None:
+    global g_num_calls, g_num_args, g_hashing_time, g_repr_time, g_num_ignored_calls
+    g_num_calls += 1
+    g_num_args += len(args)
+    hash_l = []
+    for i_arg, arg in enumerate(args):
+        start = time.time()
+        _hash = stringify_obj_type((str(i_arg),), arg)
+        g_repr_time += time.time() - start
+        if not _hash:
+            g_num_ignored_calls += 1
+            return None
+        hash_l.append(_hash)
+    start = time.time()
+    res = hash_iterable_strings(hash_l)
+    g_hashing_time += time.time() - start
+    return res
+
+
+def dump_stats() -> None:
+    print("args hasher dump stats")
+    print("total calls", g_num_calls)
+    print("ignored calls", g_num_ignored_calls)
+    print("total args", g_num_args)
+    print("hashing time", g_hashing_time)
+    print("arg representation time", g_repr_time)

gstaichi/lang/_fast_caching/config_hasher.py

@@ -0,0 +1,30 @@
+from gstaichi.lang import impl
+
+from .hash_utils import hash_iterable_strings
+
+EXCLUDE_PREFIXES = ["_", "offline_cache", "print_", "verbose_"]
+
+
+def hash_compile_config() -> str:
+    """
+    Calculates a hash string for the current compiler config.
+
+    If any value in the compiler config changes, the hash string changes too.
+
+    Though arguably we might want to blacklist certain keys, such as print_ir_debug,
+    which do not affect the compiled kernels, just stuff that gets printed during
+    the compilation process.
+    """
+    config = impl.get_runtime().prog.config()
+    config_l = []
+    for k in dir(config):
+        skip = False
+        for prefix in EXCLUDE_PREFIXES:
+            if k.startswith(prefix) or k in [""]:
+                skip = True
+        if skip:
+            continue
+        v = getattr(config, k)
+        config_l.append(f"{k}={v}")
+    config_hash = hash_iterable_strings(config_l, separator="\n")
+    return config_hash

gstaichi/lang/_fast_caching/fast_caching_types.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+
+from .._wrap_inspect import FunctionSourceInfo
+
+
+class HashedFunctionSourceInfo(BaseModel):
+    """
+    Wraps a function source info, and the hash string of that function.
+
+    By not adding the hash directly into function source info, we avoid
+    having to make hash an optional type, and checking if it's empty or not.
+
+    If you have a HashedFunctionSourceInfo object, then you are guaranteed
+    to have the hash string.
+
+    If you only have the FunctionSourceInfo object, you are guaranteed that it
+    does not have a hash string.
+    """
+
+    function_source_info: FunctionSourceInfo
+    hash: str

gstaichi/lang/_fast_caching/function_hasher.py

@@ -0,0 +1,57 @@
+import os
+from itertools import islice
+from typing import TYPE_CHECKING, Iterable
+
+from .._wrap_inspect import FunctionSourceInfo
+from .fast_caching_types import HashedFunctionSourceInfo
+from .hash_utils import hash_iterable_strings
+
+if TYPE_CHECKING:
+    from gstaichi.lang.kernel_impl import GsTaichiCallable
+
+
+def pure(fn: "GsTaichiCallable") -> "GsTaichiCallable":
+    fn.is_pure = True
+    return fn
+
+
+def _read_file(function_info: FunctionSourceInfo) -> list[str]:
+    with open(function_info.filepath) as f:
+        return list(islice(f, function_info.start_lineno, function_info.end_lineno + 1))
+
+
+def _hash_function(function_info: FunctionSourceInfo) -> str:
+    return hash_iterable_strings(_read_file(function_info))
+
+
+def hash_functions(function_infos: Iterable[FunctionSourceInfo]) -> list[HashedFunctionSourceInfo]:
+    results = []
+    for f_info in function_infos:
+        hash_ = _hash_function(f_info)
+        results.append(HashedFunctionSourceInfo(function_source_info=f_info, hash=hash_))
+    return results
+
+
+def hash_kernel(kernel_info: FunctionSourceInfo) -> str:
+    return _hash_function(kernel_info)
+
+
+def dump_stats() -> None:
+    print("function hasher dump stats")
+
+
+def _validate_hashed_function_info(hashed_function_info: HashedFunctionSourceInfo) -> bool:
+    """
+    Checks the hash
+    """
+    if not os.path.isfile(hashed_function_info.function_source_info.filepath):
+        return False
+    _hash = _hash_function(hashed_function_info.function_source_info)
+    return _hash == hashed_function_info.hash
+
+
+def validate_hashed_function_infos(function_infos: Iterable[HashedFunctionSourceInfo]) -> bool:
+    for function_info in function_infos:
+        if not _validate_hashed_function_info(function_info):
+            return False
+    return True

gstaichi/lang/_fast_caching/hash_utils.py

@@ -0,0 +1,11 @@
+import hashlib
+from typing import Iterable
+
+
+def hash_iterable_strings(strings: Iterable[str], separator: str = "_") -> str:
+    h = hashlib.sha256()
+    separator_enc = separator.encode("utf-8")
+    for v in strings:
+        h.update(v.encode("utf-8"))
+        h.update(separator_enc)
+    return h.hexdigest()

gstaichi/lang/_fast_caching/python_side_cache.py

@@ -0,0 +1,52 @@
+import os
+
+from .. import impl
+
+
+class PythonSideCache:
+    """
+    Manages a cache that is managed from the python side (we also have c++-side caches)
+
+    The cache is disk-based. When we create the PythonSideCache object, the cache
+    path is created as a sub-folder of CompileConfig.offline_cache_file_path.
+
+    Note that constructing this object is cheap, so there is no need to maintain some
+    kind of conceptual singleton instance or similar.
+
+    Each cache key value is stored to a single file, with the cache key as the filename.
+
+    No metadata is associated with the file, making management very lightweight.
+
+    We update the file date/time when we read from a particular file, so we can easily
+    implement an LRU cleaning strategy at some point in the future, based on the file
+    date/times.
+    """
+
+    def __init__(self) -> None:
+        _cache_parent_folder = impl.get_runtime().prog.config().offline_cache_file_path
+        self.cache_folder = os.path.join(_cache_parent_folder, "python_side_cache")
+        os.makedirs(self.cache_folder, exist_ok=True)
+
+    def _get_filepath(self, key: str) -> str:
+        filepath = os.path.join(self.cache_folder, f"{key}.cache.txt")
+        return filepath
+
+    def _touch(self, filepath):
+        """
+        Updates file date/time.
+        """
+        with open(filepath, "a"):
+            os.utime(filepath, None)
+
+    def store(self, key: str, value: str) -> None:
+        filepath = self._get_filepath(key)
+        with open(filepath, "w") as f:
+            f.write(value)
+
+    def try_load(self, key: str) -> str | None:
+        filepath = self._get_filepath(key)
+        if not os.path.isfile(filepath):
+            return None
+        self._touch(filepath)
+        with open(filepath) as f:
+            return f.read()

gstaichi/lang/_fast_caching/src_hasher.py

@@ -0,0 +1,75 @@
+from typing import Any, Iterable, Sequence
+
+from pydantic import BaseModel
+
+from .._wrap_inspect import FunctionSourceInfo
+from . import args_hasher, config_hasher, function_hasher
+from .fast_caching_types import HashedFunctionSourceInfo
+from .hash_utils import hash_iterable_strings
+from .python_side_cache import PythonSideCache
+
+
+def create_cache_key(kernel_source_info: FunctionSourceInfo, args: Sequence[Any]) -> str | None:
+    """
+    cache key takes into account:
+    - arg types
+    - cache value arg values
+    - kernel function (but not sub functions)
+    - compilation config (which includes arch, and debug)
+    """
+    args_hash = args_hasher.hash_args(args)
+    if args_hash is None:
+        return None
+    kernel_hash = function_hasher.hash_kernel(kernel_source_info)
+    config_hash = config_hasher.hash_compile_config()
+    cache_key = hash_iterable_strings((kernel_hash, args_hash, config_hash))
+    return cache_key
+
+
+class CacheValue(BaseModel):
+    hashed_function_source_infos: list[HashedFunctionSourceInfo]
+
+
+def store(cache_key: str, function_source_infos: Iterable[FunctionSourceInfo]) -> None:
+    """
+    Note that unlike other caches, this cache is not going to store the actual value we want.
+    This cache is only used for verification that our cache key is valid. Big picture:
+    - we have a cache key, based on args and top level kernel function
+    - we want to use this to look up LLVM IR, in C++ side cache
+    - however, before doing that, we first want to validate that the source code didn't change
+        - i.e. is our cache key still valid?
+    - the python side cache contains information we will use to verify that our cache key is valid
+        - ie the list of function source infos
+    """
+    if not cache_key:
+        return
+    cache = PythonSideCache()
+    hashed_function_source_infos = function_hasher.hash_functions(function_source_infos)
+    cache_value_obj = CacheValue(hashed_function_source_infos=list(hashed_function_source_infos))
+    cache.store(cache_key, cache_value_obj.json())
+
+
+def _try_load(cache_key: str) -> Sequence[HashedFunctionSourceInfo] | None:
+    cache = PythonSideCache()
+    maybe_cache_value_json = cache.try_load(cache_key)
+    if maybe_cache_value_json is None:
+        return None
+    cache_value_obj = CacheValue.parse_raw(maybe_cache_value_json)
+    return cache_value_obj.hashed_function_source_infos
+
+
+def validate_cache_key(cache_key: str) -> bool:
+    """
+    loads function source infos from cache, if available
+    checks the hashes against the current source code
+    """
+    maybe_hashed_function_source_infos = _try_load(cache_key)
+    if not maybe_hashed_function_source_infos:
+        return False
+    return function_hasher.validate_hashed_function_infos(maybe_hashed_function_source_infos)
+
+
+def dump_stats() -> None:
+    print("dump stats")
+    args_hasher.dump_stats()
+    function_hasher.dump_stats()