gstaichi 0.1.18.dev1__cp310-cp310-win_amd64.whl

taichi/_snode/fields_builder.py

@@ -0,0 +1,189 @@
+# type: ignore
+
+from typing import Any, Optional, Sequence, Union
+
+from taichi._lib import core as _ti_core
+from taichi._lib.core.taichi_python import SNode as SNodeCxx
+from taichi._snode.snode_tree import SNodeTree
+from taichi.lang import impl, snode
+from taichi.lang.exception import TaichiRuntimeError
+from taichi.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:`taichi.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:`taichi.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:`taichi.lang.snode.SNode.pointer`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise TaichiRuntimeError("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:`taichi.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:`taichi.lang.snode.SNode.dynamic`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise TaichiRuntimeError("Dynamic SNode is not supported on this backend.")
+
+        if dimension >= 2**31:
+            raise TaichiRuntimeError(
+                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 TaichiRuntimeError(
+                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:`taichi.lang.snode.SNode.bitmasked`"""
+        if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse):
+            raise TaichiRuntimeError("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:`taichi.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:`taichi.lang.snode.SNode.place`"""
+        self._check_not_finalized()
+        self.empty = False
+        self.root.place(*args, offset=offset)
+
+    def lazy_grad(self):
+        """Same as :func:`taichi.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:`taichi.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:`taichi.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_for_aot(self):
+        """Constructs the SNodeTree and compiles the type for AOT purpose."""
+        return self._finalize(raise_warning=False, compile_only=True)
+
+    def _finalize(self, raise_warning, compile_only):
+        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 TaichiRuntimeError("FieldsBuilder finalized")

taichi/_snode/snode_tree.py

@@ -0,0 +1,34 @@
+# type: ignore
+
+# The reason we import just the taichi.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 taichi.lang import impl
+from taichi.lang.exception import TaichiRuntimeError
+
+
+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 TaichiRuntimeError("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 TaichiRuntimeError("SNode tree has been destroyed")
+        return self.ptr.id()

taichi/_ti_module/__init__.py

@@ -0,0 +1,3 @@
+# type: ignore
+
+from taichi._ti_module.module import _main

taichi/_ti_module/cppgen.py

@@ -0,0 +1,309 @@
+# type: ignore
+
+from typing import Any, List, Optional, Set
+
+from taichi.aot.conventions.gfxruntime140 import GfxRuntime140, sr
+
+dtype2ctype = {
+    sr.DataType.f16: "half_t",
+    sr.DataType.f32: "float",
+    sr.DataType.f64: "double",
+    sr.DataType.i8: "int8_t",
+    sr.DataType.i16: "int16_t",
+    sr.DataType.i32: "int32_t",
+    sr.DataType.i64: "int64_t",
+    sr.DataType.u8: "uint8_t",
+    sr.DataType.u16: "uint16_t",
+    sr.DataType.u32: "uint32_t",
+    sr.DataType.u64: "uint64_t",
+}
+
+
+def check_arg(actual: str, expect: Any) -> List[str]:
+    out = []
+
+    expect = str(expect)
+    out += [
+        f"    if (value.{actual} != {expect}) {{",
+        f'      ti_set_last_error(TI_ERROR_INVALID_ARGUMENT, "value.{actual} != {expect}");',
+        "      return *this;",
+        "    }",
+    ]
+
+    return out
+
+
+def get_arg_dst(i: int, is_named: bool) -> str:
+    if is_named:
+        return f"args_[{i}].argument"
+    return f"args_[{i}]"
+
+
+def generate_scalar_assign(cls_name: str, i: int, arg_name: str, arg: sr.ArgumentScalar, is_named: bool) -> List[str]:
+    ctype = dtype2ctype[arg.dtype]
+
+    out = []
+
+    out += [
+        f"  {cls_name} &set_{arg_name}({ctype} value) {{",
+    ]
+
+    if is_named:
+        out += [
+            f'    args_[{i}].name = "{arg_name}";',
+        ]
+    if ctype == "float":
+        out += [
+            f"    {get_arg_dst(i, is_named)}.type = TI_ARGUMENT_TYPE_F32;",
+            f"    {get_arg_dst(i, is_named)}.value.f32 = value;",
+        ]
+    elif ctype == "int32_t":
+        out += [
+            f"    {get_arg_dst(i, is_named)}.type = TI_ARGUMENT_TYPE_I32;",
+            f"    {get_arg_dst(i, is_named)}.value.i32 = value;",
+        ]
+    else:
+        out += [
+            f"    {get_arg_dst(i, is_named)}.type = TI_ARGUMENT_TYPE_SCALAR;",
+            f"    {get_arg_dst(i, is_named)}.value.scalar.type = TI_DATA_TYPE_{arg.dtype.name.upper()};",
+            f"    *(({ctype}*)(&{get_arg_dst(i, is_named)}.value.scalar.value)) = value;",
+        ]
+        assert False, f"{ctype} is not a supported scalar type."
+
+    out += [
+        "    return *this;",
+        "  }",
+    ]
+    return out
+
+
+def generate_ndarray_assign(cls_name: str, i: int, arg_name: str, arg: sr.ArgumentNdArray, is_named: bool) -> List[str]:
+    out = []
+
+    out += [
+        f"  {cls_name} &set_{arg_name}(const TiNdArray &value) {{",
+    ]
+
+    out += check_arg("elem_type", f"TI_DATA_TYPE_{arg.dtype.name.upper()}")
+    out += check_arg("shape.dim_count", arg.ndim)
+    assert len(arg.element_shape) <= 16
+    out += check_arg("elem_shape.dim_count", len(arg.element_shape))
+    for j, dim in enumerate(arg.element_shape):
+        out += check_arg(f"elem_shape.dims[{j}]", dim)
+
+    if is_named:
+        out += [
+            f'    args_[{i}].name = "{arg_name}";',
+        ]
+    out += [
+        f"    {get_arg_dst(i, is_named)}.type = TI_ARGUMENT_TYPE_NDARRAY;",
+        f"    {get_arg_dst(i, is_named)}.value.ndarray = value;",
+        "    return *this;",
+        "  }",
+    ]
+    return out
+
+
+def generate_texture_assign(
+    cls_name: str,
+    i: int,
+    arg_name: str,
+    arg: sr.ArgumentTexture | sr.ArgumentRwTexture,
+    is_named: bool,
+) -> List[str]:
+    out = []
+
+    out += [
+        f"  {cls_name} &set_{arg_name}(const TiTexture &value) {{",
+    ]
+
+    assert arg.ndim in [1, 2, 3]
+    out += check_arg("dimension", f"TI_IMAGE_DIMENSION_{arg.ndim}D")
+    if isinstance(arg, sr.ArgumentRwTexture):
+        out += check_arg("format", f"TI_FORMAT_{arg.fmt.name.upper()}")
+
+    if is_named:
+        out += [
+            f'    args_[{i}].name = "{arg_name}";',
+        ]
+    out += [
+        f"    {get_arg_dst(i, is_named)}.type = TI_ARGUMENT_TYPE_TEXTURE;",
+        f"    {get_arg_dst(i, is_named)}.value.texture = value;",
+        "    return *this;",
+        "  }",
+    ]
+    return out
+
+
+def generate_kernel_args_builder(kernel: sr.Kernel) -> List[str]:
+    out = []
+
+    out += [
+        f"struct Kernel_{kernel.name} : public ti::Kernel {{",
+        f"  explicit Kernel_{kernel.name}(TiRuntime runtime, TiKernel kernel) :",
+        "    ti::Kernel(runtime, kernel) {",
+        f"    args_.resize({len(kernel.context.args)});",
+        "  }",
+        "",
+    ]
+
+    cls_name = f"Kernel_{kernel.name}"
+    for i, arg in enumerate(kernel.context.args):
+        arg_name = arg.name if arg.name else f"arg{i}"
+        if isinstance(arg, sr.ArgumentScalar):
+            out += generate_scalar_assign(cls_name, i, arg_name, arg, False)
+        elif isinstance(arg, sr.ArgumentNdArray):
+            out += generate_ndarray_assign(cls_name, i, arg_name, arg, False)
+        elif isinstance(arg, (sr.ArgumentTexture, sr.ArgumentRwTexture)):
+            out += generate_texture_assign(cls_name, i, arg_name, arg, False)
+        else:
+            assert False
+        out += [""]
+
+    out += [
+        "};",
+        "",
+    ]
+    return out
+
+
+def generate_graph_args_builder(graph: sr.Graph) -> List[str]:
+    out = []
+
+    out += [
+        f"struct ComputeGraph_{graph.name} : public ti::ComputeGraph {{",
+        f"  explicit ComputeGraph_{graph.name}(TiRuntime runtime, TiComputeGraph graph) :",
+        "    ti::ComputeGraph(runtime, graph) {",
+        f"    args_.resize({len(graph.args)});",
+        "  }",
+        "",
+    ]
+
+    cls_name = f"ComputeGraph_{graph.name}"
+    for i, arg in enumerate(graph.args):
+        arg_name = arg.name
+        if isinstance(arg.arg, sr.ArgumentScalar):
+            out += generate_scalar_assign(cls_name, i, arg_name, arg.arg, True)
+        elif isinstance(arg.arg, sr.ArgumentNdArray):
+            out += generate_ndarray_assign(cls_name, i, arg_name, arg.arg, True)
+        elif isinstance(arg.arg, (sr.ArgumentTexture, sr.ArgumentRwTexture)):
+            out += generate_texture_assign(cls_name, i, arg_name, arg.arg, True)
+        else:
+            assert False
+        out += [""]
+
+    out += [
+        "};",
+        "",
+    ]
+    return out
+
+
+def generate_module_content_repr(m: GfxRuntime140, module_name: str, cgraph_kernel_names: Set[str]) -> List[str]:
+    out = []
+
+    if module_name:
+        module_name = f"AotModule_{module_name}"
+    else:
+        module_name = "AotModule"
+
+    out += [
+        f"struct {module_name} : public ti::AotModule {{",
+        f"  explicit {module_name}(TiRuntime runtime, TiAotModule aot_module, bool should_destroy = true) :",
+        "    ti::AotModule(runtime, aot_module, should_destroy) {}",
+        "",
+        f"  static {module_name} load(TiRuntime runtime, const char *path) {{",
+        "    TiAotModule aot_module = ti_load_aot_module(runtime, path);",
+        f"    return {module_name}(runtime, aot_module, true);",
+        "  }",
+        f"  static {module_name} load(TiRuntime runtime, const std::string &path) {{",
+        f"    return {module_name}::load(runtime, path.c_str());",
+        "  }",
+        f"  static {module_name} create(TiRuntime runtime, const void *tcm, size_t size) {{",
+        "    TiAotModule aot_module = ti_create_aot_module(runtime, tcm, size);",
+        f"    return {module_name}(runtime, aot_module, true);",
+        "  }",
+        f"  static {module_name} create(TiRuntime runtime, const std::vector<uint8_t> &tcm) {{",
+        f"    return {module_name}::create(runtime, tcm.data(), tcm.size());",
+        "  }",
+        "",
+    ]
+    for kernel in m.metadata.kernels.values():
+        if kernel.name in cgraph_kernel_names:
+            continue
+        out += [
+            f"  Kernel_{kernel.name} get_kernel_{kernel.name}() const {{",
+            f'    return Kernel_{kernel.name}(runtime_, ti_get_aot_module_kernel(aot_module(), "{kernel.name}"));',
+            "  }",
+        ]
+    for graph in m.graphs:
+        out += [
+            f"  ComputeGraph_{graph.name} get_compute_graph_{graph.name}() const {{",
+            f'    return ComputeGraph_{graph.name}(runtime_, ti_get_aot_module_compute_graph(aot_module(), "{graph.name}"));',
+            "  }",
+        ]
+    out += [
+        "};",
+        "",
+    ]
+    return out
+
+
+def generate_module_content(m: GfxRuntime140, module_name: str) -> List[str]:
+    # This has all kernels including all the ones launched by compute graphs.
+    cgraph_kernel_names = set(dispatch.kernel.name for graph in m.graphs for dispatch in graph.dispatches)
+
+    out = []
+    for kernel in m.metadata.kernels.values():
+        if kernel.name in cgraph_kernel_names:
+            continue
+        out += generate_kernel_args_builder(kernel)
+
+    for graph in m.graphs:
+        out += generate_graph_args_builder(graph)
+
+    out += generate_module_content_repr(m, module_name, cgraph_kernel_names)
+
+    return out
+
+
+def generate_header(m: GfxRuntime140, module_name: str, namespace: str, tcm: Optional[bytes]) -> List[str]:
+    out = []
+
+    out += [
+        "// THIS IS A GENERATED HEADER; PLEASE DO NOT MODIFY.",
+        "#pragma once",
+        "#include <vector>",
+        "#include <string>",
+        "#include <taichi/cpp/taichi.hpp>",
+        "",
+    ]
+
+    if namespace:
+        out += [
+            f"namespace {namespace} {{",
+            "",
+        ]
+
+    if tcm is not None:
+        tcm_bytes = [x for x in tcm]
+
+        out += [
+            f"static const uint8_t {module_name}_tcm[{len(tcm_bytes)}] = {{",
+        ]
+
+        out += [f"  {', '.join(str(x) for x in tcm_bytes[i:i + 8])}," for i in range(0, len(tcm_bytes), 8)]
+
+        out += [
+            "};",
+            "",
+            f"static const size_t {module_name}_tcm_size = {len(tcm_bytes)};",
+            "",
+        ]
+
+    out += generate_module_content(m, module_name)
+
+    if namespace:
+        out += [f"}} // namespace {namespace}", ""]
+
+    return out

taichi/_ti_module/module.py

@@ -0,0 +1,145 @@
+# type: ignore
+
+import argparse
+import runpy
+from pathlib import Path
+from typing import List
+
+import taichi
+from taichi._ti_module.cppgen import generate_header
+from taichi.aot._export import _aot_kernels
+from taichi.aot.conventions.gfxruntime140 import GfxRuntime140
+from taichi.aot.module import Module
+from taichi.types.ndarray_type import NdarrayType
+from taichi.types.primitive_types import integer_type_ids, real_type_ids
+from taichi.types.texture_type import RWTextureType, TextureType
+
+
+def module_cppgen(parser: argparse.ArgumentParser):
+    """Generate C++ headers for Taichi modules."""
+    parser.add_argument("MODOLE", help="Path to the module directory.")
+    parser.add_argument("-n", "--namespace", type=str, help="C++ namespace if wanted.")
+    parser.add_argument(
+        "-m",
+        "--module-name",
+        type=str,
+        help="Module name to be a part of the module class. By default, it's the directory name.",
+        default=None,
+    )
+    parser.add_argument("-o", "--output", type=str, help="Output C++ header path.", default="module.h")
+    parser.add_argument(
+        "--bin2c",
+        help="Save the entire TCM archive to an in-memory buffer. This flag is ignored if the module is not a TCM archive",
+        action="store_true",
+    )
+    parser.set_defaults(func=module_cppgen_impl)
+
+
+def module_cppgen_impl(a):
+    module_path = a.MODOLE
+
+    print(f"Generating C++ header for Taichi module: {Path(module_path).absolute()}")
+
+    tcm = None
+    if a.bin2c and module_path.endswith(".tcm"):
+        with open(module_path, "rb") as f:
+            tcm = f.read()
+
+    if a.module_name:
+        module_name = a.module_name
+    else:
+        module_name = Path(module_path).name
+        if module_name.endswith(".tcm"):
+            module_name = module_name[:-4]
+
+    m = GfxRuntime140.from_module(module_path)
+
+    out = generate_header(m, module_name, a.namespace, tcm)
+
+    with open(a.output, "w") as f:
+        f.write("\n".join(out))
+
+    print(f"Module header is saved to: {Path(a.output).absolute()}")
+
+
+def module_build(parser: argparse.ArgumentParser):
+    """Build Taichi modules from python scripts."""
+    parser.add_argument("SOURCE", help="Path to the Taichi program source (Python script).")
+    parser.add_argument("-o", "--output", type=str, help="Output module path.", default=None)
+    parser.set_defaults(func=module_build_impl)
+
+
+def module_build_impl(a):
+    source_path = a.SOURCE
+    module_path = a.output
+
+    source_path = Path(source_path)
+    assert source_path.name.endswith(".py"), "Source must be a Python script."
+    if module_path is None:
+        module_path = f"{source_path.name[:-3]}.tcm"
+    module_path = Path(module_path)
+
+    print(f"Building Taichi module: {source_path}")
+    print()
+
+    d = runpy.run_path(str(source_path), run_name="__main__")
+    print()
+
+    required_caps = d["REQUIRED_CAPS"] if "REQUIRED_CAPS" in d else []
+    assert isinstance(required_caps, list), "REQUIRED_CAPS must be a list."
+
+    if required_caps:
+        print("Module requires the following capabilities:")
+        for cap in required_caps:
+            print(f"  - {cap}")
+        print()
+
+    m = Module(caps=required_caps)
+    for record in _aot_kernels:
+        print("Added kernel:", record.name)
+        template_args = None
+        if record.template_types:
+            print("  Template types:")
+            template_args = {}
+            for k, v in record.template_types.items():
+                print(f"    - {k}: {v}")
+                # TODO: (penguinliong) Remove this hack. It's not properly
+                # working with unusual numeric types like f16 or i64.
+                if isinstance(v, int) or id(v) in integer_type_ids:
+                    value = 0
+                elif isinstance(v, float) or id(v) in real_type_ids:
+                    value = 0.0
+                elif isinstance(v, NdarrayType):
+                    if v.ndim is None or v.ndim <= 0:
+                        raise ValueError("Ndarray template type must specify a non-zero dimension.")
+                    value = taichi.ndarray(v.dtype, (1,) * v.ndim)
+                elif isinstance(v, TextureType):
+                    value = taichi.Texture(taichi.Format.rgba8, (4,) * v.num_dimensions)
+                elif isinstance(v, RWTextureType):
+                    value = taichi.Texture(v.fmt, (4,) * v.num_dimensions)
+                else:
+                    raise ValueError(f"Unsupported template type: {type(v)}")
+                template_args[k] = value
+        m.add_kernel(record.kernel, template_args)
+    print()
+
+    if module_path.name.endswith(".tcm"):
+        m.archive(str(module_path))
+    else:
+        m.save(str(module_path))
+
+    print(f"Module is archive to: {module_path}")
+    print()
+
+
+def _main(arguments: List[str]):
+    """Taichi module tools."""
+    parser = argparse.ArgumentParser(prog="ti module", description=_main.__doc__)
+    subparsers = parser.add_subparsers(title="Taichi module manager commands", required=True)
+
+    cppgen_parser = subparsers.add_parser("cppgen", help=module_cppgen.__doc__)
+    build_parser = subparsers.add_parser("build", help=module_build.__doc__)
+    module_cppgen(cppgen_parser)
+    module_build(build_parser)
+    args = parser.parse_args(arguments)
+    args.func(args)

taichi/_version.py

@@ -0,0 +1 @@
+__version__ = '0.1.18.dev1'