Spaces:
Paused
Paused
| from typing import Any | |
| import warp as wp | |
| from warp.fem.types import ElementIndex, NULL_ELEMENT_INDEX | |
| from warp.fem.utils import masked_indices | |
| from warp.fem.cache import cached_arg_value, TemporaryStore, borrow_temporary | |
| from .geometry import Geometry | |
| wp.set_module_options({"enable_backward": False}) | |
| class GeometryPartition: | |
| """Base class for geometry partitions, i.e. subset of cells and sides""" | |
| class CellArg: | |
| pass | |
| class SideArg: | |
| pass | |
| def __init__(self, geometry: Geometry): | |
| self.geometry = geometry | |
| def cell_count(self) -> int: | |
| """Number of cells that are 'owned' by this partition""" | |
| raise NotImplementedError() | |
| def side_count(self) -> int: | |
| """Number of sides that are 'owned' by this partition""" | |
| raise NotImplementedError() | |
| def boundary_side_count(self) -> int: | |
| """Number of geo-boundary sides that are 'owned' by this partition""" | |
| raise NotImplementedError() | |
| def frontier_side_count(self) -> int: | |
| """Number of sides with neighbors owned by this and another partition""" | |
| raise NotImplementedError() | |
| def name(self) -> str: | |
| return f"{self.geometry.name}_{self.__class__.__name__}" | |
| def __str__(self) -> str: | |
| return self.name | |
| def cell_arg_value(self, device): | |
| raise NotImplementedError() | |
| def side_arg_value(self, device): | |
| raise NotImplementedError() | |
| def cell_index(args: CellArg, partition_cell_index: int): | |
| """Index in the geometry of a partition cell""" | |
| raise NotImplementedError() | |
| def partition_cell_index(args: CellArg, cell_index: int): | |
| """Index of a geometry cell in the partition (or ``NULL_ELEMENT_INDEX``)""" | |
| raise NotImplementedError() | |
| def side_index(args: SideArg, partition_side_index: int): | |
| """Partition side to side index""" | |
| raise NotImplementedError() | |
| def boundary_side_index(args: SideArg, boundary_side_index: int): | |
| """Boundary side to side index""" | |
| raise NotImplementedError() | |
| def frontier_side_index(args: SideArg, frontier_side_index: int): | |
| """Frontier side to side index""" | |
| raise NotImplementedError() | |
| class WholeGeometryPartition(GeometryPartition): | |
| """Trivial (NOP) partition""" | |
| def __init__( | |
| self, | |
| geometry: Geometry, | |
| ): | |
| super().__init__(geometry) | |
| self.SideArg = geometry.SideIndexArg | |
| self.side_arg_value = geometry.side_index_arg_value | |
| self.cell_index = WholeGeometryPartition._identity_element_index | |
| self.partition_cell_index = WholeGeometryPartition._identity_element_index | |
| self.side_index = WholeGeometryPartition._identity_element_index | |
| self.boundary_side_index = geometry.boundary_side_index | |
| self.frontier_side_index = WholeGeometryPartition._identity_element_index | |
| def __eq__(self, other: GeometryPartition) -> bool: | |
| # Ensures that two whole partition instances of the same geometry are considered equal | |
| return isinstance(other, WholeGeometryPartition) and self.geometry == other.geometry | |
| def cell_count(self) -> int: | |
| return self.geometry.cell_count() | |
| def side_count(self) -> int: | |
| return self.geometry.side_count() | |
| def boundary_side_count(self) -> int: | |
| return self.geometry.boundary_side_count() | |
| def frontier_side_count(self) -> int: | |
| return 0 | |
| class CellArg: | |
| pass | |
| def cell_arg_value(self, device): | |
| arg = WholeGeometryPartition.CellArg() | |
| return arg | |
| def _identity_element_index(args: Any, idx: ElementIndex): | |
| return idx | |
| def name(self) -> str: | |
| return self.geometry.name | |
| class CellBasedGeometryPartition(GeometryPartition): | |
| """Geometry partition based on a subset of cells. Interior, boundary and frontier sides are automatically categorized.""" | |
| def __init__( | |
| self, | |
| geometry: Geometry, | |
| device=None, | |
| ): | |
| super().__init__(geometry) | |
| class SideArg: | |
| partition_side_indices: wp.array(dtype=int) | |
| boundary_side_indices: wp.array(dtype=int) | |
| frontier_side_indices: wp.array(dtype=int) | |
| def side_count(self) -> int: | |
| return self._partition_side_indices.array.shape[0] | |
| def boundary_side_count(self) -> int: | |
| return self._boundary_side_indices.array.shape[0] | |
| def frontier_side_count(self) -> int: | |
| return self._frontier_side_indices.array.shape[0] | |
| def side_arg_value(self, device): | |
| arg = LinearGeometryPartition.SideArg() | |
| arg.partition_side_indices = self._partition_side_indices.array.to(device) | |
| arg.boundary_side_indices = self._boundary_side_indices.array.to(device) | |
| arg.frontier_side_indices = self._frontier_side_indices.array.to(device) | |
| return arg | |
| def side_index(args: SideArg, partition_side_index: int): | |
| """partition side to side index""" | |
| return args.partition_side_indices[partition_side_index] | |
| def boundary_side_index(args: SideArg, boundary_side_index: int): | |
| """Boundary side to side index""" | |
| return args.boundary_side_indices[boundary_side_index] | |
| def frontier_side_index(args: SideArg, frontier_side_index: int): | |
| """Frontier side to side index""" | |
| return args.frontier_side_indices[frontier_side_index] | |
| def compute_side_indices_from_cells( | |
| self, cell_arg_value: Any, cell_inclusion_test_func: wp.Function, device, temporary_store: TemporaryStore = None | |
| ): | |
| from warp.fem import cache | |
| cell_arg_type = next(iter(cell_inclusion_test_func.input_types.values())) | |
| def count_sides( | |
| geo_arg: self.geometry.SideArg, | |
| cell_arg_value: cell_arg_type, | |
| partition_side_mask: wp.array(dtype=int), | |
| boundary_side_mask: wp.array(dtype=int), | |
| frontier_side_mask: wp.array(dtype=int), | |
| ): | |
| side_index = wp.tid() | |
| inner_cell_index = self.geometry.side_inner_cell_index(geo_arg, side_index) | |
| outer_cell_index = self.geometry.side_outer_cell_index(geo_arg, side_index) | |
| inner_in = cell_inclusion_test_func(cell_arg_value, inner_cell_index) | |
| outer_in = cell_inclusion_test_func(cell_arg_value, outer_cell_index) | |
| if inner_in: | |
| # Inner neighbor in partition; count as partition side | |
| partition_side_mask[side_index] = 1 | |
| # Inner and outer element as the same -- this is a boundary side | |
| if inner_cell_index == outer_cell_index: | |
| boundary_side_mask[side_index] = 1 | |
| if inner_in != outer_in: | |
| # Exactly one neighbor in partition; count as frontier side | |
| frontier_side_mask[side_index] = 1 | |
| partition_side_mask = borrow_temporary( | |
| temporary_store, | |
| shape=(self.geometry.side_count(),), | |
| dtype=int, | |
| device=device, | |
| ) | |
| boundary_side_mask = borrow_temporary( | |
| temporary_store, | |
| shape=(self.geometry.side_count(),), | |
| dtype=int, | |
| device=device, | |
| ) | |
| frontier_side_mask = borrow_temporary( | |
| temporary_store, | |
| shape=(self.geometry.side_count(),), | |
| dtype=int, | |
| device=device, | |
| ) | |
| partition_side_mask.array.zero_() | |
| boundary_side_mask.array.zero_() | |
| frontier_side_mask.array.zero_() | |
| wp.launch( | |
| dim=partition_side_mask.array.shape[0], | |
| kernel=count_sides, | |
| inputs=[ | |
| self.geometry.side_arg_value(device), | |
| cell_arg_value, | |
| partition_side_mask.array, | |
| boundary_side_mask.array, | |
| frontier_side_mask.array, | |
| ], | |
| device=device, | |
| ) | |
| # Convert counts to indices | |
| self._partition_side_indices, _ = masked_indices(partition_side_mask.array, temporary_store=temporary_store) | |
| self._boundary_side_indices, _ = masked_indices(boundary_side_mask.array, temporary_store=temporary_store) | |
| self._frontier_side_indices, _ = masked_indices(frontier_side_mask.array, temporary_store=temporary_store) | |
| partition_side_mask.release() | |
| boundary_side_mask.release() | |
| frontier_side_mask.release() | |
| class LinearGeometryPartition(CellBasedGeometryPartition): | |
| def __init__( | |
| self, | |
| geometry: Geometry, | |
| partition_rank: int, | |
| partition_count: int, | |
| device=None, | |
| temporary_store: TemporaryStore = None, | |
| ): | |
| """Creates a geometry partition by uniformly partionning cell indices | |
| Args: | |
| geometry: the geometry to partition | |
| partition_rank: the index of the partition being created | |
| partition_count: the number of partitions that will be created over the geometry | |
| device: Warp device on which to perform and store computations | |
| """ | |
| super().__init__(geometry) | |
| total_cell_count = geometry.cell_count() | |
| cells_per_partition = (total_cell_count + partition_count - 1) // partition_count | |
| self.cell_begin = cells_per_partition * partition_rank | |
| self.cell_end = min(self.cell_begin + cells_per_partition, total_cell_count) | |
| super().compute_side_indices_from_cells( | |
| self.cell_arg_value(device), | |
| LinearGeometryPartition._cell_inclusion_test, | |
| device, | |
| temporary_store=temporary_store, | |
| ) | |
| def cell_count(self) -> int: | |
| return self.cell_end - self.cell_begin | |
| class CellArg: | |
| cell_begin: int | |
| cell_end: int | |
| def cell_arg_value(self, device): | |
| arg = LinearGeometryPartition.CellArg() | |
| arg.cell_begin = self.cell_begin | |
| arg.cell_end = self.cell_end | |
| return arg | |
| def cell_index(args: CellArg, partition_cell_index: int): | |
| """Partition cell to cell index""" | |
| return args.cell_begin + partition_cell_index | |
| def partition_cell_index(args: CellArg, cell_index: int): | |
| """Partition cell to cell index""" | |
| if cell_index > args.cell_end: | |
| return NULL_ELEMENT_INDEX | |
| partition_cell_index = cell_index - args.cell_begin | |
| if partition_cell_index < 0: | |
| return NULL_ELEMENT_INDEX | |
| return partition_cell_index | |
| def _cell_inclusion_test(arg: CellArg, cell_index: int): | |
| return cell_index >= arg.cell_begin and cell_index < arg.cell_end | |
| class ExplicitGeometryPartition(CellBasedGeometryPartition): | |
| def __init__(self, geometry: Geometry, cell_mask: "wp.array(dtype=int)", temporary_store: TemporaryStore = None): | |
| """Creates a geometry partition by uniformly partionning cell indices | |
| Args: | |
| geometry: the geometry to partition | |
| cell_mask: warp array of length ``geometry.cell_count()`` indicating which cells are selected. Array values must be either ``1`` (selected) or ``0`` (not selected). | |
| """ | |
| super().__init__(geometry) | |
| self._cell_mask = cell_mask | |
| self._cells, self._partition_cells = masked_indices(self._cell_mask, temporary_store=temporary_store) | |
| super().compute_side_indices_from_cells( | |
| self._cell_mask, | |
| ExplicitGeometryPartition._cell_inclusion_test, | |
| self._cell_mask.device, | |
| temporary_store=temporary_store, | |
| ) | |
| def cell_count(self) -> int: | |
| return self._cells.array.shape[0] | |
| class CellArg: | |
| cell_index: wp.array(dtype=int) | |
| partition_cell_index: wp.array(dtype=int) | |
| def cell_arg_value(self, device): | |
| arg = ExplicitGeometryPartition.CellArg() | |
| arg.cell_index = self._cells.array.to(device) | |
| arg.partition_cell_index = self._partition_cells.array.to(device) | |
| return arg | |
| def cell_index(args: CellArg, partition_cell_index: int): | |
| return args.cell_index[partition_cell_index] | |
| def partition_cell_index(args: CellArg, cell_index: int): | |
| return args.partition_cell_index[cell_index] | |
| def _cell_inclusion_test(mask: wp.array(dtype=int), cell_index: int): | |
| return mask[cell_index] > 0 | |