gstaichi 0.1.23.dev0__cp310-cp310-macosx_15_0_arm64.whl → 0.1.25.dev0__cp310-cp310-macosx_15_0_arm64.whl

taichi/_ti_module/cppgen.py

@@ -1,309 +0,0 @@
-# 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

@@ -1,145 +0,0 @@
-# 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

@@ -1 +0,0 @@
-__version__ = '0.1.23.dev0'

taichi/ad/__init__.py

@@ -1,3 +0,0 @@
-# type: ignore
-
-from taichi.ad._ad import *

taichi/aot/__init__.py

@@ -1,12 +0,0 @@
-# type: ignore
-
-"""Taichi's AOT (ahead of time) module.
-
-Users can use Taichi as a GPU compute shader/kernel compiler by compiling their
-Taichi kernels into an AOT module.
-"""
-
-import taichi.aot.conventions
-from taichi.aot._export import export, export_as
-from taichi.aot.conventions.gfxruntime140 import GfxRuntime140
-from taichi.aot.module import Module

taichi/aot/_export.py

@@ -1,28 +0,0 @@
-# type: ignore
-
-from typing import Any, Dict, List, Optional
-
-
-class AotExportKernel:
-    def __init__(self, f, name: str, template_types: Dict[str, Any]) -> None:
-        self.kernel = f
-        self.name = name
-        self.template_types = template_types
-
-
-_aot_kernels: List[AotExportKernel] = []
-
-
-def export_as(name: str, *, template_types: Optional[Dict[str, Any]] = None):
-    def inner(f):
-        assert hasattr(f, "_is_wrapped_kernel"), "Only Taichi kernels can be exported"
-
-        record = AotExportKernel(f, name, template_types or {})
-        _aot_kernels.append(record)
-        return f
-
-    return inner
-
-
-def export(f):
-    return export_as(f.__name__)(f)

taichi/aot/conventions/__init__.py

@@ -1,3 +0,0 @@
-# type: ignore
-
-from taichi.aot.conventions import gfxruntime140

taichi/aot/conventions/gfxruntime140/__init__.py

@@ -1,38 +0,0 @@
-# type: ignore
-
-import json
-import zipfile
-from pathlib import Path
-from typing import Any, List
-
-from taichi.aot.conventions.gfxruntime140 import dr, sr
-
-
-class GfxRuntime140:
-    def __init__(self, metadata_json: Any, graphs_json: Any) -> None:
-        metadata = dr.from_json_metadata(metadata_json)
-        graphs = [dr.from_json_graph(x) for x in graphs_json]
-        self.metadata = sr.from_dr_metadata(metadata)
-        self.graphs = [sr.from_dr_graph(self.metadata, x) for x in graphs]
-
-    @staticmethod
-    def from_module(module_path: str) -> "GfxRuntime140":
-        if Path(module_path).is_file():
-            with zipfile.ZipFile(module_path) as z:
-                with z.open("metadata.json") as f:
-                    metadata_json = json.load(f)
-                with z.open("graphs.json") as f:
-                    graphs_json = json.load(f)
-        else:
-            with open(f"{module_path}/metadata.json") as f:
-                metadata_json = json.load(f)
-            with open(f"{module_path}/graphs.json") as f:
-                graphs_json = json.load(f)
-
-        return GfxRuntime140(metadata_json, graphs_json)
-
-    def to_metadata_json(self) -> Any:
-        return dr.to_json_metadata(sr.to_dr_metadata(self.metadata))
-
-    def to_graphs_json(self) -> List[Any]:
-        return [dr.to_json_graph(sr.to_dr_graph(x)) for x in self.graphs]