warp-lang 1.7.2rc1__py3-none-macosx_10_13_universal2.whl → 1.8.1__py3-none-macosx_10_13_universal2.whl

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,

warp/fem/domain.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, Optional, Set, Union
 
 import warp as wp
@@ -27,7 +28,7 @@ from warp.fem.geometry import (
     WholeGeometryPartition,
 )
 from warp.fem.operator import Operator
-from warp.fem.types import ElementKind
+from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind
 
 GeometryOrPartition = Union[Geometry, GeometryPartition]
 
@@ -89,6 +90,9 @@ class GeometryDomain:
     element_index: wp.Function
     """Device function for retrieving an ElementIndex from a linearized index"""
 
+    element_partition_index: wp.Function
+    """Device function for retrieving linearized index in the domain's partition from an ElementIndex"""
+
     ElementArg: warp.codegen.Struct
     """Structure containing arguments to be passed to device functions computing element geometry"""
 
@@ -107,13 +111,34 @@ class GeometryDomain:
     element_normal: wp.Function
     """Device function returning the element normal at a sample point"""
 
+    element_closest_point: wp.Function
+    """Device function returning the coordinates of the closest point in a given element to a world position"""
+
+    element_coordinates: wp.Function
+    """Device function returning the coordinates corresponding to a world position in a given element reference system"""
+
     element_lookup: wp.Function
-    """Device function returning the sample point corresponding to a world position"""
+    """Device function returning the sample point in the domain's geometry corresponding to a world position"""
+
+    element_partition_lookup: wp.Function
+    """Device function returning the sample point in the domain's geometry partition corresponding to a world position"""
 
     def notify_operator_usage(self, ops: Set[Operator]):
         """Makes the Domain aware that the operators `ops` will be applied"""
         pass
 
+    @cached_property
+    def DomainArg(self):
+        return self._make_domain_arg()
+
+    def _make_domain_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class DomainArg:
+            geo: self.ElementArg
+            index: self.ElementIndexArg
+
+        return DomainArg
+
 
 class Cells(GeometryDomain):
     """A Domain containing all cells of the geometry or geometry partition"""
@@ -145,13 +170,23 @@ class Cells(GeometryDomain):
     def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
         return self.geometry_partition.cell_arg_value(device)
 
+    def fill_element_index_arg(self, arg: ElementIndexArg, device: warp.context.Devicelike):
+        self.geometry_partition.fill_cell_arg(arg, device)
+
     @property
     def element_index(self) -> wp.Function:
         return self.geometry_partition.cell_index
 
+    @property
+    def element_partition_index(self) -> wp.Function:
+        return self.geometry_partition.partition_cell_index
+
     def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
         return self.geometry.cell_arg_value(device)
 
+    def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
+        self.geometry.fill_cell_arg(arg, device)
+
     @property
     def ElementArg(self) -> warp.codegen.Struct:
         return self.geometry.CellArg
@@ -176,10 +211,46 @@ class Cells(GeometryDomain):
     def element_normal(self) -> wp.Function:
         return self.geometry.cell_normal
 
+    @property
+    def element_closest_point(self) -> wp.Function:
+        return self.geometry.cell_closest_point
+
+    @property
+    def element_coordinates(self) -> wp.Function:
+        return self.geometry.cell_coordinates
+
     @property
     def element_lookup(self) -> wp.Function:
         return self.geometry.cell_lookup
 
+    @property
+    def element_partition_lookup(self) -> wp.Function:
+        pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
+
+        @cache.dynamic_func(suffix=self.geometry_partition.name)
+        def is_in_partition(args: self.ElementIndexArg, cell_index: int):
+            return self.geometry_partition.partition_cell_index(args, cell_index) != NULL_ELEMENT_INDEX
+
+        filtered_cell_lookup = self.geometry.make_filtered_cell_lookup(filter_func=is_in_partition)
+
+        # overloads
+        filter_target = True
+        pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_partition_lookup(args: self.DomainArg, pos: pos_type, max_dist: float):
+            return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_partition_lookup(args: self.DomainArg, pos: pos_type):
+            max_dist = 0.0
+            return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)
+
+        return cell_partition_lookup
+
+    def supports_lookup(self, device):
+        return self.geometry.supports_cell_lookup(device)
+
     @property
     def domain_cell_arg(self) -> wp.Function:
         return Cells._identity_fn
@@ -200,6 +271,11 @@ class Sides(GeometryDomain):
         super().__init__(geometry)
 
         self.element_lookup = None
+        self.element_partition_lookup = None
+        self.element_filtered_lookup = None
+
+    def supports_lookup(self, device):
+        return False
 
     @property
     def element_kind(self) -> ElementKind:
