Merge branch 'GH-403' into 'main'

Negative constants evaluate to compile-time constants.

Closes GH-403

See merge request omniverse/warp!954

CHANGELOG.md

@@ -19,6 +19,7 @@
   ([GH-409](https://github.com/NVIDIA/warp/issues/409)).
 - Fix so that `wp.Tape.zero()` zeroes gradients passed via the 'grads' parameter in `wp.Tape.backward()` ([GH-407](https://github.com/NVIDIA/warp/issues/407)).
 - Fix the OpenGL renderer not working when multiple instances exist at the same time ([GH-385](https://github.com/NVIDIA/warp/issues/385)).
+- Negative constants evaluate to compile-time constants (fixes [GH-403](https://github.com/NVIDIA/warp/issues/403))
 
 ## [1.5.1] - 2025-01-02
 

warp/builtins.py

@@ -1836,6 +1836,9 @@ def tile_arange_value_func(arg_types: Mapping[str, type], arg_values: Mapping[st
     step = 1
     dtype = int
 
+    if "args" not in arg_values:
+        raise RuntimeError("wp.tile_arange() requires one or more positional arguments describing range")
+
     args = arg_values["args"]
 
     if len(args) == 1:
@@ -1876,11 +1879,30 @@ def tile_arange_dispatch_func(arg_types: Mapping[str, type], return_type: Any, a
     template_args.append(m)
     template_args.append(n)
 
-    # todo: it is somewhat redundant to create new vars here since some of start,stop,step
-    # already exist depending on which form the function was called by the user
-    start = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.start)
-    stop = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.stop)
-    step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+    if "args" not in arg_values:
+        raise RuntimeError("wp.tile_arange() requires one or more positional arguments describing range")
+
+    args = arg_values["args"]
+
+    if len(args) == 1:
+        start = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.start)
+        stop = args[0]
+        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+
+    elif len(args) == 2:
+        start = args[0]
+        stop = args[1]
+        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+
+    elif len(args) == 3:
+        start = args[0]
+        stop = args[1]
+        step = args[2]
+
+    else:
+        raise TypeError(
+            "Too many positional arguments passed to wp.tile_arange(), which accepts a maximum of 3 scalars."
+        )
 
     function_args = []
     function_args.append(start)

warp/codegen.py

@@ -1920,6 +1920,11 @@ def emit_UnaryOp(adj, node):
         # evaluate unary op arguments
         arg = adj.eval(node.operand)
 
+        # evaluate expression to a compile-time constant if arg is a constant
+        if arg.constant is not None and math.isfinite(arg.constant):
+            if isinstance(node.op, ast.USub):
+                return adj.add_constant(-arg.constant)
+
         name = builtin_operators[type(node.op)]
 
         return adj.add_builtin_call(name, [arg])

warp/tests/test_tile_reduce.py

@@ -332,23 +332,29 @@ def tile_arange_kernel(out: wp.array2d(dtype=int)):
     a = wp.tile_arange(17, dtype=int)
     b = wp.tile_arange(5, 23, dtype=int)
     c = wp.tile_arange(0, 34, 2, dtype=int)
+    d = wp.tile_arange(-1, 16, dtype=int)
+    e = wp.tile_arange(17, 0, -1, dtype=int)
 
     wp.tile_store(out, 0, 0, a)
     wp.tile_store(out, 1, 0, b)
     wp.tile_store(out, 2, 0, c)
+    wp.tile_store(out, 3, 0, d)
+    wp.tile_store(out, 4, 0, e)
 
 
 def test_tile_arange(test, device):
     N = 17
 
-    output = wp.zeros(shape=(3, N), dtype=int, device=device)
+    output = wp.zeros(shape=(5, N), dtype=int, device=device)
 
     with wp.Tape() as tape:
         wp.launch_tiled(tile_arange_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
 
     assert_np_equal(output.numpy()[0], np.arange(17))
     assert_np_equal(output.numpy()[1], np.arange(5, 22))
     assert_np_equal(output.numpy()[2], np.arange(0, 34, 2))
+    assert_np_equal(output.numpy()[3], np.arange(-1, 16))
+    assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
 devices = get_cuda_test_devices()