warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/fem/cache.py

@@ -0,0 +1,389 @@
+from typing import Callable, Optional, Union, Tuple, Dict, Any
+from copy import copy
+import bisect
+import re
+
+
+import warp as wp
+
+
+_kernel_cache = dict()
+_struct_cache = dict()
+_func_cache = dict()
+
+_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 get_func(func, suffix: str, use_qualified_name: bool = False):
+    key = _make_key(func, suffix, use_qualified_name)
+
+    if key not in _func_cache:
+        _func_cache[key] = wp.Function(
+            func=func,
+            key=key,
+            namespace="",
+            module=wp.get_module(
+                func.__module__,
+            ),
+        )
+
+    return _func_cache[key]
+
+
+def dynamic_func(suffix: str, use_qualified_name=False):
+    def wrap_func(func: Callable):
+        return get_func(func, suffix=suffix, use_qualified_name=use_qualified_name)
+
+    return wrap_func
+
+
+def get_kernel(
+    func,
+    suffix: str,
+    use_qualified_name: bool = False,
+    kernel_options: Dict[str, Any] = {},
+):
+    key = _make_key(func, suffix, use_qualified_name)
+
+    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
+        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.update(kernel_options)
+        _kernel_cache[key] = wp.Kernel(func=func, key=key, module=module)
+    return _kernel_cache[key]
+
+
+def dynamic_kernel(suffix: str, use_qualified_name=False, kernel_options: Dict[str, Any] = {}):
+    def wrap_kernel(func: Callable):
+        return get_kernel(func, suffix=suffix, use_qualified_name=use_qualified_name, 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)
+    # used in codegen
+    struct.__qualname__ = key
+
+    if key not in _struct_cache:
+        module = wp.get_module(struct.__module__)
+        _struct_cache[key] = wp.codegen.Struct(
+            cls=struct,
+            key=key,
+            module=module,
+        )
+
+    return _struct_cache[key]
+
+
+def dynamic_struct(suffix: str, use_qualified_name=False):
+    def wrap_struct(struct: type):
+        return get_struct(struct, suffix=suffix, use_qualified_name=use_qualified_name)
+
+    return wrap_struct
+
+
+def get_integrand_function(
+    integrand: "warp.fem.operator.Integrand",
+    suffix: str,
+    func=None,
+    annotations=None,
+    code_transformers=[],
+):
+    key = _make_key(integrand.func, suffix, use_qualified_name=True)
+
+    if key not in _func_cache:
+        _func_cache[key] = wp.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]
+
+
+def get_integrand_kernel(
+    integrand: "warp.fem.operator.Integrand",
+    suffix: str,
+    kernel_fn: Optional[Callable] = None,
+    kernel_options: Dict[str, Any] = {},
+    code_transformers=[],
+):
+    key = _make_key(integrand.func, suffix, use_qualified_name=True)
+
+    if key not in _kernel_cache:
+        if kernel_fn is None:
+            return None
+
+        module = wp.get_module(f"{integrand.module.name}.{integrand.name}")
+        module.options = copy(integrand.module.options)
+        module.options.update(kernel_options)
+
+        _kernel_cache[key] = wp.Kernel(func=kernel_fn, key=key, module=module, code_transformers=code_transformers)
+    return _kernel_cache[key]
+
+
+def cached_arg_value(func: Callable):
+    """Decorator to be applied to member methods assembling Arg structs, so that the result gets
+    automatically cached for the lifetime of the parent object
+    """
+
+    cache_attr = f"_{func.__name__}_cache"
+
+    def get_arg(obj, device):
+        if not hasattr(obj, cache_attr):
+            setattr(obj, cache_attr, {})
+
+        cache = getattr(obj, cache_attr, {})
+
+        device = wp.get_device(device)
+        if device.ordinal not in cache:
+            cache[device.ordinal] = func(obj, device)
+
+        return cache[device.ordinal]
+
+    return get_arg
+
+
+_cached_vec_types = {}
+_cached_mat_types = {}
+
+
+def cached_vec_type(length, dtype):
+    key = (length, dtype)
+    if key not in _cached_vec_types:
+        _cached_vec_types[key] = wp.vec(length=length, dtype=dtype)
+
+    return _cached_vec_types[key]
+
+
+def cached_mat_type(shape, dtype):
+    key = (*shape, dtype)
+    if key not in _cached_mat_types:
+        _cached_mat_types[key] = wp.mat(shape=shape, dtype=dtype)
+
+    return _cached_mat_types[key]
+
+
+class Temporary:
+    """Handle over a temporary array from a :class:`TemporaryStore`.
+
+    The array will be automatically returned to the temporary pool for reuse upon destruction of this object, unless
+    the temporary is explicitly detached from the pool using :meth:`detach`.
+    The temporary may also be explicitly returned to the pool before destruction using :meth:`release`.
+    """
+
+    def __init__(self, array: wp.array, pool: Optional["TemporaryStore.Pool"] = None, shape=None, dtype=None):
+        self._raw_array = array
+        self._array_view = array
+        self._pool = pool
+
+        if shape is not None or dtype is not None:
+            self._view_as(shape=shape, dtype=dtype)
+
+    def detach(self) -> wp.array:
+        """Detaches the temporary so it is never returned to the pool"""
+        if self._pool is not None:
+            self._pool.detach(self._raw_array)
+
+        self._pool = None
+        return self._array_view
+
+    def release(self):
+        """Returns the temporary array to the pool"""
+        if self._pool is not None:
+            self._pool.redeem(self._raw_array)
+
+        self._pool = None
+
+    @property
+    def array(self) -> wp.array:
+        """View of the array with desired shape and data type."""
+        return self._array_view
+
+    def _view_as(self, shape, dtype) -> "Temporary":
+        def _view_reshaped_truncated(array):
+            return wp.types.array(
+                ptr=array.ptr,
+                dtype=dtype,
+                shape=shape,
+                device=array.device,
+                pinned=array.pinned,
+                capacity=array.capacity,
+                copy=False,
+                owner=False,
+                grad=None if array.grad is None else _view_reshaped_truncated(array.grad),
+            )
+
+        self._array_view = _view_reshaped_truncated(self._raw_array)
+        return self
+
+    def __del__(self):
+        self.release()
+
+
+class TemporaryStore:
+    """
+    Shared pool of temporary arrays that will be persisted and reused across invocations of ``warp.fem`` functions.
+
+    A :class:`TemporaryStore` instance may either be passed explicitly to ``warp.fem`` functions that accept such an argument, for instance :func:`.integrate.integrate`,
+    or can be set globally as the default store using :func:`set_default_temporary_store`.
+
+    By default, there is no default temporary store, so that temporary allocations are not persisted.
+    """
+
+    _default_store: "TemporaryStore" = None
+
+    class Pool:
+        def __init__(self, dtype, device, pinned: bool):
+            self.dtype = dtype
+            self.device = device
+            self.pinned = pinned
+
+            self._pool = []  # Currently available arrays for borrowing, ordered by size
+            self._pool_sizes = []  # Sizes of available arrays for borrowing, ascending
+            self._allocs = {}  # All allocated arrays, including borrowed ones
+
+        def borrow(self, shape, dtype, requires_grad: bool):
+            size = 1
+            if isinstance(shape, int):
+                shape = (shape,)
+            for d in shape:
+                size *= d
+
+            index = bisect.bisect_left(
+                a=self._pool_sizes,
+                x=size,
+            )
+            if index < len(self._pool):
+                # Big enough array found, remove from pool
+                array = self._pool.pop(index)
+                self._pool_sizes.pop(index)
+                if requires_grad and array.grad is None:
+                    array.requires_grad = True
+                return Temporary(pool=self, array=array, shape=shape, dtype=dtype)
+
+            # No big enough array found, allocate new one
+            if len(self._pool) > 0:
+                grow_factor = 1.5
+                size = max(int(self._pool_sizes[-1] * grow_factor), size)
+
+            array = wp.empty(
+                shape=(size,), dtype=self.dtype, pinned=self.pinned, device=self.device, requires_grad=requires_grad
+            )
+            self._allocs[array.ptr] = array
+            return Temporary(pool=self, array=array, shape=shape, dtype=dtype)
+
+        def redeem(self, array):
+            # Insert back array into available pool
+            index = bisect.bisect_left(
+                a=self._pool_sizes,
+                x=array.size,
+            )
+            self._pool.insert(index, array)
+            self._pool_sizes.insert(index, array.size)
+
+        def detach(self, array):
+            del self._allocs[array.ptr]
+
+    def __init__(self):
+        self.clear()
+
+    def clear(self):
+        self._temporaries = {}
+
+    def borrow(self, shape, dtype, pinned: bool = False, device=None, requires_grad: bool = False) -> Temporary:
+        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)
+
+        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
+            )
+            pool = TemporaryStore.Pool(value_type, device, pinned=pinned)
+            self._temporaries[key] = pool
+
+        return pool.borrow(dtype=dtype, shape=shape, requires_grad=requires_grad)
+
+
+def set_default_temporary_store(temporary_store: Optional[TemporaryStore]):
+    """Globally sets the default :class:`TemporaryStore` instance to use for temporary allocations in ``warp.fem`` functions.
+
+    If the default temporary store is set to ``None``, temporary allocations are not persisted unless a :class:`TemporaryStore` is provided at a per-function granularity.
+    """
+
+    TemporaryStore._default_store = temporary_store
+
+
+def borrow_temporary(
+    temporary_store: Optional[TemporaryStore],
+    shape: Union[int, Tuple[int]],
+    dtype: type,
+    pinned: bool = False,
+    requires_grad: bool = False,
+    device=None,
+) -> Temporary:
+    """
+    Borrows and returns a temporary array with specified attributes from a shared pool.
+
+    If an array with sufficient capacity and matching desired attributes is already available in the pool, it will be returned.
+    Otherwise, a new allocation will be performed.
+
+    Args:
+        temporary_store: the shared pool to borrow the temporary from. If `temporary_store` is ``None``, the global default temporary store, if set, will be used.
+        shape: desired dimensions for the temporary array
+        dtype: desired data type for the temporary array
+        pinned: whether a pinned allocation is desired
+        device: device on which the memory should be allocated; if ``None``, the current device will be used.
+    """
+
+    if temporary_store is None:
+        temporary_store = TemporaryStore._default_store
+
+    if temporary_store is None:
+        return Temporary(
+            array=wp.empty(shape=shape, dtype=dtype, pinned=pinned, device=device, requires_grad=requires_grad)
+        )
+
+    return temporary_store.borrow(shape=shape, dtype=dtype, device=device, pinned=pinned, requires_grad=requires_grad)
+
+
+def borrow_temporary_like(
+    array: Union[wp.array, Temporary],
+    temporary_store: Optional[TemporaryStore],
+) -> Temporary:
+    """
+    Borrows and returns a temporary array with the same attributes as another array or temporary.
+
+    Args:
+        array: Warp or temporary array to read the desired attributes from
+        temporary_store: the shared pool to borrow the temporary from. If `temporary_store` is ``None``, the global default temporary store, if set, will be used.
+    """
+    if isinstance(array, Temporary):
+        array = array.array
+    return borrow_temporary(
+        temporary_store=temporary_store,
+        shape=array.shape,
+        dtype=array.dtype,
+        pinned=array.pinned,
+        device=array.device,
+        requires_grad=array.requires_grad,
+    )

