warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl → 1.7.1__py3-none-manylinux_2_34_aarch64.whl

warp/jax_experimental/custom_call.py

@@ -126,7 +126,14 @@ def _create_jax_warp_primitive():
 
     # Create and register the primitive.
     # TODO add default implementation that calls the kernel via warp.
-    _jax_warp_p = jax.core.Primitive("jax_warp")
+    try:
+        # newer JAX versions
+        import jax.extend
+
+        _jax_warp_p = jax.extend.core.Primitive("jax_warp")
+    except (ImportError, AttributeError):
+        # older JAX versions
+        _jax_warp_p = jax.core.Primitive("jax_warp")
     _jax_warp_p.multiple_results = True
 
     # TODO Just launch the kernel directly, but make sure the argument
@@ -262,7 +269,12 @@ def _create_jax_warp_primitive():
     capsule = PyCapsule_New(ccall_address.value, b"xla._CUSTOM_CALL_TARGET", PyCapsule_Destructor(0))
 
     # Register the callback in XLA.
-    jax.lib.xla_client.register_custom_call_target("warp_call", capsule, platform="gpu")
+    try:
+        # newer JAX versions
+        jax.ffi.register_ffi_target("warp_call", capsule, platform="gpu", api_version=0)
+    except AttributeError:
+        # older JAX versions
+        jax.lib.xla_client.register_custom_call_target("warp_call", capsule, platform="gpu")
 
     def default_layout(shape):
         return range(len(shape) - 1, -1, -1)

warp/jax_experimental/ffi.py

@@ -560,7 +560,11 @@ class FfiCallable:
 
             # call the Python function with reconstructed arguments
             with wp.ScopedStream(stream, sync_enter=False):
-                self.func(*arg_list)
+                if stream.is_capturing:
+                    with wp.ScopedCapture(stream=stream, external=True):
+                        self.func(*arg_list)
+                else:
+                    self.func(*arg_list)
 
         except Exception as e:
             print(traceback.format_exc())

warp/native/tile.h

@@ -1133,17 +1133,17 @@ struct tile_shared_t
             constexpr int lastdim = Layout::Shape::N-1;
             constexpr bool contiguous_src = Layout::Stride::dim(lastdim) == 1;
             const bool contiguous_dest = dest.data.strides[lastdim] == sizeof(T);
-            const int elements = (dest.data.shape[lastdim] - dest.offset[lastdim]);
+            const int elements = min(Layout::Shape::dim(1), (dest.data.shape[lastdim] - dest.offset[lastdim]));
             const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
            
             float4* dest128 = (float4*)&dest.data.data[dest.index_from_coord(tile_coord(0,0))];
             const bool aligned_dst = (uint64_t)(dest128)%sizeof(float4) == 0;
 
