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

gstaichi/CHANGELOG.md

@@ -0,0 +1,6 @@
+Highlights:
+   - **Miscellaneous**
+      - Add ptx caching (#130) (by **Hugh Perkins**)
+
+Full changelog:
+   - [Misc] Add ptx caching (#130) (by **Hugh Perkins**)

gstaichi/__init__.py

@@ -24,7 +24,7 @@ from gstaichi.types.primitive_types import *
 
 def __getattr__(attr):
     if attr == "cfg":
-        return None if lang.impl.get_runtime().prog is None else lang.impl.current_cfg()
+        return None if lang.impl.get_runtime()._prog is None else lang.impl.current_cfg()
     raise AttributeError(f"module '{__name__}' has no attribute '{attr}'")
 
 

gstaichi/_lib/core/gstaichi_python.pyi

@@ -4,7 +4,7 @@ gstaichi_python
 from __future__ import annotations
 import numpy
 import typing
-__all__ = ['ADJOINT', 'ADJOINT_CHECKBIT', 'AOS', 'ASTBuilder', 'Arch', 'ArgPackCxx', 'AutodiffMode', 'Axis', 'Benchmark', 'BinaryOpType', 'BitStructType', 'BitStructTypeBuilder', 'BoundaryMode', 'CC', 'CE', 'CF', 'CGd', 'CGf', 'CLAMP', 'CUCG', 'CV', 'Cell', 'CompileConfig', 'CompiledKernelData', 'Config', 'ConvType', 'CuSparseMatrix', 'CuSparseSolver', 'DUAL', 'DataTypeCxx', 'DataType_f16', 'DataType_f32', 'DataType_f64', 'DataType_gen', 'DataType_i16', 'DataType_i32', 'DataType_i64', 'DataType_i8', 'DataType_u1', 'DataType_u16', 'DataType_u32', 'DataType_u64', 'DataType_u8', 'DataType_unknown', 'DebugInfo', 'DeviceAllocation', 'DeviceCapabilityConfig', 'EC', 'EE', 'EF', 'EV', 'Edge', 'EigenSparseSolverfloat32LDLTAMD', 'EigenSparseSolverfloat32LDLTCOLAMD', 'EigenSparseSolverfloat32LLTAMD', 'EigenSparseSolverfloat32LLTCOLAMD', 'EigenSparseSolverfloat32LUAMD', 'EigenSparseSolverfloat32LUCOLAMD', 'EigenSparseSolverfloat64LDLTAMD', 'EigenSparseSolverfloat64LDLTCOLAMD', 'EigenSparseSolverfloat64LLTAMD', 'EigenSparseSolverfloat64LLTCOLAMD', 'EigenSparseSolverfloat64LUAMD', 'EigenSparseSolverfloat64LUCOLAMD', 'ExprCxx', 'ExprGroup', 'Extension', 'FC', 'FE', 'FF', 'FORWARD', 'FV', 'Face', 'Format', 'Function', 'FunctionKey', 'GsTaichiAssertionError', 'GsTaichiIndexError', 'GsTaichiRuntimeError', 'GsTaichiSyntaxError', 'GsTaichiTypeError', 'HackedSignalRegister', 'InternalOp', 'KernelCxx', 'KernelLaunchContext', 'KernelProfileTracedRecord', 'KernelProfilerQueryResult', 'Layout', 'Mesh', 'MeshElementType', 'MeshPtr', 'MeshRelationType', 'MeshTopology', 'NONE', 'NULL', 'NdarrayCxx', 'Operation', 'PRIMAL', 'Program', 'REVERSE', 'SNodeAccessFlag', 'SNodeCxx', 'SNodeGradType', 'SNodeRegistry', 'SNodeTreeCxx', 'SNodeType', 'SOA', 'SparseMatrix', 'SparseMatrixBuilder', 'SparseSolver', 'Stmt', 'Task', 'Tetrahedron', 'TextureCxx', 'TextureOpType', 'Triangle', 'Type', 'TypeFactory', 'UNSAFE', 'UnaryOpType', 'VALIDATION', 'VC', 'VE', 'VF', 'VV', 'Vector2d', 'Vector2f', 'Vector2i', 'Vector3d', 'Vector3f', 'Vector3i', 'Vector4d', 'Vector4f', 'Vector4i', 'Vertex', 'abs', 'acos', 'add', 'adstack', 'amdgpu', 'arch_from_name', 'arch_name', 'arch_uses_llvm', 'arm64', 'asin', 'assertion', 'atan2', 'bit_and', 'bit_not', 'bit_or', 'bit_sar', 'bit_shl', 'bit_shr', 'bit_struct', 'bit_xor', 'bitmasked', 'bits_cast', 'block_local', 'bls', 'cast_bits', 'cast_value', 'ceil', 'clean_offline_cache_files', 'clear_profile_info', 'clz', 'cmp_eq', 'cmp_ge', 'cmp_gt', 'cmp_le', 'cmp_lt', 'cmp_ne', 'cos', 'create_benchmark', 'create_initialized_benchmark', 'create_initialized_task', 'create_mesh', 'create_task', 'critical', 'cuda', 'cuda_version', 'dColMajor_EigenSparseMatrix', 'dRowMajor_EigenSparseMatrix', 'data64', 'data_type_name', 'data_type_size', 'debug', 'default_compile_config', 'dense', 'div', 'dynamic', 'element_order', 'element_type_name', 'error', 'exp', 'expr_abs', 'expr_acos', 'expr_add', 'expr_asin', 'expr_assume_in_range', 'expr_atan2', 'expr_atomic_add', 'expr_atomic_bit_and', 'expr_atomic_bit_or', 'expr_atomic_bit_xor', 'expr_atomic_max', 'expr_atomic_min', 'expr_atomic_mul', 'expr_atomic_sub', 'expr_bit_and', 'expr_bit_not', 'expr_bit_or', 'expr_bit_sar', 'expr_bit_shl', 'expr_bit_shr', 'expr_bit_xor', 'expr_ceil', 'expr_clz', 'expr_cmp_eq', 'expr_cmp_ge', 'expr_cmp_gt', 'expr_cmp_le', 'expr_cmp_lt', 'expr_cmp_ne', 'expr_cos', 'expr_div', 'expr_exp', 'expr_field', 'expr_floor', 'expr_floordiv', 'expr_frexp', 'expr_ifte', 'expr_inv', 'expr_log', 'expr_logic_not', 'expr_logical_and', 'expr_logical_or', 'expr_loop_unique', 'expr_matrix_field', 'expr_max', 'expr_min', 'expr_mod', 'expr_mul', 'expr_neg', 'expr_popcnt', 'expr_pow', 'expr_rcp', 'expr_round', 'expr_rsqrt', 'expr_select', 'expr_sin', 'expr_sqrt', 'expr_sub', 'expr_tan', 'expr_tanh', 'expr_truediv', 'extfunc', 'fColMajor_EigenSparseMatrix', 'fRowMajor_EigenSparseMatrix', 'finalize_snode_tree', 'floor', 'floordiv', 'flush_log', 'frexp', 'from_end_element_order', 'g2r', 'get_commit_hash', 'get_default_float_size', 'get_external_tensor_dim', 'get_external_tensor_element_dim', 'get_external_tensor_element_shape', 'get_external_tensor_element_type', 'get_external_tensor_needs_grad', 'get_external_tensor_real_func_args', 'get_external_tensor_shape_along_axis', 'get_llvm_target_support', 'get_max_num_indices', 'get_num_elements', 'get_python_package_dir', 'get_relation_access', 'get_relation_size', 'get_repo_dir', 'get_type_factory_instance', 'get_version_major', 'get_version_minor', 'get_version_patch', 'get_version_string', 'hash', 'host_arch', 'info', 'insert_internal_func_call', 'inv', 'inverse_relation', 'is_extension_supported', 'is_integral', 'is_quant', 'is_real', 'is_signed', 'is_tensor', 'is_unsigned', 'js', 'kFetchTexel', 'kLoad', 'kSampleLod', 'kStore', 'kUndefined', 'l2g', 'l2r', 'libdevice_path', 'log', 'logging_effective', 'logic_not', 'make_arg_load_expr', 'make_binary_op_expr', 'make_const_expr_bool', 'make_const_expr_fp', 'make_const_expr_int', 'make_cucg_solver', 'make_cusparse_solver', 'make_double_cg_solver', 'make_external_tensor_expr', 'make_external_tensor_grad_expr', 'make_float_cg_solver', 'make_frontend_assign_stmt', 'make_get_element_expr', 'make_global_load_stmt', 'make_global_store_stmt', 'make_rand_expr', 'make_reference', 'make_rw_texture_ptr_expr', 'make_sparse_solver', 'make_texture_ptr_expr', 'make_unary_op_expr', 'max', 'mesh', 'mesh_local', 'metal', 'min', 'mod', 'mul', 'neg', 'opencl', 'place', 'pointer', 'pop_python_print_buffer', 'popcnt', 'pow', 'print_all_units', 'print_profile_info', 'promoted_type', 'quant', 'quant_array', 'quant_basic', 'query_int64', 'rcp', 'read_only', 'relation_by_orders', 'reset_default_compile_config', 'root', 'round', 'rsqrt', 'set_core_state_python_imported', 'set_core_trigger_gdb_when_crash', 'set_index_mapping', 'set_lib_dir', 'set_logging_level', 'set_logging_level_default', 'set_num_elements', 'set_num_patches', 'set_owned_offset', 'set_patch_max_element_num', 'set_python_package_dir', 'set_relation_dynamic', 'set_relation_fixed', 'set_tmp_dir', 'set_total_offset', 'set_vulkan_visible_device', 'sgn', 'sin', 'sparse', 'sqrt', 'start_memory_monitoring', 'sub', 'subscript_with_multiple_indices', 'tan', 'tanh', 'test_cpp_exception', 'test_logging', 'test_printf', 'test_raise_error', 'test_threading', 'test_throw', 'to_end_element_order', 'toggle_python_print_buffer', 'trace', 'trigger_crash', 'trigger_sig_fpe', 'truediv', 'undefined', 'value_cast', 'vulkan', 'wait_for_debugger', 'warn', 'with_amdgpu', 'with_cuda', 'with_metal', 'with_vulkan', 'x64']
+__all__: list[str] = ['ADJOINT', 'ADJOINT_CHECKBIT', 'AOS', 'ASTBuilder', 'Arch', 'AutodiffMode', 'Axis', 'Benchmark', 'BinaryOpType', 'BitStructType', 'BitStructTypeBuilder', 'BoundaryMode', 'CC', 'CE', 'CF', 'CGd', 'CGf', 'CLAMP', 'CUCG', 'CV', 'Cell', 'CompileConfig', 'CompiledKernelData', 'Config', 'ConvType', 'CuSparseMatrix', 'CuSparseSolver', 'DUAL', 'DataTypeCxx', 'DataType_f16', 'DataType_f32', 'DataType_f64', 'DataType_gen', 'DataType_i16', 'DataType_i32', 'DataType_i64', 'DataType_i8', 'DataType_u1', 'DataType_u16', 'DataType_u32', 'DataType_u64', 'DataType_u8', 'DataType_unknown', 'DebugInfo', 'DeviceAllocation', 'DeviceCapabilityConfig', 'EC', 'EE', 'EF', 'EV', 'Edge', 'EigenSparseSolverfloat32LDLTAMD', 'EigenSparseSolverfloat32LDLTCOLAMD', 'EigenSparseSolverfloat32LLTAMD', 'EigenSparseSolverfloat32LLTCOLAMD', 'EigenSparseSolverfloat32LUAMD', 'EigenSparseSolverfloat32LUCOLAMD', 'EigenSparseSolverfloat64LDLTAMD', 'EigenSparseSolverfloat64LDLTCOLAMD', 'EigenSparseSolverfloat64LLTAMD', 'EigenSparseSolverfloat64LLTCOLAMD', 'EigenSparseSolverfloat64LUAMD', 'EigenSparseSolverfloat64LUCOLAMD', 'ExprCxx', 'ExprGroup', 'Extension', 'FC', 'FE', 'FF', 'FORWARD', 'FV', 'Face', 'Format', 'Function', 'FunctionKey', 'GsTaichiAssertionError', 'GsTaichiIndexError', 'GsTaichiRuntimeError', 'GsTaichiSyntaxError', 'GsTaichiTypeError', 'HackedSignalRegister', 'InternalOp', 'KernelCxx', 'KernelLaunchContext', 'KernelProfileTracedRecord', 'KernelProfilerQueryResult', 'Layout', 'Mesh', 'MeshElementType', 'MeshPtr', 'MeshRelationType', 'MeshTopology', 'NONE', 'NULL', 'NdarrayCxx', 'Operation', 'PRIMAL', 'Program', 'REVERSE', 'SNodeAccessFlag', 'SNodeCxx', 'SNodeGradType', 'SNodeRegistry', 'SNodeTreeCxx', 'SNodeType', 'SOA', 'SparseMatrix', 'SparseMatrixBuilder', 'SparseSolver', 'Stmt', 'Task', 'Tetrahedron', 'TextureCxx', 'TextureOpType', 'Triangle', 'Type', 'TypeFactory', 'UNSAFE', 'UnaryOpType', 'VALIDATION', 'VC', 'VE', 'VF', 'VV', 'Vector2d', 'Vector2f', 'Vector2i', 'Vector3d', 'Vector3f', 'Vector3i', 'Vector4d', 'Vector4f', 'Vector4i', 'Vertex', 'abs', 'acos', 'add', 'adstack', 'amdgpu', 'arch_from_name', 'arch_name', 'arch_uses_llvm', 'arm64', 'asin', 'assertion', 'atan2', 'bit_and', 'bit_not', 'bit_or', 'bit_sar', 'bit_shl', 'bit_shr', 'bit_struct', 'bit_xor', 'bitmasked', 'bits_cast', 'block_local', 'bls', 'cast_bits', 'cast_value', 'ceil', 'clear_profile_info', 'clz', 'cmp_eq', 'cmp_ge', 'cmp_gt', 'cmp_le', 'cmp_lt', 'cmp_ne', 'cos', 'create_benchmark', 'create_initialized_benchmark', 'create_initialized_task', 'create_mesh', 'create_task', 'critical', 'cuda', 'cuda_version', 'dColMajor_EigenSparseMatrix', 'dRowMajor_EigenSparseMatrix', 'data64', 'data_type_name', 'data_type_size', 'debug', 'default_compile_config', 'dense', 'div', 'dynamic', 'element_order', 'element_type_name', 'error', 'exp', 'expr_abs', 'expr_acos', 'expr_add', 'expr_asin', 'expr_assume_in_range', 'expr_atan2', 'expr_atomic_add', 'expr_atomic_bit_and', 'expr_atomic_bit_or', 'expr_atomic_bit_xor', 'expr_atomic_max', 'expr_atomic_min', 'expr_atomic_mul', 'expr_atomic_sub', 'expr_bit_and', 'expr_bit_not', 'expr_bit_or', 'expr_bit_sar', 'expr_bit_shl', 'expr_bit_shr', 'expr_bit_xor', 'expr_ceil', 'expr_clz', 'expr_cmp_eq', 'expr_cmp_ge', 'expr_cmp_gt', 'expr_cmp_le', 'expr_cmp_lt', 'expr_cmp_ne', 'expr_cos', 'expr_div', 'expr_exp', 'expr_field', 'expr_floor', 'expr_floordiv', 'expr_frexp', 'expr_ifte', 'expr_inv', 'expr_log', 'expr_logic_not', 'expr_logical_and', 'expr_logical_or', 'expr_loop_unique', 'expr_matrix_field', 'expr_max', 'expr_min', 'expr_mod', 'expr_mul', 'expr_neg', 'expr_popcnt', 'expr_pow', 'expr_rcp', 'expr_round', 'expr_rsqrt', 'expr_select', 'expr_sin', 'expr_sqrt', 'expr_sub', 'expr_tan', 'expr_tanh', 'expr_truediv', 'extfunc', 'fColMajor_EigenSparseMatrix', 'fRowMajor_EigenSparseMatrix', 'finalize_snode_tree', 'floor', 'floordiv', 'flush_log', 'frexp', 'from_end_element_order', 'g2r', 'get_commit_hash', 'get_default_float_size', 'get_external_tensor_dim', 'get_external_tensor_element_dim', 'get_external_tensor_element_shape', 'get_external_tensor_element_type', 'get_external_tensor_needs_grad', 'get_external_tensor_real_func_args', 'get_external_tensor_shape_along_axis', 'get_llvm_target_support', 'get_max_num_indices', 'get_num_elements', 'get_python_package_dir', 'get_relation_access', 'get_relation_size', 'get_repo_dir', 'get_type_factory_instance', 'get_version_major', 'get_version_minor', 'get_version_patch', 'get_version_string', 'hash', 'host_arch', 'info', 'insert_internal_func_call', 'inv', 'inverse_relation', 'is_extension_supported', 'is_integral', 'is_quant', 'is_real', 'is_signed', 'is_tensor', 'is_unsigned', 'js', 'kFetchTexel', 'kLoad', 'kSampleLod', 'kStore', 'kUndefined', 'l2g', 'l2r', 'libdevice_path', 'log', 'logging_effective', 'logic_not', 'make_arg_load_expr', 'make_binary_op_expr', 'make_const_expr_bool', 'make_const_expr_fp', 'make_const_expr_int', 'make_cucg_solver', 'make_cusparse_solver', 'make_double_cg_solver', 'make_external_tensor_expr', 'make_external_tensor_grad_expr', 'make_float_cg_solver', 'make_frontend_assign_stmt', 'make_get_element_expr', 'make_global_load_stmt', 'make_global_store_stmt', 'make_rand_expr', 'make_reference', 'make_rw_texture_ptr_expr', 'make_sparse_solver', 'make_texture_ptr_expr', 'make_unary_op_expr', 'max', 'mesh', 'mesh_local', 'metal', 'min', 'mod', 'mul', 'neg', 'opencl', 'place', 'pointer', 'pop_python_print_buffer', 'popcnt', 'pow', 'print_all_units', 'print_profile_info', 'promoted_type', 'quant', 'quant_array', 'quant_basic', 'query_int64', 'rcp', 'read_only', 'relation_by_orders', 'reset_default_compile_config', 'root', 'round', 'rsqrt', 'set_core_state_python_imported', 'set_core_trigger_gdb_when_crash', 'set_index_mapping', 'set_lib_dir', 'set_logging_level', 'set_logging_level_default', 'set_num_elements', 'set_num_patches', 'set_owned_offset', 'set_patch_max_element_num', 'set_python_package_dir', 'set_relation_dynamic', 'set_relation_fixed', 'set_tmp_dir', 'set_total_offset', 'set_vulkan_visible_device', 'sgn', 'sin', 'sparse', 'sqrt', 'start_memory_monitoring', 'sub', 'subscript_with_multiple_indices', 'tan', 'tanh', 'test_cpp_exception', 'test_logging', 'test_printf', 'test_raise_error', 'test_threading', 'test_throw', 'to_end_element_order', 'toggle_python_print_buffer', 'trace', 'trigger_crash', 'trigger_sig_fpe', 'truediv', 'undefined', 'value_cast', 'vulkan', 'wait_for_debugger', 'warn', 'with_amdgpu', 'with_cuda', 'with_metal', 'with_vulkan', 'x64']
 class ASTBuilder:
     def begin_frontend_if(self, arg0: ..., arg1: DebugInfo) -> None:
         ...
@@ -167,26 +167,6 @@ class Arch:
     @property
     def value(self) -> int:
         ...
-class ArgPackCxx:
-    def data_type(self) -> DataTypeCxx:
-        ...
-    def device_allocation(self) -> DeviceAllocation:
-        ...
-    def device_allocation_ptr(self) -> int:
-        ...
-    def nelement(self) -> int:
-        ...
-    def set_arg_float(self, arg0: tuple[int, ...], arg1: float) -> None:
-        ...
-    def set_arg_int(self, arg0: tuple[int, ...], arg1: int) -> None:
-        ...
-    def set_arg_nested_argpack(self, arg0: int, arg1: ArgPackCxx) -> None:
-        ...
-    def set_arg_uint(self, arg0: tuple[int, ...], arg1: int) -> None:
-        ...
-    @property
-    def dtype(self) -> DataTypeCxx:
-        ...
 class AutodiffMode:
     """
     Members:
@@ -1213,8 +1193,6 @@ class KernelCxx:
         ...
     def finalize_rets(self) -> None:
         ...
-    def insert_argpack_param_and_push(self, arg0: str) -> tuple[int, ...]:
-        ...
     def insert_arr_param(self, arg0: DataTypeCxx, arg1: int, arg2: tuple[int, ...], arg3: str) -> tuple[int, ...]:
         ...
     def insert_ndarray_param(self, arg0: DataTypeCxx, arg1: int, arg2: str, arg3: bool) -> tuple[int, ...]:
@@ -1233,8 +1211,6 @@ class KernelCxx:
         ...
     def no_activate(self, arg0: SNodeCxx) -> None:
         ...
-    def pop_argpack_stack(self) -> None:
-        ...
 class KernelLaunchContext:
     def get_struct_ret_float(self, arg0: tuple[int, ...]) -> float:
         ...
@@ -1242,8 +1218,6 @@ class KernelLaunchContext:
         ...
     def get_struct_ret_uint(self, arg0: tuple[int, ...]) -> int:
         ...
-    def set_arg_argpack(self, arg0: tuple[int, ...], arg1: ArgPackCxx) -> None:
-        ...
     def set_arg_external_array_with_shape(self, arg0: tuple[int, ...], arg1: int, arg2: int, arg3: tuple[int, ...], arg4: int) -> None:
         ...
     def set_arg_float(self, arg0: tuple[int, ...], arg1: float) -> None:
@@ -1558,8 +1532,6 @@ class Program:
         ...
     def config(self) -> CompileConfig:
         ...
-    def create_argpack(self, dt: DataTypeCxx) -> ...:
-        ...
     def create_function(self, arg0: ...) -> Function:
         ...
     def create_kernel(self, arg0: typing.Callable[[...], None], arg1: str, arg2: AutodiffMode) -> KernelCxx:
@@ -1570,10 +1542,10 @@ class Program:
         ...
     def create_texture(self, fmt: ..., shape: tuple[int, ...] = ()) -> ...:
         ...
-    def delete_argpack(self, arg0: ...) -> None:
-        ...
     def delete_ndarray(self, arg0: ...) -> None:
         ...
+    def dump_cache_data_to_disk(self) -> None:
+        ...
     def fill_float(self, arg0: ..., arg1: float) -> None:
         ...
     def fill_int(self, arg0: ..., arg1: int) -> None:
@@ -1604,6 +1576,8 @@ class Program:
         ...
     def launch_kernel(self, arg0: CompiledKernelData, arg1: ...) -> None:
         ...
+    def load_fast_cache(self, arg0: str, arg1: str, arg2: CompileConfig, arg3: DeviceCapabilityConfig) -> CompiledKernelData:
+        ...
     def make_id_expr(self, arg0: str) -> ...:
         ...
     def make_sparse_matrix_from_ndarray(self, arg0: ..., arg1: ...) -> None:
@@ -1618,6 +1592,8 @@ class Program:
         ...
     def set_kernel_profiler_toolkit(self, arg0: str) -> bool:
         ...
+    def store_fast_cache(self, arg0: str, arg1: ..., arg2: CompileConfig, arg3: DeviceCapabilityConfig, arg4: CompiledKernelData) -> None:
+        ...
     def sync_kernel_profiler(self) -> None:
         ...
     def synchronize(self) -> None:
@@ -2013,8 +1989,6 @@ class Type:
     def to_string(self) -> str:
         ...
 class TypeFactory:
-    def get_argpack_type(self, arg0: list[tuple[DataTypeCxx, str]]) -> DataTypeCxx:
-        ...
     def get_ndarray_struct_type(self, dt: DataTypeCxx, ndim: int, needs_grad: bool) -> Type:
         ...
     def get_quant_fixed_type(self, digits_type: Type, compute_type: Type, scale: float) -> Type:
@@ -2027,8 +2001,6 @@ class TypeFactory:
         ...
     def get_struct_type(self, arg0: list[tuple[DataTypeCxx, str]]) -> DataTypeCxx:
         ...
-    def get_struct_type_for_argpack_ptr(self, dt: DataTypeCxx, layout: str = 'none') -> Type:
-        ...
     def get_tensor_type(self, shape: tuple[int, ...], element_type: DataType) -> DataTypeCxx:
         ...
 class UnaryOpType:
@@ -2442,8 +2414,6 @@ def arch_uses_llvm(arg0: Arch) -> bool:
     ...
 def bits_cast(arg0: ExprCxx, arg1: DataTypeCxx) -> ExprCxx:
     ...
-def clean_offline_cache_files(arg0: str) -> int:
-    ...
 def clear_profile_info() -> None:
     ...
 def create_benchmark(arg0: str) -> ...:
@@ -2670,7 +2640,7 @@ def libdevice_path() -> str:
     ...
 def logging_effective(arg0: str) -> bool:
     ...
-def make_arg_load_expr(arg_id: tuple[int, ...], dt: DataTypeCxx, is_ptr: bool = False, create_load: bool = True, arg_depth: int = 0, dbg_info: DebugInfo = ...) -> ExprCxx:
+def make_arg_load_expr(arg_id: tuple[int, ...], dt: DataTypeCxx, is_ptr: bool = False, create_load: bool = True, dbg_info: DebugInfo = ...) -> ExprCxx:
     ...
 def make_binary_op_expr(arg0: BinaryOpType, arg1: ExprCxx, arg2: ExprCxx) -> ExprCxx:
     ...
@@ -2686,7 +2656,7 @@ def make_cusparse_solver(arg0: DataTypeCxx, arg1: str, arg2: str) -> SparseSolve
     ...
 def make_double_cg_solver(arg0: SparseMatrix, arg1: int, arg2: float, arg3: bool) -> CGd:
     ...
-def make_external_tensor_expr(arg0: DataTypeCxx, arg1: int, arg2: tuple[int, ...], arg3: bool, arg4: int, arg5: BoundaryMode) -> ExprCxx:
+def make_external_tensor_expr(arg0: DataTypeCxx, arg1: int, arg2: tuple[int, ...], arg3: bool, arg4: BoundaryMode) -> ExprCxx:
     ...
 def make_external_tensor_grad_expr(arg0: ExprCxx) -> ExprCxx:
     ...
@@ -2704,11 +2674,11 @@ def make_rand_expr(arg0: DataTypeCxx, arg1: DebugInfo) -> ExprCxx:
     ...
 def make_reference(arg0: ExprCxx, arg1: DebugInfo) -> ExprCxx:
     ...
-def make_rw_texture_ptr_expr(arg0: tuple[int, ...], arg1: int, arg2: int, arg3: Format, arg4: int, arg5: DebugInfo) -> ExprCxx:
+def make_rw_texture_ptr_expr(arg0: tuple[int, ...], arg1: int, arg2: Format, arg3: int, arg4: DebugInfo) -> ExprCxx:
     ...
 def make_sparse_solver(arg0: DataTypeCxx, arg1: str, arg2: str) -> SparseSolver:
     ...
-def make_texture_ptr_expr(arg0: tuple[int, ...], arg1: int, arg2: int, arg3: DebugInfo) -> ExprCxx:
+def make_texture_ptr_expr(arg0: tuple[int, ...], arg1: int, arg2: DebugInfo) -> ExprCxx:
     ...
 def make_unary_op_expr(arg0: UnaryOpType, arg1: ExprCxx) -> ExprCxx:
     ...

gstaichi/_test_tools/__init__.py

@@ -0,0 +1,18 @@
+import gstaichi as ti
+
+from . import textwrap2
+
+
+def ti_init_same_arch(**options) -> None:
+    """
+    Used in tests to call ti.init, passing in the same arch as currently
+    configured. Since it's fairly fiddly to do that, extracting this out
+    to this helper function.
+    """
+    assert ti.cfg is not None
+    options = dict(options)
+    options["arch"] = getattr(ti, ti.cfg.arch.name)
+    ti.init(**options)
+
+
+__all__ = ["textwrap2"]

gstaichi/_test_tools/dataclass_test_tools.py

@@ -0,0 +1,36 @@
+import dataclasses
+from typing import Any, cast
+
+import gstaichi as ti
+
+
+def _make_child_obj(obj_type: Any) -> Any:
+    if isinstance(obj_type, ti.types.NDArray):
+        ndarray_type = cast(ti.types.ndarray, obj_type)
+        assert ndarray_type.ndim is not None
+        shape = tuple([10] * ndarray_type.ndim)
+        child_obj = ti.ndarray(ndarray_type.dtype, shape=shape)
+    elif dataclasses.is_dataclass(obj_type):
+        child_obj = build_struct(obj_type)
+    elif isinstance(obj_type, ti.Template) or obj_type == ti.Template:
+        child_obj = ti.field(ti.i32, (10,))
+    else:
+        raise Exception("unknown type ", obj_type)
+    return child_obj
+
+
+def build_struct(struct_type: Any) -> Any:
+    member_objects = {}
+    for field in dataclasses.fields(struct_type):
+        child_obj = _make_child_obj(field.type)
+        member_objects[field.name] = child_obj
+    dataclass_object = struct_type(**member_objects)
+    return dataclass_object
+
+
+def build_obj_tuple_from_type_dict(name_to_type: dict[str, Any]) -> tuple[Any, ...]:
+    obj_l = []
+    for _name, param_type in name_to_type.items():
+        child_obj = _make_child_obj(param_type)
+        obj_l.append(child_obj)
+    return tuple(obj_l)

gstaichi/_test_tools/textwrap2.py

@@ -0,0 +1,6 @@
+def dedent(indent_size: int, v: str) -> str:
+    lines = []
+    for line in v.split("\n"):
+        line = line[indent_size:]
+        lines.append(line)
+    return "\n".join(lines)

gstaichi/_version.py

@@ -1 +1 @@
-__version__ = '0.1.25.dev0'
+__version__ = '2.0.0'

gstaichi/examples/minimal.py

@@ -2,7 +2,7 @@ import gstaichi as ti
 
 
 @ti.kernel
-def lcg_ti(B: int, lcg_its: int, a: ti.types.ndarray(ti.i32, 1)) -> None:
+def lcg_ti(B: int, lcg_its: int, a: ti.types.NDArray[ti.i32, 1]) -> None:
     """
     Linear congruential generator https://en.wikipedia.org/wiki/Linear_congruential_generator
     """

gstaichi/lang/__init__.py

@@ -1,10 +1,10 @@
 # type: ignore
 
 from gstaichi.lang import impl, simt
+from gstaichi.lang._fast_caching.function_hasher import pure
 from gstaichi.lang._ndarray import *
 from gstaichi.lang._ndrange import ndrange
 from gstaichi.lang._texture import Texture
-from gstaichi.lang.argpack import *
 from gstaichi.lang.exception import *
 from gstaichi.lang.field import *
 from gstaichi.lang.impl import *

gstaichi/lang/_dataclass_util.py

@@ -0,0 +1,31 @@
+def create_flat_name(basename: str, child_name: str) -> str:
+    """
+    Appends child_name to basename, separated by __ti_.
+    If basename does not start with __ti_ then prefix the resulting string
+    with __ti_.
+
+    Note that we want to avoid adding prefix __ti_ if already included in `basename`,
+    to avoid duplicating said delimiter.
+
+    We'll use this when expanding py dataclass members, e.g.
+
+    @dataclasses.dataclass
+    def Foo:
+        a: int
+        b: int
+
+    foo = Foo(a=5, b=3)
+
+    When we expand out foo, we'll replace foo with the following names instead:
+    - __ti_foo__ti_a
+    - __ti_foo__ti_b
+
+    We use the __ti_ to ensure that it's easy to ensure no collision with existing user-defined
+    names. We require the user to not create any fields or variables which themselves are prefixed
+    with __ti_, and given this constraint, the names we create will not conflict with user-generated
+    names.
+    """
+    full_name = f"{basename}__ti_{child_name}"
+    if not full_name.startswith("__ti_"):
+        full_name = f"__ti_{full_name}"
+    return full_name

gstaichi/lang/_fast_caching/__init__.py

@@ -0,0 +1,3 @@
+from .args_hasher import FIELD_METADATA_CACHE_VALUE
+
+__all__ = ["FIELD_METADATA_CACHE_VALUE"]

gstaichi/lang/_fast_caching/args_hasher.py

@@ -0,0 +1,110 @@
+import dataclasses
+import enum
+import numbers
+import time
+from typing import Any, Sequence
+
+import numpy as np
+
+from .._ndarray import ScalarNdarray
+from ..field import ScalarField
+from ..matrix import MatrixField, MatrixNdarray, VectorNdarray
+from ..util import is_data_oriented
+from .hash_utils import hash_iterable_strings
+
+g_num_calls = 0
+g_num_args = 0
+g_hashing_time = 0
+g_repr_time = 0
+g_num_ignored_calls = 0
+
+
+FIELD_METADATA_CACHE_VALUE = "add_value_to_cache_key"
+
+
+def dataclass_to_repr(path: tuple[str, ...], arg: Any) -> str:
+    repr_l = []
+    for field in dataclasses.fields(arg):
+        child_value = getattr(arg, field.name)
+        _repr = stringify_obj_type(path + (field.name,), child_value)
+        full_repr = f"{field.name}: ({_repr})"
+        if field.metadata.get(FIELD_METADATA_CACHE_VALUE, False):
+            full_repr += f" = {child_value}"
+        repr_l.append(full_repr)
+    return "[" + ",".join(repr_l) + "]"
+
+
+def stringify_obj_type(path: tuple[str, ...], obj: Any) -> str | None:
+    """
+    Convert an object into a string representation that only depends on its type.
+
+    String should somehow represent the type of obj. Doesnt have to be hashed, nor does it have
+    to be the actual python type string, just a string that is representative of the type, and won't collide
+    with different (allowed) types.
+
+    `path` is used during debugging.
+    """
+    # TODO: We should have a way of printing this without having to hack the code really. Using logger perhaps?
+    # (I have another PR that addreses this https://github.com/Genesis-Embodied-AI/gstaichi/pull/144/files)
+    arg_type = type(obj)
+    if isinstance(obj, ScalarNdarray):
+        return f"[nd-{obj.dtype}-{len(obj.shape)}]"
+    if isinstance(obj, VectorNdarray):
+        return f"[ndv-{obj.n}-{obj.dtype}-{len(obj.shape)}]"
+    if isinstance(obj, ScalarField):
+        return f"[f-{obj.snode._id}-{obj.dtype}-{obj.shape}]"
+    if isinstance(obj, MatrixNdarray):
+        return f"[ndm-{obj.m}-{obj.n}-{obj.dtype}-{len(obj.shape)}]"
+    if "torch.Tensor" in str(arg_type):
+        return f"[pt-{obj.dtype}-{obj.ndim}]"
+    if isinstance(obj, np.ndarray):
+        return f"[np-{obj.dtype}-{obj.ndim}]"
+    if isinstance(obj, MatrixField):
+        return f"[fm-{obj.m}-{obj.n}-{obj.snode._id}-{obj.dtype}-{obj.shape}]"
+    if dataclasses.is_dataclass(obj):
+        return dataclass_to_repr(path, obj)
+    if is_data_oriented(obj):
+        child_repr_l = []
+        for k, v in obj.__dict__.items():
+            _child_repr = stringify_obj_type((*path, k), v)
+            if _child_repr is None:
+                print("not representable child", k, type(v), "path", path)
+                return None
+            child_repr_l.append(f"{k}: {_child_repr}")
+        return ", ".join(child_repr_l)
+    if issubclass(arg_type, (numbers.Number, np.number)):
+        return str(arg_type)
+    if arg_type is np.bool_:
+        # np is deprecating bool. Treat specially/carefully
+        return "np.bool_"
+    if isinstance(obj, enum.Enum):
+        return f"enum-{obj.name}-{obj.value}"
+    return None
+
+
+def hash_args(args: Sequence[Any]) -> str | None:
+    global g_num_calls, g_num_args, g_hashing_time, g_repr_time, g_num_ignored_calls
+    g_num_calls += 1
+    g_num_args += len(args)
+    hash_l = []
+    for i_arg, arg in enumerate(args):
+        start = time.time()
+        _hash = stringify_obj_type((str(i_arg),), arg)
+        g_repr_time += time.time() - start
+        if not _hash:
+            g_num_ignored_calls += 1
+            return None
+        hash_l.append(_hash)
+    start = time.time()
+    res = hash_iterable_strings(hash_l)
+    g_hashing_time += time.time() - start
+    return res
+
+
+def dump_stats() -> None:
+    print("args hasher dump stats")
+    print("total calls", g_num_calls)
+    print("ignored calls", g_num_ignored_calls)
+    print("total args", g_num_args)
+    print("hashing time", g_hashing_time)
+    print("arg representation time", g_repr_time)

gstaichi/lang/_fast_caching/config_hasher.py

@@ -0,0 +1,30 @@
+from gstaichi.lang import impl
+
+from .hash_utils import hash_iterable_strings
+
+EXCLUDE_PREFIXES = ["_", "offline_cache", "print_", "verbose_"]
+
+
+def hash_compile_config() -> str:
+    """
+    Calculates a hash string for the current compiler config.
+
+    If any value in the compiler config changes, the hash string changes too.
+
+    Though arguably we might want to blacklist certain keys, such as print_ir_debug,
+    which do not affect the compiled kernels, just stuff that gets printed during
+    the compilation process.
+    """
+    config = impl.get_runtime().prog.config()
+    config_l = []
+    for k in dir(config):
+        skip = False
+        for prefix in EXCLUDE_PREFIXES:
+            if k.startswith(prefix) or k in [""]:
+                skip = True
+        if skip:
+            continue
+        v = getattr(config, k)
+        config_l.append(f"{k}={v}")
+    config_hash = hash_iterable_strings(config_l, separator="\n")
+    return config_hash

gstaichi/lang/_fast_caching/fast_caching_types.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+
+from .._wrap_inspect import FunctionSourceInfo
+
+
+class HashedFunctionSourceInfo(BaseModel):
+    """
+    Wraps a function source info, and the hash string of that function.
+
+    By not adding the hash directly into function source info, we avoid
+    having to make hash an optional type, and checking if it's empty or not.
+
+    If you have a HashedFunctionSourceInfo object, then you are guaranteed
+    to have the hash string.
+
+    If you only have the FunctionSourceInfo object, you are guaranteed that it
+    does not have a hash string.
+    """
+
+    function_source_info: FunctionSourceInfo
+    hash: str

gstaichi/lang/_fast_caching/function_hasher.py

@@ -0,0 +1,57 @@
+import os
+from itertools import islice
+from typing import TYPE_CHECKING, Iterable
+
+from .._wrap_inspect import FunctionSourceInfo
+from .fast_caching_types import HashedFunctionSourceInfo
+from .hash_utils import hash_iterable_strings
+
+if TYPE_CHECKING:
+    from gstaichi.lang.kernel_impl import GsTaichiCallable
+
+
+def pure(fn: "GsTaichiCallable") -> "GsTaichiCallable":
+    fn.is_pure = True
+    return fn
+
+
+def _read_file(function_info: FunctionSourceInfo) -> list[str]:
+    with open(function_info.filepath) as f:
+        return list(islice(f, function_info.start_lineno, function_info.end_lineno + 1))
+
+
+def _hash_function(function_info: FunctionSourceInfo) -> str:
+    return hash_iterable_strings(_read_file(function_info))
+
+
+def hash_functions(function_infos: Iterable[FunctionSourceInfo]) -> list[HashedFunctionSourceInfo]:
+    results = []
+    for f_info in function_infos:
+        hash_ = _hash_function(f_info)
+        results.append(HashedFunctionSourceInfo(function_source_info=f_info, hash=hash_))
+    return results
+
+
+def hash_kernel(kernel_info: FunctionSourceInfo) -> str:
+    return _hash_function(kernel_info)
+
+
+def dump_stats() -> None:
+    print("function hasher dump stats")
+
+
+def _validate_hashed_function_info(hashed_function_info: HashedFunctionSourceInfo) -> bool:
+    """
+    Checks the hash
+    """
+    if not os.path.isfile(hashed_function_info.function_source_info.filepath):
+        return False
+    _hash = _hash_function(hashed_function_info.function_source_info)
+    return _hash == hashed_function_info.hash
+
+
+def validate_hashed_function_infos(function_infos: Iterable[HashedFunctionSourceInfo]) -> bool:
+    for function_info in function_infos:
+        if not _validate_hashed_function_info(function_info):
+            return False
+    return True

gstaichi/lang/_fast_caching/hash_utils.py

@@ -0,0 +1,11 @@
+import hashlib
+from typing import Iterable
+
+
+def hash_iterable_strings(strings: Iterable[str], separator: str = "_") -> str:
+    h = hashlib.sha256()
+    separator_enc = separator.encode("utf-8")
+    for v in strings:
+        h.update(v.encode("utf-8"))
+        h.update(separator_enc)
+    return h.hexdigest()

gstaichi/lang/_fast_caching/python_side_cache.py

@@ -0,0 +1,52 @@
+import os
+
+from .. import impl
+
+
+class PythonSideCache:
+    """
+    Manages a cache that is managed from the python side (we also have c++-side caches)
+
+    The cache is disk-based. When we create the PythonSideCache object, the cache
+    path is created as a sub-folder of CompileConfig.offline_cache_file_path.
+
+    Note that constructing this object is cheap, so there is no need to maintain some
+    kind of conceptual singleton instance or similar.
+
+    Each cache key value is stored to a single file, with the cache key as the filename.
+
+    No metadata is associated with the file, making management very lightweight.
+
+    We update the file date/time when we read from a particular file, so we can easily
+    implement an LRU cleaning strategy at some point in the future, based on the file
+    date/times.
+    """
+
+    def __init__(self) -> None:
+        _cache_parent_folder = impl.get_runtime().prog.config().offline_cache_file_path
+        self.cache_folder = os.path.join(_cache_parent_folder, "python_side_cache")
+        os.makedirs(self.cache_folder, exist_ok=True)
+
+    def _get_filepath(self, key: str) -> str:
+        filepath = os.path.join(self.cache_folder, f"{key}.cache.txt")
+        return filepath
+
+    def _touch(self, filepath):
+        """
+        Updates file date/time.
+        """
+        with open(filepath, "a"):
+            os.utime(filepath, None)
+
+    def store(self, key: str, value: str) -> None:
+        filepath = self._get_filepath(key)
+        with open(filepath, "w") as f:
+            f.write(value)
+
+    def try_load(self, key: str) -> str | None:
+        filepath = self._get_filepath(key)
+        if not os.path.isfile(filepath):
+            return None
+        self._touch(filepath)
+        with open(filepath) as f:
+            return f.read()