gstaichi 0.1.18.dev1__cp310-cp310-win_amd64.whl

taichi/graph/_graph.py

@@ -0,0 +1,292 @@
+# type: ignore
+
+import warnings
+from typing import Any, Dict, List
+
+from taichi._lib import core as _ti_core
+from taichi.aot.utils import produce_injected_args
+from taichi.lang import impl, kernel_impl
+from taichi.lang._ndarray import Ndarray
+from taichi.lang._texture import Texture
+from taichi.lang.exception import TaichiRuntimeError
+from taichi.lang.matrix import Matrix, MatrixType
+from taichi.types import enums
+from taichi.types.texture_type import FORMAT2TY_CH, TY_CH2FORMAT
+
+ArgKind = _ti_core.ArgKind
+
+
+def gen_cpp_kernel(kernel_fn, args):
+    kernel = kernel_fn._primal
+    assert isinstance(kernel, kernel_impl.Kernel)
+    injected_args = produce_injected_args(kernel, symbolic_args=args)
+    key = kernel.ensure_compiled(*injected_args)
+    return kernel.compiled_kernels[key]
+
+
+def flatten_args(args):
+    unzipped_args = []
+    # Tuple for matrix args
+    # FIXME remove this when native Matrix type is ready
+    for arg in args:
+        if isinstance(arg, list):
+            for sublist in arg:
+                unzipped_args.extend(sublist)
+        else:
+            unzipped_args.append(arg)
+    return unzipped_args
+
+
+class Sequential:
+    def __init__(self, seq):
+        self.seq_ = seq
+
+    def dispatch(self, kernel_fn, *args):
+        kernel_cpp = gen_cpp_kernel(kernel_fn, args)
+        unzipped_args = flatten_args(args)
+        self.seq_.dispatch(kernel_cpp, unzipped_args)
+
+
+class GraphBuilder:
+    def __init__(self):
+        self._graph_builder = _ti_core.GraphBuilder()
+
+    def dispatch(self, kernel_fn, *args):
+        kernel_cpp = gen_cpp_kernel(kernel_fn, args)
+        unzipped_args = flatten_args(args)
+        self._graph_builder.dispatch(kernel_cpp, unzipped_args)
+
+    def create_sequential(self):
+        return Sequential(self._graph_builder.create_sequential())
+
+    def append(self, node):
+        # TODO: support appending dispatch node as well.
+        assert isinstance(node, Sequential)
+        self._graph_builder.seq().append(node.seq_)
+
+    def compile(self):
+        return Graph(self._graph_builder.compile())
+
+
+class Graph:
+    def __init__(self, compiled_graph) -> None:
+        self._compiled_graph = compiled_graph
+
+    def run(self, args):
+        # Support native python numerical types (int, float), Ndarray.
+        # Taichi Matrix types are flattened into (int, float) arrays.
+        # TODO diminish the flatten behavior when Matrix becomes a Taichi native type.
+        flattened = {}
+        for k, v in args.items():
+            if isinstance(v, Ndarray):
+                flattened[k] = v.arr
+            elif isinstance(v, Texture):
+                flattened[k] = v.tex
+            elif isinstance(v, Matrix):
+                flattened[k] = v.entries
+            elif isinstance(v, (int, float)):
+                flattened[k] = v
+            else:
+                raise TaichiRuntimeError(
+                    f"Only python int, float, ti.Matrix and ti.Ndarray are supported as runtime arguments but got {type(v)}"
+                )
+        self._compiled_graph.jit_run(impl.get_runtime().prog.config(), flattened)
+
+
+def _deprecate_arg_args(kwargs: Dict[str, Any]):
+    if "field_dim" in kwargs:
+        warnings.warn(
+            "The field_dim argument for ndarray will be deprecated in v1.6.0, use ndim instead.",
+            DeprecationWarning,
+        )
+        if "ndim" in kwargs:
+            raise TaichiRuntimeError(
+                "field_dim is deprecated, please do not specify field_dim and ndim at the same time."
+            )
+        kwargs["ndim"] = kwargs["field_dim"]
+        del kwargs["field_dim"]
+    tag = kwargs["tag"]
+
+    if tag == ArgKind.SCALAR:
+        if "element_shape" in kwargs:
+            raise TaichiRuntimeError(
+                "The element_shape argument for scalar is deprecated in v1.6.0, and is removed in v1.7.0. "
+                "Please remove them."
+            )
+
+    if tag == ArgKind.NDARRAY:
+        if "element_shape" not in kwargs:
+            if "dtype" in kwargs:
+                dtype = kwargs["dtype"]
+                if isinstance(dtype, MatrixType):
+                    kwargs["dtype"] = dtype.dtype
+                    kwargs["element_shape"] = dtype.get_shape()
+                else:
+                    kwargs["element_shape"] = ()
+        else:
+            raise TaichiRuntimeError(
+                "The element_shape argument for ndarray is deprecated in v1.6.0, and it is removed in v1.7.0. "
+                "Please use vector or matrix data type instead."
+            )
+
+    if tag == ArgKind.RWTEXTURE or tag == ArgKind.TEXTURE:
+        if "dtype" in kwargs:
+            warnings.warn(
+                "The dtype argument for texture will be deprecated in v1.6.0, use format instead.",
+                DeprecationWarning,
+            )
+            del kwargs["dtype"]
+
+        if "shape" in kwargs:
+            raise TaichiRuntimeError(
+                "The shape argument for texture is deprecated in v1.6.0, and it is removed in v1.7.0. "
+                "Please use ndim instead. (Note that you no longer need the exact texture size.)"
+            )
+
+        if "channel_format" in kwargs or "num_channels" in kwargs:
+            if "fmt" in kwargs:
+                raise TaichiRuntimeError(
+                    "channel_format and num_channels are deprecated, please do not specify channel_format/num_channels and fmt at the same time."
+                )
+            if tag == ArgKind.RWTEXTURE:
+                fmt = TY_CH2FORMAT[(kwargs["channel_format"], kwargs["num_channels"])]
+                kwargs["fmt"] = fmt
+                raise TaichiRuntimeError(
+                    "The channel_format and num_channels arguments for texture are deprecated in v1.6.0, "
+                    "and they are removed in v1.7.0. Please use fmt instead."
+                )
+            else:
+                raise TaichiRuntimeError(
+                    "The channel_format and num_channels arguments are no longer required for non-RW textures "
+                    "since v1.6.0, and they are removed in v1.7.0. Please remove them."
+                )
+
+
+def _check_args(kwargs: Dict[str, Any], allowed_kwargs: List[str]):
+    for k, v in kwargs.items():
+        if k not in allowed_kwargs:
+            raise TaichiRuntimeError(
+                f"Invalid argument: {k}, you can only create a graph argument with: {allowed_kwargs}"
+            )
+        if k == "tag":
+            if not isinstance(v, ArgKind):
+                raise TaichiRuntimeError(f"tag must be a ArgKind variant, but found {type(v)}.")
+        if k == "name":
+            if not isinstance(v, str):
+                raise TaichiRuntimeError(f"name must be a string, but found {type(v)}.")
+
+
+def _make_arg_scalar(kwargs: Dict[str, Any]):
+    allowed_kwargs = [
+        "tag",
+        "name",
+        "dtype",
+    ]
+    _check_args(kwargs, allowed_kwargs)
+    name = kwargs["name"]
+    dtype = kwargs["dtype"]
+    if isinstance(dtype, MatrixType):
+        raise TaichiRuntimeError(f"Tag ArgKind.SCALAR must specify a scalar type, but found {type(dtype)}.")
+    return _ti_core.Arg(ArgKind.SCALAR, name, dtype, 0, [])
+
+
+def _make_arg_ndarray(kwargs: Dict[str, Any]):
+    allowed_kwargs = [
+        "tag",
+        "name",
+        "dtype",
+        "ndim",
+        "element_shape",
+    ]
+    _check_args(kwargs, allowed_kwargs)
+    name = kwargs["name"]
+    ndim = kwargs["ndim"]
+    dtype = kwargs["dtype"]
+    element_shape = kwargs["element_shape"]
+    if isinstance(dtype, MatrixType):
+        raise TaichiRuntimeError(f"Tag ArgKind.NDARRAY must specify a scalar type, but found {dtype}.")
+    return _ti_core.Arg(ArgKind.NDARRAY, name, dtype, ndim, element_shape)
+
+
+def _make_arg_matrix(kwargs: Dict[str, Any]):
+    allowed_kwargs = [
+        "tag",
+        "name",
+        "dtype",
+    ]
+    _check_args(kwargs, allowed_kwargs)
+    name = kwargs["name"]
+    dtype = kwargs["dtype"]
+    if not isinstance(dtype, MatrixType):
+        raise TaichiRuntimeError(f"Tag ArgKind.MATRIX must specify matrix type, but got {dtype}.")
+    return _ti_core.Arg(ArgKind.MATRIX, f"{name}", dtype.dtype, 0, [dtype.n, dtype.m])
+
+
+def _make_arg_texture(kwargs: Dict[str, Any]):
+    allowed_kwargs = [
+        "tag",
+        "name",
+        "ndim",
+    ]
+    _check_args(kwargs, allowed_kwargs)
+    name = kwargs["name"]
+    ndim = kwargs["ndim"]
+    return _ti_core.Arg(ArgKind.TEXTURE, name, impl.f32, 4, [2] * ndim)
+
+
+def _make_arg_rwtexture(kwargs: Dict[str, Any]):
+    allowed_kwargs = [
+        "tag",
+        "name",
+        "ndim",
+        "fmt",
+    ]
+    _check_args(kwargs, allowed_kwargs)
+    name = kwargs["name"]
+    ndim = kwargs["ndim"]
+    fmt = kwargs["fmt"]
+    if fmt == enums.Format.unknown:
+        raise TaichiRuntimeError(f"Tag ArgKind.RWTEXTURE must specify a valid color format, but found {fmt}.")
+    channel_format, num_channels = FORMAT2TY_CH[fmt]
+    return _ti_core.Arg(ArgKind.RWTEXTURE, name, channel_format, num_channels, [2] * ndim)
+
+
+def _make_arg(kwargs: Dict[str, Any]):
+    assert "tag" in kwargs
+    _deprecate_arg_args(kwargs)
+    proc = {
+        ArgKind.SCALAR: _make_arg_scalar,
+        ArgKind.NDARRAY: _make_arg_ndarray,
+        ArgKind.MATRIX: _make_arg_matrix,
+        ArgKind.TEXTURE: _make_arg_texture,
+        ArgKind.RWTEXTURE: _make_arg_rwtexture,
+    }
+    tag = kwargs["tag"]
+    return proc[tag](kwargs)
+
+
+def _kwarg_rewriter(args, kwargs):
+    for i, arg in enumerate(args):
+        rewrite_map = {
+            0: "tag",
+            1: "name",
+            2: "dtype",
+            3: "ndim",
+            4: "field_dim",
+            5: "element_shape",
+            6: "channel_format",
+            7: "shape",
+            8: "num_channels",
+        }
+        if i in rewrite_map:
+            kwargs[rewrite_map[i]] = arg
+        else:
+            raise TaichiRuntimeError(f"Unexpected {i}th positional argument")
+
+
+def Arg(*args, **kwargs):
+    _kwarg_rewriter(args, kwargs)
+    return _make_arg(kwargs)
+
+
+__all__ = ["GraphBuilder", "Graph", "Arg", "ArgKind"]

