gstaichi 2.1.1rc3__cp313-cp313-macosx_11_0_arm64.whl

gstaichi/lang/_wrap_inspect.py

@@ -0,0 +1,215 @@
+# type: ignore
+
+# GsTaichi's custom inspect module.
+# This module is used by GsTaichi's ast transformer to parse the source code.
+# Currently this module is aimed for working in the following modes:
+# 1. Usual Python/IPython mode, e.g. python script.py
+#    In this case we mainly rely on the built-in `inspect` module, except
+#    we need some hacks when we are in IPython mode and there is a cell magic.
+# 2. Blender's scripting mode, e.g. Users write GsTaichi code in the scripting
+#    window in Blender and press the run button. In this case we need to
+#    retrieve the source using Blender's `bpy.data.texts` and write it to a temp
+#    file so that the inspect module can parse.
+# 3. The interactive shell mode, e.g. Users directly type their code in the
+#    interactive shell. In this case we use `dill` to get the source.
+#
+# NB: Running GsTaichi in other modes are likely not supported.
+
+import atexit
+import inspect
+import os
+import tempfile
+from typing import Callable
+
+import dill
+from pydantic import BaseModel
+
+_builtin_getfile = inspect.getfile
+_builtin_findsource = inspect.findsource
+
+
+def _find_source_with_custom_getfile_func(func, obj):
+    """Use a custom function `func` to replace inspect's `getfile`, return the
+    source found by the new routine and restore the original `getfile` back.
+    """
+    inspect.getfile = func  # replace with our custom func
+    source = inspect.findsource(obj)
+    inspect.getfile = _builtin_getfile  # restore
+    return source
+
+
+def _blender_get_text_name(filename: str):
+    """Extract filename from path in the Blender mode."""
+    # In Blender's scripting mode, unsaved files are named
+    # like `/Text`, `/Text.001`, `/test.py`, etc.
+    # We simply remove this path seperator.
+    if filename.startswith(os.path.sep) and filename.count(os.path.sep) == 1:
+        return filename[1:]  # "/Text.001" --> "Text.001"
+
+    # Saved text files are named like `some-path/xxx.blend/Text` or
+    # `some-path/xxx.blend/test.py`
+    # We drop the path and extract the filename with extension.
+    index = filename.rfind(".blend" + os.path.sep)
+    if index != -1:
+        return filename[index + 7 :]  # "xxx.blend/test.py" --> "test.py"
+
+    return None
+
+
+def _blender_findsource(obj):
+    try:
+        import bpy  # pylint: disable=import-outside-toplevel
+    except:
+        raise ImportError("Not in Blender environment!")
+
+    # Inspect's built-in `getfile` returns the filename like
+    # `/Text`, `/Text.001`, `some-path/xxx.blend/test.py`
+    # This filename may not be a full valid path.
+    filename = _builtin_getfile(obj)
+    # Extract the text name without path
+    text_name = _blender_get_text_name(filename)
+    if text_name is None:
+        raise IOError("Object `{obj.__name__}` is not defined in a .blend file!")
+    # Get the lines of code via text_name
+    lines = bpy.data.texts[text_name].as_string()
+    # Now we have found the lines of code.
+    # We first check if they are already cached, to avoid file io in each query.
+    try:
+        filename = _blender_findsource._saved_inspect_cache[lines]  # pylint: disable=no-member
+    except KeyError:
+        # Save the code to a valid path.
+        fd, filename = tempfile.mkstemp(prefix="_Blender_", suffix=f"_{text_name}.py")
+        os.close(fd)
+
+        with open(filename, "w") as f:
+            f.write(lines)
+
+        _blender_findsource._saved_inspect_cache[lines] = filename  # pylint: disable=no-member
+        atexit.register(os.unlink, filename)  # Remove file when program exits
+
+    # Our custom getfile function
+    def wrapped_getfile(ob):
+        if id(ob) == id(obj):
+            return filename
+
+        return _builtin_getfile(ob)
+
+    return _find_source_with_custom_getfile_func(wrapped_getfile, obj)
+
+
+_blender_findsource._saved_inspect_cache = {}
+
+
+def _Python_IPython_findsource(obj):
+    try:
+        # In Python and IPython the builtin inspect would suffice in most cases
+        return _builtin_findsource(obj)
+    except IOError:
+        # Except that the cell has a magic command like %%time or %%timeit
+        # In this case the filename returned by the built-in's getfile is wrong,
+        # it becomes something like `<timed exec>` or `<magic-timeit>`.
+        filename = _builtin_getfile(obj)
+        if filename in {"<timed exec>", "<magic-timeit>"}:
+            try:
+                ip = get_ipython()
+                if ip is not None:
+                    # So we are in IPython's cell magic
+                    session_id = ip.history_manager.get_last_session_id()
+                    fd, filename = tempfile.mkstemp(prefix="_IPython_", suffix=f"_{session_id}.py")
+                    os.close(fd)
+                    # The latest lines of code can be retrived from here
+                    lines = ip.history_manager._i00
+
+                    # `lines` is a string that also contains the cell magic
+                    # command, we need to remove the magic command
+                    # (and spaces/sep around it) to obtain a valid Python code
+                    # snippet before saving it to a file
+                    index = lines.find("%time")
+                    lines_stripped = lines[index:]
+                    lines_stripped = lines_stripped.split(maxsplit=1)[1]
+
+                    with open(filename, "w") as f:
+                        f.write(lines_stripped)
+
+                    atexit.register(os.unlink, filename)  # Remove the file after the program exits
+                    func = lambda obj: filename
+                    return _find_source_with_custom_getfile_func(func, obj)
+
+            except ImportError:
+                pass
+        raise IOError(
+            f"Cannot find source code for Object: {obj}, it's likely \
+you are not running GsTaichi from command line or IPython."
+        )
+
+
+def _REPL_findsource(obj):
+    """Findsource in the interactive shell mode."""
+    return dill.source.findsource(obj)
+
+
+def _custom_findsource(obj):
+    try:
+        return _Python_IPython_findsource(obj)
+    except IOError:
+        try:
+            return _REPL_findsource(obj)
+        except:
+            try:
+                return _blender_findsource(obj)
+            except:
+                raise IOError(
+                    f"Cannot find source code for Object: {obj}, this \
+is possibly because of you are running GsTaichi in an environment that GsTaichi's own \
+inspect module cannot find the source. Please report an issue to help us fix: \
+https://github.com/taichi-dev/gstaichi/issues"
+                )
+
+
+class _InspectContextManager:
+    def __enter__(self):
+        inspect.findsource = _custom_findsource
+        return self
+
+    def __exit__(self, *_):
+        inspect.findsource = _builtin_findsource
+
+
+def getsourcelines(obj):
+    with _InspectContextManager():
+        return inspect.getsourcelines(obj)
+
+
+def getsourcefile(obj):
+    with _InspectContextManager():
+        ret = inspect.getsourcefile(obj)
+        if ret is None:
+            ret = inspect.getfile(obj)
+        return ret
+
+
+class FunctionSourceInfo(BaseModel):
+    function_name: str
+    filepath: str
+    start_lineno: int
+    end_lineno: int
+
+    class Config:
+        frozen = True
+
+
+def get_source_info_and_src(func: Callable) -> tuple[FunctionSourceInfo, list[str]]:
+    file = getsourcefile(func)
+    name = func.__name__
+    src, start_lineno = getsourcelines(func)
+    end_lineno = start_lineno + len(src) - 1
+    func_info = FunctionSourceInfo(
+        function_name=name,
+        filepath=file,
+        start_lineno=start_lineno,
+        end_lineno=end_lineno,
+    )
+    return (func_info, src)
+
+
+__all__ = ["getsourcelines", "getsourcefile", "get_source_info_and_src"]

