warp-lang 0.15.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/__init__.py

@@ -73,6 +73,7 @@ from warp.context import is_peer_access_supported, is_peer_access_enabled, set_p
 from warp.tape import Tape
 from warp.utils import ScopedTimer, ScopedDevice, ScopedStream
 from warp.utils import ScopedMempool, ScopedMempoolAccess, ScopedPeerAccess
+from warp.utils import ScopedCapture
 from warp.utils import transform_expand, quat_between_vectors
 
 from warp.torch import from_torch, to_torch

warp/codegen.py

@@ -919,12 +919,16 @@ class Adjoint:
                     break
 
         # if it is a user-function then build it recursively
-        if not func.is_builtin():
+        if not func.is_builtin() and func not in adj.builder.functions:
             adj.builder.build_function(func)
+            # add custom grad, replay functions to the list of functions
+            # to be built later (invalid code could be generated if we built them now)
+            # so that they are not missed when only the forward function is imported
+            # from another module
             if func.custom_grad_func:
-                adj.builder.build_function(func.custom_grad_func)
+                adj.builder.deferred_functions.append(func.custom_grad_func)
             if func.custom_replay_func:
-                adj.builder.build_function(func.custom_replay_func)
+                adj.builder.deferred_functions.append(func.custom_replay_func)
 
         # evaluate the function type based on inputs
         arg_types = [strip_reference(a.type) for a in args if not isinstance(a, warp.context.Function)]

warp/config.py

@@ -5,7 +5,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-version = "0.15.0"
+version = "1.0.0"
 
 verify_fp = False  # verify inputs and outputs are finite after each launch
 verify_cuda = False  # if true will check CUDA errors after each kernel launch / memory operation
@@ -13,6 +13,7 @@ print_launches = False  # if true will print out launch information
 
 mode = "release"
 verbose = False  # print extra informative messages
+verbose_warnings = False  # whether file and line info gets included in Warp warnings
 quiet = False  # suppress all output except errors and warnings
 
 cache_kernels = True

warp/constants.py

@@ -26,6 +26,8 @@ __all__ = [
     "phi",
     "PI",
     "pi",
+    "HALF_PI",
+    "half_pi",
     "TAU",
     "tau",
 ]
@@ -37,6 +39,7 @@ LN2 = ln2 = constant(0.69314718055994530942)  # ln(2)
 LN10 = ln10 = constant(2.30258509299404568402)  # ln(10)
 PHI = phi = constant(1.61803398874989484820)  # golden constant
 PI = pi = constant(3.14159265358979323846)  # pi
+HALF_PI = half_pi = constant(1.57079632679489661923)  # half pi
 TAU = tau = constant(6.28318530717958647692)  # 2 * pi
 
 INF = inf = constant(math.inf)

warp/context.py

@@ -7,6 +7,7 @@
 
 import ast
 import ctypes
+import functools
 import gc
 import hashlib
 import inspect
@@ -643,8 +644,10 @@ def func(f):
         func=f, key=name, namespace="", module=m, value_func=None
     )  # value_type not known yet, will be inferred during Adjoint.build()
 
-    # return the top of the list of overloads for this key
-    return m.functions[name]
+    # use the top of the list of overloads for this key
+    g = m.functions[name]
+    # copy over the function attributes, including docstring
+    return functools.update_wrapper(g, f)
 
 
 def func_native(snippet, adj_snippet=None, replay_snippet=None):
@@ -665,8 +668,9 @@ def func_native(snippet, adj_snippet=None, replay_snippet=None):
             adj_native_snippet=adj_snippet,
             replay_snippet=replay_snippet,
         )  # cuda snippets do not have a return value_type
-
-        return m.functions[name]
+        g = m.functions[name]
+        # copy over the function attributes, including docstring
+        return functools.update_wrapper(g, f)
 
     return snippet_func
 
@@ -710,7 +714,11 @@ def func_grad(forward_fn):
         def match_function(f):
             # check whether the function overload f matches the signature of the provided gradient function
             if not hasattr(f.adj, "return_var"):
-                f.adj.build(None, f.module.options)
+                # we have to temporarily build this function to figure out its return type(s);
+                # note that we do not have a ModuleBuilder instance here at this wrapping stage, hence we
+                # have to create a dummy builder
+                builder = ModuleBuilder(Module("dummy", None), f.module.options)
+                f.adj.build(builder)
             expected_args = list(f.input_types.items())
             if f.adj.return_var is not None:
                 expected_args += [(f"adj_ret_{var.label}", var.type) for var in f.adj.return_var]
@@ -745,7 +753,7 @@ def func_grad(forward_fn):
                     continue
                 if match_function(f):
                     add_custom_grad(f)