-            if (contiguous_dest && contiguous_src && aligned_size && aligned_dst)
-            {                    
-                constexpr int M = Layout::Shape::dim(0);
-                constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);            
+            constexpr int M = Layout::Shape::dim(0);
+            constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
 
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_dst && N)
+            {                               
                 // alias of shared tile with 128bit type
                 using SrcLayout = tile_layout_strided_t<tile_shape_t<M, N>>;
                 tile_shared_t<float4, SrcLayout> src128((float4*)data.ptr);
@@ -1222,17 +1222,17 @@ struct tile_shared_t
             constexpr int lastdim = Layout::Shape::N-1;
             constexpr bool contiguous_dest = Layout::Stride::dim(lastdim) == 1;
             const bool contiguous_src = src.data.strides[lastdim] == sizeof(T);
-            const int elements = (src.data.shape[lastdim] - src.offset[lastdim]);
+            const int elements = min(Layout::Shape::dim(1), (src.data.shape[lastdim] - src.offset[lastdim]));
             const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
-
+           
             float4* src128 = (float4*)&src.data.data[src.index_from_coord(tile_coord(0,0))];
             const bool aligned_src = (uint64_t)(src128)%sizeof(float4) == 0;
 
-            if (contiguous_dest && contiguous_src && aligned_size && aligned_src)
-            {            
-                constexpr int M = Layout::Shape::dim(0);
-                constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
+            constexpr int M = Layout::Shape::dim(0);
+            constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
 
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_src && N)
+            {
                 // alias of shared tile with 128bit type
                 using DestLayout = tile_layout_strided_t<tile_shape_t<M, N>>;
                 tile_shared_t<float4, DestLayout> dest128((float4*)data.ptr);                

warp/native/warp.cu

@@ -3027,7 +3027,7 @@ size_t cuda_compile_program(const char* cuda_src, const char* program_name, int
                     fprintf(stderr, "Warp error: num_ltoirs > 0 but ltoir_input_types, ltoirs or ltoir_sizes are NULL\n");
                     return size_t(-1);
                 }
-                nvJitLinkHandle handle;
+                nvJitLinkHandle handle = nullptr;
                 std::vector<const char *> lopts = {"-dlto", arch_opt_lto};
                 if (use_ptx) {
                     lopts.push_back("-ptx");

warp/render/render_opengl.py

@@ -13,11 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import ctypes
 import sys
 import time
 from collections import defaultdict
-from typing import List, Optional, Tuple, Union
+from typing import List, Union
 
 import numpy as np
 
@@ -1500,16 +1502,16 @@ class OpenGLRenderer:
 
     def setup_tiled_rendering(
         self,
-        instances: List[List[int]],
+        instances: list[list[int]],
         rescale_window: bool = False,
-        tile_width: Optional[int] = None,
-        tile_height: Optional[int] = None,
-        tile_ncols: Optional[int] = None,
-        tile_nrows: Optional[int] = None,
-        tile_positions: Optional[List[Tuple[int]]] = None,
-        tile_sizes: Optional[List[Tuple[int]]] = None,
-        projection_matrices: Optional[List[Mat44]] = None,
-        view_matrices: Optional[List[Mat44]] = None,
+        tile_width: int | None = None,
+        tile_height: int | None = None,
+        tile_ncols: int | None = None,
+        tile_nrows: int | None = None,
+        tile_positions: list[tuple[int]] | None = None,
+        tile_sizes: list[tuple[int]] | None = None,
+        projection_matrices: list[Mat44] | None = None,
+        view_matrices: list[Mat44] | None = None,
     ):
         """
         Set up tiled rendering where the render buffer is split into multiple tiles that can visualize
@@ -1602,11 +1604,11 @@ class OpenGLRenderer:
     def update_tile(
         self,
         tile_id,
-        instances: Optional[List[int]] = None,
-        projection_matrix: Optional[Mat44] = None,
-        view_matrix: Optional[Mat44] = None,
-        tile_size: Optional[Tuple[int]] = None,
-        tile_position: Optional[Tuple[int]] = None,
+        instances: list[int] | None = None,
+        projection_matrix: Mat44 | None = None,
+        view_matrix: Mat44 | None = None,
+        tile_size: tuple[int] | None = None,
+        tile_position: tuple[int] | None = None,
     ):
         """
         Update the shape instances, projection matrix, view matrix, tile size, or tile position
@@ -1806,7 +1808,7 @@ class OpenGLRenderer:
 
         return np.array((scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 1), dtype=np.float32)
 
-    def update_model_matrix(self, model_matrix: Optional[Mat44] = None):
+    def update_model_matrix(self, model_matrix: Mat44 | None = None):
         gl = OpenGLRenderer.gl
 
         self._switch_context()
@@ -3092,7 +3094,7 @@ Instances: {len(self._instances)}"""
         parent_body: str = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: Tuple[float, float, float] = None,
+        color: tuple[float, float, float] = None,
     ):
         """Add a arrow for visualization
 

warp/render/render_usd.py

@@ -13,10 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import numpy as np
 
 import warp as wp
 
+UP_AXIS_TOKEN = ("X", "Y", "Z")
+UP_AXIS_VEC = (
+    np.array((1.0, 0.0, 0.0), dtype=float),
+    np.array((0.0, 1.0, 0.0), dtype=float),
+    np.array((0.0, 0.0, 1.0), dtype=float),
+)
+
 
 def _usd_add_xform(prim):
     from pxr import UsdGeom
@@ -29,7 +38,13 @@ def _usd_add_xform(prim):
     prim.AddScaleOp()
 
 
-def _usd_set_xform(xform, pos: tuple, rot: tuple, scale: tuple, time):
+def _usd_set_xform(
+    xform,
+    pos: tuple | None = None,
+    rot: tuple | None = None,
+    scale: tuple | None = None,
+    time: float = 0.0,
+):
     from pxr import Gf, UsdGeom
 
     xform = UsdGeom.Xform(xform)
@@ -108,7 +123,7 @@ class UsdRenderer:
         self.stage.SetDefaultPrim(self.root.GetPrim())
         self.stage.SetStartTimeCode(0.0)
         self.stage.SetEndTimeCode(0.0)
-        self.stage.SetTimeCodesPerSecond(self.fps)
+        self.stage.SetFramesPerSecond(self.fps)
 
         if up_axis == "X":
             UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.x)
@@ -622,7 +637,82 @@ class UsdRenderer:
 
         return prim_path
 
-    def render_line_list(self, name, vertices, indices, color, radius):
+    def render_arrow(
+        self,
+        name: str,
+        pos: tuple,
+        rot: tuple,
+        base_radius: float,
+        base_height: float,
+        cap_radius: float = None,
+        cap_height: float = None,
+        parent_body: str = None,
+        is_template: bool = False,
+        up_axis: int = 1,
+        color: tuple[float, float, float] = None,
+        visible: bool = True,
+    ):
+        from pxr import Gf, Sdf, UsdGeom
+
+        if is_template:
+            prim_path = self._resolve_path(name, parent_body, is_template)
+            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
+            blueprint_prim = blueprint.GetPrim()
+            blueprint_prim.SetInstanceable(True)
+            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
+            arrow_path = prim_path.AppendChild("arrow")
+        else:
+            arrow_path = self._resolve_path(name, parent_body)
+            prim_path = arrow_path
+
+        arrow = UsdGeom.Xform.Get(self.stage, arrow_path)
+        if not arrow:
+            arrow = UsdGeom.Xform.Define(self.stage, arrow_path)
+            _usd_add_xform(arrow)
+
+        base_path = arrow_path.AppendChild("base")
+        base = UsdGeom.Xform.Get(self.stage, base_path)
+        if not base:
+            base = UsdGeom.Cylinder.Define(self.stage, base_path)
+            _usd_add_xform(base)
+
+        base.GetRadiusAttr().Set(float(base_radius))
+        base.GetHeightAttr().Set(float(base_height))
+        base.GetAxisAttr().Set(UP_AXIS_TOKEN[up_axis])
+        _usd_set_xform(base, UP_AXIS_VEC[up_axis] * base_height * 0.5)
+
+        cap_path = arrow_path.AppendChild("cap")
+        cap = UsdGeom.Xform.Get(self.stage, cap_path)
+        if not cap:
+            cap = UsdGeom.Cone.Define(self.stage, arrow_path.AppendChild("cap"))
+            _usd_add_xform(cap)
+
+        cap.GetRadiusAttr().Set(float(cap_radius))
+        cap.GetHeightAttr().Set(float(cap_height))
+        cap.GetAxisAttr().Set(UP_AXIS_TOKEN[up_axis])
+        _usd_set_xform(cap, UP_AXIS_VEC[up_axis] * (base_height + cap_height * 0.5))
+
+        if color is not None:
+            base.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+            cap.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)
+
+        self._shape_constructors[name] = UsdGeom.Xform
+
+        if not is_template:
+            _usd_set_xform(arrow, pos, rot, (1.0, 1.0, 1.0), self.time)
+
+        arrow.GetVisibilityAttr().Set("inherited" if visible else "invisible", self.time)
+        return prim_path
+
+    def render_line_list(
+        self,
+        name: str,
+        vertices,
+        indices,
+        color: tuple = None,
+        radius: float = 0.01,
+        visible: bool = True,
+    ):
         """Debug helper to add a line list as a set of capsules
 
         Args:

warp/sim/collide.py

@@ -614,9 +614,9 @@ def volume_grad(volume: wp.uint64, p: wp.vec3):
 @wp.func
 def counter_increment(counter: wp.array(dtype=int), counter_index: int, tids: wp.array(dtype=int), tid: int):
     # increment counter, remember which thread received which counter value
-    next_count = wp.atomic_add(counter, counter_index, 1)
-    tids[tid] = next_count
-    return next_count
+    count = wp.atomic_add(counter, counter_index, 1)
+    tids[tid] = count
+    return count
 
 
 @wp.func_replay(counter_increment)
@@ -629,10 +629,10 @@ def limited_counter_increment(
     counter: wp.array(dtype=int), counter_index: int, tids: wp.array(dtype=int), tid: int, index_limit: int
 ):
     # increment counter but only if it is smaller than index_limit, remember which thread received which counter value
-    next_count = wp.atomic_add(counter, counter_index, 1)
-    if next_count < index_limit or index_limit < 0:
-        tids[tid] = next_count
-        return next_count
+    count = wp.atomic_add(counter, counter_index, 1)
+    if count < index_limit or index_limit < 0:
+        tids[tid] = count
+        return count
     tids[tid] = -1
     return -1
 
@@ -1547,6 +1547,8 @@ def handle_contact_pairs(
                 # reached contact point limit
                 return
         index = counter_increment(contact_count, 0, contact_tids, tid)
+        if index == -1:
+            return
         contact_shape0[index] = shape_a
         contact_shape1[index] = shape_b
         # transform from world into body frame (so the contact point includes the shape transform)
@@ -1690,7 +1692,7 @@ def collide(
                 model.rigid_contact_normal = wp.empty_like(model.rigid_contact_normal)
                 model.rigid_contact_thickness = wp.empty_like(model.rigid_contact_thickness)
                 model.rigid_contact_count = wp.zeros_like(model.rigid_contact_count)
-                model.rigid_contact_tids = wp.zeros_like(model.rigid_contact_tids)
+                model.rigid_contact_tids = wp.full_like(model.rigid_contact_tids, -1)
                 model.rigid_contact_shape0 = wp.empty_like(model.rigid_contact_shape0)
                 model.rigid_contact_shape1 = wp.empty_like(model.rigid_contact_shape1)
 
@@ -1698,7 +1700,7 @@ def collide(
                     model.rigid_contact_pairwise_counter = wp.zeros_like(model.rigid_contact_pairwise_counter)
             else:
                 model.rigid_contact_count.zero_()
-                model.rigid_contact_tids.zero_()
+                model.rigid_contact_tids.fill_(-1)
 
                 if model.rigid_contact_pairwise_counter is not None:
                     model.rigid_contact_pairwise_counter.zero_()