gstaichi/lang/any_array.py

@@ -0,0 +1,99 @@
+# type: ignore
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.expr import Expr, make_expr_group
+from gstaichi.lang.util import gstaichi_scope
+from gstaichi.types.enums import Layout
+from gstaichi.types.ndarray_type import NdarrayTypeMetadata
+
+
+class AnyArray:
+    """Class for arbitrary arrays in Python AST.
+
+    Args:
+        ptr (gstaichi_python.Expr): A gstaichi_python.Expr wrapping a gstaichi_python.ExternalTensorExpression.
+        element_shape (Tuple[Int]): () if scalar elements (default), (n) if vector elements, and (n, m) if matrix elements.
+        layout (Layout): Memory layout.
+    """
+
+    def __init__(self, ptr):
+        assert ptr.is_external_tensor_expr()
+        self.ptr = ptr
+        self.ptr.type_check(impl.get_runtime().prog.config())
+
+    def element_shape(self):
+        return _ti_core.get_external_tensor_element_shape(self.ptr)
+
+    def layout(self):
+        # 0: scalar; 1: vector (SOA); 2: matrix (SOA); -1: vector
+        # (AOS); -2: matrix (AOS)
+        element_dim = _ti_core.get_external_tensor_element_dim(self.ptr)
+        if element_dim == 1 or element_dim == 2:
+            return Layout.SOA
+        return Layout.AOS
+
+    def get_type(self):
+        return NdarrayTypeMetadata(
+            _ti_core.get_external_tensor_element_type(self.ptr), None, _ti_core.get_external_tensor_needs_grad(self.ptr)
+        )  # AnyArray can take any shape
+
+    @property
+    @gstaichi_scope
+    def grad(self):
+        """Returns the gradient of this array."""
+        return AnyArray(_ti_core.make_external_tensor_grad_expr(self.ptr))
+
+    @property
+    @gstaichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        return [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr, i, dbg_info)) for i in range(dim)]
+
+    @gstaichi_scope
+    def _loop_range(self):
+        """Gets the corresponding gstaichi_python.Expr to serve as loop range.
+
+        Returns:
+            gstaichi_python.Expr: See above.
+        """
+        return self.ptr
+
+
+class AnyArrayAccess:
+    """Class for first-level access to AnyArray with Vector/Matrix elements in Python AST.
+
+    Args:
+        arr (AnyArray): See above.
+        indices_first (Tuple[Int]): Indices of first-level access.
+    """
+
+    def __init__(self, arr, indices_first):
+        self.arr = arr
+        self.indices_first = indices_first
+
+    @gstaichi_scope
+    def subscript(self, i, j):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+
+        indices_second = (i,) if len(self.arr.element_shape()) == 1 else (i, j)
+        if self.arr.layout() == Layout.SOA:
+            indices = indices_second + self.indices_first
+        else:
+            indices = self.indices_first + indices_second
+        return Expr(
+            ast_builder.expr_subscript(
+                self.arr.ptr,
+                make_expr_group(*indices),
+                _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()),
+            )
+        )
+
+
+__all__ = []

gstaichi/lang/ast/__init__.py

@@ -0,0 +1,5 @@
+# type: ignore
+
+from gstaichi.lang.ast.ast_transformer_utils import ASTTransformerContext
+from gstaichi.lang.ast.checkers import KernelSimplicityASTChecker
+from gstaichi.lang.ast.transform import transform_tree