gstaichi 0.1.18.dev1__cp310-cp310-macosx_15_0_arm64.whl

taichi/algorithms/_algorithms.py

@@ -0,0 +1,117 @@
+# type: ignore
+
+from taichi._kernels import (
+    blit_from_field_to_field,
+    scan_add_inclusive,
+    sort_stage,
+    uniform_add,
+    warp_shfl_up_i32,
+)
+from taichi.lang.impl import current_cfg, field
+from taichi.lang.kernel_impl import data_oriented
+from taichi.lang.misc import cuda, vulkan
+from taichi.lang.runtime_ops import sync
+from taichi.lang.simt import subgroup
+from taichi.types.primitive_types import i32
+
+
+def parallel_sort(keys, values=None):
+    """Odd-even merge sort
+
+    References:
+        https://developer.nvidia.com/gpugems/gpugems2/part-vi-simulation-and-numerical-algorithms/chapter-46-improved-gpu-sorting
+        https://en.wikipedia.org/wiki/Batcher_odd%E2%80%93even_mergesort
+    """
+    N = keys.shape[0]
+
+    num_stages = 0
+    p = 1
+    while p < N:
+        k = p
+        while k >= 1:
+            invocations = int((N - k - k % p) / (2 * k)) + 1
+            if values is None:
+                sort_stage(keys, 0, keys, N, p, k, invocations)
+            else:
+                sort_stage(keys, 1, values, N, p, k, invocations)
+            num_stages += 1
+            sync()
+            k = int(k / 2)
+        p = int(p * 2)
+
+
+@data_oriented
+class PrefixSumExecutor:
+    """Parallel Prefix Sum (Scan) Helper
+
+    Use this helper to perform an inclusive in-place's parallel prefix sum.
+
+    References:
+        https://developer.download.nvidia.com/compute/cuda/1.1-Beta/x86_website/projects/scan/doc/scan.pdf
+        https://github.com/NVIDIA/cuda-samples/blob/master/Samples/2_Concepts_and_Techniques/shfl_scan/shfl_scan.cu
+    """
+
+    def __init__(self, length):
+        self.sorting_length = length
+
+        BLOCK_SZ = 64
+        GRID_SZ = int((length + BLOCK_SZ - 1) / BLOCK_SZ)
+
+        # Buffer position and length
+        # This is a single buffer implementation for ease of aot usage
+        ele_num = length
+        self.ele_nums = [ele_num]
+        start_pos = 0
+        self.ele_nums_pos = [start_pos]
+
+        while ele_num > 1:
+            ele_num = int((ele_num + BLOCK_SZ - 1) / BLOCK_SZ)
+            self.ele_nums.append(ele_num)
+            start_pos += BLOCK_SZ * ele_num
+            self.ele_nums_pos.append(start_pos)
+
+        self.large_arr = field(i32, shape=start_pos)
+
+    def run(self, input_arr):
+        length = self.sorting_length
+        ele_nums = self.ele_nums
+        ele_nums_pos = self.ele_nums_pos
+
+        if input_arr.dtype != i32:
+            raise RuntimeError("Only ti.i32 type is supported for prefix sum.")
+
+        if current_cfg().arch == cuda:
+            inclusive_add = warp_shfl_up_i32
+        elif current_cfg().arch == vulkan:
+            inclusive_add = subgroup.inclusive_add
+        else:
+            raise RuntimeError(f"{str(current_cfg().arch)} is not supported for prefix sum.")
+
+        blit_from_field_to_field(self.large_arr, input_arr, 0, length)
+
+        # Kogge-Stone construction
+        for i in range(len(ele_nums) - 1):
+            if i == len(ele_nums) - 2:
+                scan_add_inclusive(
+                    self.large_arr,
+                    ele_nums_pos[i],
+                    ele_nums_pos[i + 1],
+                    True,
+                    inclusive_add,
+                )
+            else:
+                scan_add_inclusive(
+                    self.large_arr,
+                    ele_nums_pos[i],
+                    ele_nums_pos[i + 1],
+                    False,
+                    inclusive_add,
+                )
+
+        for i in range(len(ele_nums) - 3, -1, -1):
+            uniform_add(self.large_arr, ele_nums_pos[i], ele_nums_pos[i + 1])
+
+        blit_from_field_to_field(input_arr, self.large_arr, 0, length)
+
+
+__all__ = ["parallel_sort", "PrefixSumExecutor"]

