gstaichi 2.1.1rc3__cp310-cp310-macosx_11_0_arm64.whl

gstaichi/lang/runtime_ops.py

@@ -0,0 +1,13 @@
+# type: ignore
+
+from gstaichi.lang import impl
+
+
+def sync():
+    """Blocks the calling thread until all the previously
+    launched GsTaichi kernels have completed.
+    """
+    impl.get_runtime().sync()
+
+
+__all__ = ["sync"]

gstaichi/lang/shell.py

@@ -0,0 +1,35 @@
+# type: ignore
+
+import functools
+import os
+import sys
+
+from gstaichi._lib import core as _ti_core
+from gstaichi._logging import info
+
+pybuf_enabled = False
+_env_enable_pybuf = os.environ.get("TI_ENABLE_PYBUF", "1")
+if not _env_enable_pybuf or int(_env_enable_pybuf):
+    # When using in Jupyter / IDLE, the sys.stdout will be their wrapped ones.
+    # While sys.__stdout__ should always be the raw console stdout.
+    pybuf_enabled = sys.stdout is not sys.__stdout__
+
+_ti_core.toggle_python_print_buffer(pybuf_enabled)
+
+
+def _shell_pop_print(old_call):
+    if not pybuf_enabled:
+        # zero-overhead!
+        return old_call
+
+    info("Graphical python shell detected, using wrapped sys.stdout")
+
+    @functools.wraps(old_call)
+    def new_call(*args, **kwargs):
+        ret = old_call(*args, **kwargs)
+        # print's in kernel won't take effect until ti.sync(), discussion:
+        # https://github.com/taichi-dev/gstaichi/pull/1303#discussion_r444897102
+        print(_ti_core.pop_python_print_buffer(), end="")
+        return ret
+
+    return new_call

gstaichi/lang/simt/__init__.py

@@ -0,0 +1,5 @@
+# type: ignore
+
+from gstaichi.lang.simt import block, grid, subgroup, warp
+
+__all__ = ["warp", "subgroup", "block", "grid"]

gstaichi/lang/simt/block.py

@@ -0,0 +1,94 @@
+# type: ignore
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.expr import make_expr_group
+from gstaichi.lang.util import gstaichi_scope
+
+
+def arch_uses_spv(arch):
+    return arch == _ti_core.vulkan or arch == _ti_core.metal
+
+
+def sync():
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda or arch == _ti_core.amdgpu:
+        return impl.call_internal("block_barrier", with_runtime_context=False)
+    if arch_uses_spv(arch):
+        return impl.call_internal("workgroupBarrier", with_runtime_context=False)
+    raise ValueError(f"ti.block.shared_array is not supported for arch {arch}")
+
+
+def sync_all_nonzero(predicate):
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda:
+        return impl.call_internal("block_barrier_and_i32", predicate, with_runtime_context=False)
+    raise ValueError(f"ti.block.sync_all_nonzero is not supported for arch {arch}")
+
+
+def sync_any_nonzero(predicate):
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda:
+        return impl.call_internal("block_barrier_or_i32", predicate, with_runtime_context=False)
+    raise ValueError(f"ti.block.sync_any_nonzero is not supported for arch {arch}")
+
+
+def sync_count_nonzero(predicate):
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda:
+        return impl.call_internal("block_barrier_count_i32", predicate, with_runtime_context=False)
+    raise ValueError(f"ti.block.sync_count_nonzero is not supported for arch {arch}")
+
+
+def mem_sync():
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda:
+        return impl.call_internal("block_barrier", with_runtime_context=False)
+    if arch_uses_spv(arch):
+        return impl.call_internal("workgroupMemoryBarrier", with_runtime_context=False)
+    raise ValueError(f"ti.block.mem_sync is not supported for arch {arch}")
+
+
+def thread_idx():
+    arch = impl.get_runtime().prog.config().arch
+    if arch_uses_spv(arch):
+        return impl.call_internal("localInvocationId", with_runtime_context=False)
+    raise ValueError(f"ti.block.thread_idx is not supported for arch {arch}")
+
+
+def global_thread_idx():
+    arch = impl.get_runtime().prog.config().arch
+    if arch == _ti_core.cuda or _ti_core.amdgpu:
+        return impl.get_runtime().compiling_callable.ast_builder().insert_thread_idx_expr()
+    if arch_uses_spv(arch):
+        return impl.call_internal("globalInvocationId", with_runtime_context=False)
+    raise ValueError(f"ti.block.global_thread_idx is not supported for arch {arch}")
+
+
+class SharedArray:
+    _is_gstaichi_class = True
+
+    def __init__(self, shape, dtype):
+        if isinstance(shape, int):
+            self.shape = (shape,)
+        elif (isinstance(shape, tuple) or isinstance(shape, list)) and all(isinstance(s, int) for s in shape):
+            self.shape = shape
+        else:
+            raise ValueError(
+                f"ti.simt.block.shared_array shape must be an integer or a tuple of integers, but got {shape}"
+            )
+        if isinstance(dtype, impl.MatrixType):
+            dtype = dtype.tensor_type
+        self.dtype = dtype
+        self.shared_array_proxy = impl.expr_init_shared_array(self.shape, dtype)
+
+    @gstaichi_scope
+    def subscript(self, *indices):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        return impl.Expr(
+            ast_builder.expr_subscript(
+                self.shared_array_proxy,
+                make_expr_group(*indices),
+                _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()),
+            )
+        )

