warp-lang 1.1.0__py3-none-manylinux2014_x86_64.whl → 1.2.0__py3-none-manylinux2014_x86_64.whl

warp/fem/space/shape/__init__.py

@@ -1,3 +1,9 @@
+from enum import Enum
+from typing import Optional
+
+from warp.fem.geometry import element as _element
+from warp.fem.polynomial import Polynomial
+
 from .cube_shape_function import (
     CubeNonConformingPolynomialShapeFunctions,
     CubeSerendipityShapeFunctions,
@@ -11,3 +17,74 @@ from .square_shape_function import (
 )
 from .tet_shape_function import TetrahedronNonConformingPolynomialShapeFunctions, TetrahedronPolynomialShapeFunctions
 from .triangle_shape_function import Triangle2DNonConformingPolynomialShapeFunctions, Triangle2DPolynomialShapeFunctions
+
+
+class ElementBasis(Enum):
+    """Choice of basis function to equip individual elements"""
+
+    LAGRANGE = 0
+    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = 1
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = 2
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+
+
+def get_shape_function(
+    element: _element.Element,
+    space_dimension: int,
+    degree: int,
+    element_basis: ElementBasis,
+    family: Optional[Polynomial] = None,
+):
+    """
+    Equips a reference element with a shape function basis.
+
+    Args:
+        element: the reference element on which to build the shape function
+        space_dimension: the dimension of the embedding space
+        degree: polynomial degree of the per-element shape functions
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the corresponding shape function
+    """
+
+    if degree == 0:
+        return ConstantShapeFunction(element, space_dimension)
+
+    if family is None:
+        family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+    if isinstance(element, _element.Square):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return SquareNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return SquareSerendipityShapeFunctions(degree=degree, family=family)
+
+        return SquareBipolynomialShapeFunctions(degree=degree, family=family)
+    if isinstance(element, _element.Triangle):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Triangle2DNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Triangle elements")
+
+        return Triangle2DPolynomialShapeFunctions(degree=degree)
+
+    if isinstance(element, _element.Cube):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return CubeNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return CubeSerendipityShapeFunctions(degree=degree, family=family)
+
+        return CubeTripolynomialShapeFunctions(degree=degree, family=family)
+    if isinstance(element, _element.Tetrahedron):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TetrahedronNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Tet elements")
+
+        return TetrahedronPolynomialShapeFunctions(degree=degree)
+
+    return NotImplementedError("Unrecognized element type")

warp/fem/space/shape/cube_shape_function.py

@@ -10,16 +10,6 @@ from warp.fem.types import Coords
 
 from .tet_shape_function import TetrahedronPolynomialShapeFunctions
 
-_CUBE_EDGE_INDICES = wp.constant(
-    wp.mat(shape=(3, 4), dtype=int)(
-        [
-            [0, 4, 2, 6],
-            [3, 1, 7, 5],
-            [8, 11, 9, 10],
-        ]
-    )
-)
-
 
 class CubeTripolynomialShapeFunctions:
     VERTEX = 0
@@ -96,13 +86,13 @@ class CubeTripolynomialShapeFunctions:
                         return CubeTripolynomialShapeFunctions.VERTEX, type_instance, 0
 
                     # z edge
-                    type_instance = _CUBE_EDGE_INDICES[2, mi * 2 + mj]
+                    type_instance = 8 + mi * 2 + mj
                     type_index = k - 1
                     return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index
 
                 if zk + mk == 1:
                     # y edge
-                    type_instance = _CUBE_EDGE_INDICES[1, mk * 2 + mi]
+                    type_instance = 4 + mk * 2 + mi
                     type_index = j - 1
                     return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index
 
@@ -114,7 +104,7 @@ class CubeTripolynomialShapeFunctions:
             if zj + mj == 1:
                 if zk + mk == 1:
                     # x edge
-                    type_instance = _CUBE_EDGE_INDICES[0, mj * 2 + mk]
+                    type_instance = mj * 2 + mk
                     type_index = i - 1
                     return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index
 
@@ -399,7 +389,7 @@ class CubeSerendipityShapeFunctions:
     def _edge_coords(type_index: int):
         index_in_side = type_index // 4
         side_offset = type_index - 4 * index_in_side
-        return (wp.vec3i(index_in_side + 1, side_offset // 2, side_offset % 2),)
+        return wp.vec3i(index_in_side + 1, side_offset // 2, side_offset & 1)
 
     @wp.func
     def _edge_axis(node_type: int):
@@ -410,7 +400,7 @@ class CubeSerendipityShapeFunctions:
         index_in_side = type_index // 4
         side_offset = type_index - 4 * index_in_side
 
-        return _CUBE_EDGE_INDICES[node_type - CubeSerendipityShapeFunctions.EDGE_X, side_offset], index_in_side
+        return 4 * (node_type - CubeSerendipityShapeFunctions.EDGE_X) + side_offset, index_in_side
 
     def _get_node_lobatto_indices(self):
         ORDER = self.ORDER

warp/fem/space/tetmesh_function_space.py

@@ -1,16 +1,13 @@
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Tetmesh
-from warp.fem.types import Coords, ElementIndex
+from warp.fem.types import ElementIndex
 
-from .basis_space import ShapeBasisSpace, TraceBasisSpace
 from .shape import (
-    ConstantShapeFunction,
     ShapeFunction,
-    TetrahedronNonConformingPolynomialShapeFunctions,
     TetrahedronPolynomialShapeFunctions,
 )
-from .topology import DiscontinuousSpaceTopologyMixin, SpaceTopology, forward_base_topology
+from .topology import SpaceTopology, forward_base_topology
 
 
 @wp.struct
@@ -130,44 +127,6 @@ class TetmeshSpaceTopology(SpaceTopology):
             tet_face_indices[t1, t1_face] = e
 
 
-class TetmeshDiscontinuousSpaceTopology(
-    DiscontinuousSpaceTopologyMixin,
-    SpaceTopology,
-):
-    def __init__(self, mesh: Tetmesh, shape: ShapeFunction):
-        super().__init__(mesh, shape.NODES_PER_ELEMENT)
-
-
-class TetmeshBasisSpace(ShapeBasisSpace):
-    def __init__(self, topology: TetmeshSpaceTopology, shape: ShapeFunction):
-        super().__init__(topology, shape)
-
-        self._mesh: Tetmesh = topology.geometry
-
-
-class TetmeshPiecewiseConstantBasis(TetmeshBasisSpace):
-    def __init__(self, mesh: Tetmesh):
-        shape = ConstantShapeFunction(mesh.reference_cell(), space_dimension=3)
-        topology = TetmeshDiscontinuousSpaceTopology(mesh, shape)
-        super().__init__(shape=shape, topology=topology)
-
-    class Trace(TraceBasisSpace):
-        @wp.func
-        def _node_coords_in_element(
-            side_arg: Tetmesh.SideArg,
-            basis_arg: TetmeshBasisSpace.BasisArg,
-            element_index: ElementIndex,
-            node_index_in_element: int,
-        ):
-            return Coords(1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0)
-
-        def make_node_coords_in_element(self):
-            return self._node_coords_in_element
-
-    def trace(self):
-        return TetmeshPiecewiseConstantBasis.Trace(self)
-
-
 class TetmeshPolynomialSpaceTopology(TetmeshSpaceTopology):
     def __init__(self, mesh: Tetmesh, shape: TetrahedronPolynomialShapeFunctions):
         super().__init__(mesh, shape, need_tet_edge_indices=shape.ORDER >= 2, need_tet_face_indices=shape.ORDER >= 3)
@@ -258,37 +217,8 @@ class TetmeshPolynomialSpaceTopology(TetmeshSpaceTopology):
         return element_node_index
 
 
-class TetmeshPolynomialBasisSpace(TetmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Tetmesh,
-        degree: int,
-    ):
-        shape = TetrahedronPolynomialShapeFunctions(degree)
-        topology = forward_base_topology(TetmeshPolynomialSpaceTopology, mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class TetmeshDGPolynomialBasisSpace(TetmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Tetmesh,
-        degree: int,
-    ):
-        shape = TetrahedronPolynomialShapeFunctions(degree)
-        topology = TetmeshDiscontinuousSpaceTopology(mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class TetmeshNonConformingPolynomialBasisSpace(TetmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Tetmesh,
-        degree: int,
-    ):
-        shape = TetrahedronNonConformingPolynomialShapeFunctions(degree)
-        topology = TetmeshDiscontinuousSpaceTopology(mesh, shape)
+def make_tetmesh_space_topology(mesh: Tetmesh, shape: ShapeFunction):
+    if isinstance(shape, TetrahedronPolynomialShapeFunctions):
+        return forward_base_topology(TetmeshPolynomialSpaceTopology, mesh, shape)
 
-        super().__init__(topology, shape)
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/trimesh_2d_function_space.py

@@ -1,16 +1,13 @@
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Trimesh2D
-from warp.fem.types import Coords, ElementIndex
+from warp.fem.types import ElementIndex
 
-from .basis_space import ShapeBasisSpace, TraceBasisSpace
 from .shape import (
-    ConstantShapeFunction,
     ShapeFunction,
-    Triangle2DNonConformingPolynomialShapeFunctions,
     Triangle2DPolynomialShapeFunctions,
 )
-from .topology import DiscontinuousSpaceTopologyMixin, SpaceTopology, forward_base_topology
+from .topology import SpaceTopology, forward_base_topology
 
 
 @wp.struct
@@ -95,44 +92,6 @@ class Trimesh2DSpaceTopology(SpaceTopology):
             tri_edge_indices[t1, t1_edge] = e
 
 
-class Trimesh2DDiscontinuousSpaceTopology(
-    DiscontinuousSpaceTopologyMixin,
-    SpaceTopology,
-):
-    def __init__(self, mesh: Trimesh2D, shape: ShapeFunction):
-        super().__init__(mesh, shape.NODES_PER_ELEMENT)
-
-
-class Trimesh2DBasisSpace(ShapeBasisSpace):
-    def __init__(self, topology: Trimesh2DSpaceTopology, shape: ShapeFunction):
-        super().__init__(topology, shape)
-
-        self._mesh: Trimesh2D = topology.geometry
-
-
-class Trimesh2DPiecewiseConstantBasis(Trimesh2DBasisSpace):
-    def __init__(self, mesh: Trimesh2D):
-        shape = ConstantShapeFunction(mesh.reference_cell(), space_dimension=2)
-        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
-        super().__init__(shape=shape, topology=topology)
-
-    class Trace(TraceBasisSpace):
-        @wp.func
-        def _node_coords_in_element(
-            side_arg: Trimesh2D.SideArg,
-            basis_arg: Trimesh2DBasisSpace.BasisArg,
-            element_index: ElementIndex,
-            node_index_in_element: int,
-        ):
-            return Coords(0.5, 0.0, 0.0)
-
-        def make_node_coords_in_element(self):
-            return self._node_coords_in_element
-
-    def trace(self):
-        return Trimesh2DPiecewiseConstantBasis.Trace(self)
-
-
 class Trimesh2DPolynomialSpaceTopology(Trimesh2DSpaceTopology):
     def __init__(self, mesh: Trimesh2D, shape: Triangle2DPolynomialShapeFunctions):
         super().__init__(mesh, shape)
@@ -187,37 +146,8 @@ class Trimesh2DPolynomialSpaceTopology(Trimesh2DSpaceTopology):
         return element_node_index
 
 
-class Trimesh2DPolynomialBasisSpace(Trimesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Trimesh2D,
-        degree: int,
-    ):
-        shape = Triangle2DPolynomialShapeFunctions(degree)
-        topology = forward_base_topology(Trimesh2DPolynomialSpaceTopology, mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class Trimesh2DDGPolynomialBasisSpace(Trimesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Trimesh2D,
-        degree: int,
-    ):
-        shape = Triangle2DPolynomialShapeFunctions(degree)
-        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class Trimesh2DNonConformingPolynomialBasisSpace(Trimesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Trimesh2D,
-        degree: int,
-    ):
-        shape = Triangle2DNonConformingPolynomialShapeFunctions(degree)
-        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
+def make_trimesh_2d_space_topology(mesh: Trimesh2D, shape: ShapeFunction):
+    if isinstance(shape, Triangle2DPolynomialShapeFunctions):
+        return forward_base_topology(Trimesh2DPolynomialSpaceTopology, mesh, shape)
 
-        super().__init__(topology, shape)
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/native/array.h

@@ -489,7 +489,10 @@ CUDA_CALLABLE inline array_t<T> view(array_t<T>& src, int i)
     assert(i >= 0 && i < src.shape[0]);
 
     array_t<T> a;
-    a.data = data_at_byte_offset(src, byte_offset(src, i));
+    size_t offset = byte_offset(src, i);
+    a.data = data_at_byte_offset(src, offset);
+    if (src.grad)
+        a.grad = grad_at_byte_offset(src, offset);
     a.shape[0] = src.shape[1];
     a.shape[1] = src.shape[2];
     a.shape[2] = src.shape[3];
@@ -509,7 +512,10 @@ CUDA_CALLABLE inline array_t<T> view(array_t<T>& src, int i, int j)
     assert(j >= 0 && j < src.shape[1]);
 
     array_t<T> a;
-    a.data = data_at_byte_offset(src, byte_offset(src, i, j));
+    size_t offset = byte_offset(src, i, j);
+    a.data = data_at_byte_offset(src, offset);
+    if (src.grad)
+        a.grad = grad_at_byte_offset(src, offset);
     a.shape[0] = src.shape[2];
     a.shape[1] = src.shape[3];
     a.strides[0] = src.strides[2];
@@ -528,7 +534,10 @@ CUDA_CALLABLE inline array_t<T> view(array_t<T>& src, int i, int j, int k)
     assert(k >= 0 && k < src.shape[2]);
 
     array_t<T> a;
-    a.data = data_at_byte_offset(src, byte_offset(src, i, j, k));
+    size_t offset = byte_offset(src, i, j, k);
+    a.data = data_at_byte_offset(src, offset);
+    if (src.grad)
+        a.grad = grad_at_byte_offset(src, offset);
     a.shape[0] = src.shape[3];
     a.strides[0] = src.strides[3];
     a.ndim = src.ndim-3;

warp/native/builtin.h

@@ -84,8 +84,17 @@ struct half
 
     unsigned short u;
 
-    CUDA_CALLABLE inline bool operator==(const half& h) const { return u == h.u; }
-    CUDA_CALLABLE inline bool operator!=(const half& h) const { return u != h.u; }
+    CUDA_CALLABLE inline bool operator==(const half& h) const
+    {
+        // Use float32 to get IEEE 754 behavior in case of a NaN
+        return float32(h) == float32(*this);
+    }
+
+    CUDA_CALLABLE inline bool operator!=(const half& h) const
+    {
+        // Use float32 to get IEEE 754 behavior in case of a NaN
+        return float32(h) != float32(*this);
+    }
     CUDA_CALLABLE inline bool operator>(const half& h) const { return half_to_float(*this) > half_to_float(h); }
     CUDA_CALLABLE inline bool operator>=(const half& h) const { return half_to_float(*this) >= half_to_float(h); }
     CUDA_CALLABLE inline bool operator<(const half& h) const { return half_to_float(*this) < half_to_float(h); }
@@ -293,7 +302,9 @@ inline CUDA_CALLABLE T bit_xor(T a, T b) { return a^b; } \
 inline CUDA_CALLABLE T lshift(T a, T b) { return a<<b; } \
 inline CUDA_CALLABLE T rshift(T a, T b) { return a>>b; } \
 inline CUDA_CALLABLE T invert(T x) { return ~x; } \
-inline CUDA_CALLABLE bool isfinite(T x) { return true; } \
+inline CUDA_CALLABLE bool isfinite(T x) { return ::isfinite(double(x)); } \
+inline CUDA_CALLABLE bool isnan(T x) { return ::isnan(double(x)); } \
+inline CUDA_CALLABLE bool isinf(T x) { return ::isinf(double(x)); } \
 inline CUDA_CALLABLE void adj_mul(T a, T b, T& adj_a, T& adj_b, T adj_ret) { } \
 inline CUDA_CALLABLE void adj_div(T a, T b, T ret, T& adj_a, T& adj_b, T adj_ret) { } \
 inline CUDA_CALLABLE void adj_add(T a, T b, T& adj_a, T& adj_b, T adj_ret) { } \
@@ -313,7 +324,10 @@ inline CUDA_CALLABLE void adj_bit_or(T a, T b, T& adj_a, T& adj_b, T adj_ret) {
 inline CUDA_CALLABLE void adj_bit_xor(T a, T b, T& adj_a, T& adj_b, T adj_ret) { } \
 inline CUDA_CALLABLE void adj_lshift(T a, T b, T& adj_a, T& adj_b, T adj_ret) { } \
 inline CUDA_CALLABLE void adj_rshift(T a, T b, T& adj_a, T& adj_b, T adj_ret) { } \
-inline CUDA_CALLABLE void adj_invert(T x, T adj_x, T& adj_ret) { }
+inline CUDA_CALLABLE void adj_invert(T x, T adj_x, T& adj_ret) { } \
+inline CUDA_CALLABLE void adj_isnan(const T&, T&, bool) { } \
+inline CUDA_CALLABLE void adj_isinf(const T&, T&, bool) { } \
+inline CUDA_CALLABLE void adj_isfinite(const T&, T&, bool) { }
 
 inline CUDA_CALLABLE int8 abs(int8 x) { return ::abs(x); }
 inline CUDA_CALLABLE int16 abs(int16 x) { return ::abs(x); }
@@ -354,7 +368,7 @@ inline CUDA_CALLABLE uint32 sign(uint32 x) { return 1; }
 inline CUDA_CALLABLE uint64 sign(uint64 x) { return 1; }
 
 
-// Catch-all for non-float types
+// Catch-all for non-float, non-integer types
 template<typename T>
 inline bool CUDA_CALLABLE isfinite(const T&)
 {
@@ -374,6 +388,32 @@ inline bool CUDA_CALLABLE isfinite(double x)
     return ::isfinite(x);
 }
 
+inline bool CUDA_CALLABLE isnan(half x)
+{
+    return ::isnan(float(x));
+}
+inline bool CUDA_CALLABLE isnan(float x)
+{
+    return ::isnan(x);
+}
+inline bool CUDA_CALLABLE isnan(double x)
+{
+    return ::isnan(x);
+}
+
+inline bool CUDA_CALLABLE isinf(half x)
+{
+    return ::isinf(float(x));
+}
+inline bool CUDA_CALLABLE isinf(float x)
+{
+    return ::isinf(x);
+}
+inline bool CUDA_CALLABLE isinf(double x)
+{
+    return ::isinf(x);
+}
+
 template<typename T>
 inline CUDA_CALLABLE void print(const T&)
 {
@@ -466,6 +506,9 @@ inline CUDA_CALLABLE void adj_div(T a, T b, T ret, T& adj_a, T& adj_b, T adj_ret
         assert(0);\
     })\
 }\
+inline CUDA_CALLABLE void adj_isnan(const T&, T&, bool) { }\
+inline CUDA_CALLABLE void adj_isinf(const T&, T&, bool) { }\
+inline CUDA_CALLABLE void adj_isfinite(const T&, T&, bool) { }
 
 DECLARE_FLOAT_OPS(float16)
 DECLARE_FLOAT_OPS(float32)

warp/native/bvh.cpp

@@ -339,6 +339,19 @@ void bvh_rem_descriptor(uint64_t id)
 }
 
 
+// create in-place given existing descriptor
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, BVH& bvh)
+{
+    memset(&bvh, 0, sizeof(BVH));
+
+    bvh.item_lowers = lowers;
+    bvh.item_uppers = uppers;
+    bvh.num_items = num_items;
+
+    MedianBVHBuilder builder;
+    builder.build(bvh, lowers, uppers, num_items);
+}
+
 void bvh_destroy_host(BVH& bvh)
 {
     delete[] bvh.node_lowers;
@@ -360,16 +373,7 @@ void bvh_destroy_host(BVH& bvh)
 uint64_t bvh_create_host(vec3* lowers, vec3* uppers, int num_items)
 {
     BVH* bvh = new BVH();
-    memset(bvh, 0, sizeof(BVH));
-
-    bvh->context = NULL;
-
-    bvh->item_lowers = lowers;
-    bvh->item_uppers = uppers;
-    bvh->num_items = num_items;
-
-    MedianBVHBuilder builder;
-    builder.build(*bvh, lowers, uppers, num_items);
+    wp::bvh_create_host(lowers, uppers, num_items, *bvh);
 
     return (uint64_t)bvh;
 }

warp/native/bvh.cu

@@ -455,7 +455,28 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
 
 }
 
-void bvh_destroy_device(wp::BVH& bvh)
+// create in-place given existing descriptor
+void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, BVH& bvh_host)
+{
+    ContextGuard guard(context);
+
+    bvh_host.num_items = num_items;
+    bvh_host.max_nodes = 2*num_items;
+    bvh_host.node_lowers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
+    bvh_host.node_uppers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
+    bvh_host.node_parents = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
+    bvh_host.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
+    bvh_host.root = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int));
+    bvh_host.item_lowers = lowers;
+    bvh_host.item_uppers = uppers;
+
+    bvh_host.context = context ? context : cuda_context_get_current();
+
+    LinearBVHBuilderGPU builder;
+    builder.build(bvh_host, lowers, uppers, num_items, NULL);
+}
+
+void bvh_destroy_device(BVH& bvh)
 {
     ContextGuard guard(bvh.context);
 
@@ -485,20 +506,7 @@ uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, in
     ContextGuard guard(context);
 
     wp::BVH bvh_host;
-    bvh_host.num_items = num_items;
-    bvh_host.max_nodes = 2*num_items;
-    bvh_host.node_lowers = (wp::BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_host.node_uppers = (wp::BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_host.node_parents = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-    bvh_host.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-    bvh_host.root = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int));
-    bvh_host.item_lowers = lowers;
-    bvh_host.item_uppers = uppers;
-
-    bvh_host.context = context ? context : cuda_context_get_current();
-
-    wp::LinearBVHBuilderGPU builder;
-    builder.build(bvh_host, lowers, uppers, num_items, NULL);
+    bvh_create_device(WP_CURRENT_CONTEXT, lowers, uppers, num_items, bvh_host);
 
     // create device-side BVH descriptor
     wp::BVH* bvh_device = (wp::BVH*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVH));

warp/native/bvh.h

@@ -417,6 +417,7 @@ CUDA_CALLABLE void bvh_rem_descriptor(uint64_t id);
 
 #if !__CUDA_ARCH__
 
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, BVH& bvh);
 void bvh_destroy_host(wp::BVH& bvh);
 void bvh_refit_host(wp::BVH& bvh);
 

warp/native/clang/clang.cpp

@@ -413,6 +413,8 @@ WP_API int load_obj(const char* object_file, const char* module_name)
             SYMBOL(memcpy), SYMBOL(memset), SYMBOL(memmove),
             SYMBOL(_wp_assert),
             SYMBOL(_wp_isfinite),
+            SYMBOL(_wp_isnan),
+            SYMBOL(_wp_isinf),
         #if defined(_WIN64)
             // For functions with large stack frames the compiler will emit a call to
             // __chkstk() to linearly touch each memory page. This grows the stack without
@@ -492,4 +494,3 @@ WP_API uint64_t lookup(const char* dll_name, const char* function_name)
 }  // extern "C"
 
 }  // namespace wp
