Fix Explicit release of temporaries due to ref-cycle in fem.borrow_temporary

warp/_src/fem/adaptivity.py

@@ -65,6 +65,7 @@ def adaptive_nanogrid_from_hierarchy(
             device=device,
             inputs=[l, voxel_offsets[l], grid_voxels, merged_ijks],
         )
+        grid_voxels.release()
 
     # Allocate merged grid
     grid_info = grids[0].get_grid_info()
@@ -74,6 +75,7 @@ def adaptive_nanogrid_from_hierarchy(
         translation=grid_info.translation,
         device=device,
     )
+    merged_ijks.release()
 
     # Get unique voxel and corresponding level
     cell_count = cell_grid.get_voxel_count()
@@ -89,6 +91,7 @@ def adaptive_nanogrid_from_hierarchy(
         dim=cell_count,
         inputs=[level_count, cell_grid_ids, cell_ijk, cell_level],
     )
+    cell_ijk.release()
 
     cell_grid, cell_level = enforce_nanogrid_grading(
         cell_grid, cell_level, level_count=level_count, grading=grading, temporary_store=temporary_store
@@ -170,10 +173,17 @@ def adaptive_nanogrid_from_field(
                 fine_level,
             ],
         )
+        cell_refinement.release()
+
+        prev_cell_ijk = cell_ijk
+        prev_cell_level = cell_level
 
         # Fine is the new coarse
         cell_ijk = fine_ijk
         cell_level = fine_level
+        prev_cell_ijk.release()
+        prev_cell_level.release()
+    fine_count.release()
 
     wp.launch(_adjust_refined_ijk, dim=fine_shape, device=device, inputs=[cell_ijk, cell_level])
 
@@ -195,6 +205,8 @@ def adaptive_nanogrid_from_field(
         device=device,
         inputs=[fine_grid.id, cell_ijk, cell_level, fine_level],
     )