taichi/aot/__init__.py

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

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

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

taichi/aot/conventions/gfxruntime140/__init__.py

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

taichi/aot/conventions/gfxruntime140/dr.py

@@ -0,0 +1,244 @@
+# type: ignore
+
+"""
+Data representation of all JSON data structures following the GfxRuntime140
+convention.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from taichi.aot.utils import dump_json_data_model, json_data_model
+
+
+@json_data_model
+class FieldAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        dtype = j["dtype"]
+        dtype_name = j["dtype_name"]
+        element_shape = j["element_shape"]
+        field_name = j["field_name"]
+        is_scalar = j["is_scalar"]
+        mem_offset_in_parent = j["mem_offset_in_parent"]
+        shape = j["shape"]
+
+        self.dtype: int = int(dtype)
+        self.dtype_name: str = str(dtype_name)
+        self.element_shape: List[int] = [int(x) for x in element_shape]
+        self.field_name: str = str(field_name)
+        self.is_scalar: bool = bool(is_scalar)
+        self.mem_offset_in_parent: int = int(mem_offset_in_parent)
+        self.shape: List[int] = [int(x) for x in shape]
+
+
+@json_data_model
+class ArgumentAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        index = j["key"][0]
+        dtype = j["value"]["dtype"]
+        element_shape = j["value"]["element_shape"]
+        field_dim = j["value"]["field_dim"]
+        fmt = j["value"]["format"]
+        is_array = j["value"]["is_array"]
+        offset_in_mem = j["value"]["offset_in_mem"]
+        stride = j["value"]["stride"]
+        # (penguinliong) Note that the name field is optional for kernels.
+        # Kernels are always launched by indexed arguments and this is for
+        # debugging and header generation only.
+        name = j["value"]["name"] if "name" in j["value"] and len(j["value"]["name"]) > 0 else None
+        ptype = j["value"]["ptype"] if "ptype" in j["value"] else None
+
+        self.dtype: int = int(dtype)
+        self.element_shape: List[int] = [int(x) for x in element_shape]
+        self.field_dim: int = int(field_dim)
+        self.format: int = int(fmt)
+        self.index: int = int(index)
+        self.is_array: bool = bool(is_array)
+        self.offset_in_mem: int = int(offset_in_mem)
+        self.stride: int = int(stride)
+        self.name: Optional[str] = str(name) if name is not None else None
+        self.ptype: Optional[int] = int(ptype) if ptype is not None else None
+
+
+@json_data_model
+class ContextAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        arg_attribs_vec_ = j["arg_attribs_vec_"]
+        args_bytes_ = j["args_bytes_"]
+        arr_access = j["arr_access"]
+        ret_attribs_vec_ = j["ret_attribs_vec_"]
+        rets_bytes_ = j["rets_bytes_"]
+
+        self.arg_attribs_vec_: List[ArgumentAttributes] = [ArgumentAttributes(x) for x in arg_attribs_vec_]
+        self.arg_attribs_vec_.sort(key=lambda x: x.index)
+        self.args_bytes_: int = int(args_bytes_)
+        self.arr_access: List[int] = [int(x["value"]) for x in arr_access]
+        self.ret_attribs_vec_: List[ArgumentAttributes] = [ArgumentAttributes(x) for x in ret_attribs_vec_]
+        self.rets_bytes_: int = int(rets_bytes_)
+
+
+@json_data_model
+class Buffer:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        root_id = j["root_id"][0]
+        ty = j["type"]
+
+        self.root_id: int = int(root_id)
+        self.type: int = int(ty)
+
+
+@json_data_model
+class BufferBinding:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        binding = j["binding"]
+        buffer = j["buffer"]
+
+        self.binding: int = int(binding)
+        self.buffer: Buffer = Buffer(buffer)
+
+
+@json_data_model
+class TextureBinding:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        arg_id = j["arg_id"]
+        binding = j["binding"]
+        is_storage = j["is_storage"]
+
+        self.arg_id: int = int(arg_id)
+        self.binding: int = int(binding)
+        self.is_storage: bool = bool(is_storage)
+
+
+@json_data_model
+class RangeForAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        begin = j["begin"]
+        const_begin = j["const_begin"]
+        const_end = j["const_end"]
+        end = j["end"]
+
+        self.begin: int = int(begin)
+        self.const_begin: bool = bool(const_begin)
+        self.const_end: bool = bool(const_end)
+        self.end: int = int(end)
+
+
+@json_data_model
+class TaskAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        advisory_num_threads_per_group = j["advisory_num_threads_per_group"]
+        advisory_total_num_threads = j["advisory_total_num_threads"]
+        buffer_binds = j["buffer_binds"]
+        name = j["name"]
+        range_for_attribs = j["range_for_attribs"] if "range_for_attribs" in j else None
+        task_type = j["task_type"]
+        texture_binds = j["texture_binds"]
+
+        self.advisory_num_threads_per_group: int = int(advisory_num_threads_per_group)
+        self.advisory_total_num_threads: int = int(advisory_total_num_threads)
+        self.buffer_binds: List[BufferBinding] = [BufferBinding(x) for x in buffer_binds]
+        self.name: str = str(name)
+        self.range_for_attribs: Optional[RangeForAttributes] = (
+            RangeForAttributes(range_for_attribs) if range_for_attribs is not None else None
+        )
+        self.task_type: int = int(task_type)
+        self.texture_binds: List[TextureBinding] = [TextureBinding(x) for x in texture_binds]
+
+
+@json_data_model
+class DeviceCapabilityLevel:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        key = j["key"]
+        value = j["value"]
+
+        self.key: str = str(key)
+        self.value: int = int(value)
+
+
+@json_data_model
+class KernelAttributes:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        ctx_attribs = j["ctx_attribs"]
+        is_jit_evaluator = j["is_jit_evaluator"]
+        name = j["name"]
+        tasks_attribs = j["tasks_attribs"]
+
+        self.ctx_attribs: ContextAttributes = ContextAttributes(ctx_attribs)
+        self.is_jit_evaluator: bool = bool(is_jit_evaluator)
+        self.name: str = str(name)
+        self.tasks_attribs: List[TaskAttributes] = [TaskAttributes(x) for x in tasks_attribs]
+
+
+@json_data_model
+class Metadata:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        fields = j["fields"]
+        kernels = j["kernels"]
+        required_caps = j["required_caps"]
+        root_buffer_size = j["root_buffer_size"]
+
+        self.fields: List[FieldAttributes] = [FieldAttributes(x) for x in fields]
+        self.kernels: List[KernelAttributes] = [KernelAttributes(x) for x in kernels]
+        self.required_caps: List[DeviceCapabilityLevel] = [DeviceCapabilityLevel(x) for x in required_caps]
+        self.root_buffer_size: int = int(root_buffer_size)
+
+
+def from_json_metadata(j: Dict[str, Any]) -> Metadata:
+    return Metadata(j)
+
+
+def to_json_metadata(meta_data: Metadata) -> Dict[str, Any]:
+    return dump_json_data_model(meta_data)
+
+
+@json_data_model
+class SymbolicArgument:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        dtype_id = j["dtype_id"]
+        element_shape = j["element_shape"]
+        field_dim = j["field_dim"]
+        name = j["name"]
+        num_channels = j["num_channels"]
+        tag = j["tag"]
+
+        self.dtype_id: int = int(dtype_id)
+        self.element_shape: List[int] = [int(x) for x in element_shape]
+        self.field_dim: int = int(field_dim)
+        self.name: str = str(name)
+        self.num_channels: int = int(num_channels)
+        self.tag: int = int(tag)
+
+
+@json_data_model
+class Dispatch:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        kernel_name = j["kernel_name"]
+        symbolic_args = j["symbolic_args"]
+
+        self.kernel_name: str = str(kernel_name)
+        self.symbolic_args: List[SymbolicArgument] = [SymbolicArgument(x) for x in symbolic_args]
+
+
+@json_data_model
+class GraphData:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        dispatches = j["dispatches"]
+
+        self.dispatches = [Dispatch(x) for x in dispatches]
+
+
+@json_data_model
+class Graph:
+    def __init__(self, j: Dict[str, Any]) -> None:
+        key = j["key"]
+        value = j["value"]
+
+        self.key = str(key)
+        self.value = GraphData(value)
+
+
+def from_json_graph(j: Dict[str, Any]) -> Graph:
+    return Graph(j)
+
+
+def to_json_graph(graph: Graph) -> Dict[str, Any]:
+    return dump_json_data_model(graph)