gstaichi/lang/simt/grid.py

@@ -0,0 +1,7 @@
+# type: ignore
+
+from gstaichi.lang import impl
+
+
+def memfence():
+    return impl.call_internal("grid_memfence", with_runtime_context=False)

gstaichi/lang/simt/subgroup.py

@@ -0,0 +1,191 @@
+# type: ignore
+
+from gstaichi.lang import impl
+
+
+def barrier():
+    return impl.call_internal("subgroupBarrier", with_runtime_context=False)
+
+
+def memory_barrier():
+    return impl.call_internal("subgroupMemoryBarrier", with_runtime_context=False)
+
+
+def elect():
+    return impl.call_internal("subgroupElect", with_runtime_context=False)
+
+
+def all_true(cond):
+    # TODO
+    pass
+
+
+def any_true(cond):
+    # TODO
+    pass
+
+
+def all_equal(value):
+    # TODO
+    pass
+
+
+def broadcast_first(value):
+    # TODO
+    pass
+
+
+def broadcast(value, index):
+    return impl.call_internal("subgroupBroadcast", value, index, with_runtime_context=False)
+
+
+def group_size():
+    return impl.call_internal("subgroupSize", with_runtime_context=False)
+
+
+def invocation_id():
+    return impl.call_internal("subgroupInvocationId", with_runtime_context=False)
+
+
+def reduce_add(value):
+    return impl.call_internal("subgroupAdd", value, with_runtime_context=False)
+
+
+def reduce_mul(value):
+    return impl.call_internal("subgroupMul", value, with_runtime_context=False)
+
+
+def reduce_min(value):
+    return impl.call_internal("subgroupMin", value, with_runtime_context=False)
+
+
+def reduce_max(value):
+    return impl.call_internal("subgroupMax", value, with_runtime_context=False)
+
+
+def reduce_and(value):
+    return impl.call_internal("subgroupAnd", value, with_runtime_context=False)
+
+
+def reduce_or(value):
+    return impl.call_internal("subgroupOr", value, with_runtime_context=False)
+
+
+def reduce_xor(value):
+    return impl.call_internal("subgroupXor", value, with_runtime_context=False)
+
+
+def inclusive_add(value):
+    return impl.call_internal("subgroupInclusiveAdd", value, with_runtime_context=False)
+
+
+def inclusive_mul(value):
+    return impl.call_internal("subgroupInclusiveMul", value, with_runtime_context=False)
+
+
+def inclusive_min(value):
+    return impl.call_internal("subgroupInclusiveMin", value, with_runtime_context=False)
+
+
+def inclusive_max(value):
+    return impl.call_internal("subgroupInclusiveMax", value, with_runtime_context=False)
+
+
+def inclusive_and(value):
+    return impl.call_internal("subgroupInclusiveAnd", value, with_runtime_context=False)
+
+
+def inclusive_or(value):
+    return impl.call_internal("subgroupInclusiveOr", value, with_runtime_context=False)
+
+
+def inclusive_xor(value):
+    return impl.call_internal("subgroupInclusiveXor", value, with_runtime_context=False)
+
+
+def exclusive_add(value):
+    # TODO
+    pass
+
+
+def exclusive_mul(value):
+    # TODO
+    pass
+
+
+def exclusive_min(value):
+    # TODO
+    pass
+
+
+def exclusive_max(value):
+    # TODO
+    pass
+
+
+def exclusive_and(value):
+    # TODO
+    pass
+
+
+def exclusive_or(value):
+    # TODO
+    pass
+
+
+def exclusive_xor(value):
+    # TODO
+    pass
+
+
+def shuffle(value, index):
+    return impl.call_internal("subgroupShuffle", value, index, with_runtime_context=False)
+
+
+def shuffle_xor(value, mask):
+    # TODO
+    pass
+
+
+def shuffle_up(value, offset):
+    return impl.call_internal("subgroupShuffleUp", value, offset, with_runtime_context=False)
+
+
+def shuffle_down(value, offset):
+    return impl.call_internal("subgroupShuffleDown", value, offset, with_runtime_context=False)
+
+
+__all__ = [
+    "barrier",
+    "memory_barrier",
+    "elect",
+    "all_true",
+    "any_true",
+    "all_equal",
+    "broadcast_first",
+    "reduce_add",
+    "reduce_mul",
+    "reduce_min",
+    "reduce_max",
+    "reduce_and",
+    "reduce_or",
+    "reduce_xor",
+    "inclusive_add",
+    "inclusive_mul",
+    "inclusive_min",
+    "inclusive_max",
+    "inclusive_and",
+    "inclusive_or",
+    "inclusive_xor",
+    "exclusive_add",
+    "exclusive_mul",
+    "exclusive_min",
+    "exclusive_max",
+    "exclusive_and",
+    "exclusive_or",
+    "exclusive_xor",
+    "shuffle",
+    "shuffle_xor",
+    "shuffle_up",
+    "shuffle_down",
+]

