warp-lang 1.7.2__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/fabric.py

@@ -22,12 +22,12 @@ from warp.types import *
 
 
 class fabricbucket_t(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("index_start", ctypes.c_size_t),
         ("index_end", ctypes.c_size_t),
         ("ptr", ctypes.c_void_p),
         ("lengths", ctypes.c_void_p),
-    ]
+    )
 
     def __init__(self, index_start=0, index_end=0, ptr=None, lengths=None):
         self.index_start = index_start
@@ -37,11 +37,11 @@ class fabricbucket_t(ctypes.Structure):
 
 
 class fabricarray_t(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("buckets", ctypes.c_void_p),  # array of fabricbucket_t on the correct device
         ("nbuckets", ctypes.c_size_t),
         ("size", ctypes.c_size_t),
-    ]
+    )
 
     def __init__(self, buckets=None, nbuckets=0, size=0):
         self.buckets = ctypes.c_void_p(buckets)
@@ -50,11 +50,11 @@ class fabricarray_t(ctypes.Structure):
 
 
 class indexedfabricarray_t(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("fa", fabricarray_t),
         ("indices", ctypes.c_void_p),
         ("size", ctypes.c_size_t),
-    ]
+    )
 
     def __init__(self, fa=None, indices=None):
         if fa is None:
@@ -121,7 +121,7 @@ class fabricarray(noncontiguous_array_base[T]):
     _vars = None
 
     def __new__(cls, *args, **kwargs):
-        instance = super(fabricarray, cls).__new__(cls)
+        instance = super().__new__(cls)
         instance.deleter = None
         return instance
 

warp/fem/__init__.py