warp/fem/dirichlet.py

@@ -0,0 +1,181 @@
+from typing import Any, Optional
+
+import warp as wp
+
+from warp.types import type_length, type_is_matrix
+from warp.sparse import BsrMatrix, bsr_copy, bsr_mv, bsr_mm, bsr_assign, bsr_axpy
+
+from .utils import array_axpy
+
+
+def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None):
+    """
+    Scale projector so that it becomes idempotent, and apply the same scaling to fixed_value if provided
+    """
+
+    if projector_matrix.nrow < projector_matrix.nnz or projector_matrix.ncol != projector_matrix.nrow:
+        raise ValueError("Projector must be a square diagonal matrix, with at most one non-zero block per row")
+
+    # Cast blocks to matrix type if necessary
+    projector_values = projector_matrix.values
+    if not type_is_matrix(projector_values.dtype):
+        projector_values = wp.array(
+            data=None,
+            ptr=projector_values.ptr,
+            capacity=projector_values.capacity,
+            owner=False,
+            device=projector_values.device,
+            dtype=wp.mat(shape=projector_matrix.block_shape, dtype=projector_matrix.scalar_type),
+            shape=projector_values.shape[0],
+        )
+
+    if fixed_value is None:
+        wp.launch(
+            kernel=_normalize_dirichlet_projector_kernel,
+            dim=projector_matrix.nrow,
+            device=projector_values.device,
+            inputs=[projector_matrix.offsets, projector_matrix.columns, projector_values],
+        )
+
+    else:
+        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:
+            # array of scalars, convert to 1d array of vectors
+            fixed_value = wp.array(
+                data=None,
+                ptr=fixed_value.ptr,
+                capacity=fixed_value.capacity,
+                owner=False,
+                device=fixed_value.device,
+                dtype=wp.vec(length=projector_matrix.block_shape[0], dtype=projector_matrix.scalar_type),
+                shape=fixed_value.shape[0],
+            )
+
+        wp.launch(
+            kernel=_normalize_dirichlet_projector_and_values_kernel,
+            dim=projector_matrix.nrow,
+            device=projector_values.device,
+            inputs=[projector_matrix.offsets, projector_matrix.columns, projector_values, fixed_value],
+        )
+
+
+def project_system_rhs(
+    system_matrix: BsrMatrix, system_rhs: wp.array, projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None
+):
+    """Projects the right-hand-side of a linear system to enforce Dirichlet boundary conditions
+
+    ``rhs = (I - projector) * ( rhs - system * projector * fixed_value) + projector * fixed_value``
+    """
+
+    rhs_tmp = wp.empty_like(system_rhs)
+    rhs_tmp.assign(system_rhs)
+
+    if fixed_value is None:
+        system_rhs.zero_()
+    else:
+        bsr_mv(A=projector_matrix, x=fixed_value, y=system_rhs, alpha=1.0, beta=0.0)
+
+    bsr_mv(A=system_matrix, x=system_rhs, y=rhs_tmp, alpha=-1.0, beta=1.0)
+
+    # here rhs_tmp = system_rhs - system_matrix * projector * fixed_value
+    # system_rhs = projector * fixed_value
+    array_axpy(x=rhs_tmp, y=system_rhs, alpha=1.0, beta=1.0)
+    bsr_mv(A=projector_matrix, x=rhs_tmp, y=system_rhs, alpha=-1.0, beta=1.0)
+
+
+def project_system_matrix(system_matrix: BsrMatrix, projector_matrix: BsrMatrix):
+    """Projects the right-hand-side of a linear system to enforce Dirichlet boundary conditions
+
+    ``system = (I - projector) * system * (I - projector) + projector``
+    """
+
+    complement_system = bsr_copy(system_matrix)
+    bsr_mm(x=projector_matrix, y=system_matrix, z=complement_system, alpha=-1.0, beta=1.0)
+
+    bsr_assign(dest=system_matrix, src=complement_system)
+    bsr_axpy(x=projector_matrix, y=system_matrix)
+    bsr_mm(x=complement_system, y=projector_matrix, z=system_matrix, alpha=-1.0, beta=1.0)
+
+
+def project_linear_system(
+    system_matrix: BsrMatrix,
+    system_rhs: wp.array,
+    projector_matrix: BsrMatrix,
+    fixed_value: Optional[wp.array] = None,
+    normalize_projector=True,
+):
+    """
+    Projects both the left-hand-side and right-hand-side of a linear system to enforce Dirichlet boundary conditions
+
+    If normalize_projector is True, first apply scaling so that the projector_matrix is idempotent
+    """
+    if normalize_projector:
+        normalize_dirichlet_projector(projector_matrix, fixed_value)
+
+    project_system_rhs(system_matrix, system_rhs, projector_matrix, fixed_value)
+    project_system_matrix(system_matrix, projector_matrix)
+
+
+@wp.kernel
+def _normalize_dirichlet_projector_kernel(
+    offsets: wp.array(dtype=int),
+    columns: wp.array(dtype=int),
+    block_values: wp.array(dtype=Any),
+):
+    row = wp.tid()
+
+    beg = offsets[row]
+    end = offsets[row + 1]
+
+    if beg == end:
+        return
+
+    diag = wp.lower_bound(columns, beg, end, row)
+
+    if diag < end and columns[diag] == row:
+        P = block_values[diag]
+
+        P_sq = P * P
+        trace_P = wp.trace(P)
+        trace_P_sq = wp.trace(P_sq)
+
+        if wp.nonzero(trace_P_sq):
+            scale = trace_P / trace_P_sq
+            block_values[diag] = scale * P
+        else:
+            block_values[diag] = P - P
+
+
+@wp.kernel
+def _normalize_dirichlet_projector_and_values_kernel(
+    offsets: wp.array(dtype=int),
+    columns: wp.array(dtype=int),
+    block_values: wp.array(dtype=Any),
+    fixed_values: wp.array(dtype=Any),
+):
+    row = wp.tid()
+
+    beg = offsets[row]
+    end = offsets[row + 1]
+
+    if beg == end:
+        return
+
+    diag = wp.lower_bound(columns, beg, end, row)
+
+    if diag < end and columns[diag] == row:
+        P = block_values[diag]
+
+        P_sq = P * P
+        trace_P = wp.trace(P)
+        trace_P_sq = wp.trace(P_sq)
+
+        if wp.nonzero(trace_P_sq):
+            scale = trace_P / trace_P_sq
+            block_values[diag] = scale * P
+            fixed_values[row] = scale * fixed_values[row]
+        else:
+            block_values[diag] = P - P
+            fixed_values[row] = fixed_values[row] - fixed_values[row]