-                    return
+                    return grad_fn
             raise RuntimeError(
                 f"No function overload found for gradient function {grad_fn.__qualname__} for function {forward_fn.key}"
             )
@@ -767,6 +775,8 @@ def func_grad(forward_fn):
                     f"\n{', '.join(map(lambda nt: f'{nt[0]}: {nt[1].__name__}', expected_args))}"
                 )
 
+        return grad_fn
+
     return wrapper
 
 
@@ -810,6 +820,7 @@ def func_replay(forward_fn):
             skip_adding_overload=True,
             code_transformers=f.adj.transformers,
         )
+        return replay_fn
 
     return wrapper
 
@@ -830,6 +841,7 @@ def kernel(f=None, *, enable_backward=None):
             module=m,
             options=options,
         )
+        k = functools.update_wrapper(k, f)
         return k
 
     if f is None:
@@ -843,7 +855,7 @@ def kernel(f=None, *, enable_backward=None):
 def struct(c):
     m = get_module(c.__module__)
     s = warp.codegen.Struct(cls=c, key=warp.codegen.make_full_qualified_name(c), module=m)
-
+    s = functools.update_wrapper(s, c)
     return s
 
 
@@ -1159,6 +1171,7 @@ class ModuleBuilder:
         self.structs = {}
         self.options = options
         self.module = module
+        self.deferred_functions = []
 
         # build all functions declared in the module
         for func in module.functions.values():
@@ -1175,6 +1188,10 @@ class ModuleBuilder:
                 for k in kernel.overloads.values():
                     self.build_kernel(k)
 
+        # build all functions outside this module which are called from functions or kernels in this module
+        for func in self.deferred_functions:
+            self.build_function(func)
+
     def build_struct_recursive(self, struct: warp.codegen.Struct):
         structs = []
 
@@ -1667,7 +1684,7 @@ class Module:
                     if cuda_module is not None:
                         self.cuda_modules[device.context] = cuda_module
                     else:
-                        raise Exception("Failed to load CUDA module")
+                        raise Exception(f"Failed to load CUDA module '{self.name}'")
 
                 except Exception as e:
                     self.cuda_build_failed = True
@@ -2871,6 +2888,15 @@ class Runtime:
         # initialize kernel cache
         warp.build.init_kernel_cache(warp.config.kernel_cache_dir)
 
+        devices_without_uva = []
+        devices_without_mempool = []
+        for cuda_device in self.cuda_devices:
+            if cuda_device.is_primary:
+                if not cuda_device.is_uva:
+                    devices_without_uva.append(cuda_device)
+                if not cuda_device.is_mempool_supported:
+                    devices_without_mempool.append(cuda_device)
+
         # print device and version information
         if not warp.config.quiet:
             greeting = []
@@ -2893,16 +2919,12 @@ class Runtime:
             alias_str = f'"{self.cpu_device.alias}"'
             name_str = f'"{self.cpu_device.name}"'
             greeting.append(f"     {alias_str:10s} : {name_str}")
-            devices_without_uva = []
-            devices_without_mempool = []
             for cuda_device in self.cuda_devices:
                 alias_str = f'"{cuda_device.alias}"'
                 if cuda_device.is_primary:
                     name_str = f'"{cuda_device.name}"'
                     arch_str = f"sm_{cuda_device.arch}"
                     mem_str = f"{cuda_device.total_memory / 1024 / 1024 / 1024:.0f} GiB"
-                    if not cuda_device.is_uva:
-                        devices_without_uva.append(cuda_device)
                     if cuda_device.is_mempool_supported:
                         if cuda_device.is_mempool_enabled:
                             mempool_str = "mempool enabled"
@@ -2910,7 +2932,6 @@ class Runtime:
                             mempool_str = "mempool supported"
                     else:
                         mempool_str = "mempool not supported"
-                        devices_without_mempool.append(cuda_device)
                     greeting.append(f"     {alias_str:10s} : {name_str} ({mem_str}, {arch_str}, {mempool_str})")
                 else:
                     primary_alias_str = f'"{self.cuda_primary_devices[cuda_device.ordinal].alias}"'
@@ -2971,14 +2992,16 @@ class Runtime:
             # --quick flag.  The consequences of running with an older driver can be obscure and severe,
             # so make sure we print a very visible warning.
             if self.driver_version < self.toolkit_version and not self.core.is_cuda_compatibility_enabled():
