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

warp/fem/space/partition.py

@@ -0,0 +1,350 @@
+from typing import Any, Optional, Union
+
+import warp as wp
+from warp.fem.cache import (
+    TemporaryStore,
+    borrow_temporary,
+    borrow_temporary_like,
+    cached_arg_value,
+)
+from warp.fem.geometry import GeometryPartition, WholeGeometryPartition
+from warp.fem.types import NULL_NODE_INDEX
+from warp.fem.utils import _iota_kernel, compress_node_indices
+
+from .function_space import FunctionSpace
+from .topology import SpaceTopology
+
+wp.set_module_options({"enable_backward": False})
+
+
+class SpacePartition:
+    class PartitionArg:
+        pass
+
+    def __init__(self, space_topology: SpaceTopology, geo_partition: GeometryPartition):
+        self.space_topology = space_topology
+        self.geo_partition = geo_partition
+
+    def node_count(self):
+        """Returns number of nodes in this partition"""
+
+    def owned_node_count(self) -> int:
+        """Returns number of nodes in this partition, excluding exterior halo"""
+
+    def interior_node_count(self) -> int:
+        """Returns number of interior nodes in this partition"""
+
+    def space_node_indices(self) -> wp.array:
+        """Return the global function space indices for nodes in this partition"""
+
+    def partition_arg_value(self, device):
+        pass
+
+    @staticmethod
+    def partition_node_index(args: "PartitionArg", space_node_index: int):
+        """Returns the index in the partition of a function space node, or -1 if it does not exist"""
+
+    def __str__(self) -> str:
+        return self.name
+
+    @property
+    def name(self) -> str:
+        return f"{self.__class__.__name__}"
+
+
+class WholeSpacePartition(SpacePartition):
+    @wp.struct
+    class PartitionArg:
+        pass
+
+    def __init__(self, space_topology: SpaceTopology):
+        super().__init__(space_topology, WholeGeometryPartition(space_topology.geometry))
+        self._node_indices = None
+
+    def node_count(self):
+        """Returns number of nodes in this partition"""
+        return self.space_topology.node_count()
+
+    def owned_node_count(self) -> int:
+        """Returns number of nodes in this partition, excluding exterior halo"""
+        return self.space_topology.node_count()
+
+    def interior_node_count(self) -> int:
+        """Returns number of interior nodes in this partition"""
+        return self.space_topology.node_count()
+
+    def space_node_indices(self):
+        """Return the global function space indices for nodes in this partition"""
+        if self._node_indices is None:
+            self._node_indices = borrow_temporary(temporary_store=None, shape=(self.node_count(),), dtype=int)
+            wp.launch(kernel=_iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array, 1])
+        return self._node_indices.array
+
+    def partition_arg_value(self, device):
+        return WholeSpacePartition.PartitionArg()
+
+    @wp.func
+    def partition_node_index(args: Any, space_node_index: int):
+        return space_node_index
+
+    def __eq__(self, other: SpacePartition) -> bool:
+        return isinstance(other, SpacePartition) and self.space_topology == other.space_topology
+
+    @property
+    def name(self) -> str:
+        return "Whole"
+
+
+class NodeCategory:
+    OWNED_INTERIOR = wp.constant(0)
+    """Node is touched exclusively by this partition, not touched by frontier side"""
+    OWNED_FRONTIER = wp.constant(1)
+    """Node is touched by a frontier side, but belongs to an element of this partition"""
+    HALO_LOCAL_SIDE = wp.constant(2)
+    """Node belongs to an element of another partition, but is touched by one of our frontier side"""
+    HALO_OTHER_SIDE = wp.constant(3)
+    """Node belongs to an element of another partition, and is not touched by one of our frontier side"""
+    EXTERIOR = wp.constant(4)
+    """Node is never referenced by this partition"""
+
+    COUNT = 5
+
+
+class NodePartition(SpacePartition):
+    @wp.struct
+    class PartitionArg:
+        space_to_partition: wp.array(dtype=int)
+
+    def __init__(
+        self,
+        space_topology: SpaceTopology,
+        geo_partition: GeometryPartition,
+        with_halo: bool = True,
+        device=None,
+        temporary_store: TemporaryStore = None,
+    ):
+        super().__init__(space_topology=space_topology, geo_partition=geo_partition)
+
+        self._compute_node_indices_from_sides(device, with_halo, temporary_store)
+
+    def node_count(self) -> int:
+        """Returns number of nodes referenced by this partition, including exterior halo"""
+        return int(self._category_offsets.array.numpy()[NodeCategory.HALO_OTHER_SIDE + 1])
+
+    def owned_node_count(self) -> int:
+        """Returns number of nodes in this partition, excluding exterior halo"""
+        return int(self._category_offsets.array.numpy()[NodeCategory.OWNED_FRONTIER + 1])
+
+    def interior_node_count(self) -> int:
+        """Returns number of interior nodes in this partition"""
+        return int(self._category_offsets.array.numpy()[NodeCategory.OWNED_INTERIOR + 1])
+
+    def space_node_indices(self):
+        """Return the global function space indices for nodes in this partition"""
+        return self._node_indices.array
+
+    @cached_arg_value
+    def partition_arg_value(self, device):
+        arg = NodePartition.PartitionArg()
+        arg.space_to_partition = self._space_to_partition.array.to(device)
+        return arg
+
+    @wp.func
+    def partition_node_index(args: PartitionArg, space_node_index: int):
+        return args.space_to_partition[space_node_index]
+
+    def _compute_node_indices_from_sides(self, device, with_halo: bool, temporary_store: TemporaryStore):
+        from warp.fem import cache
+
+        trace_topology = self.space_topology.trace()
+        NODES_PER_CELL = self.space_topology.NODES_PER_ELEMENT
+        NODES_PER_SIDE = trace_topology.NODES_PER_ELEMENT
+
+        @cache.dynamic_kernel(suffix=f"{self.geo_partition.name}_{self.space_topology.name}")
+        def node_category_from_cells_kernel(
+            geo_arg: self.geo_partition.geometry.CellArg,
+            geo_partition_arg: self.geo_partition.CellArg,
+            space_arg: self.space_topology.TopologyArg,
+            node_mask: wp.array(dtype=int),
+        ):
+            partition_cell_index = wp.tid()
+
+            cell_index = self.geo_partition.cell_index(geo_partition_arg, partition_cell_index)
+
+            for n in range(NODES_PER_CELL):
+                space_nidx = self.space_topology.element_node_index(geo_arg, space_arg, cell_index, n)
+                node_mask[space_nidx] = NodeCategory.OWNED_INTERIOR
+
+        @cache.dynamic_kernel(suffix=f"{self.geo_partition.name}_{self.space_topology.name}")
+        def node_category_from_owned_sides_kernel(
+            geo_arg: self.geo_partition.geometry.SideArg,
+            geo_partition_arg: self.geo_partition.SideArg,
+            space_arg: trace_topology.TopologyArg,
+            node_mask: wp.array(dtype=int),
+        ):
+            partition_side_index = wp.tid()
+
+            side_index = self.geo_partition.side_index(geo_partition_arg, partition_side_index)
+
+            for n in range(NODES_PER_SIDE):
+                space_nidx = trace_topology.element_node_index(geo_arg, space_arg, side_index, n)
+
+                if node_mask[space_nidx] == NodeCategory.EXTERIOR:
+                    node_mask[space_nidx] = NodeCategory.HALO_LOCAL_SIDE
+
+        @cache.dynamic_kernel(suffix=f"{self.geo_partition.name}_{self.space_topology.name}")
+        def node_category_from_frontier_sides_kernel(
+            geo_arg: self.geo_partition.geometry.SideArg,
+            geo_partition_arg: self.geo_partition.SideArg,
+            space_arg: trace_topology.TopologyArg,
+            node_mask: wp.array(dtype=int),
+        ):
+            frontier_side_index = wp.tid()
+
+            side_index = self.geo_partition.frontier_side_index(geo_partition_arg, frontier_side_index)
+
+            for n in range(NODES_PER_SIDE):
+                space_nidx = trace_topology.element_node_index(geo_arg, space_arg, side_index, n)
+                if node_mask[space_nidx] == NodeCategory.EXTERIOR:
+                    node_mask[space_nidx] = NodeCategory.HALO_OTHER_SIDE
+                elif node_mask[space_nidx] == NodeCategory.OWNED_INTERIOR:
+                    node_mask[space_nidx] = NodeCategory.OWNED_FRONTIER
+
+        node_category = borrow_temporary(
+            temporary_store,
+            shape=(self.space_topology.node_count(),),
+            dtype=int,
+            device=device,
+        )
+        node_category.array.fill_(value=NodeCategory.EXTERIOR)
+
+        wp.launch(
+            dim=self.geo_partition.cell_count(),
+            kernel=node_category_from_cells_kernel,
+            inputs=[
+                self.geo_partition.geometry.cell_arg_value(device),
+                self.geo_partition.cell_arg_value(device),
+                self.space_topology.topo_arg_value(device),
+                node_category.array,
+            ],
+            device=device,
+        )
+
+        if with_halo:
+            wp.launch(
+                dim=self.geo_partition.side_count(),
+                kernel=node_category_from_owned_sides_kernel,
+                inputs=[
+                    self.geo_partition.geometry.side_arg_value(device),
+                    self.geo_partition.side_arg_value(device),
+                    self.space_topology.topo_arg_value(device),
+                    node_category.array,
+                ],
+                device=device,
+            )
+
+            wp.launch(
+                dim=self.geo_partition.frontier_side_count(),
+                kernel=node_category_from_frontier_sides_kernel,
+                inputs=[
+                    self.geo_partition.geometry.side_arg_value(device),
+                    self.geo_partition.side_arg_value(device),
+                    self.space_topology.topo_arg_value(device),
+                    node_category.array,
+                ],
+                device=device,
+            )
+
+        self._finalize_node_indices(node_category.array, temporary_store)
+
+        node_category.release()
+
+    def _finalize_node_indices(self, node_category: wp.array(dtype=int), temporary_store: TemporaryStore):
+        category_offsets, node_indices, _, __ = compress_node_indices(NodeCategory.COUNT, node_category)
+
+        # Copy offsets to cpu
+        device = node_category.device
+        self._category_offsets = borrow_temporary(
+            temporary_store,
+            shape=category_offsets.array.shape,
+            dtype=category_offsets.array.dtype,
+            pinned=device.is_cuda,
+            device="cpu",
+        )
+        wp.copy(src=category_offsets.array, dest=self._category_offsets.array)
+
+        if device.is_cuda:
+            # TODO switch to synchronize_event once available
+            wp.synchronize_stream(wp.get_stream(device))
+
+        category_offsets.release()
+
+        # Compute global to local indices
+        self._space_to_partition = borrow_temporary_like(node_indices, temporary_store)
+        wp.launch(
+            kernel=NodePartition._scatter_partition_indices,
+            dim=self.space_topology.node_count(),
+            device=device,
+            inputs=[self.node_count(), node_indices.array, self._space_to_partition.array],
+        )
+
+        # Copy to shrinked-to-fit array
+        self._node_indices = borrow_temporary(temporary_store, shape=(self.node_count()), dtype=int, device=device)
+        wp.copy(dest=self._node_indices.array, src=node_indices.array, count=self.node_count())
+
+        node_indices.release()
+
+    @wp.kernel
+    def _scatter_partition_indices(
+        local_node_count: int,
+        node_indices: wp.array(dtype=int),
+        space_to_partition_indices: wp.array(dtype=int),
+    ):
+        local_idx = wp.tid()
+        space_idx = node_indices[local_idx]
+
+        if local_idx < local_node_count:
+            space_to_partition_indices[space_idx] = local_idx
+        else:
+            space_to_partition_indices[space_idx] = NULL_NODE_INDEX
+
+
+def make_space_partition(
+    space: Optional[FunctionSpace] = None,
+    geometry_partition: Optional[GeometryPartition] = None,
+    space_topology: Optional[SpaceTopology] = None,
+    with_halo: bool = True,
+    device=None,
+    temporary_store: TemporaryStore = None,
+) -> SpacePartition:
+    """Computes the subset of nodes from a function space topology that touch a geometry partition
+
+    Either `space_topology` or `space` must be provided (and will be considered in that order).
+
+    Args:
+        space: (deprecated) the function space defining the topology if `space_topology` is ``None``.
+        geometry_partition: The subset of the space geometry.  If not provided, use the whole geometry.
+        space_topology: the topology of the function space to consider. If ``None``, deduced from `space`.
+        with_halo: if True, include the halo nodes (nodes from exterior frontier cells to the partition)
+        device: Warp device on which to perform and store computations
+
+    Returns:
+        the resulting space partition
+    """
+
+    if space_topology is None:
+        space_topology = space.topology
+
+    space_topology = space_topology.full_space_topology()
+
+    if geometry_partition is not None:
+        if geometry_partition.cell_count() < geometry_partition.geometry.cell_count():
+            return NodePartition(
+                space_topology=space_topology,
+                geo_partition=geometry_partition,
+                with_halo=with_halo,
+                device=device,
+                temporary_store=temporary_store,
+            )
+
+    return WholeSpacePartition(space_topology)