-

warp/native/crt.cpp

@@ -17,6 +17,16 @@ extern "C" WP_API int _wp_isfinite(double x)
     return std::isfinite(x);
 }
 
+extern "C" WP_API int _wp_isnan(double x)
+{
+    return std::isnan(x);
+}
+
+extern "C" WP_API int _wp_isinf(double x)
+{
+    return std::isinf(x);
+}
+
 extern "C" WP_API void _wp_assert(const char* expression, const char* file, unsigned int line)
 {
     fflush(stdout);
@@ -29,4 +39,4 @@ extern "C" WP_API void _wp_assert(const char* expression, const char* file, unsi
     // Now invoke the standard assert(), which may abort the program or break
     // into the debugger as decided by the runtime environment.
     assert(false && "assert() failed");
-}
+}

warp/native/crt.h

@@ -38,6 +38,12 @@ extern "C" WP_API void _wp_assert(const char* message, const char* file, unsigne
 // Helper for implementing isfinite()
 extern "C" WP_API int _wp_isfinite(double);
 
+// Helper for implementing isnan()
+extern "C" WP_API int _wp_isnan(double);
+
+// Helper for implementing isinf()
+extern "C" WP_API int _wp_isinf(double);
+
 #endif  // !__CUDA_ARCH__
 
 #if !defined(WP_NO_CRT)
@@ -124,6 +130,7 @@ enum {
 #define INFINITY   ((float)(DBL_MAX * DBL_MAX))
 #define HUGE_VAL   ((double)INFINITY)
 #define HUGE_VALF  ((float)INFINITY)
+#define NAN        ((float)(0.0 / 0.0))
 
 /// stdint.h
 typedef signed char      int8_t;
@@ -320,6 +327,16 @@ inline bool isfinite(double x)
     return _wp_isfinite(x);
 }
 
+inline bool isnan(double x)
+{
+    return _wp_isnan(x);
+}
+
+inline bool isinf(double x)
+{
+    return _wp_isinf(x);
+}
+
 // assert.h
 #ifdef NDEBUG
     #define assert(expression) ((void)0)
@@ -332,4 +349,4 @@ inline bool isfinite(double x)
 
 #endif  // !__CUDACC__
 
-#endif // WP_NO_CRT
+#endif // WP_NO_CRT