warp-lang 1.4.1__py3-none-macosx_10_13_universal2.whl → 1.5.0__py3-none-macosx_10_13_universal2.whl

warp/fem/space/trimesh_function_space.py

@@ -0,0 +1,179 @@
+import warp as wp
+from warp.fem import cache
+from warp.fem.geometry import Trimesh
+from warp.fem.types import ElementIndex
+
+from .shape import TriangleShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+
+@wp.struct
+class TrimeshTopologyArg:
+    edge_vertex_indices: wp.array(dtype=wp.vec2i)
+    tri_edge_indices: wp.array2d(dtype=int)
+
+    vertex_count: int
+    edge_count: int
+
+
+class TrimeshSpaceTopology(SpaceTopology):
+    TopologyArg = TrimeshTopologyArg
+
+    def __init__(self, mesh: Trimesh, shape: TriangleShapeFunction):
+        self._shape = shape
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+
+        self._compute_tri_edge_indices()
+        self.element_node_index = self._make_element_node_index()
+        self.element_node_sign = self._make_element_node_sign()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    @cache.cached_arg_value
+    def topo_arg_value(self, device):
+        arg = TrimeshTopologyArg()
+        arg.tri_edge_indices = self._tri_edge_indices.to(device)
+        arg.edge_vertex_indices = self._mesh.edge_vertex_indices.to(device)
+
+        arg.vertex_count = self._mesh.vertex_count()
+        arg.edge_count = self._mesh.side_count()
+        return arg
+
+    def _compute_tri_edge_indices(self):
+        self._tri_edge_indices = wp.empty(
+            dtype=int, device=self._mesh.tri_vertex_indices.device, shape=(self._mesh.cell_count(), 3)
+        )
+
+        wp.launch(
+            kernel=TrimeshSpaceTopology._compute_tri_edge_indices_kernel,
+            dim=self._mesh.edge_tri_indices.shape,
+            device=self._mesh.tri_vertex_indices.device,
+            inputs=[
+                self._mesh.edge_tri_indices,
+                self._mesh.edge_vertex_indices,
+                self._mesh.tri_vertex_indices,
+                self._tri_edge_indices,
+            ],
+        )
+
+    @wp.func
+    def _find_edge_index_in_tri(
+        edge_vtx: wp.vec2i,
+        tri_vtx: wp.vec3i,
+    ):
+        for k in range(2):
+            if (edge_vtx[0] == tri_vtx[k] and edge_vtx[1] == tri_vtx[k + 1]) or (
+                edge_vtx[1] == tri_vtx[k] and edge_vtx[0] == tri_vtx[k + 1]
+            ):
+                return k
+        return 2
+
+    @wp.kernel
+    def _compute_tri_edge_indices_kernel(
+        edge_tri_indices: wp.array(dtype=wp.vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
+        tri_vertex_indices: wp.array2d(dtype=int),
+        tri_edge_indices: wp.array2d(dtype=int),
+    ):
+        e = wp.tid()
+
+        edge_vtx = edge_vertex_indices[e]
+        edge_tris = edge_tri_indices[e]
+
+        t0 = edge_tris[0]
+        t0_vtx = wp.vec3i(tri_vertex_indices[t0, 0], tri_vertex_indices[t0, 1], tri_vertex_indices[t0, 2])
+        t0_edge = TrimeshSpaceTopology._find_edge_index_in_tri(edge_vtx, t0_vtx)
+        tri_edge_indices[t0, t0_edge] = e
+
+        t1 = edge_tris[1]
+        if t1 != t0:
+            t1_vtx = wp.vec3i(tri_vertex_indices[t1, 0], tri_vertex_indices[t1, 1], tri_vertex_indices[t1, 2])
+            t1_edge = TrimeshSpaceTopology._find_edge_index_in_tri(edge_vtx, t1_vtx)
+            tri_edge_indices[t1, t1_edge] = e
+
+    def node_count(self) -> int:
+        return (
+            self._mesh.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self._mesh.side_count() * self._shape.EDGE_NODE_COUNT
+            + self._mesh.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        INTERIOR_NODES_PER_SIDE = self._shape.EDGE_NODE_COUNT
+        INTERIOR_NODES_PER_CELL = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: TrimeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == TriangleShapeFunction.VERTEX:
+                    vertex = type_index // VERTEX_NODE_COUNT
+                    vertex_node = type_index - VERTEX_NODE_COUNT * vertex
+                    return geo_arg.topology.tri_vertex_indices[element_index][vertex] * VERTEX_NODE_COUNT + vertex_node
+
+            global_offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(INTERIOR_NODES_PER_SIDE > 0):
+                if node_type == TriangleShapeFunction.EDGE:
+                    edge = type_index // INTERIOR_NODES_PER_SIDE
+                    edge_node = type_index - INTERIOR_NODES_PER_SIDE * edge
+
+                    global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
+
+                    if (
+                        topo_arg.edge_vertex_indices[global_edge_index][0]
+                        != geo_arg.topology.tri_vertex_indices[element_index][edge]
+                    ):
+                        edge_node = INTERIOR_NODES_PER_SIDE - 1 - edge_node
+
+                    return global_offset + INTERIOR_NODES_PER_SIDE * global_edge_index + edge_node
+
+                global_offset += INTERIOR_NODES_PER_SIDE * topo_arg.edge_count
+
+            return global_offset + INTERIOR_NODES_PER_CELL * element_index + type_index
+
+        return element_node_index
+
+    def _make_element_node_sign(self):
+        INTERIOR_NODES_PER_SIDE = self._shape.EDGE_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_sign(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: TrimeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if node_type == TriangleShapeFunction.EDGE:
+                edge = type_index // INTERIOR_NODES_PER_SIDE
+
+                global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
+                return wp.select(
+                    topo_arg.edge_vertex_indices[global_edge_index][0]
+                    == geo_arg.topology.tri_vertex_indices[element_index][edge],
+                    -1.0,
+                    1.0,
+                )
+
+            return 1.0
+
+        return element_node_sign
+
+
+def make_trimesh_space_topology(mesh: Trimesh, shape: TriangleShapeFunction):
+    if isinstance(shape, TriangleShapeFunction):
+        return forward_base_topology(TrimeshSpaceTopology, mesh, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/utils.py

@@ -1,323 +1,18 @@
-from typing import Any, Tuple, Union
+from typing import Tuple, Union
 
 import numpy as np
 
 import warp as wp
 import warp.fem.cache as cache
+from warp.fem.linalg import (  # noqa: F401 (for backward compatibility, not part of public API but used in examples)
+    array_axpy,
+    inverse_qr,
+    symmetric_eigenvalues_qr,
+)
 from warp.fem.types import NULL_NODE_INDEX
 from warp.utils import array_scan, radix_sort_pairs, runlength_encode
 
 
-@wp.func
-def generalized_outer(x: Any, y: Any):
-    """Generalized outer product allowing for the first argument to be a scalar"""
-    return wp.outer(x, y)
-
-
-@wp.func
-def generalized_outer(x: wp.float32, y: wp.vec2):
-    return x * y
-
-
-@wp.func
-def generalized_outer(x: wp.float32, y: wp.vec3):
-    return x * y
-
-
-@wp.func
-def generalized_inner(x: Any, y: Any):
-    """Generalized inner product allowing for the first argument to be a tensor"""
-    return wp.dot(x, y)
-
-
-@wp.func
-def generalized_inner(x: wp.mat22, y: wp.vec2):
-    return x[0] * y[0] + x[1] * y[1]
-
-
-@wp.func
-def generalized_inner(x: wp.mat33, y: wp.vec3):
-    return x[0] * y[0] + x[1] * y[1] + x[2] * y[2]
-
-
-@wp.func
-def unit_element(template_type: Any, coord: int):
-    """Returns a instance of `template_type` with a single coordinate set to 1 in the canonical basis"""
-
-    t = type(template_type)(0.0)
-    t[coord] = 1.0
-    return t
-
-
-@wp.func
-def unit_element(template_type: wp.float32, coord: int):
-    return 1.0
-
-
-@wp.func
-def unit_element(template_type: wp.mat22, coord: int):
-    t = wp.mat22(0.0)
-    row = coord // 2
-    col = coord - 2 * row
-    t[row, col] = 1.0
-    return t
-
-
-@wp.func
-def unit_element(template_type: wp.mat33, coord: int):
-    t = wp.mat33(0.0)
-    row = coord // 3
-    col = coord - 3 * row
-    t[row, col] = 1.0
-    return t
-
-
-@wp.func
-def symmetric_part(x: Any):
-    """Symmetric part of a square tensor"""
-    return 0.5 * (x + wp.transpose(x))
-
-
-@wp.func
-def skew_part(x: wp.mat22):
-    """Skew part of a 2x2 tensor as corresponding rotation angle"""
-    return 0.5 * (x[1, 0] - x[0, 1])
-
-
-@wp.func
-def skew_part(x: wp.mat33):
-    """Skew part of a 3x3 tensor as the corresponding rotation vector"""
-    a = 0.5 * (x[2, 1] - x[1, 2])
-    b = 0.5 * (x[0, 2] - x[2, 0])
-    c = 0.5 * (x[1, 0] - x[0, 1])
-    return wp.vec3(a, b, c)
-
-
-@wp.func
-def householder_qr_decomposition(A: Any):
-    """
-    QR decomposition of a square matrix using Householder reflections
-
-    Returns Q and R such that Q R = A, Q orthonormal (such that QQ^T = Id), R upper triangular
-    """
-
-    x = type(A[0])()
-    Q = wp.identity(n=type(x).length, dtype=A.dtype)
-
-    zero = x.dtype(0.0)
-    two = x.dtype(2.0)
-
-    for i in range(type(x).length):
-        for k in range(type(x).length):
-            x[k] = wp.select(k < i, A[k, i], zero)
-
-        alpha = wp.length(x) * wp.sign(x[i])
-        x[i] += alpha
-        two_over_x_sq = wp.select(alpha == zero, two / wp.length_sq(x), zero)
-
-        A -= wp.outer(two_over_x_sq * x, x * A)
-        Q -= wp.outer(Q * x, two_over_x_sq * x)
-
-    return Q, A
-
-
-@wp.func
-def householder_make_hessenberg(A: Any):
-    """Transforms a square matrix to Hessenberg form (single lower diagonal) using Householder reflections
-
-    Returns:
-        Q and H such that Q H Q^T = A, Q orthonormal, H under Hessenberg form
-        If A is symmetric, H will be tridiagonal
-    """
-
-    x = type(A[0])()
-    Q = wp.identity(n=type(x).length, dtype=A.dtype)
-
-    zero = x.dtype(0.0)
-    two = x.dtype(2.0)
-
-    for i in range(1, type(x).length):
-        for k in range(type(x).length):
-            x[k] = wp.select(k < i, A[k, i - 1], zero)
-
-        alpha = wp.length(x) * wp.sign(x[i])
-        x[i] += alpha
-        two_over_x_sq = wp.select(alpha == zero, two / wp.length_sq(x), zero)
-
-        # apply on both sides
-        A -= wp.outer(two_over_x_sq * x, x * A)
-        A -= wp.outer(A * x, two_over_x_sq * x)
-        Q -= wp.outer(Q * x, two_over_x_sq * x)
-
-    return Q, A
-
-
-@wp.func
-def solve_triangular(R: Any, b: Any):
-    """Solves for R x = b where R is an upper triangular matrix
-
-    Returns x
-    """
-    zero = b.dtype(0)
-    x = type(b)(b.dtype(0))
-    for i in range(b.length, 0, -1):
-        j = i - 1
-        r = b[j] - wp.dot(R[j], x)
-        x[j] = wp.select(R[j, j] == zero, r / R[j, j], zero)
-
-    return x
-
-
-@wp.func
-def inverse_qr(A: Any):
-    # Computes a square matrix inverse using QR factorization
-
-    Q, R = householder_qr_decomposition(A)
-
-    A_inv = type(A)()
-    for i in range(type(A[0]).length):
-        A_inv[i] = solve_triangular(R, Q[i])  # ith column of Q^T
-
-    return wp.transpose(A_inv)
-
-
-@wp.func
-def _wilkinson_shift(a: Any, b: Any, c: Any, tol: Any):
-    # Wilkinson shift: estimate eigenvalue of 2x2 symmetric matrix [a, c, c, b]
-    d = (a - b) * type(tol)(0.5)
-    return b + d - wp.sign(d) * wp.sqrt(d * d + c * c)
-
-
-@wp.func
-def _givens_rotation(a: Any, b: Any):
-    # Givens rotation [[c -s], [s c]] such that sa+cb =0
-    zero = type(a)(0.0)
-    one = type(a)(1.0)
-
-    b2 = b * b
-    if b2 == zero:
-        # id rotation
-        return one, zero
-
-    scale = one / wp.sqrt(a * a + b2)
-    return a * scale, -b * scale
-
-
-@wp.func
-def tridiagonal_symmetric_eigenvalues_qr(D: Any, L: Any, Q: Any, tol: Any):
-    """
-    Computes the eigenvalues and eigen vectors of a symmetric tridiagonal matrix using the
-    Symmetric tridiagonal QR algorithm with implicit Wilkinson shift
-
-    Args:
-        D: Main diagonal of the matrix
-        L: Lower diagonal of the matrix, indexed such that L[i] = A[i+1, i]
-        Q: Initialization for the eigenvectors, useful if a pre-transformation has been applied, otherwise may be identity
-        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
-
-    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
-
-
-    Ref: Arbenz P, Numerical Methods for Solving Large Scale Eigenvalue Problems, Chapter 4 (QR algorithm, Mar 13, 2018)
-    """
-
-    two = D.dtype(2.0)
-
-    # so that we can use the type length in expressions
-    # this will prevent unrolling by warp, but should be ok for native code
-    m = int(0)
-    for _ in range(type(D).length):
-        m += 1
-
-    start = int(0)
-    y = D.dtype(0.0)  # moving buldge
-    x = D.dtype(0.0)  # coeff atop buldge
-
-    for _ in range(32 * m):  # failsafe, usually converges faster than that
-        # Iterate over all independent (deflated) blocks
-        end = int(-1)
-
-        for k in range(m - 1):
-            if k >= end:
-                # Check if new block is starting
-                if k == end or wp.abs(L[k]) <= tol * (wp.abs(D[k]) + wp.abs(D[k + 1])):
-                    continue
-
-                # Find end of block
-                start = k
-                end = start + 1
-                while end + 1 < m:
-                    if wp.abs(L[end]) <= tol * (wp.abs(D[end + 1]) + wp.abs(D[end])):
-                        break
-                    end += 1
-
-                # Wilkinson shift (an eigenvalue of the last 2x2 block)
-                shift = _wilkinson_shift(D[end - 1], D[end], L[end - 1], tol)
-
-                # start with eliminating lower diag of first column of shifted matrix
-                # (i.e. first step of excplit QR factorization)
-                # Then all further steps eliminate the buldge (second diag) of the non-shifted matrix
-                x = D[start] - shift
-                y = L[start]
-
-            c, s = _givens_rotation(x, y)
-
-            # Apply Givens rotation on both sides of tridiagonal matrix
-
-            # middle block
-            d = D[k] - D[k + 1]
-            z = (two * c * L[k] + d * s) * s
-            D[k] -= z
-            D[k + 1] += z
-            L[k] = d * c * s + (c * c - s * s) * L[k]
-
-            if k > start:
-                L[k - 1] = c * x - s * y
-
-            x = L[k]
-            y = -s * L[k + 1]  # new buldge
-            L[k + 1] *= c
-
-            # apply givens rotation on left of Q
-            # note: Q is transposed compared to usual impls, as Warp makes it easier to index rows
-            Qk0 = Q[k]
-            Qk1 = Q[k + 1]
-            Q[k] = c * Qk0 - s * Qk1
-            Q[k + 1] = c * Qk1 + s * Qk0
-
-        if end <= 0:
-            # We did nothing, so diagonalization must have been achieved
-            break
-
-    return D, Q
-
-
-@wp.func
-def symmetric_eigenvalues_qr(A: Any, tol: Any):
-    """
-    Computes the eigenvalues and eigen vectors of a square symmetric matrix A using the QR algorithm
-
-    Args:
-        A: square symmetric matrix
-        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
-
-    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
-    """
-
-    # Put A under Hessenberg form (tridiagonal)
-    Q, H = householder_make_hessenberg(A)
-
-    # tridiagonal storage for H
-    D = wp.get_diag(H)
-    L = type(D)(A.dtype(0.0))
-    for i in range(1, type(D).length):
-        L[i - 1] = H[i, i - 1]
-
-    Qt = wp.transpose(Q)
-    ev, P = tridiagonal_symmetric_eigenvalues_qr(D, L, Qt, tol)
-    return ev, P
-
-
 def compress_node_indices(
     node_count: int,
     node_indices: wp.array(dtype=int),
@@ -458,20 +153,6 @@ def masked_indices(
     return indices_temp, offsets_temp
 
 
-def array_axpy(x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 1.0):
-    """Performs y = alpha*x + beta*y"""
-
-    dtype = wp.types.type_scalar_type(x.dtype)
-
-    alpha = dtype(alpha)
-    beta = dtype(beta)
-
-    if not wp.types.types_equal(x.dtype, y.dtype) or x.shape != y.shape or x.device != y.device:
-        raise ValueError("x and y arrays must have same dat atype, shape and device")
-
-    wp.launch(kernel=_array_axpy_kernel, dim=x.shape, device=x.device, inputs=[x, y, alpha, beta])
-
-
 @wp.kernel
 def _iota_kernel(indices: wp.array(dtype=int), divisor: int):
     indices[wp.tid()] = wp.tid() // divisor
@@ -515,12 +196,6 @@ def _masked_indices_kernel(
         masked_to_global[masked_idx] = i
 
 
-@wp.kernel
-def _array_axpy_kernel(x: wp.array(dtype=Any), y: wp.array(dtype=Any), alpha: Any, beta: Any):
-    i = wp.tid()
-    y[i] = beta * y[i] + alpha * x[i]
-
-
 def grid_to_tris(Nx: int, Ny: int):
     """Constructs a triangular mesh topology by dividing each cell of a dense 2D grid into two triangles.
 

warp/jax_experimental.py

@@ -102,7 +102,9 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
     assert hooks.forward, "Failed to find kernel entry point"
 
     # Launch the kernel.
-    wp.context.runtime.core.cuda_launch_kernel(device.context, hooks.forward, bounds.size, 0, kernel_params, stream)
+    wp.context.runtime.core.cuda_launch_kernel(
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
+    )
 
 
 # TODO: is there a simpler way of getting the Jax "current" device?