warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/geometry/partition.py

@@ -4,6 +4,7 @@ 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
 
@@ -12,9 +13,14 @@ 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
 
@@ -41,6 +47,37 @@ class GeometryPartition:
     def __str__(self) -> str:
         return self.name
 
+    def cell_arg_value(self, device):
+        raise NotImplementedError()
+
+    def side_arg_value(self, device):
+        raise NotImplementedError()
+
+    @staticmethod
+    def cell_index(args: CellArg, partition_cell_index: int):
+        """Index in the geometry of a partition cell"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def partition_cell_index(args: CellArg, cell_index: int):
+        """Index of a geometry cell in the partition (or ``NULL_ELEMENT_INDEX``)"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def side_index(args: SideArg, partition_side_index: int):
+        """Partition side to side index"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def boundary_side_index(args: SideArg, boundary_side_index: int):
+        """Boundary side to side index"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def frontier_side_index(args: SideArg, frontier_side_index: int):
+        """Frontier side to side index"""
+        raise NotImplementedError()
+
 
 class WholeGeometryPartition(GeometryPartition):
     """Trivial (NOP) partition"""
@@ -89,6 +126,10 @@ class WholeGeometryPartition(GeometryPartition):
     def _identity_element_index(args: Any, idx: ElementIndex):
         return idx
 
+    @property
+    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."""
@@ -107,19 +148,20 @@ class CellBasedGeometryPartition(GeometryPartition):
         frontier_side_indices: wp.array(dtype=int)
 
     def side_count(self) -> int:
-        return self._partition_side_indices.shape[0]
+        return self._partition_side_indices.array.shape[0]
 
     def boundary_side_count(self) -> int:
-        return self._boundary_side_indices.shape[0]
+        return self._boundary_side_indices.array.shape[0]
 
     def frontier_side_count(self) -> int:
-        return self._frontier_side_indices.shape[0]
+        return self._frontier_side_indices.array.shape[0]
 
+    @cached_arg_value
     def side_arg_value(self, device):
         arg = LinearGeometryPartition.SideArg()
-        arg.partition_side_indices = self._partition_side_indices.to(device)
-        arg.boundary_side_indices = self._boundary_side_indices.to(device)
-        arg.frontier_side_indices = self._frontier_side_indices.to(device)
+        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
 
     @wp.func
@@ -138,16 +180,16 @@ class CellBasedGeometryPartition(GeometryPartition):
         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,
+        self, cell_arg_value: Any, cell_inclusion_test_func: wp.Function, device, temporary_store: TemporaryStore = None
     ):
         from warp.fem import cache
 
-        def count_side_fn(
+        cell_arg_type = next(iter(cell_inclusion_test_func.input_types.values()))
+
+        @cache.dynamic_kernel(suffix=f"{self.geometry.name}_{cell_inclusion_test_func.key}")
+        def count_sides(
             geo_arg: self.geometry.SideArg,
-            cell_arg_value: Any,
+            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),
@@ -171,44 +213,50 @@ class CellBasedGeometryPartition(GeometryPartition):
                 # Exactly one neighbor in partition; count as frontier side
                 frontier_side_mask[side_index] = 1
 
-        count_sides = cache.get_kernel(
-            count_side_fn,
-            suffix=f"{self.geometry.name}_{cell_inclusion_test_func.key}",
-        )
-
-        partition_side_mask = wp.zeros(
+        partition_side_mask = borrow_temporary(
+            temporary_store,
             shape=(self.geometry.side_count(),),
             dtype=int,
             device=device,
         )
-        boundary_side_mask = wp.zeros(
+        boundary_side_mask = borrow_temporary(
+            temporary_store,
             shape=(self.geometry.side_count(),),
             dtype=int,
             device=device,
         )
-        frontier_side_mask = wp.zeros(
+        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.shape[0],
+            dim=partition_side_mask.array.shape[0],
             kernel=count_sides,
             inputs=[
                 self.geometry.side_arg_value(device),
                 cell_arg_value,
-                partition_side_mask,
-                boundary_side_mask,
-                frontier_side_mask,
+                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)
-        self._boundary_side_indices, _ = masked_indices(boundary_side_mask)
-        self._frontier_side_indices, _ = masked_indices(frontier_side_mask)
+        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):
@@ -218,6 +266,7 @@ class LinearGeometryPartition(CellBasedGeometryPartition):
         partition_rank: int,
         partition_count: int,
         device=None,
+        temporary_store: TemporaryStore = None,
     ):
         """Creates a geometry partition by uniformly partionning cell indices
 
@@ -239,6 +288,7 @@ class LinearGeometryPartition(CellBasedGeometryPartition):
             self.cell_arg_value(device),
             LinearGeometryPartition._cell_inclusion_test,
             device,
+            temporary_store=temporary_store,
         )
 
     def cell_count(self) -> int:
@@ -278,7 +328,7 @@ class LinearGeometryPartition(CellBasedGeometryPartition):
 
 
 class ExplicitGeometryPartition(CellBasedGeometryPartition):
-    def __init__(self, geometry: Geometry, cell_mask: "wp.array(dtype=int)"):
+    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:
@@ -289,26 +339,28 @@ class ExplicitGeometryPartition(CellBasedGeometryPartition):
         super().__init__(geometry)
 
         self._cell_mask = cell_mask
-        self._cells, self._partition_cells = masked_indices(self._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.shape[0]
+        return self._cells.array.shape[0]
 
     @wp.struct
     class CellArg:
         cell_index: wp.array(dtype=int)
         partition_cell_index: wp.array(dtype=int)
 
+    @cached_arg_value
     def cell_arg_value(self, device):
         arg = ExplicitGeometryPartition.CellArg()
-        arg.cell_index = self._cells.to(device)
-        arg.partition_cell_index = self._partition_cells.to(device)
+        arg.cell_index = self._cells.array.to(device)
+        arg.partition_cell_index = self._partition_cells.array.to(device)
         return arg
 
     @wp.func