gstaichi 0.1.18.dev1__cp310-cp310-macosx_15_0_arm64.whl

taichi/aot/module.py

@@ -0,0 +1,253 @@
+# type: ignore
+
+import datetime
+import os
+import warnings
+from contextlib import contextmanager
+from glob import glob
+from pathlib import Path, PurePosixPath
+from shutil import rmtree
+from tempfile import mkdtemp
+from zipfile import ZipFile
+
+import taichi
+from taichi.aot.utils import produce_injected_args, produce_injected_args_from_template
+from taichi.lang import impl, kernel_impl
+from taichi.lang.field import ScalarField
+from taichi.lang.matrix import MatrixField
+from taichi.types.annotations import template
+
+
+class KernelTemplate:
+    def __init__(self, kernel_fn, aot_module):
+        self._kernel_fn = kernel_fn
+        self._aot_module = aot_module
+
+    @staticmethod
+    def keygen(v, key_p, fields):
+        if isinstance(v, (int, float, bool)):
+            key_p += "=" + str(v) + ","
+            return key_p
+        for ky, val in fields:
+            if val is v:
+                key_p += "=" + ky + ","
+                return key_p
+        raise RuntimeError(
+            "Arg type must be of type int/float/boolean" f" or taichi field. Type {str(type(v))}" " is not supported"
+        )
+
+    def instantiate(self, **kwargs):
+        name = self._kernel_fn.__name__
+        kernel = self._kernel_fn._primal
+        assert isinstance(kernel, kernel_impl.Kernel)
+        injected_args = []
+        key_p = ""
+        anno_index = 0
+        template_args = {}
+
+        for index, (key, value) in enumerate(kwargs.items()):
+            template_args[index] = (key, value)
+
+        for arg in kernel.arguments:
+            if arg.annotation == template or isinstance(arg.annotation, template):
+                (k, v) = template_args[anno_index]
+                key_p += k
+                key_p = self.keygen(v, key_p, self._aot_module._fields.items())
+                injected_args.append(v)
+                anno_index += 1
+            else:
+                injected_args.append(0)
+        kernel.ensure_compiled(*injected_args)
+        self._aot_module._aot_builder.add_kernel_template(name, key_p, kernel.kernel_cpp)
+
+        # kernel AOT
+        self._aot_module._kernels.append(kernel)
+
+
+class Module:
+    """An AOT module to save and load Taichi kernels.
+
+    This module serializes the Taichi kernels for a specific arch. The
+    serialized module can later be loaded to run on that backend, without the
+    Python environment.
+
+    Example:
+      Usage::
+
+        m = ti.aot.Module(ti.metal)
+        m.add_kernel(foo)
+        m.add_kernel(bar)
+
+        m.save('/path/to/module')
+
+        # Now the module file '/path/to/module' contains the Metal kernels
+        # for running ``foo`` and ``bar``.
+    """
+
+    def __init__(self, arch=None, caps=None):
+        """Creates a new AOT module instance
+
+        Args:
+          arch: Target backend architecture. Default to the one initialized in :func:`~taichi.lang.init` if not specified.
+          caps (List[str]): Enabled device capabilities.
+        """
+        if caps is None:
+            caps = []
+        curr_arch = impl.current_cfg().arch
+        if arch is None:
+            arch = curr_arch
+        elif arch != curr_arch:
+            # TODO: we'll support this eventually but not yet...
+            warnings.warn(
+                f"AOT compilation to a different arch than the current one is not yet supported, switching to {curr_arch}"
+            )
+            arch = curr_arch
+
+        self._arch = arch
+        self._kernels = []
+        self._fields = {}
+        rtm = impl.get_runtime()
+        rtm._finalize_root_fb_for_aot()
+        self._aot_builder = rtm.prog.make_aot_module_builder(arch, caps)
+        self._content = []
+
+    def add_field(self, name, field):
+        """Add a taichi field to the AOT module.
+
+        Args:
+          name: name of taichi field
+          field: taichi field
+
+        Example::
+
+            >>> a = ti.field(ti.f32, shape=(4,4))
+            >>> b = ti.field("something")
+            >>>
+            >>> m.add_field(a)
+            >>> m.add_field(b)
+            >>>
+            >>> # Must add in sequence
+        """
+        is_scalar = True
+        self._fields[name] = field
+        column_num = 1
+        row_num = 1
+        if isinstance(field, MatrixField):
+            is_scalar = False
+            row_num = field.m
+            column_num = field.n
+        else:
+            assert isinstance(field, ScalarField)
+        self._aot_builder.add_field(
+            name,
+            field.snode.ptr,
+            is_scalar,
+            field.dtype,
+            field.snode.shape,
+            row_num,
+            column_num,
+        )
+
+    def add_kernel(self, kernel_fn, template_args=None, name=None):
+        """Add a taichi kernel to the AOT module.
+
+        Args:
+          kernel_fn (Function): the function decorated by taichi `kernel`.
+          template_args (Dict[str, Any]): a dict where key is the template
+            parameter name, and value is the instantiating arg. Note that this
+            works for both :class:`~taichi.types.template` and for
+            `:class:`~taichi.types.ndarray`.
+          name (str): Name to identify this kernel in the module. If not
+            provided, uses the built-in ``__name__`` attribute of `kernel_fn`.
+
+        """
+        kernel_name = name or kernel_fn.__name__
+        kernel = kernel_fn._primal
+        assert isinstance(kernel, kernel_impl.Kernel)
+        if template_args is not None:
+            injected_args = produce_injected_args_from_template(kernel, template_args)
+        else:
+            injected_args = produce_injected_args(kernel)
+        kernel.ensure_compiled(*injected_args)
+        self._aot_builder.add(kernel_name, kernel.kernel_cpp)
+
+        # kernel AOT
+        self._kernels.append(kernel)
+
+        self._content += ["kernel:" + kernel_name]
+
+    def add_graph(self, name, graph):
+        self._aot_builder.add_graph(name, graph._compiled_graph)
+        self._content += ["cgraph:" + name]
+
+    @contextmanager
+    def add_kernel_template(self, kernel_fn):
+        """Add a taichi kernel (with template parameters) to the AOT module.
+
+        Args:
+          kernel_fn (Function): the function decorated by taichi `kernel`.
+
+        Example::
+
+            >>> @ti.kernel
+            >>> def bar_tmpl(a: ti.template()):
+            >>>   x = a
+            >>>   # or y = a
+            >>>   # do something with `x` or `y`
+            >>>
+            >>> m = ti.aot.Module(arch)
+            >>> with m.add_kernel_template(bar_tmpl) as kt:
+            >>>   kt.instantiate(a=x)
+            >>>   kt.instantiate(a=y)
+            >>>
+            >>> @ti.kernel
+            >>> def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
+            >>>   x = a
+            >>>   y = b
+            >>>   # do something with `x` and `y`
+            >>>
+            >>> with m.add_kernel_template(bar_tmpl) as kt:
+            >>>   kt.instantiate(a=x, b=y)
+
+        TODO:
+          * Support external array
+        """
+        kt = KernelTemplate(kernel_fn, self)
+        yield kt
+
+    def save(self, filepath):
+        """
+        Args:
+          filepath (str): path to a folder to store aot files.
+        """
+        filepath = str(PurePosixPath(Path(filepath)))
+        self._aot_builder.dump(filepath, "")
+        with open(f"{filepath}/__content__", "w") as f:
+            f.write("\n".join(self._content))
+        with open(f"{filepath}/__version__", "w") as f:
+            f.write(".".join(str(x) for x in taichi.__version__))
+
+    def archive(self, filepath: str):
+        """
+        Args:
+          filepath (str): path to the stored archive of aot artifacts, MUST
+            end with `.tcm`.
+        """
+        assert filepath.endswith(".tcm"), "AOT module artifact archive must ends with .tcm"
+        tcm_path = Path(filepath).absolute()
+        assert tcm_path.parent.exists(), "Output directory doesn't exist"
+
+        temp_dir = mkdtemp(prefix="tcm_")
+        # Save first as usual.
+        self.save(temp_dir)
+
+        fixed_time = datetime.datetime(2000, 12, 1).timestamp()
+
+        # Package all artifacts into a zip archive and attach contend data.
+        with ZipFile(tcm_path, "w") as z:
+            for path in glob(f"{temp_dir}/*", recursive=True):
+                os.utime(path, (fixed_time, fixed_time))
+                z.write(path, Path.relative_to(Path(path), temp_dir))
+
+        # Remove cached files
+        rmtree(temp_dir)

