warp-lang 1.7.0__py3-none-macosx_10_13_universal2.whl → 1.7.2__py3-none-macosx_10_13_universal2.whl

warp/context.py

@@ -457,6 +457,24 @@ class Function:
         return f"<Function {self.key}({inputs_str})>"
 
 
+def get_builtin_type(return_type: type) -> type:
+    # The return_type might just be vector_t(length=3,dtype=wp.float32), so we've got to match that
+    # in the list of hard coded types so it knows it's returning one of them:
+    if hasattr(return_type, "_wp_generic_type_hint_"):
+        return_type_match = tuple(
+            x
+            for x in generic_vtypes
+            if x._wp_generic_type_hint_ == return_type._wp_generic_type_hint_
+            and x._wp_type_params_ == return_type._wp_type_params_
+        )
+        if not return_type_match:
+            raise RuntimeError("No match")
+
+        return return_type_match[0]
+
+    return return_type
+
+
 def call_builtin(func: Function, *params: Any) -> Tuple[bool, Any]:
     uses_non_warp_array_type = False
 
@@ -769,6 +787,12 @@ class Kernel:
 
         return f"{self.key}_{hash_suffix}"
 
+    def __call__(self, *args, **kwargs):
+        # we implement this function only to ensure Kernel is a callable object
+        # so that we can document Warp kernels in the same way as Python functions
+        # annotated by @wp.kernel (see functools.update_wrapper())
+        raise NotImplementedError("Kernel.__call__() is not implemented, please use wp.launch() instead")
+
 
 # ----------------------
 
