warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/fem/space/partition.py

@@ -19,7 +19,7 @@ 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 _iota_kernel, compress_node_indices
+from warp.fem.utils import compress_node_indices
 
 from .function_space import FunctionSpace
 from .topology import SpaceTopology
@@ -87,7 +87,7 @@ class WholeSpacePartition(SpacePartition):
         """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=_iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array, 1])
+            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):
@@ -104,6 +104,10 @@ class WholeSpacePartition(SpacePartition):
     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)

warp/fem/space/quadmesh_function_space.py

@@ -166,10 +166,10 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
             if wp.static(EDGE_NODE_COUNT > 0):
                 # EDGE_X, EDGE_Y
-                side_start = wp.select(
+                side_start = wp.where(
                     node_type == SquareShapeFunction.EDGE_X,
-                    wp.select(type_instance == 0, 1, 3),
-                    wp.select(type_instance == 0, 2, 0),
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
                 )
 
                 side_index = topo_arg.quad_edge_indices[element_index, side_start]
@@ -178,7 +178,7 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
                 # Flip indexing direction
                 flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
-                index_in_side = wp.select(flipped, type_index, EDGE_NODE_COUNT - 1 - type_index)
+                index_in_side = wp.where(flipped, EDGE_NODE_COUNT - 1 - type_index, type_index)
 
                 return global_offset + EDGE_NODE_COUNT * side_index + index_in_side
 
@@ -197,10 +197,10 @@ class QuadmeshSpaceTopology(SpaceTopology):
             node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
 
             if node_type == SquareShapeFunction.EDGE_X or node_type == SquareShapeFunction.EDGE_Y:
-                side_start = wp.select(
+                side_start = wp.where(
                     node_type == SquareShapeFunction.EDGE_X,
-                    wp.select(type_instance == 0, 1, 3),
-                    wp.select(type_instance == 0, 2, 0),
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
                 )
 
                 side_index = topo_arg.quad_edge_indices[element_index, side_start]
@@ -209,7 +209,7 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
                 # Flip indexing direction
                 flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
-                return wp.select(flipped, 1.0, -1.0)
+                return wp.where(flipped, -1.0, 1.0)
 
             return 1.0
 

warp/fem/space/shape/cube_shape_function.py

@@ -63,7 +63,7 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
 
         self.ORDER = wp.constant(degree)
         self.NODES_PER_ELEMENT = wp.constant((degree + 1) ** 3)
-        self.NODES_PER_EDGE = wp.constant(degree + 1)
+        self.NODES_PER_SIDE = wp.constant((degree + 1) ** 2)
 
         if is_closed(self.family):
             self.VERTEX_NODE_COUNT = wp.constant(1)
@@ -152,13 +152,13 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
         ):
             i, j, k = self._node_ijk(node_index_in_elt)
 
-            zi = wp.select(i == 0, 0, 1)
-            zj = wp.select(j == 0, 0, 1)
-            zk = wp.select(k == 0, 0, 1)
+            zi = wp.where(i == 0, 1, 0)
+            zj = wp.where(j == 0, 1, 0)
+            zk = wp.where(k == 0, 1, 0)
 
-            mi = wp.select(i == ORDER, 0, 1)
-            mj = wp.select(j == ORDER, 0, 1)
-            mk = wp.select(k == ORDER, 0, 1)
+            mi = wp.where(i == ORDER, 1, 0)
+            mj = wp.where(j == ORDER, 1, 0)
+            mk = wp.where(k == ORDER, 1, 0)
 
             if zi + mi == 1:
                 if zj + mj == 1:
@@ -504,7 +504,7 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
 
         self.ORDER = wp.constant(degree)
         self.NODES_PER_ELEMENT = wp.constant(8 + 12 * (degree - 1))
-        self.NODES_PER_EDGE = wp.constant(degree + 1)
+        self.NODES_PER_SIDE = wp.constant(4 * degree)
 
         self.VERTEX_NODE_COUNT = wp.constant(1)
         self.EDGE_NODE_COUNT = wp.constant(degree - 1)
@@ -634,9 +634,9 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             if node_type == CubeSerendipityShapeFunctions.VERTEX:
                 node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
 
-                cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
-                cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2])
+                cx = wp.where(node_ijk[0] == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_ijk[1] == 0, 1.0 - coords[1], coords[1])
+                cz = wp.where(node_ijk[2] == 0, 1.0 - coords[2], coords[2])
 
                 w = cx * cy * cz
 
@@ -659,8 +659,8 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             local_coords = Grid3D._world_to_local(axis, coords)
 
             w = float(1.0)
-            w *= wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1])
-            w *= wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2])
+            w *= wp.where(node_all[1] == 0, 1.0 - local_coords[1], local_coords[1])
+            w *= wp.where(node_all[2] == 0, 1.0 - local_coords[2], local_coords[2])
 
             for k in range(ORDER_PLUS_ONE):
                 if k != node_all[0]:
@@ -690,13 +690,13 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             if node_type == CubeSerendipityShapeFunctions.VERTEX:
                 node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
 
-                cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
-                cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2])
+                cx = wp.where(node_ijk[0] == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_ijk[1] == 0, 1.0 - coords[1], coords[1])
+                cz = wp.where(node_ijk[2] == 0, 1.0 - coords[2], coords[2])
 
-                gx = wp.select(node_ijk[0] == 0, 1.0, -1.0)
-                gy = wp.select(node_ijk[1] == 0, 1.0, -1.0)
-                gz = wp.select(node_ijk[2] == 0, 1.0, -1.0)
+                gx = wp.where(node_ijk[0] == 0, -1.0, 1.0)
+                gy = wp.where(node_ijk[1] == 0, -1.0, 1.0)
+                gz = wp.where(node_ijk[2] == 0, -1.0, 1.0)
 
                 if wp.static(ORDER == 2):
                     w = cx + cy + cz - 3.0 + LOBATTO_COORDS[1]
@@ -728,11 +728,11 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
 
             local_coords = Grid3D._world_to_local(axis, coords)
 
-            w_long = wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1])
-            w_lat = wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2])
+            w_long = wp.where(node_all[1] == 0, 1.0 - local_coords[1], local_coords[1])
+            w_lat = wp.where(node_all[2] == 0, 1.0 - local_coords[2], local_coords[2])
 
-            g_long = wp.select(node_all[1] == 0, 1.0, -1.0)
-            g_lat = wp.select(node_all[2] == 0, 1.0, -1.0)
+            g_long = wp.where(node_all[1] == 0, -1.0, 1.0)
+            g_lat = wp.where(node_all[2] == 0, -1.0, 1.0)
 
             w_alt = LAGRANGE_SCALE[node_all[0]]
             g_alt = float(0.0)

warp/fem/space/shape/square_shape_function.py

@@ -461,8 +461,8 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
             node_i, node_j = self._node_lobatto_indices(node_type, type_instance, type_index)
 
             if node_type == SquareSerendipityShapeFunctions.VERTEX:
-                cx = wp.select(node_i == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_j == 0, coords[1], 1.0 - coords[1])
+                cx = wp.where(node_i == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_j == 0, 1.0 - coords[1], coords[1])
 
                 w = cx * cy
 
@@ -475,7 +475,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
 
             w = float(1.0)
             if node_type == SquareSerendipityShapeFunctions.EDGE_Y:
-                w *= wp.select(node_i == 0, coords[0], 1.0 - coords[0])
+                w *= wp.where(node_i == 0, 1.0 - coords[0], coords[0])
             else:
                 for k in range(ORDER_PLUS_ONE):
                     if k != node_i:
@@ -484,7 +484,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
                 w *= LAGRANGE_SCALE[node_i]
 
             if node_type == SquareSerendipityShapeFunctions.EDGE_X:
-                w *= wp.select(node_j == 0, coords[1], 1.0 - coords[1])
+                w *= wp.where(node_j == 0, 1.0 - coords[1], coords[1])
             else:
                 for k in range(ORDER_PLUS_ONE):
                     if k != node_j:
@@ -513,11 +513,11 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
             node_i, node_j = self._node_lobatto_indices(node_type, type_instance, type_index)
 
             if node_type == SquareSerendipityShapeFunctions.VERTEX:
-                cx = wp.select(node_i == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_j == 0, coords[1], 1.0 - coords[1])
+                cx = wp.where(node_i == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_j == 0, 1.0 - coords[1], coords[1])
 
-                gx = wp.select(node_i == 0, 1.0, -1.0)
-                gy = wp.select(node_j == 0, 1.0, -1.0)
+                gx = wp.where(node_i == 0, -1.0, 1.0)
+                gy = wp.where(node_j == 0, -1.0, 1.0)
 
                 if ORDER == 2:
                     w = cx + cy - 2.0 + LOBATTO_COORDS[1]
@@ -537,7 +537,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
                     return wp.vec2(grad_x, grad_y) * DEGREE_3_CIRCLE_SCALE
 
             if node_type == SquareSerendipityShapeFunctions.EDGE_X:
-                prefix_x = wp.select(node_j == 0, coords[1], 1.0 - coords[1])
+                prefix_x = wp.where(node_j == 0, 1.0 - coords[1], coords[1])
             else:
                 prefix_x = LAGRANGE_SCALE[node_j]
                 for k in range(ORDER_PLUS_ONE):
@@ -545,7 +545,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
                         prefix_x *= coords[1] - LOBATTO_COORDS[k]
 
             if node_type == SquareSerendipityShapeFunctions.EDGE_Y:
-                prefix_y = wp.select(node_i == 0, coords[0], 1.0 - coords[0])
+                prefix_y = wp.where(node_i == 0, 1.0 - coords[0], coords[0])
             else:
                 prefix_y = LAGRANGE_SCALE[node_i]
                 for k in range(ORDER_PLUS_ONE):
@@ -553,7 +553,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
                         prefix_y *= coords[0] - LOBATTO_COORDS[k]
 
             if node_type == SquareSerendipityShapeFunctions.EDGE_X:
-                grad_y = wp.select(node_j == 0, 1.0, -1.0) * prefix_y
+                grad_y = wp.where(node_j == 0, -1.0, 1.0) * prefix_y
             else:
                 prefix_y *= LAGRANGE_SCALE[node_j]
                 grad_y = float(0.0)
@@ -564,7 +564,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
                         prefix_y *= delta_y
 
             if node_type == SquareSerendipityShapeFunctions.EDGE_Y:
-                grad_x = wp.select(node_i == 0, 1.0, -1.0) * prefix_x
+                grad_x = wp.where(node_i == 0, -1.0, 1.0) * prefix_x
             else:
                 prefix_x *= LAGRANGE_SCALE[node_i]
                 grad_x = float(0.0)

warp/fem/space/shape/triangle_shape_function.py

@@ -196,7 +196,7 @@ class TrianglePolynomialShapeFunctions(TriangleShapeFunction):
         def trace_node_quadrature_weight(node_index_in_element: int):
             node_type, type_index = self.node_type_and_type_index(node_index_in_element)
 
-            return wp.select(node_type == TrianglePolynomialShapeFunctions.VERTEX, EDGE_WEIGHT, VERTEX_WEIGHT)
+            return wp.where(node_type == TrianglePolynomialShapeFunctions.VERTEX, VERTEX_WEIGHT, EDGE_WEIGHT)
 
         return trace_node_quadrature_weight
 

warp/fem/space/tetmesh_function_space.py

@@ -244,10 +244,10 @@ class TetmeshSpaceTopology(SpaceTopology):
                     edge = type_index // INTERIOR_NODES_PER_EDGE
                     c1, c2 = TetrahedronShapeFunction.edge_vidx(edge)
 
-                    return wp.select(
+                    return wp.where(
                         geo_arg.tet_vertex_indices[element_index][c1] > geo_arg.tet_vertex_indices[element_index][c2],
-                        1.0,
                         -1.0,
+                        1.0,
                     )
 
             if wp.static(INTERIOR_NODES_PER_FACE > 0):
@@ -257,7 +257,7 @@ class TetmeshSpaceTopology(SpaceTopology):
                     global_face_index = topo_arg.tet_face_indices[element_index][face]
                     inner = topo_arg.face_tet_indices[global_face_index][0]
 
-                    return wp.select(inner == element_index, -1.0, 1.0)
+                    return wp.where(inner == element_index, 1.0, -1.0)
 
             return 1.0
 

warp/fem/space/trimesh_function_space.py

@@ -175,11 +175,11 @@ class TrimeshSpaceTopology(SpaceTopology):
                 edge = type_index // INTERIOR_NODES_PER_SIDE
 
                 global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
-                return wp.select(
+                return wp.where(
                     topo_arg.edge_vertex_indices[global_edge_index][0]
                     == geo_arg.topology.tri_vertex_indices[element_index][edge],
-                    -1.0,
                     1.0,
+                    -1.0,
                 )
 
             return 1.0

warp/fem/utils.py

@@ -56,13 +56,11 @@ def compress_node_indices(
         sorted_node_indices = sorted_node_indices_temp.array
         sorted_array_indices = sorted_array_indices_temp.array
 
-        wp.copy(dest=sorted_node_indices, src=node_indices, count=index_count)
-
         indices_per_element = 1 if node_indices.ndim == 1 else node_indices.shape[-1]
         wp.launch(
-            kernel=_iota_kernel,
+            kernel=_prepare_node_sort_kernel,
             dim=index_count,
-            inputs=[sorted_array_indices, indices_per_element],
+            inputs=[node_indices.flatten(), sorted_node_indices, sorted_array_indices, indices_per_element],
         )
 
         # Sort indices
@@ -169,8 +167,16 @@ def masked_indices(
 
 
 @wp.kernel
-def _iota_kernel(indices: wp.array(dtype=int), divisor: int):
-    indices[wp.tid()] = wp.tid() // divisor
+def _prepare_node_sort_kernel(
+    node_indices: wp.array(dtype=int),
+    sort_keys: wp.array(dtype=int),
+    sort_values: wp.array(dtype=int),
+    divisor: int,
+):
+    i = wp.tid()
+    node = node_indices[i]
+    sort_keys[i] = wp.where(node >= 0, node, NULL_NODE_INDEX)
+    sort_values[i] = i // divisor
 
 
 @wp.kernel

warp/jax.py

@@ -58,6 +58,19 @@ def device_from_jax(jax_device) -> warp.context.Device:
         raise RuntimeError(f"Unsupported Jax device platform '{jax_device.platform}'")
 
 
+def get_jax_device():
+    """Get the current Jax device."""
+    import jax
+
+    # TODO: is there a simpler way of getting the Jax "current" device?
+    # check if jax.default_device() context manager is active
+    device = jax.config.jax_default_device
+    # if default device is not set, use first device
+    if device is None:
+        device = jax.local_devices()[0]
+    return device
+
+
 def dtype_to_jax(warp_dtype):
     """Return the Jax dtype corresponding to a Warp dtype.
 