@@ -225,6 +301,9 @@ class Sides(GeometryDomain):
     def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
         return self.geometry_partition.side_arg_value(device)
 
+    def fill_element_index_arg(self, arg: "ElementIndexArg", device: warp.context.Devicelike):
+        self.geometry_partition.fill_side_arg(arg, device)
+
     @property
     def element_index(self) -> wp.Function:
         return self.geometry_partition.side_index
@@ -236,6 +315,9 @@ class Sides(GeometryDomain):
     def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
         return self.geometry.side_arg_value(device)
 
+    def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
+        self.geometry.fill_side_arg(arg, device)
+
     @property
     def element_position(self) -> wp.Function:
         return self.geometry.side_position
@@ -256,6 +338,14 @@ class Sides(GeometryDomain):
     def element_normal(self) -> wp.Function:
         return self.geometry.side_normal
 
+    @property
+    def element_closest_point(self) -> wp.Function:
+        return self.geometry.side_closest_point
+
+    @property
+    def element_coordinates(self) -> wp.Function:
+        return self.geometry.side_coordinates
+
     @property
     def element_inner_cell_index(self) -> wp.Function:
         return self.geometry.side_inner_cell_index
@@ -276,9 +366,18 @@ class Sides(GeometryDomain):
     def cell_to_element_coords(self) -> wp.Function:
         return self.geometry.side_from_cell_coords
 
-    @property
+    @cached_property
     def domain_cell_arg(self) -> wp.Function:
-        return self.geometry.side_to_cell_arg
+        CellDomainArg = self.cell_domain().DomainArg
+
+        @cache.dynamic_func(suffix=self.name)
+        def domain_cell_arg(x: self.DomainArg):
+            return CellDomainArg(
+                self.geometry.side_to_cell_arg(x.geo),
+                self.geometry_partition.side_to_cell_arg(x.index),
+            )
+
+        return domain_cell_arg
 
     def cell_domain(self):
         return Cells(self.geometry_partition)
@@ -359,11 +458,13 @@ class Subdomain(GeometryDomain):
         self.geometry_element_count = self._domain.geometry_element_count
         self.reference_element = self._domain.reference_element
         self.element_arg_value = self._domain.element_arg_value
+        self.fill_element_arg = self._domain.fill_element_arg
         self.element_measure = self._domain.element_measure
         self.element_measure_ratio = self._domain.element_measure_ratio
         self.element_position = self._domain.element_position
         self.element_deformation_gradient = self._domain.element_deformation_gradient
         self.element_lookup = self._domain.element_lookup
+        self.element_partition_lookup = self._domain.element_partition_lookup
         self.element_normal = self._domain.element_normal
 
     @property
@@ -399,13 +500,42 @@ class Subdomain(GeometryDomain):
     @cache.cached_arg_value
     def element_index_arg_value(self, device: warp.context.Devicelike):
         arg = self.ElementIndexArg()
-        arg.domain_arg = self._domain.element_index_arg_value(device)
-        arg.element_indices = self._element_indices.to(device)
+        self.fill_element_index_arg(arg, device)
         return arg
 
+    def fill_element_index_arg(self, arg: "GeometryDomain.ElementIndexArg", device: warp.context.Devicelike):
+        self._domain.fill_element_index_arg(arg.domain_arg, device)
+        arg.element_indices = self._element_indices.to(device)
+
     def _make_element_index(self) -> wp.Function:
         @cache.dynamic_func(suffix=self.name)
         def element_index(arg: self.ElementIndexArg, index: int):
             return self._domain.element_index(arg.domain_arg, arg.element_indices[index])
 
         return element_index
+
+    def _make_element_partition_index(self) -> wp.Function:
+        @cache.dynamic_func(suffix=self.name)
+        def element_partition_index(arg: self.ElementIndexArg, element_index: int):
+            return self._domain.element_partition_index(arg.domain_arg, element_index)
+
+        return element_partition_index
+
+    def supports_lookup(self, device):
+        return self._domain.supports_lokup(device)
+
+    def cell_domain(self):
+        return self._domain.cell_domain()
+
+    @cached_property
+    def domain_cell_arg(self) -> wp.Function:
+        CellDomainArg = self.cell_domain().DomainArg
+
+        @cache.dynamic_func(suffix=self.name)
+        def domain_cell_arg(x: self.DomainArg):
+            return CellDomainArg(
+                self.geometry.side_to_cell_arg(x.geo),
+                self.geometry_partition.side_to_cell_arg(x.index.domain_arg),
+            )
+
+        return domain_cell_arg