-                print("******************************************************************\n"
-                      "* WARNING:                                                       *\n"
-                      "*   Warp was compiled without CUDA compatibility support         *\n"
-                      "*   (quick build).  The CUDA Toolkit version used to build       *\n"
-                      "*   Warp is not fully supported by the current driver.           *\n"
-                      "*   Some CUDA functionality may not work correctly!              *\n"
-                      "*   Update the driver or rebuild Warp without the --quick flag.  *\n"
-                      "******************************************************************\n")
+                print(
+                    "******************************************************************\n"
+                    "* WARNING:                                                       *\n"
+                    "*   Warp was compiled without CUDA compatibility support         *\n"
+                    "*   (quick build).  The CUDA Toolkit version used to build       *\n"
+                    "*   Warp is not fully supported by the current driver.           *\n"
+                    "*   Some CUDA functionality may not work correctly!              *\n"
+                    "*   Update the driver or rebuild Warp without the --quick flag.  *\n"
+                    "******************************************************************\n"
+                )
 
             # ensure initialization did not change the initial context (e.g. querying available memory)
             self.core.cuda_context_set_current(initial_context)

warp/examples/assets/cartpole.urdf

@@ -0,0 +1,110 @@
+<?xml version="1.0"?>
+<robot name="cartpole">
+
+  <link name="slider">
+    <visual>
+      <geometry>
+        <box size="0.03 8 0.03"/>
+      </geometry>
+      <material name="slider_mat">
+        <color rgba="0.9 0.6 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.03 8 0.03"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <joint name="slider_to_cart" type="prismatic">
+    <axis xyz="0 1 0"/>
+    <origin xyz="0 0 0"/>
+    <parent link="slider"/>
+    <child link="cart"/>
+    <limit effort="1000.0" lower="-4" upper="4" velocity="100"/>
+  </joint>
+  
+  <link name="cart">
+    <visual>
+      <geometry>
+        <box size="0.2 0.25 0.2"/>
+      </geometry>
+      <material name="cart_mat">
+        <color rgba="0.3 0.5 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+          <box size="0.2 0.25 0.2"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <mass value="1"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <joint name="cart_to_pole1" type="continuous">
+    <axis xyz="1 0 0"/>
+    <origin xyz="0.12 0 0"/>
+    <parent link="cart"/>
+    <child link="pole1"/>
+    <limit effort="1000.0" velocity="8"/>
+  </joint>
+  
+  <link name="pole1">
+    <visual>
+      <geometry>
+        <box size="0.04 0.06 1.0"/>	
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+      <material name="pole_mat">
+        <color rgba="0.1 0.1 0.3 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.04 0.06 1.0"/>	
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+    </collision>
+    <inertial>
+      <origin xyz="0 0 0.47"/>
+      <mass value="0.25"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <link name="pole2">
+    <visual>
+      <geometry>
+        <box size="0.04 0.06 1.0"/> 
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+      <material name="pole_mat">
+        <color rgba="0.1 0.1 0.3 1"/>
+      </material>
+    </visual>
+    <collision>
+      <geometry>
+        <box size="0.04 0.06 1.0"/> 
+      </geometry>
+      <origin xyz="0 0 0.47"/>
+    </collision>
+    <inertial>
+      <origin xyz="0 0 0.47"/>
+      <mass value="0.25"/>
+      <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+    </inertial>
+  </link>
+
+  <joint name="pole1_to_pole2" type="continuous">
+    <axis xyz="1 0 0"/>
+    <origin xyz="0.0 0 1.0"/>
+    <parent link="pole1"/>
+    <child link="pole2"/>
+    <limit effort="1000.0" velocity="8"/>
+  </joint>
+
+</robot>

warp/examples/assets/cube.usda