@@ -156,7 +169,7 @@ def to_jax(warp_array):
     """
     import jax.dlpack
 
-    return jax.dlpack.from_dlpack(warp.to_dlpack(warp_array))
+    return jax.dlpack.from_dlpack(warp_array)
 
 
 def from_jax(jax_array, dtype=None) -> warp.array:

warp/jax_experimental/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 .custom_call import jax_kernel

warp/{jax_experimental.py → jax_experimental/custom_call.py}

@@ -15,10 +15,9 @@
 
 import ctypes
 
-import jax
-
 import warp as wp
 from warp.context import type_str
+from warp.jax import get_jax_device
 from warp.types import array_t, launch_bounds_t, strides_from_shape
 
 _jax_warp_p = None
@@ -29,35 +28,33 @@ _registered_kernels = [None]
 _registered_kernel_to_id = {}
 
 
-def jax_kernel(wp_kernel, launch_dims=None):
+def jax_kernel(kernel, launch_dims=None):
     """Create a Jax primitive from a Warp kernel.
 
     NOTE: This is an experimental feature under development.
 
     Args:
-        wp_kernel: The Warp kernel to be wrapped.
+        kernel: The Warp kernel to be wrapped.
         launch_dims: Optional. Specify the kernel launch dimensions. If None,
                      dimensions are inferred from the shape of the first argument.
                      This option when set will specify the output dimensions.
 
