warp-lang 1.6.1__py3-none-win_amd64.whl → 1.7.0__py3-none-win_amd64.whl

warp/fem/space/shape/square_shape_function.py

@@ -1,3 +1,18 @@
+# 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.
+
 import math
 
 import numpy as np
@@ -446,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
 
@@ -460,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:
@@ -469,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:
@@ -498,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]
@@ -522,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):
@@ -530,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):
@@ -538,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)
@@ -549,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/tet_shape_function.py

@@ -1,3 +1,18 @@
+# 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.
+
 import numpy as np
 
 import warp as wp

warp/fem/space/shape/triangle_shape_function.py

@@ -1,3 +1,18 @@
+# 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.
+
 import numpy as np
 
 import warp as wp
@@ -181,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

@@ -1,3 +1,18 @@
+# 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.
+
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Tetmesh
@@ -229,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):
@@ -242,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/topology.py

@@ -1,3 +1,18 @@
+# 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 Optional, Tuple, Type
 
 import warp as wp

warp/fem/space/trimesh_function_space.py

@@ -1,3 +1,18 @@
+# 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.
+
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Trimesh
@@ -160,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/types.py

@@ -1,3 +1,18 @@
+# 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 enum import Enum
 
 import warp as wp

warp/fem/utils.py

@@ -1,3 +1,18 @@
+# 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 Tuple, Union
 
 import numpy as np
@@ -41,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
@@ -154,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

@@ -1,9 +1,17 @@
-# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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.
 
 import warp
 
@@ -50,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.
 
@@ -148,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}

@@ -1,16 +1,23 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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.
 
 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
@@ -21,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)
@@ -94,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.
@@ -107,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
 
@@ -280,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")