warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/fem/space/nanogrid_function_space.py

@@ -0,0 +1,201 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Union
+
+import warp as wp
+from warp.fem import cache
+from warp.fem.geometry import AdaptiveNanogrid, Nanogrid
+from warp.fem.types import ElementIndex
+
+from .shape import CubeShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+
+@wp.struct
+class NanogridTopologyArg:
+    vertex_grid: wp.uint64
+    face_grid: wp.uint64
+    edge_grid: wp.uint64
+
+    vertex_count: int
+    edge_count: int
+    face_count: int
+
+
+class NanogridSpaceTopology(SpaceTopology):
+    TopologyArg = NanogridTopologyArg
+
+    def __init__(
+        self,
+        grid: Union[Nanogrid, AdaptiveNanogrid],
+        shape: CubeShapeFunction,
+    ):
+        self._shape = shape
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._grid = grid
+
+        self._vertex_grid = grid.vertex_grid.id
+
+        need_edge_indices = shape.EDGE_NODE_COUNT > 0
+        need_face_indices = shape.FACE_NODE_COUNT > 0
+
+        if isinstance(grid, Nanogrid):
+            self._edge_grid = grid.edge_grid.id if need_edge_indices else -1
+            self._face_grid = grid.face_grid.id if need_face_indices else -1
+            self._edge_count = grid.edge_count() if need_edge_indices else 0
+            self._face_count = grid.side_count() if need_face_indices else 0
+        else:
+            self._edge_grid = grid.stacked_edge_grid.id if need_edge_indices else -1
+            self._face_grid = grid.stacked_face_grid.id if need_face_indices else -1
+            self._edge_count = grid.stacked_edge_count() if need_edge_indices else 0
+            self._face_count = grid.stacked_face_count() if need_face_indices else 0
+
+        self.element_node_index = self._make_element_node_index()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    @cache.cached_arg_value
+    def topo_arg_value(self, device):
+        arg = NanogridTopologyArg()
+
+        arg.vertex_grid = self._vertex_grid
+        arg.face_grid = self._face_grid
+        arg.edge_grid = self._edge_grid
+
+        arg.vertex_count = self._grid.vertex_count()
+        arg.face_count = self._face_count
+        arg.edge_count = self._edge_count
+        return arg
+
+    def _make_element_node_index(self):
+        element_node_index_generic = self._make_element_node_index_generic()
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: Nanogrid.CellArg,
+            topo_arg: NanogridTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            ijk = geo_arg.cell_ijk[element_index]
+            return element_node_index_generic(topo_arg, element_index, node_index_in_elt, ijk, 0)
+
+        if isinstance(self._grid, Nanogrid):
+            return element_node_index
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index_adaptive(
+            geo_arg: AdaptiveNanogrid.CellArg,
+            topo_arg: NanogridTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            ijk = geo_arg.cell_ijk[element_index]
+            level = int(geo_arg.cell_level[element_index])
+            return element_node_index_generic(topo_arg, element_index, node_index_in_elt, ijk, level)
+
+        return element_node_index_adaptive
+
+    def node_count(self) -> int:
+        return (
+            self._grid.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self._edge_count * self._shape.EDGE_NODE_COUNT
+            + self._face_count * self._shape.FACE_NODE_COUNT
+            + self._grid.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index_generic(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        FACE_NODE_COUNT = self._shape.FACE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index_generic(
+            topo_arg: NanogridTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+            ijk: wp.vec3i,
+            level: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.VERTEX:
+                    n_ijk = _cell_vertex_coord(ijk, level, type_instance)
+                    return (
+                        wp.volume_lookup_index(topo_arg.vertex_grid, n_ijk[0], n_ijk[1], n_ijk[2]) * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.EDGE:
+                    axis = type_instance >> 2
+                    node_offset = type_instance & 3
+
+                    n_ijk = _cell_edge_coord(ijk, level, axis, node_offset)
+
+                    edge_index = wp.volume_lookup_index(topo_arg.edge_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+                    return offset + EDGE_NODE_COUNT * edge_index + type_index
+
+                offset += EDGE_NODE_COUNT * topo_arg.edge_count
+
+            if wp.static(FACE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.FACE:
+                    axis = type_instance >> 1
+                    node_offset = type_instance & 1
+
+                    n_ijk = _cell_face_coord(ijk, level, axis, node_offset)
+
+                    face_index = wp.volume_lookup_index(topo_arg.face_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+                    return offset + FACE_NODE_COUNT * face_index + type_index
+
+                offset += FACE_NODE_COUNT * topo_arg.face_count
+
+            return offset + INTERIOR_NODE_COUNT * element_index + type_index
+
+        return element_node_index_generic
+
+
+@wp.func
+def _cell_vertex_coord(cell_ijk: wp.vec3i, cell_level: int, n: int):
+    return cell_ijk + AdaptiveNanogrid.fine_ijk(wp.vec3i((n & 4) >> 2, (n & 2) >> 1, n & 1), cell_level)
+
+
+@wp.func
+def _cell_edge_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
+    e_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
+    e_ijk[(axis + 1) % 3] += offset >> 1
+    e_ijk[(axis + 2) % 3] += offset & 1
+    return AdaptiveNanogrid.encode_axis_and_level(e_ijk, axis, cell_level)
+
+
+@wp.func
+def _cell_face_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
+    f_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
+    f_ijk[axis] += offset
+    return AdaptiveNanogrid.encode_axis_and_level(f_ijk, axis, cell_level)
+
+
+def make_nanogrid_space_topology(grid: Union[Nanogrid, AdaptiveNanogrid], shape: CubeShapeFunction):
+    if isinstance(shape, CubeShapeFunction):
+        return forward_base_topology(NanogridSpaceTopology, grid, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/partition.py

@@ -0,0 +1,367 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional
+
+import warp as wp
+import warp.fem.cache as cache
+from warp.fem.geometry import GeometryPartition, WholeGeometryPartition
+from warp.fem.types import NULL_NODE_INDEX
+from warp.fem.utils import 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 ``NULL_NODE_INDEX`` 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 = cache.borrow_temporary(temporary_store=None, shape=(self.node_count(),), dtype=int)
+            wp.launch(kernel=self._iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array])
+        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"
+
+    @wp.kernel
+    def _iota_kernel(indices: wp.array(dtype=int)):
+        indices[wp.tid()] = wp.tid()
+
+
+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: cache.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
+
+    @cache.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: cache.TemporaryStore):
+        from warp.fem import cache
+
+        trace_topology = self.space_topology.trace()
+
+        @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)
+
+            cell_node_count = self.space_topology.element_node_count(geo_arg, space_arg, cell_index)
+            for n in range(cell_node_count):
+                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)
+
+            side_node_count = trace_topology.element_node_count(geo_arg, space_arg, side_index)
+            for n in range(side_node_count):
+                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)
+
+            side_node_count = trace_topology.element_node_count(geo_arg, space_arg, side_index)
+            for n in range(side_node_count):
+                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 = cache.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: cache.TemporaryStore):
+        category_offsets, node_indices = compress_node_indices(
+            NodeCategory.COUNT, node_category, temporary_store=temporary_store
+        )
+
+        # Copy offsets to cpu
+        device = node_category.device
+        with wp.ScopedDevice(device):
+            self._category_offsets = cache.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)
+            copy_event = cache.capture_event()
+
+            # Compute global to local indices
+            self._space_to_partition = cache.borrow_temporary_like(node_indices, temporary_store)
+            wp.launch(
+                kernel=NodePartition._scatter_partition_indices,
+                dim=self.space_topology.node_count(),
+                device=device,
+                inputs=[category_offsets.array, node_indices.array, self._space_to_partition.array],
+            )
+
+            # Copy to shrinked-to-fit array
+            cache.synchronize_event(copy_event)  # Transfer to host must be finished to access node_count()
+            self._node_indices = cache.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(
+        category_offsets: wp.array(dtype=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]
+
+        local_node_count = category_offsets[NodeCategory.EXTERIOR]  # all but exterior nodes
+        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: cache.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)