-    Current limitations:
-    - All kernel arguments must be arrays.
-    - If launch_dims is not provided, kernel launch dimensions are inferred from the shape of the first argument.
-    - Input arguments are followed by output arguments in the Warp kernel definition.
-    - There must be at least one input argument and at least one output argument.
-    - All arrays must be contiguous.
-    - Only the CUDA backend is supported.
+    Limitations:
+        - All kernel arguments must be contiguous arrays.
+        - Input arguments are followed by output arguments in the Warp kernel definition.
+        - There must be at least one input argument and at least one output argument.
+        - Only the CUDA backend is supported.
     """
 
     if _jax_warp_p is None:
         # Create and register the primitive
         _create_jax_warp_primitive()
-    if wp_kernel not in _registered_kernel_to_id:
+    if kernel not in _registered_kernel_to_id:
         id = len(_registered_kernels)
-        _registered_kernels.append(wp_kernel)
-        _registered_kernel_to_id[wp_kernel] = id
+        _registered_kernels.append(kernel)
+        _registered_kernel_to_id[kernel] = id
     else:
-        id = _registered_kernel_to_id[wp_kernel]
+        id = _registered_kernel_to_id[kernel]
 
     def bind(*args):
         return _jax_warp_p.bind(*args, kernel=id, launch_dims=launch_dims)
@@ -102,7 +99,7 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
         kernel_params[i + 1] = arg_ptr
 
     # Get current device.
-    device = wp.device_from_jax(_get_jax_device())
+    device = wp.device_from_jax(get_jax_device())
 
     # Get kernel hooks.
     # Note: module was loaded during jit lowering.
@@ -115,16 +112,6 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
     )
 
 
-# TODO: is there a simpler way of getting the Jax "current" device?
-def _get_jax_device():
-    # check if jax.default_device() context manager is active
-    device = jax.config.jax_default_device
-    # if default device is not set, use first device
-    if device is None:
-        device = jax.local_devices()[0]
-    return device
-
-
 def _create_jax_warp_primitive():
     from functools import reduce
 
@@ -288,7 +275,7 @@ def _create_jax_warp_primitive():
         # TODO This may not be necessary, but it is perhaps better not to be
         # mucking with kernel loading while already running the workload.
         module = wp_kernel.module
-        device = wp.device_from_jax(_get_jax_device())
+        device = wp.device_from_jax(get_jax_device())
         if not module.load(device):
             raise Exception("Could not load kernel on device")