@@ -55,6 +55,8 @@ from .operator import (
     degree,
     div,
     div_outer,
+    element_closest_point,
+    element_coordinates,
     grad,
     grad_average,
     grad_jump,
@@ -65,8 +67,11 @@ from .operator import (
     lookup,
     measure,
     measure_ratio,
+    node_count,
+    node_index,
     normal,
     outer,
+    partition_lookup,
     position,
     to_cell_side,
     to_inner_cell,

warp/fem/adaptivity.py

@@ -50,7 +50,7 @@ def adaptive_nanogrid_from_hierarchy(
     # Concatenate voxels for each grid
     voxel_counts = [grid.get_voxel_count() for grid in grids]
 
-    voxel_offsets = np.cumsum(np.array([0] + voxel_counts))
+    voxel_offsets = np.cumsum(np.array([0, *voxel_counts]))
     merged_ijks = cache.borrow_temporary(temporary_store, dtype=wp.vec3i, shape=int(voxel_offsets[-1]), device=device)
     for l in range(level_count):
         voxel_count = voxel_counts[l]

warp/fem/cache.py

@@ -15,13 +15,15 @@
 
 import ast
 import bisect
+import hashlib
 import re
 import weakref
-from copy import copy
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import warp as wp
+from warp.codegen import get_annotations
 from warp.fem.operator import Integrand
+from warp.fem.types import Domain, Field
 
 _kernel_cache = {}
 _struct_cache = {}
@@ -30,31 +32,88 @@ _func_cache = {}
 _key_re = re.compile("[^0-9a-zA-Z_]+")
 
 
-def _make_key(obj, suffix: str, use_qualified_name):
-    base_name = f"{obj.__module__}.{obj.__qualname__}" if use_qualified_name else obj.__name__
-    return _key_re.sub("", f"{base_name}_{suffix}")
+def _make_key(obj, suffix: str, options: Optional[Dict[str, Any]] = None):
+    # human-readable part
+    key = _key_re.sub("", f"{obj.__name__}_{suffix}")
 
+    sorted_opts = sorted(options.items()) if options is not None else ()
+    opts_str = "".join(
+        (
+            obj.__module__,
+            obj.__qualname__,
+            suffix,
+            *(opt[0] for opt in sorted_opts),
+            *(str(opt[1]) for opt in sorted_opts),
+        )
+    )
+    uid = hashlib.blake2b(bytes(opts_str, encoding="utf-8"), digest_size=4).hexdigest()
 
-def get_func(func, suffix: str, use_qualified_name: bool = False, code_transformers=None):
-    key = _make_key(func, suffix, use_qualified_name)
+    # avoid long keys, issues on win
+    key = f"{key[:64]}_{uid}"
 
-    if key not in _func_cache:
-        _func_cache[key] = wp.Function(
-            func=func,
-            key=key,
-            namespace="",
-            module=wp.get_module(
-                func.__module__,
-            ),
+    return key
+
+
+def _arg_type_name(arg_type):
+    if isinstance(arg_type, str):
+        return arg_type
+    if arg_type in (Field, Domain):
+        return ""
+    return wp.types.get_type_code(wp.types.type_to_warp(arg_type))
+
+
+def _make_cache_key(func, key, argspec=None):
+    if argspec is None:
+        annotations = get_annotations(func)
+    else:
+        annotations = argspec.annotations
+
+    sig_key = (key, tuple((k, _arg_type_name(v)) for k, v in annotations.items()))
+    return sig_key
+
+
+def _register_function(
+    func,
+    key,
+    module,
+    **kwargs,
+):
+    # wp.Function will override existing func for a given key...
+    # manually add back our overloads
+    existing = module.functions.get(key)
+    new_fn = wp.Function(
+        func=func,
+        key=key,
+        namespace="",
+        module=module,
+        **kwargs,
+    )
+
+    if existing:
+        existing.add_overload(new_fn)
+        module.functions[key] = existing
+    return module.functions[key]
+
+
+def get_func(func, suffix: str, code_transformers=None):
+    key = _make_key(func, suffix)
+    cache_key = _make_cache_key(func, key)
+
+    if cache_key not in _func_cache:
+        module = wp.get_module(func.__module__)
+        _func_cache[cache_key] = _register_function(
+            func,
+            key,
+            module,
             code_transformers=code_transformers,
         )
 
-    return _func_cache[key]
+    return _func_cache[cache_key]
 
 
-def dynamic_func(suffix: str, use_qualified_name=False, code_transformers=None):
+def dynamic_func(suffix: str, code_transformers=None):
     def wrap_func(func: Callable):
-        return get_func(func, suffix=suffix, use_qualified_name=use_qualified_name, code_transformers=code_transformers)
+        return get_func(func, suffix=suffix, code_transformers=code_transformers)
 
     return wrap_func
 
@@ -62,38 +121,35 @@ def dynamic_func(suffix: str, use_qualified_name=False, code_transformers=None):
 def get_kernel(
     func,
     suffix: str,
-    use_qualified_name: bool = False,
-    kernel_options: Dict[str, Any] = None,
+    kernel_options: Optional[Dict[str, Any]] = None,
 ):
     if kernel_options is None:
         kernel_options = {}
 
-    key = _make_key(func, suffix, use_qualified_name)
+    key = _make_key(func, suffix, kernel_options)
+    cache_key = _make_cache_key(func, key)
 
-    if key not in _kernel_cache:
-        # Avoid creating too long file names -- can lead to issues on Windows
-        # We could hash the key, but prefer to keep it human-readable
+    if cache_key not in _kernel_cache:
         module_name = f"{func.__module__}.dyn.{key}"
-        module_name = module_name[:128] if len(module_name) > 128 else module_name
         module = wp.get_module(module_name)
-        module.options = copy(wp.get_module(func.__module__).options)
+        module.options = dict(wp.get_module(func.__module__).options)
         module.options.update(kernel_options)
-        _kernel_cache[key] = wp.Kernel(func=func, key=key, module=module)
-    return _kernel_cache[key]
+        _kernel_cache[cache_key] = wp.Kernel(func=func, key=key, module=module, options=kernel_options)
+    return _kernel_cache[cache_key]
 
 
-def dynamic_kernel(suffix: str, use_qualified_name=False, kernel_options: Dict[str, Any] = None):
+def dynamic_kernel(suffix: str, kernel_options: Optional[Dict[str, Any]] = None):
     if kernel_options is None:
         kernel_options = {}
 
     def wrap_kernel(func: Callable):
-        return get_kernel(func, suffix=suffix, use_qualified_name=use_qualified_name, kernel_options=kernel_options)
+        return get_kernel(func, suffix=suffix, kernel_options=kernel_options)
 
     return wrap_kernel
 
 
-def get_struct(struct: type, suffix: str, use_qualified_name: bool = False):
-    key = _make_key(struct, suffix, use_qualified_name)
+def get_struct(struct: type, suffix: str):
+    key = _make_key(struct, suffix)
     # used in codegen
     struct.__qualname__ = key
 
@@ -108,9 +164,9 @@ def get_struct(struct: type, suffix: str, use_qualified_name: bool = False):
     return _struct_cache[key]
 
 
-def dynamic_struct(suffix: str, use_qualified_name=False):
+def dynamic_struct(suffix: str):
     def wrap_struct(struct: type):
-        return get_struct(struct, suffix=suffix, use_qualified_name=use_qualified_name)
+        return get_struct(struct, suffix=suffix)
 
     return wrap_struct
 
@@ -125,35 +181,36 @@ def get_argument_struct(arg_types: Dict[str, type]):
         setattr(Args, name, None)
         annotations[name] = arg_type
 
-    def arg_type_name(arg_type):
-        return wp.types.get_type_code(wp.types.type_to_warp(arg_type))
-
     try:
         Args.__annotations__ = annotations
     except AttributeError:
         Args.__dict__.__annotations__ = annotations
 
-    suffix = "_".join([f"{name}_{arg_type_name(arg_type)}" for name, arg_type in annotations.items()])
+    suffix = "_".join([f"{name}_{_arg_type_name(arg_type)}" for name, arg_type in annotations.items()])
 
     return get_struct(Args, suffix=suffix)
 
 
-def populate_argument_struct(Args: wp.codegen.Struct, values: Dict[str, Any], func_name: str):
+def populate_argument_struct(
+    Args: wp.codegen.Struct, values: Dict[str, Any], func_name: str, value_struct_values: Optional = None
+):
     if values is None:
         values = {}
 
-    value_struct_values = Args()
-    for k, v in values.items():
-        try:
+    if value_struct_values is None:
+        value_struct_values = Args()
+
+    try:
+        for k, v in values.items():
             setattr(value_struct_values, k, v)
-        except Exception as err:
-            if k not in Args.vars:
-                raise ValueError(
-                    f"Passed value argument '{k}' does not match any of the function '{func_name}' parameters"
-                ) from err
+    except Exception as err:
+        if k not in Args.vars:
             raise ValueError(
-                f"Passed value argument '{k}' of type '{wp.types.type_repr(v)}' is incompatible with the function '{func_name}' parameter of type '{wp.types.type_repr(Args.vars[k].type)}'"
+                f"Passed value argument '{k}' does not match any of the function '{func_name}' parameters"
             ) from err
+        raise ValueError(
+            f"Passed value argument '{k}' of type '{wp.types.type_repr(v)}' is incompatible with the function '{func_name}' parameter of type '{wp.types.type_repr(Args.vars[k].type)}'"
+        ) from err
 
     missing_values = Args.vars.keys() - values.keys()
     if missing_values:
@@ -208,26 +265,26 @@ def get_integrand_function(
     annotations=None,
     code_transformers=None,
 ):
-    key = _make_key(integrand.func, suffix, use_qualified_name=True)
+    key = _make_key(integrand.func, suffix)
+    cache_key = _make_cache_key(integrand.func, key, integrand.argspec)
 
-    if key not in _func_cache:
-        _func_cache[key] = wp.Function(
+    if cache_key not in _func_cache:
+        _func_cache[cache_key] = _register_function(
             func=integrand.func if func is None else func,
             key=key,
-            namespace="",
             module=integrand.module,
             overloaded_annotations=annotations,
             code_transformers=code_transformers,
         )
 
-    return _func_cache[key]
+    return _func_cache[cache_key]
 
 
 def get_integrand_kernel(
     integrand: Integrand,
     suffix: str,
     kernel_fn: Optional[Callable] = None,
-    kernel_options: Dict[str, Any] = None,
+    kernel_options: Optional[Dict[str, Any]] = None,
     code_transformers=None,
 ):
     options = integrand.module.options.copy()
@@ -235,15 +292,15 @@ def get_integrand_kernel(
     if kernel_options is not None:
         options.update(kernel_options)
 
-    kernel_key = _make_key(integrand.func, suffix, use_qualified_name=True)
-    opts_key = "".join([f"{k}:{v}" for k, v in sorted(options.items())])
-    cache_key = kernel_key + opts_key
+    kernel_key = _make_key(integrand.func, suffix, options=options)
+    cache_key = _make_cache_key(integrand, kernel_key, integrand.argspec)
 
     if cache_key not in _kernel_cache:
         if kernel_fn is None:
             return None
 
-        module = wp.get_module(f"{integrand.module.name}.{integrand.name}")
+        module = wp.get_module(f"{integrand.module.name}.{kernel_key}")
+        module.options = options
         _kernel_cache[cache_key] = wp.Kernel(
             func=kernel_fn, key=kernel_key, module=module, code_transformers=code_transformers, options=options
         )
@@ -272,6 +329,40 @@ def cached_arg_value(func: Callable):
     return get_arg
 
 
+def setup_dynamic_attributes(
+    obj,
+    cls: Optional[type] = None,
+    constructors: Optional[Dict[str, Callable]] = None,
+    key: Optional[str] = None,
+):
+    if cls is None:
+        cls = type(obj)
+
+    if key is None:
+        key = obj.name
+
+    if constructors is None:
+        constructors = cls._dynamic_attribute_constructors
+
+    key = (key, frozenset(constructors.keys()))
+
+    if not hasattr(cls, "_cached_dynamic_attrs"):
+        cls._cached_dynamic_attrs = {}
+
+    attrs = cls._cached_dynamic_attrs.get(key)
+    if attrs is None:
+        attrs = {}
+        # create attributes one-by-one, as some may depend on previous ones
+        for k, v in constructors.items():
+            attr = v(obj)
+            attrs[k] = attr
+            setattr(obj, k, attr)
+        cls._cached_dynamic_attrs[key] = attrs
+    else:
+        for k, v in attrs.items():
+            setattr(obj, k, v)
+
+
 _cached_vec_types = {}
 _cached_mat_types = {}
 
@@ -301,7 +392,7 @@ class Temporary:
     """
 
     def __new__(cls, *args, **kwargs):
-        instance = super(Temporary, cls).__new__(cls)
+        instance = super().__new__(cls)
         instance._pool = None
         return instance
 
@@ -447,15 +538,13 @@ class TemporaryStore:
         dtype = wp.types.type_to_warp(dtype)
         device = wp.get_device(device)
 
-        type_length = wp.types.type_length(dtype)
-        key = (dtype._type_, type_length, pinned, device.ordinal)
+        type_size = wp.types.type_size(dtype)
+        key = (dtype._type_, type_size, pinned, device.ordinal)
 
         pool = self._temporaries.get(key, None)
         if pool is None:
             value_type = (
-                cached_vec_type(length=type_length, dtype=wp.types.type_scalar_type(dtype))
-                if type_length > 1
-                else dtype
+                cached_vec_type(length=type_size, dtype=wp.types.type_scalar_type(dtype)) if type_size > 1 else dtype
             )
             pool = TemporaryStore.Pool(value_type, device, pinned=pinned)
             self._temporaries[key] = pool

warp/fem/dirichlet.py

@@ -18,7 +18,7 @@ from typing import Any, Optional
 import warp as wp
 from warp.fem.linalg import array_axpy, symmetric_eigenvalues_qr
 from warp.sparse import BsrMatrix, bsr_assign, bsr_axpy, bsr_copy, bsr_mm, bsr_mv
-from warp.types import type_is_matrix, type_length
+from warp.types import type_is_matrix, type_size
 
 
 def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None):
@@ -53,7 +53,7 @@ def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Opti
         if fixed_value.shape[0] != projector_matrix.nrow:
             raise ValueError("Fixed value array must be of length equal to the number of rows of blocks")
 
-        if type_length(fixed_value.dtype) == 1:
+        if type_size(fixed_value.dtype) == 1:
             # array of scalars, convert to 1d array of vectors
             fixed_value = wp.array(
                 data=None,