gstaichi/lang/simt/warp.py

@@ -0,0 +1,96 @@
+# type: ignore
+
+from gstaichi.lang import impl
+
+
+def all_nonzero(mask, predicate):
+    return impl.call_internal("cuda_all_sync_i32", mask, predicate, with_runtime_context=False)
+
+
+def any_nonzero(mask, predicate):
+    return impl.call_internal("cuda_any_sync_i32", mask, predicate, with_runtime_context=False)
+
+
+def unique(mask, predicate):
+    return impl.call_internal("cuda_uni_sync_i32", mask, predicate, with_runtime_context=False)
+
+
+def ballot(predicate):
+    return impl.call_internal("cuda_ballot_i32", predicate, with_runtime_context=False)
+
+
+def shfl_sync_i32(mask, val, offset):
+    # lane offset is 31 for warp size 32
+    return impl.call_internal("cuda_shfl_sync_i32", mask, val, offset, 31, with_runtime_context=False)
+
+
+def shfl_sync_f32(mask, val, offset):
+    # lane offset is 31 for warp size 32
+    return impl.call_internal("cuda_shfl_sync_f32", mask, val, offset, 31, with_runtime_context=False)
+
+
+def shfl_up_i32(mask, val, offset):
+    # lane offset is 0 for warp size 32
+    return impl.call_internal("cuda_shfl_up_sync_i32", mask, val, offset, 0, with_runtime_context=False)
+
+
+def shfl_up_f32(mask, val, offset):
+    # lane offset is 0 for warp size 32
+    return impl.call_internal("cuda_shfl_up_sync_f32", mask, val, offset, 0, with_runtime_context=False)
+
+
+def shfl_down_i32(mask, val, offset):
+    # lane offset is 31 for warp size 32
+    return impl.call_internal("cuda_shfl_down_sync_i32", mask, val, offset, 31, with_runtime_context=False)
+
+
+def shfl_down_f32(mask, val, offset):
+    # lane offset is 31 for warp size 32
+    return impl.call_internal("cuda_shfl_down_sync_f32", mask, val, offset, 31, with_runtime_context=False)
+
+
+def shfl_xor_i32(mask, val, offset):
+    return impl.call_internal("cuda_shfl_xor_sync_i32", mask, val, offset, 31, with_runtime_context=False)
+
+
+def match_any(mask, value):
+    # These intrinsics are only available on compute_70 or higher
+    # https://docs.nvidia.com/cuda/pdf/NVVM_IR_Specification.pdf
+    if impl.get_cuda_compute_capability() < 70:
+        raise AssertionError("match_any intrinsic only available on compute_70 or higher")
+    return impl.call_internal("cuda_match_any_sync_i32", mask, value, with_runtime_context=False)
+
+
+def match_all(mask, val):
+    # These intrinsics are only available on compute_70 or higher
+    # https://docs.nvidia.com/cuda/pdf/NVVM_IR_Specification.pdf
+    if impl.get_cuda_compute_capability() < 70:
+        raise AssertionError("match_all intrinsic only available on compute_70 or higher")
+    return impl.call_internal("cuda_match_all_sync_i32", mask, val, with_runtime_context=False)
+
+
+def active_mask():
+    return impl.call_internal("cuda_active_mask", with_runtime_context=False)
+
+
+def sync(mask):
+    return impl.call_internal("warp_barrier", mask, with_runtime_context=False)
+
+
+__all__ = [
+    "all_nonzero",
+    "any_nonzero",
+    "unique",
+    "ballot",
+    "shfl_sync_i32",
+    "shfl_sync_f32",
+    "shfl_up_i32",
+    "shfl_up_f32",
+    "shfl_down_i32",
+    "shfl_down_f32",
+    "shfl_xor_i32",
+    "match_any",
+    "match_all",
+    "active_mask",
+    "sync",
+]