Return a constant value from len() where possible

Author: christophercrouzet

warp/builtins.py

@@ -76,6 +76,14 @@ def extract_tuple(arg, as_constant=False):
     return out
 
 
+def static_len_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    if arg_types is None:
+        return int
+
+    length = warp.types.type_length(arg_types["a"])
+    return Var(None, type=int, constant=length)
+
+
 # ---------------------------------
 # Scalar Math
 
@@ -7720,7 +7728,7 @@ def static(expr):
 add_builtin(
     "len",
     input_types={"a": vector(length=Any, dtype=Scalar)},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of elements in a vector.",
     group="Utility",
     export=False,
@@ -7729,7 +7737,7 @@ def static(expr):
 add_builtin(
     "len",
     input_types={"a": quaternion(dtype=Scalar)},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of elements in a quaternion.",
     group="Utility",
     export=False,
@@ -7738,7 +7746,7 @@ def static(expr):
 add_builtin(
     "len",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar)},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of rows in a matrix.",
     group="Utility",
     export=False,
@@ -7747,7 +7755,7 @@ def static(expr):
 add_builtin(
     "len",
     input_types={"a": transformation(dtype=Float)},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of elements in a transformation.",
     group="Utility",
     export=False,
@@ -7765,7 +7773,7 @@ def static(expr):
 add_builtin(
     "len",
     input_types={"a": tile(dtype=Any, shape=Tuple[int, ...])},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of rows in a tile.",
     group="Utility",
     export=False,
@@ -7835,10 +7843,11 @@ def tuple_extract_dispatch_func(input_types: Mapping[str, type], return_type: An
     missing_grad=True,
 )
 
+
 add_builtin(
     "len",
     input_types={"a": Tuple},
-    value_type=int,
+    value_func=static_len_value_func,
     doc="Return the number of elements in a tuple.",
     group="Utility",
     export=False,

warp/codegen.py

@@ -1423,6 +1423,13 @@ def add_call(adj, func, args, kwargs, type_args, min_outputs=None):
             bound_arg_values,
         )
 
+        # Handle the special case where a Var instance is returned from the `value_func`
+        # callback, in which case we replace the call with a reference to that variable.
+        if isinstance(return_type, Var):
+            return adj.register_var(return_type)
+        elif isinstance(return_type, Sequence) and all(isinstance(x, Var) for x in return_type):
+            return tuple(adj.register_var(x) for x in return_type)
+
         if get_origin(return_type) is tuple:
             types = get_args(return_type)
             return_type = warp.types.tuple_t(types=types, values=(None,) * len(types))
@@ -2334,7 +2341,6 @@ def emit_Call(adj, node):
         args = tuple(adj.resolve_arg(x) for x in node.args)
         kwargs = {x.arg: adj.resolve_arg(x.value) for x in node.keywords}
 
-        # add the call and build the callee adjoint if needed (func.adj)
         out = adj.add_call(func, args, kwargs, type_args, min_outputs=min_outputs)
 
         if warp.config.verify_autograd_array_access:
@@ -3049,29 +3055,10 @@ def evaluate_static_expression(adj, node) -> tuple[Any, str]:
 
         # Replace all constant `len()` expressions with their value.
         if "len" in static_code:
-
-            def eval_len(obj):
-                if type_is_vector(obj):
-                    return obj._length_
-                elif type_is_quaternion(obj):
-                    return obj._length_
-                elif type_is_matrix(obj):
-                    return obj._shape_[0]
-                elif type_is_transformation(obj):
-                    return obj._length_
-                elif is_tuple(obj):
-                    return len(obj.types)
-                elif is_tile(obj):
-                    return obj.shape[0]
-                elif get_origin(obj) is tuple:
-                    return len(get_args(obj))
-
-                return len(obj)
-
             len_expr_ctx = vars_dict.copy()
             constant_types = {k: v.type for k, v in adj.symbols.items() if isinstance(v, Var) and v.type is not None}
             len_expr_ctx.update(constant_types)
-            len_expr_ctx.update({"len": eval_len})
+            len_expr_ctx.update({"len": warp.types.type_length})
 
             # We want to replace the expression code in-place,
             # so reparse it to get the correct column info.

warp/tests/test_tuple.py

@@ -136,14 +136,17 @@ def loop_user_func(values: Tuple[int, int, int]):
     for i in range(wp.static(len(values))):
         out += values[i]
 
+    for i in range(len(values)):
+        out += values[i] * 2
+
     return out
 
 
 @wp.kernel
 def test_loop():
     t = (1, 2, 3)
     res = loop_user_func(t)
-    wp.expect_eq(res, 6)
+    wp.expect_eq(res, 18)
 
 
 @wp.func
@@ -152,26 +155,29 @@ def loop_variadic_any_user_func(values: Any):
     for i in range(wp.static(len(values))):
         out += values[i]
 
+    for i in range(len(values)):
+        out += values[i] * 2
+
     return out
 
 
 @wp.kernel
 def test_loop_variadic_any():
     t1 = (1,)
     res = loop_variadic_any_user_func(t1)
-    wp.expect_eq(res, 1)
+    wp.expect_eq(res, 3)
 
     t2 = (2, 3)
     res = loop_variadic_any_user_func(t2)
-    wp.expect_eq(res, 5)
+    wp.expect_eq(res, 15)
 
     t3 = (3, 4, 5)
     res = loop_variadic_any_user_func(t3)
-    wp.expect_eq(res, 12)
+    wp.expect_eq(res, 36)
 
     t4 = (4, 5, 6, 7)
     res = loop_variadic_any_user_func(t4)
-    wp.expect_eq(res, 22)
+    wp.expect_eq(res, 66)
 
 
 @wp.func

warp/types.py

@@ -1388,6 +1388,21 @@ def type_ctype(dtype):
         return dtype._type_
 
 
+def type_length(obj):
+    if is_tile(obj):
+        return obj.shape[0]
+    elif is_tuple(obj):
+        return len(obj.types)
+    elif get_origin(obj) is tuple:
+        return len(get_args(obj))
+    elif hasattr(obj, "_shape_"):
+        return obj._shape_[0]
+    elif hasattr(obj, "_length_"):
+        return obj._length_
+
+    return len(obj)
+
+
 def type_size(dtype):
     if dtype == float or dtype == int or isinstance(dtype, warp.codegen.Struct):
         return 1