gstaichi 0.1.25.dev0__cp311-cp311-macosx_15_0_arm64.whl → 2.0.0__cp311-cp311-macosx_15_0_arm64.whl

gstaichi/lang/impl.py

@@ -8,6 +8,7 @@ from gstaichi._lib import core as _ti_core
 from gstaichi._lib.core.gstaichi_python import (
     DataTypeCxx,
     Function,
+    KernelCxx,
     Program,
 )
 from gstaichi._snode.fields_builder import FieldsBuilder
@@ -70,17 +71,14 @@ from gstaichi.types.primitive_types import (
 
 @gstaichi_scope
 def expr_init_shared_array(shape, element_type):
-    compiling_callable = get_runtime().compiling_callable
-    assert compiling_callable is not None
-    return compiling_callable.ast_builder().expr_alloca_shared_array(
-        shape, element_type, _ti_core.DebugInfo(get_runtime().get_current_src_info())
-    )
+    ast_builder = get_runtime().compiling_callable.ast_builder()
+    debug_info = _ti_core.DebugInfo(get_runtime().get_current_src_info())
+    return ast_builder.expr_alloca_shared_array(shape, element_type, debug_info)
 
 
 @gstaichi_scope
 def expr_init(rhs):
     compiling_callable = get_runtime().compiling_callable
-    assert compiling_callable is not None
     if rhs is None:
         return Expr(
             compiling_callable.ast_builder().expr_alloca(_ti_core.DebugInfo(get_runtime().get_current_src_info()))
@@ -167,7 +165,7 @@ def _calc_slice(index, default_stop):
                 "GsTaichi does not support variables in slice now, please use constant instead of it."
             )
 
-    check_validity(start), check_validity(stop), check_validity(step)
+    _ = check_validity(start), check_validity(stop), check_validity(step)
     return [_ for _ in range(start, stop, step)]
 
 
@@ -194,9 +192,7 @@ def validate_subscript_index(value, index):
 @gstaichi_scope
 def subscript(ast_builder, value, *_indices, skip_reordered=False):
     dbg_info = _ti_core.DebugInfo(get_runtime().get_current_src_info())
-    compiling_callable = get_runtime().compiling_callable
-    assert compiling_callable is not None
-    ast_builder = compiling_callable.ast_builder()
+    ast_builder = get_runtime().compiling_callable.ast_builder()
     # Directly evaluate in Python for non-GsTaichi types
     if not isinstance(
         value,
@@ -337,8 +333,8 @@ class PyGsTaichi:
         self._prog: Program | None = None
         self.src_info_stack = []
         self.inside_kernel: bool = False
-        self.compiling_callable: Kernel | Function | None = None  # pointer to instance of lang::Kernel/Function
-        self._current_kernel: Kernel | None = None
+        self._compiling_callable: KernelCxx | Kernel | Function | None = None
+        self._current_kernel: "Kernel | None" = None
         self.global_vars = []
         self.grad_vars = []
         self.dual_vars = []
@@ -350,10 +346,18 @@ class PyGsTaichi:
         self.target_tape = None
         self.fwd_mode_manager = None
         self.grad_replaced = False
-        self.kernels = kernels or []
+        self.kernels: list[Kernel] = kernels or []
         self._signal_handler_registry = None
         self.unfinalized_fields_builder = {}
 
+    @property
+    def compiling_callable(self) -> KernelCxx | Kernel | Function:
+        if self._compiling_callable is None:
+            raise GsTaichiRuntimeError(
+                "_compiling_callable attribute not initialized. Maybe you forgot to call `ti.init()` first?"
+            )
+        return self._compiling_callable
+
     @property
     def prog(self) -> Program:
         if self._prog is None:
@@ -364,7 +368,7 @@ class PyGsTaichi:
     def current_kernel(self) -> Kernel:
         if self._current_kernel is None:
             raise GsTaichiRuntimeError(
-                "_pr_current_kernelog attribute not initialized. Maybe you forgot to call `ti.init()` first?"
+                "_current_kernel attribute not initialized. Maybe you forgot to call `ti.init()` first?"
             )
         return self._current_kernel
 
@@ -373,7 +377,7 @@ class PyGsTaichi:
 
     def clear_compiled_functions(self):
         for k in self.kernels:
-            k.compiled_kernels.clear()
+            k.materialized_kernels.clear()
 
     def finalize_fields_builder(self, builder):
         self.unfinalized_fields_builder.pop(builder)
@@ -390,7 +394,7 @@ class PyGsTaichi:
     def get_num_compiled_functions(self):
         count = 0
         for k in self.kernels:
-            count += len(k.compiled_kernels)
+            count += len(k.materialized_kernels)
         return count
 
     def src_info_guard(self, info):
@@ -962,11 +966,9 @@ def ti_print(*_vars, sep=" ", end="\n"):
 
     _vars = add_separators(_vars)
     contents, formats = ti_format_list_to_content_entries(_vars)
-    compiling_callable = get_runtime().compiling_callable
-    assert compiling_callable is not None
-    compiling_callable.ast_builder().create_print(
-        contents, formats, _ti_core.DebugInfo(get_runtime().get_current_src_info())
-    )
+    ast_builder = get_runtime().compiling_callable.ast_builder()
+    debug_info = _ti_core.DebugInfo(get_runtime().get_current_src_info())
+    ast_builder.create_print(contents, formats, debug_info)
 
 
 @gstaichi_scope
@@ -996,9 +998,8 @@ def ti_format(*args):
 def ti_assert(cond, msg, extra_args, dbg_info):
     # Mostly a wrapper to help us convert from Expr (defined in Python) to
     # _ti_core.Expr (defined in C++)
-    compiling_callable = get_runtime().compiling_callable
-    assert compiling_callable is not None
-    compiling_callable.ast_builder().create_assert_stmt(Expr(cond).ptr, msg, extra_args, dbg_info)
+    ast_builder = get_runtime().compiling_callable.ast_builder()
+    ast_builder.create_assert_stmt(Expr(cond).ptr, msg, extra_args, dbg_info)
 
 
 @gstaichi_scope

gstaichi/lang/kernel_arguments.py

@@ -4,6 +4,10 @@ import inspect
 
 import gstaichi.lang
 from gstaichi._lib import core as _ti_core
+from gstaichi._lib.core.gstaichi_python import (
+    BoundaryMode,
+    DataTypeCxx,
+)
 from gstaichi.lang import impl, ops
 from gstaichi.lang._texture import RWTextureAccessor, TextureSampler
 from gstaichi.lang.any_array import AnyArray
@@ -15,11 +19,18 @@ from gstaichi.types.compound_types import CompoundType
 from gstaichi.types.primitive_types import RefType, u64
 
 
-class KernelArgument:
-    def __init__(self, _annotation, _name, _default=inspect.Parameter.empty):
-        self.annotation = _annotation
-        self.name = _name
-        self.default = _default
+class ArgMetadata:
+    """
+    Metadata about an argument to a function
+    """
+
+    def __init__(self, annotation, name, default=inspect.Parameter.empty):
+        self.annotation = annotation
+        self.name = name
+        self.default = default
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(annotation={self.annotation}, name={self.name}, default={self.default})"
 
 
 class SparseMatrixEntry:
@@ -48,7 +59,7 @@ class SparseMatrixProxy:
         return SparseMatrixEntry(self.ptr, i, j, self.dtype)
 
 
-def decl_scalar_arg(dtype, name, arg_depth):
+def decl_scalar_arg(dtype, name):
     is_ref = False
     if isinstance(dtype, RefType):
         is_ref = True
@@ -60,9 +71,7 @@ def decl_scalar_arg(dtype, name, arg_depth):
         arg_id = impl.get_runtime().compiling_callable.insert_scalar_param(dtype, name)
 
     argload_di = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
-    return Expr(
-        _ti_core.make_arg_load_expr(arg_id, dtype, is_ref, create_load=True, arg_depth=arg_depth, dbg_info=argload_di)
-    )
+    return Expr(_ti_core.make_arg_load_expr(arg_id, dtype, is_ref, create_load=True, dbg_info=argload_di))
 
 
 def get_type_for_kernel_args(dtype, name):
@@ -86,35 +95,22 @@ def get_type_for_kernel_args(dtype, name):
     return dtype
 
 
-def decl_matrix_arg(matrixtype, name, arg_depth):
+def decl_matrix_arg(matrixtype, name):
     arg_type = get_type_for_kernel_args(matrixtype, name)
     arg_id = impl.get_runtime().compiling_callable.insert_scalar_param(arg_type, name)
     argload_di = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
-    arg_load = Expr(
-        _ti_core.make_arg_load_expr(arg_id, arg_type, create_load=False, arg_depth=arg_depth, dbg_info=argload_di)
-    )
+    arg_load = Expr(_ti_core.make_arg_load_expr(arg_id, arg_type, create_load=False, dbg_info=argload_di))
     return matrixtype.from_gstaichi_object(arg_load)
 
 
-def decl_struct_arg(structtype, name, arg_depth):
+def decl_struct_arg(structtype, name):
     arg_type = get_type_for_kernel_args(structtype, name)
     arg_id = impl.get_runtime().compiling_callable.insert_scalar_param(arg_type, name)
     argload_di = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
-    arg_load = Expr(
-        _ti_core.make_arg_load_expr(arg_id, arg_type, create_load=False, arg_depth=arg_depth, dbg_info=argload_di)
-    )
+    arg_load = Expr(_ti_core.make_arg_load_expr(arg_id, arg_type, create_load=False, dbg_info=argload_di))
     return structtype.from_gstaichi_object(arg_load)
 
 
-def push_argpack_arg(name):
-    impl.get_runtime().compiling_callable.insert_argpack_param_and_push(name)
-
-
-def decl_argpack_arg(argpacktype, member_dict):
-    impl.get_runtime().compiling_callable.pop_argpack_stack()
-    return argpacktype.from_gstaichi_object(member_dict)
-
-
 def decl_sparse_matrix(dtype, name):
     value_type = cook_dtype(dtype)
     ptr_type = cook_dtype(u64)
@@ -126,16 +122,18 @@ def decl_sparse_matrix(dtype, name):
     )
 
 
-def decl_ndarray_arg(element_type, ndim, name, needs_grad, boundary):
+def decl_ndarray_arg(
+    element_type: DataTypeCxx, ndim: int, name: str, needs_grad: bool, boundary: BoundaryMode
+) -> AnyArray:
     arg_id = impl.get_runtime().compiling_callable.insert_ndarray_param(element_type, ndim, name, needs_grad)
-    return AnyArray(_ti_core.make_external_tensor_expr(element_type, ndim, arg_id, needs_grad, 0, boundary))
+    return AnyArray(_ti_core.make_external_tensor_expr(element_type, ndim, arg_id, needs_grad, boundary))
 
 
 def decl_texture_arg(num_dimensions, name):
     # FIXME: texture_arg doesn't have element_shape so better separate them
     arg_id = impl.get_runtime().compiling_callable.insert_texture_param(num_dimensions, name)
     dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
-    return TextureSampler(_ti_core.make_texture_ptr_expr(arg_id, num_dimensions, 0, dbg_info), num_dimensions)
+    return TextureSampler(_ti_core.make_texture_ptr_expr(arg_id, num_dimensions, dbg_info), num_dimensions)
 
 
 def decl_rw_texture_arg(num_dimensions, buffer_format, lod, name):
@@ -143,7 +141,7 @@ def decl_rw_texture_arg(num_dimensions, buffer_format, lod, name):
     arg_id = impl.get_runtime().compiling_callable.insert_rw_texture_param(num_dimensions, buffer_format, name)
     dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
     return RWTextureAccessor(
-        _ti_core.make_rw_texture_ptr_expr(arg_id, num_dimensions, 0, buffer_format, lod, dbg_info), num_dimensions
+        _ti_core.make_rw_texture_ptr_expr(arg_id, num_dimensions, buffer_format, lod, dbg_info), num_dimensions
     )