Merge branch 'ccrouzet/fix-iter-reverse' into 'main'

Fix `iter_reverse()` For Ranges With Steps Other than 1

Closes GH-311

See merge request omniverse/warp!770

CHANGELOG.md

@@ -14,6 +14,14 @@
 - Fix robustness to very low desired tolerance in `wp.fem.utils.symmetric_eigenvalues_qr`
 - Fix invalid code generation error messages when nesting dynamic and static for-loops
 - Fix caching of kernels with static expressions
+- Expose a `reversed()` built-in for iterators.
+
+### Changed
+
+### Fixed
+
+- Fix `iter_reverse()` not working as expected for ranges with steps other than 1 ([GH-311](https://github.com/NVIDIA/warp/issues/311)).
+
 
 ## [1.4.0] - 2024-10-01
 

docs/modules/functions.rst

@@ -818,6 +818,11 @@ Utility
            All matrices are assumed to be stored in flattened row-major memory layout (NumPy default).
 
 
+.. py:function:: reversed(range: range_t) -> range_t
+
+    Returns the range in reversed order.
+
+
 .. py:function:: printf(fmt: str, *args: Any) -> None
 
     Allows printing formatted strings using C-style format specifiers.

warp/builtins.py

@@ -2387,6 +2387,16 @@ def spatial_vector_dispatch_func(input_types: Mapping[str, type], return_type: A
     "iter_next", input_types={"query": mesh_query_aabb_t}, value_type=int, group="Utility", export=False, hidden=True
 )
 
+add_builtin(
+    "reversed",
+    input_types={"range": range_t},
+    value_type=range_t,
+    native_func="iter_reverse",
+    group="Utility",
+    doc="""Returns the range in reversed order.""",
+    export=False,
+)
+
 # ---------------------------------
 # Volumes
 

warp/native/bvh.h

@@ -404,6 +404,10 @@ CUDA_CALLABLE inline bvh_query_t iter_reverse(const bvh_query_t& query)
     return query;
 }
 
+CUDA_CALLABLE inline void adj_iter_reverse(const bvh_query_t& query, bvh_query_t& adj_query, bvh_query_t& adj_ret)
+{
+}
+
 
 // stub
 CUDA_CALLABLE inline void adj_bvh_query_next(bvh_query_t& query, int& index, bvh_query_t&, int&, bool&) 

warp/native/hashgrid.h

@@ -209,6 +209,10 @@ CUDA_CALLABLE inline hash_grid_query_t iter_reverse(const hash_grid_query_t& que
     return query;
 }
 
+CUDA_CALLABLE inline void adj_iter_reverse(const hash_grid_query_t& query, hash_grid_query_t& adj_query, hash_grid_query_t& adj_ret)
+{
+}
+
 
 
 CUDA_CALLABLE inline int hash_grid_point_id(uint64_t id, int& index)

warp/native/mesh.h

@@ -1693,6 +1693,10 @@ CUDA_CALLABLE inline mesh_query_aabb_t iter_reverse(const mesh_query_aabb_t& que
     return query;
 }
 
+CUDA_CALLABLE inline void adj_iter_reverse(const mesh_query_aabb_t& query, mesh_query_aabb_t& adj_query, mesh_query_aabb_t& adj_ret)
+{
+}
+
 
 // stub
 CUDA_CALLABLE inline void adj_mesh_query_aabb_next(mesh_query_aabb_t& query, int& index, mesh_query_aabb_t&, int&, bool&) 

warp/native/range.h

@@ -97,13 +97,26 @@ CUDA_CALLABLE inline range_t iter_reverse(const range_t& r)
 {
     // generates a reverse range, equivalent to reversed(range())
     range_t rev;
-    rev.start = r.end-1;
-    rev.end = r.start-1;
+
+    if (r.step > 0)
+    {
+        rev.start = r.start + int((r.end - r.start - 1) / r.step) * r.step;
+    }
+    else
+    {
+        rev.start = r.start + int((r.end - r.start + 1) / r.step) * r.step;
+    }
+
+    rev.end = r.start - r.step;
     rev.step = -r.step;
 
     rev.i = rev.start;
 
     return rev;
 }
 
+CUDA_CALLABLE inline void adj_iter_reverse(const range_t& r, range_t& adj_r, range_t& adj_ret)
+{
+}
+
 } // namespace wp

warp/stubs.py

@@ -1160,6 +1160,12 @@ def closest_point_edge_edge(p1: vec3f, q1: vec3f, p2: vec3f, q2: vec3f, epsilon:
     ...
 
 
+@over
+def reversed(range: range_t) -> range_t:
+    """Returns the range in reversed order."""
+    ...
+
+
 @over
 def volume_sample(id: uint64, uvw: vec3f, sampling_mode: int32, dtype: Any) -> Any:
     """Sample the volume of type `dtype` given by ``id`` at the volume local-space point ``uvw``.

warp/tests/test_iter.py

@@ -0,0 +1,68 @@
+# 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.
+
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def reversed_kernel(
+    start: wp.int32,
+    end: wp.int32,
+    step: wp.int32,
+    out_count: wp.array(dtype=wp.int32),
+    out_values: wp.array(dtype=wp.int32),
+):
+    count = wp.int32(0)
+    for i in reversed(range(start, end, step)):
+        out_values[count] = i
+        count += 1
+
+    out_count[0] = count
+
+
+def test_reversed(test, device):
+    count = wp.empty(1, dtype=wp.int32)
+    values = wp.empty(32, dtype=wp.int32)
+
+    start, end, step = (-2, 8, 3)
+    wp.launch(
+        reversed_kernel,
+        dim=1,
+        inputs=(start, end, step),
+        outputs=(count, values),
+    )
+    expected = tuple(reversed(range(start, end, step)))
+    assert count.numpy()[0] == len(expected)
+    assert_np_equal(values.numpy()[: len(expected)], expected)
+
+    start, end, step = (9, -3, -2)
+    wp.launch(
+        reversed_kernel,
+        dim=1,
+        inputs=(start, end, step),
+        outputs=(count, values),
+    )
+    expected = tuple(reversed(range(start, end, step)))
+    assert count.numpy()[0] == len(expected)
+    assert_np_equal(values.numpy()[: len(expected)], expected)
+
+
+devices = get_test_devices()
+
+
+class TestIter(unittest.TestCase):
+    pass
+
+
+add_function_test(TestIter, "test_reversed", test_reversed, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)