@@ -0,0 +1,42 @@
+#usda 1.0
+(
+    doc = "Blender v2.92.0"
+    metersPerUnit = 1
+    upAxis = "Z"
+    defaultPrim = "Cube"
+)
+
+def Xform "Cube"
+{
+    matrix4d xformOp:transform = ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) )
+    uniform token[] xformOpOrder = ["xformOp:transform"]
+
+    def Mesh "Cube_001"
+    {
+        int[] faceVertexCounts = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
+        int[] faceVertexIndices = [1, 2, 0, 3, 6, 2, 7, 4, 6, 5, 0, 4, 6, 0, 2, 3, 5, 7, 1, 3, 2, 3, 7, 6, 7, 5, 4, 5, 1, 0, 6, 4, 0, 3, 1, 5]
+        normal3f[] normals = [(-1, 0, 0), (-1, 0, 0), (-1, 0, 0), (0, 1, -0), (0, 1, -0), (0, 1, -0), (1, 0, -0), (1, 0, -0), (1, 0, -0), (0, -1, 0), (0, -1, 0), (0, -1, 0), (0, 0, -1), (0, 0, -1), (0, 0, -1), (0, 0, 1), (0, 0, 1), (0, 0, 1), (-1, 0, 0), (-1, 0, 0), (-1, 0, 0), (0, 1, -0), (0, 1, -0), (0, 1, -0), (1, 0, 0), (1, 0, 0), (1, 0, 0), (0, -1, 0), (0, -1, 0), (0, -1, 0), (0, 0, -1), (0, 0, -1), (0, 0, -1), (0, -0, 1), (0, -0, 1), (0, -0, 1)] (
+            interpolation = "faceVarying"
+        )
+        point3f[] points = [(-1, -1, -1), (-1, -1, 1), (-1, 1, -1), (-1, 1, 1), (1, -1, -1), (1, -1, 1), (1, 1, -1), (1, 1, 1)]
+        uniform token subdivisionScheme = "none"
+    }
+}
+
+def Xform "Camera"
+{
+    matrix4d xformOp:transform = ( (0.6859206557273865, 0.7276763319969177, 0, 0), (-0.32401347160339355, 0.305420845746994, 0.8953956365585327, 0), (0.6515582203865051, -0.6141703724861145, 0.44527140259742737, 0), (7.358891487121582, -6.925790786743164, 4.958309173583984, 1) )
+    uniform token[] xformOpOrder = ["xformOp:transform"]
+
+    def Camera "Camera"
+    {
+        float2 clippingRange = (0.1, 100)
+        float focalLength = 50
+        float horizontalAperture = 36
+        float horizontalApertureOffset = 0
+        token projection = "perspective"
+        float verticalAperture = 20.25
+        float verticalApertureOffset = 0
+    }
+}
+

warp/examples/assets/nv_ant.xml

@@ -0,0 +1,92 @@
+<mujoco model="ant">
+  <custom>
+    <numeric data="0.0 0.0 0.55 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0" name="init_qpos"/>
+  </custom>
+
+  <default>
+    <joint armature="0.01" damping="0.1" limited="true"/>
+    <geom condim="3" density="5.0" friction="1.5 0.1 0.1" margin="0.01" rgba="0.97 0.38 0.06 1"/>
+  </default>
+
+  <compiler inertiafromgeom="true" angle="degree"/>
+
+  <option timestep="0.016" iterations="50" tolerance="1e-10" solver="Newton" jacobian="dense" cone="pyramidal"/>
+
+  <size nconmax="50" njmax="200" nstack="10000"/>
+  <visual>
+      <map force="0.1" zfar="30"/>
+      <rgba haze="0.15 0.25 0.35 1"/>
+      <quality shadowsize="2048"/>
+      <global offwidth="800" offheight="800"/>
+  </visual>
+
+  <asset>
+      <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="512"/> 
+      <texture name="texplane" type="2d" builtin="checker" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2" width="512" height="512" mark="cross" markrgb=".8 .8 .8"/>
+      <texture name="texgeom" type="cube" builtin="flat" mark="cross" width="127" height="1278" 
+          rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1" random="0.01"/>  
+
+      <material name="matplane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true"/>
+      <material name="matgeom" texture="texgeom" texuniform="true" rgba="0.8 0.6 .4 1"/>
+  </asset>
+
+  <worldbody>
+    <geom name="floor" pos="0 0 0" size="0 0 .25" type="plane" material="matplane" condim="3"/>
+
+    <light directional="false" diffuse=".2 .2 .2" specular="0 0 0" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
+    <light mode="targetbodycom" target="torso" directional="false" diffuse=".8 .8 .8" specular="0.3 0.3 0.3" pos="0 0 4.0" dir="0 0 -1"/>
+
+    <body name="torso" pos="0 0 0.75">
+      <geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere"/>
+      <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+      <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule"/>
+      <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule"/>
+      <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule" rgba=".999 .2 .02 1"/>
+
+      <joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free"/>
+      <body name="front_left_leg" pos="0.2 0.2 0">
+        <joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        <body pos="0.2 0.2 0" name="front_left_foot">
+          <joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 100" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        </body>
+      </body>
+      <body name="front_right_leg" pos="-0.2 0.2 0">
+        <joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule"/>
+        <body pos="-0.2 0.2 0" name="front_right_foot">
+          <joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-100 -30" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule"/>
+        </body>
+      </body>
+      <body name="left_back_leg" pos="-0.2 -0.2 0">
+        <joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule"/>
+        <body pos="-0.2 -0.2 0" name="left_back_foot">
+          <joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-100 -30" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule"/>
+        </body>
+      </body>
+      <body name="right_back_leg" pos="0.2 -0.2 0">
+        <joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-40 40" type="hinge"/>
+        <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        <body pos="0.2 -0.2 0" name="right_back_foot">
+          <joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 100" type="hinge"/>
+          <geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule" rgba=".999 .2 .1 1"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="15"/>
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="15"/>
+  </actuator>
+</mujoco>