taichi/lang/__init__.py

@@ -0,0 +1,50 @@
+# type: ignore
+
+from taichi.lang import impl, simt
+from taichi.lang._ndarray import *
+from taichi.lang._ndrange import ndrange
+from taichi.lang._texture import Texture
+from taichi.lang.argpack import *
+from taichi.lang.exception import *
+from taichi.lang.field import *
+from taichi.lang.impl import *
+from taichi.lang.kernel_impl import *
+from taichi.lang.matrix import *
+from taichi.lang.mesh import *
+from taichi.lang.misc import *  # pylint: disable=W0622
+from taichi.lang.ops import *  # pylint: disable=W0622
+from taichi.lang.runtime_ops import *
+from taichi.lang.snode import *
+from taichi.lang.source_builder import *
+from taichi.lang.struct import *
+from taichi.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",
+    ]
+]

taichi/lang/_ndarray.py

@@ -0,0 +1,348 @@
+# type: ignore
+
+import numpy as np
+
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.exception import TaichiIndexError
+from taichi.lang.util import cook_dtype, get_traceback, python_scope, to_numpy_type
+from taichi.types import primitive_types
+from taichi.types.enums import Layout
+from taichi.types.ndarray_type import NdarrayTypeMetadata
+from taichi.types.utils import is_real, is_signed
+
+
+class Ndarray:
+    """Taichi ndarray class.
+
+    Args:
+        dtype (DataType): Data type of each value.
+        shape (Tuple[int]): Shape of the Ndarray.
+    """
+
+    def __init__(self):
+        self.host_accessor = None
+        self.shape = None
+        self.element_type = None
+        self.dtype = None
+        self.arr = None
+        self.layout = Layout.AOS
+        self.grad = None
+
+    def get_type(self):
+        return NdarrayTypeMetadata(self.element_type, self.shape, self.grad is not None)
+
+    @property
+    def element_shape(self):
+        """Gets ndarray element shape.
+
+        Returns:
+            Tuple[Int]: Ndarray element shape.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def __setitem__(self, key, value):
+        """Sets ndarray element in Python scope.
+
+        Args:
+            key (Union[List[int], int, None]): Coordinates of the ndarray element.
+            value (element type): Value to set.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def __getitem__(self, key):
+        """Gets ndarray element in Python scope.
+
+        Args:
+            key (Union[List[int], int, None]): Coordinates of the ndarray element.
+
+        Returns:
+            element type: Value retrieved.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def fill(self, val):
+        """Fills ndarray with a specific scalar value.
+
+        Args:
+            val (Union[int, float]): Value to fill.
+        """
+        if impl.current_cfg().arch != _ti_core.Arch.cuda and impl.current_cfg().arch != _ti_core.Arch.x64:
+            self._fill_by_kernel(val)
+        elif _ti_core.is_tensor(self.element_type):
+            self._fill_by_kernel(val)
+        elif self.dtype == primitive_types.f32:
+            impl.get_runtime().prog.fill_float(self.arr, val)
+        elif self.dtype == primitive_types.i32:
+            impl.get_runtime().prog.fill_int(self.arr, val)
+        elif self.dtype == primitive_types.u32:
+            impl.get_runtime().prog.fill_uint(self.arr, val)
+        else:
+            self._fill_by_kernel(val)
+
+    @python_scope
+    def _ndarray_to_numpy(self):
+        """Converts ndarray to a numpy array.
+
+        Returns:
+            numpy.ndarray: The result numpy array.
+        """
+        arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype))
+        from taichi._kernels import ndarray_to_ext_arr  # pylint: disable=C0415
+
+        ndarray_to_ext_arr(self, arr)
+        impl.get_runtime().sync()
+        return arr
+
+    @python_scope
+    def _ndarray_matrix_to_numpy(self, as_vector):
+        """Converts matrix ndarray to a numpy array.
+
+        Returns:
+            numpy.ndarray: The result numpy array.
+        """
+        arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype))
+        from taichi._kernels import ndarray_matrix_to_ext_arr  # pylint: disable=C0415
+
+        layout_is_aos = 1
+        ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector)
+        impl.get_runtime().sync()
+        return arr
+
+    @python_scope
+    def _ndarray_from_numpy(self, arr):
+        """Loads all values from a numpy array.
+
+        Args:
+            arr (numpy.ndarray): The source numpy array.
+        """
+        if not isinstance(arr, np.ndarray):
+            raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided")
+        if tuple(self.arr.total_shape()) != tuple(arr.shape):
+            raise ValueError(f"Mismatch shape: {tuple(self.arr.shape)} expected, but {tuple(arr.shape)} provided")
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+
+        from taichi._kernels import ext_arr_to_ndarray  # pylint: disable=C0415
+
+        ext_arr_to_ndarray(arr, self)
+        impl.get_runtime().sync()
+
+    @python_scope
+    def _ndarray_matrix_from_numpy(self, arr, as_vector):
+        """Loads all values from a numpy array.
+
+        Args:
+            arr (numpy.ndarray): The source numpy array.
+        """
+        if not isinstance(arr, np.ndarray):
+            raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided")
+        if tuple(self.arr.total_shape()) != tuple(arr.shape):
+            raise ValueError(
+                f"Mismatch shape: {tuple(self.arr.total_shape())} expected, but {tuple(arr.shape)} provided"
+            )
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+
+        from taichi._kernels import ext_arr_to_ndarray_matrix  # pylint: disable=C0415
+
+        layout_is_aos = 1
+        ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector)
+        impl.get_runtime().sync()
+
+    @python_scope
+    def _get_element_size(self):
+        """Returns the size of one element in bytes.
+
+        Returns:
+            Size in bytes.
+        """
+        return self.arr.element_size()
+
+    @python_scope
+    def _get_nelement(self):
+        """Returns the total number of elements.
+
+        Returns:
+            Total number of elements.
+        """
+        return self.arr.nelement()
+
+    @python_scope
+    def copy_from(self, other):
+        """Copies all elements from another ndarray.
+
+        The shape of the other ndarray needs to be the same as `self`.
+
+        Args:
+            other (Ndarray): The source ndarray.
+        """
+        assert isinstance(other, Ndarray)
+        assert tuple(self.arr.shape) == tuple(other.arr.shape)
+        from taichi._kernels import ndarray_to_ndarray  # pylint: disable=C0415
+
+        ndarray_to_ndarray(self, other)
+        impl.get_runtime().sync()
+
+    def _set_grad(self, grad):
+        """Sets the gradient ndarray.
+
+        Args:
+            grad (Ndarray): The gradient ndarray.
+        """
+        self.grad = grad
+
+    def __deepcopy__(self, memo=None):
+        """Copies all elements to a new ndarray.
+
+        Returns:
+            Ndarray: The result ndarray.
+        """
+        raise NotImplementedError()
+
+    def _fill_by_kernel(self, val):
+        """Fills ndarray with a specific scalar value using a ti.kernel.
+
+        Args:
+            val (Union[int, float]): Value to fill.
+        """
+        raise NotImplementedError()
+
+    @python_scope
+    def _pad_key(self, key):
+        if key is None:
+            key = ()
+        if not isinstance(key, (tuple, list)):
+            key = (key,)
+        if len(key) != len(self.arr.total_shape()):
+            raise TaichiIndexError(f"{len(self.arr.total_shape())}d ndarray indexed with {len(key)}d indices: {key}")
+        return key
+
+    @python_scope
+    def _initialize_host_accessor(self):
+        if self.host_accessor:
+            return
+        impl.get_runtime().materialize()
+        self.host_accessor = NdarrayHostAccessor(self.arr)
+
+
+class ScalarNdarray(Ndarray):
+    """Taichi ndarray with scalar elements.
+
+    Args:
+        dtype (DataType): Data type of each value.
+        shape (Tuple[int]): Shape of the ndarray.
+    """
+
+    def __init__(self, dtype, arr_shape):
+        super().__init__()
+        self.dtype = cook_dtype(dtype)
+        self.arr = impl.get_runtime().prog.create_ndarray(
+            self.dtype, arr_shape, layout=Layout.NULL, zero_fill=True, dbg_info=_ti_core.DebugInfo(get_traceback())
+        )
+        self.shape = tuple(self.arr.shape)
+        self.element_type = dtype
+
+    def __del__(self):
+        if (
+            impl is not None
+            and impl.get_runtime is not None
+            and impl.get_runtime() is not None
+            and impl.get_runtime().prog is not None
+        ):
+            impl.get_runtime().prog.delete_ndarray(self.arr)
+
+    @property
+    def element_shape(self):
+        return ()
+
+    @python_scope
+    def __setitem__(self, key, value):
+        self._initialize_host_accessor()
+        self.host_accessor.setter(value, *self._pad_key(key))
+
+    @python_scope
+    def __getitem__(self, key):
+        self._initialize_host_accessor()
+        return self.host_accessor.getter(*self._pad_key(key))
+
+    @python_scope
+    def to_numpy(self):
+        return self._ndarray_to_numpy()
+
+    @python_scope
+    def from_numpy(self, arr):
+        self._ndarray_from_numpy(arr)
+
+    def __deepcopy__(self, memo=None):
+        ret_arr = ScalarNdarray(self.dtype, self.shape)
+        ret_arr.copy_from(self)
+        return ret_arr
+
+    def _fill_by_kernel(self, val):
+        from taichi._kernels import fill_ndarray  # pylint: disable=C0415
+
+        fill_ndarray(self, val)
+
+    def __repr__(self):
+        return "<ti.ndarray>"
+
+
+class NdarrayHostAccessor:
+    def __init__(self, ndarray):
+        dtype = ndarray.element_data_type()
+        if is_real(dtype):
+
+            def getter(*key):
+                return ndarray.read_float(key)
+
+            def setter(value, *key):
+                ndarray.write_float(key, value)
+
+        else:
+            if is_signed(dtype):
+
+                def getter(*key):
+                    return ndarray.read_int(key)
+
+            else:
+
+                def getter(*key):
+                    return ndarray.read_uint(key)
+
+            def setter(value, *key):
+                ndarray.write_int(key, value)
+
+        self.getter = getter
+        self.setter = setter
+
+
+class NdarrayHostAccess:
+    """Class for accessing VectorNdarray/MatrixNdarray in Python scope.
+    Args:
+        arr (Union[VectorNdarray, MatrixNdarray]): See above.
+        indices_first (Tuple[Int]): Indices of first-level access (coordinates in the field).
+        indices_second (Tuple[Int]): Indices of second-level access (indices in the vector/matrix).
+    """
+
+    def __init__(self, arr, indices_first, indices_second):
+        self.ndarr = arr
+        self.arr = arr.arr
+        self.indices = indices_first + indices_second
+
+        def getter():
+            self.ndarr._initialize_host_accessor()
+            return self.ndarr.host_accessor.getter(*self.ndarr._pad_key(self.indices))
+
+        def setter(value):
+            self.ndarr._initialize_host_accessor()
+            self.ndarr.host_accessor.setter(value, *self.ndarr._pad_key(self.indices))
+
+        self.getter = getter
+        self.setter = setter
+
+
+__all__ = ["Ndarray", "ScalarNdarray"]