@@ -1068,7 +1092,7 @@ def kernel(
 # decorator to register struct, @struct
 def struct(c: type):
     m = get_module(c.__module__)
-    s = warp.codegen.Struct(cls=c, key=warp.codegen.make_full_qualified_name(c), module=m)
+    s = warp.codegen.Struct(key=warp.codegen.make_full_qualified_name(c), cls=c, module=m)
     s = functools.update_wrapper(s, c)
     return s
 
@@ -1439,6 +1463,24 @@ def register_api_function(
     """
     function.group = group
     function.hidden = hidden
+
+    # Update the docstring to mark these functions as being available from kernels and Python's runtime.
+    assert function.__doc__.startswith("\n")
+    leading_space_count = sum(1 for _ in itertools.takewhile(str.isspace, function.__doc__[1:]))
+    assert leading_space_count % 4 == 0
+    indent_level = leading_space_count // 4
+    indent = "    "
+    function.__doc__ = (
+        f"\n"
+        f"{indent * indent_level}.. hlist::\n"
+        f"{indent * (indent_level + 1)}:columns: 8\n"
+        f"\n"
+        f"{indent * (indent_level + 1)}* Kernel\n"
+        f"{indent * (indent_level + 1)}* Python\n"
+        f"{indent * (indent_level + 1)}* Differentiable\n"
+        f"{function.__doc__}"
+    )
+
     builtin_functions[function.key] = function
 
 
@@ -6504,12 +6546,46 @@ def type_str(t):
     return t.__name__
 
 
-def print_function(f, file, noentry=False):  # pragma: no cover
+def ctype_ret_str(t):
+    return get_builtin_type(t).__name__
+
+
+def resolve_exported_function_sig(f):
+    if not f.export or f.generic:
+        return None
+
+    # only export simple types that don't use arrays or templated types
+    if not f.is_simple():
+        return None
+
+    # Runtime arguments that are to be passed to the function, not its template signature.
+    if f.export_func is not None:
+        func_args = f.export_func(f.input_types)
+    else:
+        func_args = f.input_types
+
+    # todo: construct a default value for each of the functions args
+    # so we can generate the return type for overloaded functions
+    return_type = f.value_func(func_args, None)
+
+    try:
+        return_type_str = ctype_ret_str(return_type)
+    except Exception:
+        return None
+
+    if return_type_str.startswith("Tuple"):
+        return None
+
+    return (func_args, return_type)
+
+
+def print_function(f, file, is_exported, noentry=False):  # pragma: no cover
     """Writes a function definition to a file for use in reST documentation
 
     Args:
         f: The function being written
         file: The file object for output
+        is_exported: Whether the function is available in Python's runtime
         noentry: If True, then the :noindex: and :nocontentsentry: directive
           options will be added
 
@@ -6537,11 +6613,21 @@ def print_function(f, file, noentry=False):  # pragma: no cover
         print("    :nocontentsentry:", file=file)
     print("", file=file)
 
+    print("    .. hlist::", file=file)
+    print("       :columns: 8", file=file)
+    print("", file=file)
+    print("       * Kernel", file=file)
+
+    if is_exported:
+        print("       * Python", file=file)
+
+    if not f.missing_grad:
+        print("       * Differentiable", file=file)
+
+    print("", file=file)
+
     if f.doc != "":
-        if not f.missing_grad:
-            print(f"    {f.doc}", file=file)
-        else:
-            print(f"    {f.doc} [1]_", file=file)
+        print(f"    {f.doc}", file=file)
         print("", file=file)
 
     print(file=file)
@@ -6557,8 +6643,10 @@ def export_functions_rst(file):  # pragma: no cover
         ".. functions:\n"
         ".. currentmodule:: warp\n"
         "\n"
-        "Kernel Reference\n"
-        "================"
+        "Built-Ins Reference\n"
+        "===================\n"
+        "This section lists the Warp types and functions available to use from Warp kernels and optionally also from the Warp Python runtime API.\n"
+        "For a listing of the API that is exclusively intended to be used at the *Python Scope* and run inside the CPython interpreter, see the :doc:`runtime` section.\n"
     )
 
     print(header, file=file)
@@ -6603,9 +6691,12 @@ def export_functions_rst(file):  # pragma: no cover
         if hasattr(f, "overloads"):
             # append all overloads to the group
             for o in f.overloads:
-                groups[f.group].append(o)
+                sig = resolve_exported_function_sig(f)
+                is_exported = sig is not None
+                groups[f.group].append((o, is_exported))
         else:
-            groups[f.group].append(f)
+            is_exported = False
+            groups[f.group].append((f, is_exported))
 
     # Keep track of what function and query types have been written
     written_functions = set()
@@ -6624,7 +6715,7 @@ def export_functions_rst(file):  # pragma: no cover
         print(k, file=file)
         print("---------------", file=file)
 
-        for f in g:
+        for f, is_exported in g:
             if f.func:
                 # f is a Warp function written in Python, we can use autofunction
                 print(f".. autofunction:: {f.func.__module__}.{f.key}", file=file)
@@ -6637,15 +6728,11 @@ def export_functions_rst(file):  # pragma: no cover
 
             if f.key in written_functions:
                 # Add :noindex: + :nocontentsentry: since Sphinx gets confused
-                print_function(f, file=file, noentry=True)
+                print_function(f, file, is_exported, noentry=True)
             else:
-                if print_function(f, file=file):
+                if print_function(f, file, is_exported):
                     written_functions.add(f.key)
 
-    # footnotes
-    print(".. rubric:: Footnotes", file=file)
-    print(".. [1] Function gradients have not been implemented for backpropagation.", file=file)
-
 
 def export_stubs(file):  # pragma: no cover
     """Generates stub file for auto-complete of builtin functions"""
@@ -6745,14 +6832,6 @@ def export_builtins(file: io.TextIOBase):  # pragma: no cover
         else:
             return t.__name__
 
-    def ctype_ret_str(t):
-        if isinstance(t, int):
-            return "int"
-        elif isinstance(t, float):
-            return "float"
-        else:
-            return t.__name__
-
     file.write("namespace wp {\n\n")
     file.write('extern "C" {\n\n')
 
@@ -6760,40 +6839,24 @@ def export_builtins(file: io.TextIOBase):  # pragma: no cover
         if not hasattr(g, "overloads"):
             continue
         for f in g.overloads:
-            if not f.export or f.generic:
+            sig = resolve_exported_function_sig(f)
+            if sig is None:
                 continue
 
-            # only export simple types that don't use arrays
-            # or templated types
-            if not f.is_simple():
-                continue
-
-            try:
-                # todo: construct a default value for each of the functions args
-                # so we can generate the return type for overloaded functions
-                return_type = ctype_ret_str(f.value_func(None, None))
-            except Exception:
-                continue
-
-            if return_type.startswith("Tuple"):
-                continue
-
-            # Runtime arguments that are to be passed to the function, not its template signature.
-            if f.export_func is not None:
-                func_args = f.export_func(f.input_types)
-            else:
-                func_args = f.input_types
+            func_args, return_type = sig
 
             args = ", ".join(f"{ctype_arg_str(v)} {k}" for k, v in func_args.items())
             params = ", ".join(func_args.keys())
 
+            return_str = ctype_ret_str(return_type)
+
             if args == "":
-                file.write(f"WP_API void {f.mangled_name}({return_type}* ret) {{ *ret = wp::{f.key}({params}); }}\n")
-            elif return_type == "None":
+                file.write(f"WP_API void {f.mangled_name}({return_str}* ret) {{ *ret = wp::{f.key}({params}); }}\n")
+            elif return_type is None:
                 file.write(f"WP_API void {f.mangled_name}({args}) {{ wp::{f.key}({params}); }}\n")
             else:
                 file.write(
-                    f"WP_API void {f.mangled_name}({args}, {return_type}* ret) {{ *ret = wp::{f.key}({params}); }}\n"
+                    f"WP_API void {f.mangled_name}({args}, {return_str}* ret) {{ *ret = wp::{f.key}({params}); }}\n"
                 )
 
     file.write('\n}  // extern "C"\n\n')

warp/examples/benchmarks/benchmark_cloth.py

@@ -160,7 +160,7 @@ def run_benchmark(mode, dim, timers, render=False):
         stage = Usd.Stage.CreateNew("benchmark.usd")
         stage.SetStartTimeCode(0.0)
         stage.SetEndTimeCode(sim_duration * sim_fps)
-        stage.SetTimeCodesPerSecond(sim_fps)
+        stage.SetFramesPerSecond(sim_fps)
 
         grid = UsdGeom.Mesh.Define(stage, "/root")
         grid.GetPointsAttr().Set(cloth.positions, 0.0)