taichi/aot/utils.py

@@ -0,0 +1,151 @@
+# type: ignore
+
+from typing import Any
+
+from taichi.lang._ndarray import ScalarNdarray
+from taichi.lang._texture import Texture
+from taichi.lang.exception import TaichiCompilationError
+from taichi.lang.matrix import (
+    Matrix,
+    MatrixNdarray,
+    MatrixType,
+    VectorNdarray,
+    VectorType,
+)
+from taichi.lang.util import cook_dtype
+from taichi.types.annotations import template
+from taichi.types.enums import Format
+from taichi.types.ndarray_type import NdarrayType
+from taichi.types.texture_type import RWTextureType, TextureType
+
+template_types = (NdarrayType, TextureType, template)
+
+
+def check_type_match(lhs, rhs):
+    if isinstance(lhs, MatrixType) and isinstance(rhs, MatrixType):
+        return lhs.n == rhs.n and lhs.m == rhs.m and (lhs.dtype == rhs.dtype or lhs.dtype is None or rhs.dtype is None)
+    if isinstance(lhs, MatrixType) or isinstance(rhs, MatrixType):
+        return False
+
+    return cook_dtype(lhs) == cook_dtype(rhs)
+
+
+def produce_injected_args_from_template(kernel, template_args):
+    injected_args = []
+    num_template_args = len([arg.annotation for arg in kernel.arguments if isinstance(arg.annotation, template_types)])
+    assert num_template_args == len(
+        template_args
+    ), f"Need {num_template_args} inputs to instantiate the template parameters, got {len(template_args)}"
+    for arg in kernel.arguments:
+        anno = arg.annotation
+        if isinstance(anno, template_types):
+            injected_args.append(template_args[arg.name])
+        elif isinstance(anno, RWTextureType):
+            texture_shape = (2,) * anno.num_dimensions
+            fmt = anno.fmt
+            injected_args.append(Texture(fmt, texture_shape))
+        else:
+            injected_args.append(0)
+    return injected_args
+
+
+def produce_injected_args(kernel, symbolic_args=None):
+    injected_args = []
+    for i, arg in enumerate(kernel.arguments):
+        anno = arg.annotation
+        if isinstance(anno, NdarrayType):
+            if symbolic_args is not None:
+                # TODO: reconstruct dtype to be TensorType from taichi_core instead of the Python ones
+                element_dim = len(symbolic_args[i].element_shape)
+                if element_dim == 0 or symbolic_args[i].element_shape == (1,):
+                    dtype = symbolic_args[i].dtype()
+                elif element_dim == 1:
+                    dtype = VectorType(symbolic_args[i].element_shape[0], symbolic_args[i].dtype())
+                elif element_dim == 2:
+                    dtype = MatrixType(
+                        symbolic_args[i].element_shape[0],
+                        symbolic_args[i].element_shape[1],
+                        2,
+                        symbolic_args[i].dtype(),
+                    )
+                else:
+                    raise TaichiCompilationError("Not supported")
+                ndim = symbolic_args[i].field_dim
+            else:
+                ndim = anno.ndim
+                dtype = anno.dtype
+
+            if anno.ndim is not None and ndim != anno.ndim:
+                raise TaichiCompilationError(
+                    f"{ndim} from Arg {arg.name} doesn't match kernel's annotated ndim={anno.ndim}"
+                )
+
+            if anno.dtype is not None and not check_type_match(dtype, anno.dtype):
+                raise TaichiCompilationError(
+                    f" Arg {arg.name}'s dtype {dtype.to_string()} doesn't match kernel's annotated dtype={anno.dtype.to_string()}"
+                )
+
+            if isinstance(dtype, VectorType):
+                injected_args.append(VectorNdarray(dtype.n, dtype=dtype.dtype, shape=(2,) * ndim))
+            elif isinstance(dtype, MatrixType):
+                injected_args.append(MatrixNdarray(dtype.n, dtype.m, dtype=dtype.dtype, shape=(2,) * ndim))
+            else:
+                injected_args.append(ScalarNdarray(dtype, (2,) * ndim))
+        elif isinstance(anno, RWTextureType):
+            texture_shape = (2,) * anno.num_dimensions
+            fmt = anno.fmt
+            injected_args.append(Texture(fmt, texture_shape))
+        elif isinstance(anno, TextureType):
+            texture_shape = (2,) * anno.num_dimensions
+            injected_args.append(Texture(Format.rgba8, texture_shape))
+        elif isinstance(anno, MatrixType):
+            if symbolic_args is not None:
+                symbolic_mat_n = symbolic_args[i].element_shape[0]
+                symbolic_mat_m = symbolic_args[i].element_shape[1]
+
+                if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
+                    raise RuntimeError(
+                        f"Matrix dimension mismatch, expected ({anno.n}, {anno.m}) "
+                        f"but dispatched shape ({symbolic_mat_n}, {symbolic_mat_m})."
+                    )
+            injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
+        else:
+            if symbolic_args is not None:
+                dtype = symbolic_args[i].dtype()
+            else:
+                dtype = anno
+
+            if not check_type_match(dtype, anno):
+                raise TaichiCompilationError(
+                    f" Arg {arg.name}'s dtype {dtype.to_string()} doesn't match kernel's annotated dtype={anno.to_string()}"
+                )
+            # For primitive types, we can just inject a dummy value.
+            injected_args.append(0)
+    return injected_args
+
+
+def json_data_model(f):
+    """
+    Decorates a JSON data model. A JSON data model MUST NOT have any member
+    functions and it MUST be constructible from a JSON object.
+
+    This is merely a marker.
+    """
+    f._is_json_data_model = True
+    return f
+
+
+def is_json_data_model(cls) -> bool:
+    return hasattr(cls, "_is_json_data_model")
+
+
+def dump_json_data_model(x: object) -> Any:
+    if isinstance(x, (int, float, str, bool, type(None))):
+        return x
+    if isinstance(x, (list, tuple)):
+        return [dump_json_data_model(e) for e in x]
+    if isinstance(x, dict):
+        return {k: dump_json_data_model(v) for k, v in x.items()}
+    if is_json_data_model(x):
+        return {k: dump_json_data_model(v) for k, v in x.__dict__.items() if k != "_is_json_data_model"}
+    return x

taichi/assets/.git

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

taichi/examples/minimal.py

@@ -0,0 +1,28 @@
+import taichi 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()

taichi/experimental.py

@@ -0,0 +1,16 @@
+# type: ignore
+
+import warnings
+
+from taichi.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"]

taichi/graph/__init__.py

@@ -0,0 +1,3 @@
+# type: ignore
+
+from ._graph import *