+    cell_ijk.release()
+    cell_level.release()
 
     fine_grid, fine_level = enforce_nanogrid_grading(
         fine_grid, fine_level, level_count=level_count, grading=grading, temporary_store=temporary_store
@@ -262,6 +274,8 @@ def enforce_nanogrid_grading(
         # Add new coordinates
         fine_shape = int(fine_count.numpy()[0])
         if fine_shape == cell_count:
+            cell_ijk.release()
+            refinement.release()
             break
 
         fine_ijk = cache.borrow_temporary(temporary_store, shape=fine_shape, dtype=wp.vec3i, device=device)
@@ -280,6 +294,8 @@ def enforce_nanogrid_grading(
                 fine_level,
             ],
         )
+        cell_ijk.release()
+        refinement.release()
 
         # Rebuild grid and levels
         cell_grid = wp.Volume.allocate_by_voxels(
@@ -292,7 +308,10 @@ def enforce_nanogrid_grading(
             device=device,
             inputs=[cell_grid.id, fine_ijk, fine_level, cell_level],
         )
+        fine_ijk.release()
+        fine_level.release()
 
+    fine_count.release()
     return cell_grid, cell_level
 
 

warp/_src/fem/domain.py

@@ -455,13 +455,21 @@ def __init__(
         if element_indices is None:
             if element_mask is None:
                 raise ValueError("Either 'element_mask' or 'element_indices' should be provided")
-            element_indices, _ = utils.masked_indices(mask=element_mask, temporary_store=temporary_store)
+            element_indices, element_global_to_local = utils.masked_indices(
+                mask=element_mask, temporary_store=temporary_store
+            )
             element_indices = element_indices.detach()
+            element_global_to_local.release()
+            owns_element_indices = True
         elif element_mask is not None:
             raise ValueError("Only one of 'element_mask' and 'element_indices' should be provided")
+        else:
+            # If Temporary are passed, then they are not owned by the class, hence cannot be released by the class
+            owns_element_indices = False
 
         self._domain = domain
         self._element_indices = element_indices
+        self._owns_element_indices = owns_element_indices
         self.ElementIndexArg = self._make_element_index_arg()
         self.element_index = self._make_element_index()
 
@@ -479,6 +487,13 @@ def __init__(
         self.element_partition_lookup = self._domain.element_partition_lookup
         self.element_normal = self._domain.element_normal
 
+    def __del__(self):
+        if getattr(self, "_owns_element_indices", False):
+            element_indices = getattr(self, "_element_indices", None)
+            if element_indices is not None and hasattr(element_indices, "release"):
+                element_indices.release()
+                self._element_indices = None
+
     @property
     def name(self) -> str:
         return f"{self._domain.name}_Subdomain"

warp/_src/fem/geometry/adaptive_nanogrid.py

@@ -417,8 +417,10 @@ def _build_face_grid(self, temporary_store: Optional[cache.TemporaryStore] = Non
                 boundary_face_mask,
             ],
         )
-        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask)
-        self._boundary_face_indices = boundary_face_indices.detach()
+        boundary_face_indices, boundary_face_global_to_local = utils.masked_indices(boundary_face_mask)
+        self._replace_boundary_face_indices(boundary_face_indices.detach())
+        boundary_face_global_to_local.release()
+        boundary_face_mask.release()
 
     def _ensure_stacked_edge_grid(self):
         if self._stacked_edge_grid is None:
@@ -565,6 +567,7 @@ def _build_node_grid(cell_ijk, cell_level, cell_grid: wp.Volume, temporary_store
     node_grid = wp.Volume.allocate_by_voxels(
         cell_nodes.flatten(), voxel_size=cell_grid.get_voxel_size()[0], device=cell_ijk.device
     )
+    cell_nodes.release()
 
     return node_grid
 
@@ -577,6 +580,7 @@ def _build_cell_face_grid(cell_ijk, cell_level, grid: wp.Volume, temporary_store
     face_grid = wp.Volume.allocate_by_voxels(
         cell_faces.flatten(), voxel_size=grid.get_voxel_size()[0], device=cell_ijk.device
     )
+    cell_faces.release()
 
     return face_grid
 
@@ -634,6 +638,9 @@ def _build_completed_face_grid(
     face_grid = wp.Volume.allocate_by_voxels(
         cat_face_ijk.flatten(), voxel_size=cell_face_grid.get_voxel_size(), device=device
     )
+    cat_face_ijk.release()
+    cell_face_ijk.release()
+    additional_face_count.release()
 
     return face_grid
 
@@ -651,6 +658,7 @@ def _build_stacked_face_grid(cell_ijk, cell_level, grid: wp.Volume, temporary_st
     face_grid = wp.Volume.allocate_by_voxels(
         cell_faces.flatten(), voxel_size=grid.get_voxel_size()[0], device=cell_ijk.device
     )
+    cell_faces.release()
 
     return face_grid
 
@@ -668,6 +676,7 @@ def _build_stacked_edge_grid(cell_ijk, cell_level, grid: wp.Volume, temporary_st
     edge_grid = wp.Volume.allocate_by_voxels(
         cell_edges.flatten(), voxel_size=grid.get_voxel_size()[0], device=cell_ijk.device
     )
+    cell_edges.release()
 
     return edge_grid
 

warp/_src/fem/geometry/hexmesh.py

@@ -157,8 +157,9 @@ def __init__(
         self._face_vertex_indices: wp.array = None
         self._face_hex_indices: wp.array = None
         self._face_hex_face_orientation: wp.array = None
-        self._vertex_hex_offsets: wp.array = None
-        self._vertex_hex_indices: wp.array = None
+        self._vertex_hex_offsets: wp.array = None  # owned temporary reused between rebuilds
+        self._vertex_hex_indices: wp.array = None  # owned temporary reused between rebuilds
+        self._boundary_face_indices: wp.array = None  # owned temporary reused between rebuilds
         self._hex_edge_indices: wp.array = None
         self._edge_count = 0
         self._build_topology(temporary_store)
@@ -183,6 +184,9 @@ def __init__(
         if build_bvh:
             self.build_bvh(self.positions.device)
 
+    def __del__(self):
+        self._release_owned_temporaries()
+
     def cell_count(self):
         return self.hex_vertex_indices.shape[0]
 
@@ -447,8 +451,8 @@ def _build_topology(self, temporary_store: TemporaryStore):
         vertex_hex_offsets, vertex_hex_indices = compress_node_indices(
             self.vertex_count(), self.hex_vertex_indices, temporary_store=temporary_store
         )
-        self._vertex_hex_offsets = vertex_hex_offsets.detach()
-        self._vertex_hex_indices = vertex_hex_indices.detach()
+        self._replace_owned_array("_vertex_hex_offsets", vertex_hex_offsets.detach())
+        self._replace_owned_array("_vertex_hex_indices", vertex_hex_indices.detach())
 
         vertex_start_face_count = borrow_temporary(temporary_store, dtype=int, device=device, shape=self.vertex_count())
         vertex_start_face_count.zero_()
@@ -546,8 +550,24 @@ def _build_topology(self, temporary_store: TemporaryStore):
             ],
         )
 
-        boundary_face_indices, _ = masked_indices(boundary_mask)
-        self._boundary_face_indices = boundary_face_indices.detach()
+        boundary_face_indices, boundary_face_global_to_local = masked_indices(boundary_mask)
+        self._replace_owned_array("_boundary_face_indices", boundary_face_indices.detach())
+        boundary_face_global_to_local.release()
+        boundary_mask.release()
+
+    def _replace_owned_array(self, attr_name: str, new_value):
+        if hasattr(self, attr_name):
+            old_value = getattr(self, attr_name)
+            if old_value is not None and old_value is not new_value and hasattr(old_value, "release"):
+                old_value.release()
+        setattr(self, attr_name, new_value)
+
+    def _release_owned_temporaries(self):
+        for attr in ("_vertex_hex_offsets", "_vertex_hex_indices", "_boundary_face_indices"):
+            value = getattr(self, attr, None)
+            if value is not None and hasattr(value, "release"):
+                value.release()
+                setattr(self, attr, None)
 
     def _compute_hex_edges(self, temporary_store: Optional[TemporaryStore] = None):
         from warp._src.fem.utils import host_read_at_index

warp/_src/fem/geometry/nanogrid.py

@@ -59,6 +59,9 @@ def __init__(
         self._cell_grid_info = cell_grid.get_grid_info()
         self._init_transform()
 
+    def __del__(self):
+        self._release_owned_temporaries()
+
     def reference_cell(self) -> Element:
         return Element.CUBE
 
@@ -126,6 +129,17 @@ def fill_side_index_arg(self, arg: SideIndexArg, device):
     def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
         return args.boundary_face_indices[boundary_side_index]
 
+    def _replace_boundary_face_indices(self, new_value):
+        if self._boundary_face_indices is not None and self._boundary_face_indices is not new_value:
+            if hasattr(self._boundary_face_indices, "release"):
+                self._boundary_face_indices.release()
+        self._boundary_face_indices = new_value
+
+    def _release_owned_temporaries(self):
+        if self._boundary_face_indices is not None and hasattr(self._boundary_face_indices, "release"):
+            self._boundary_face_indices.release()
+            self._boundary_face_indices = None
+
     def make_filtered_cell_lookup(grid_geo, filter_func: wp.Function = None):
         suffix = f"{grid_geo.name}{filter_func.key if filter_func is not None else ''}"
 
@@ -539,8 +553,10 @@ def _build_face_grid(self, temporary_store: Optional[cache.TemporaryStore] = Non
             device=device,
             inputs=[self._cell_grid.id, self._face_ijk, self._face_flags, boundary_face_mask],
         )
-        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask)
-        self._boundary_face_indices = boundary_face_indices.detach()
+        boundary_face_indices, boundary_face_global_to_local = utils.masked_indices(boundary_face_mask)
+        self._replace_boundary_face_indices(boundary_face_indices.detach())
+        boundary_face_global_to_local.release()
+        boundary_face_mask.release()
 
     def _build_edge_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
         self._edge_grid = _build_edge_grid(self._cell_ijk, self._cell_grid, temporary_store)
@@ -608,6 +624,7 @@ def _build_node_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
     node_grid = wp.Volume.allocate_by_voxels(
         cell_nodes.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
+    cell_nodes.release()
 
     return node_grid
 
@@ -620,6 +637,7 @@ def _build_face_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
     face_grid = wp.Volume.allocate_by_voxels(
         cell_faces.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
+    cell_faces.release()
 
     return face_grid
 
@@ -632,6 +650,7 @@ def _build_edge_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
     edge_grid = wp.Volume.allocate_by_voxels(
         cell_edges.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
+    cell_edges.release()
 
     return edge_grid