warp-lang 1.6.0__py3-none-manylinux2014_aarch64.whl → 1.6.1__py3-none-manylinux2014_aarch64.whl

warp/examples/sim/example_cloth_self_contact.py

@@ -29,17 +29,20 @@ from warp.sim.model import PARTICLE_FLAG_ACTIVE
 
 @wp.kernel
 def initialize_rotation(
-    particle_indices_to_rot: wp.array(dtype=wp.int32),
+    # input
+    vertex_indices_to_rot: wp.array(dtype=wp.int32),
     pos: wp.array(dtype=wp.vec3),
     rot_centers: wp.array(dtype=wp.vec3),
     rot_axes: wp.array(dtype=wp.vec3),
+    t: wp.array(dtype=float),
+    # output
     roots: wp.array(dtype=wp.vec3),
     roots_to_ps: wp.array(dtype=wp.vec3),
 ):
     tid = wp.tid()
-    particle_index = particle_indices_to_rot[wp.tid()]
+    v_index = vertex_indices_to_rot[wp.tid()]
 
-    p = pos[particle_index]
+    p = pos[v_index]
     rot_center = rot_centers[tid]
     rot_axis = rot_axes[tid]
     op = p - rot_center
@@ -51,21 +54,31 @@ def initialize_rotation(
     roots[tid] = root
     roots_to_ps[tid] = root_to_p
 
+    if tid == 0:
+        t[0] = 0.0
+
 
 @wp.kernel
 def apply_rotation(
-    time: float,
-    angular_velocity: float,
-    particle_indices_to_rot: wp.array(dtype=wp.int32),
-    rot_centers: wp.array(dtype=wp.vec3),
+    # input
+    vertex_indices_to_rot: wp.array(dtype=wp.int32),
     rot_axes: wp.array(dtype=wp.vec3),
     roots: wp.array(dtype=wp.vec3),
     roots_to_ps: wp.array(dtype=wp.vec3),
+    t: wp.array(dtype=float),
+    angular_velocity: float,
+    dt: float,
+    end_time: float,
+    # output
     pos_0: wp.array(dtype=wp.vec3),
     pos_1: wp.array(dtype=wp.vec3),
 ):
+    cur_t = t[0]
+    if cur_t > end_time:
+        return
+
     tid = wp.tid()
-    particle_index = particle_indices_to_rot[wp.tid()]
+    v_index = vertex_indices_to_rot[wp.tid()]
 
     rot_axis = rot_axes[tid]
 
@@ -73,7 +86,7 @@ def apply_rotation(
     uy = rot_axis[1]
     uz = rot_axis[2]
 
-    theta = time * angular_velocity
+    theta = cur_t * angular_velocity
 
     R = wp.mat33(
         wp.cos(theta) + ux * ux * (1.0 - wp.cos(theta)),
@@ -92,24 +105,29 @@ def apply_rotation(
     root_to_p_rot = R * root_to_p
     p_rot = root + root_to_p_rot
 
-    pos_0[particle_index] = p_rot
-    pos_1[particle_index] = p_rot
+    pos_0[v_index] = p_rot
+    pos_1[v_index] = p_rot
+
+    if tid == 0:
+        t[0] = cur_t + dt
 
 
 class Example:
-    def __init__(self, stage_path="example_cloth_self_contact.usd", num_frames=1500):
+    def __init__(self, stage_path="example_cloth_self_contact.usd", num_frames=600):
         fps = 60
         self.frame_dt = 1.0 / fps
+        # must be an even number when using CUDA Graph
         self.num_substeps = 10
-        self.iterations = 10
+        self.iterations = 4
         self.dt = self.frame_dt / self.num_substeps
 
         self.num_frames = num_frames
         self.sim_time = 0.0
         self.profiler = {}
 
-        self.rot_angular_velocity = math.pi / 6
-        self.rot_end_time = 21
+        self.rot_angular_velocity = math.pi / 3
+        self.rot_end_time = 10
+        self.use_cuda_graph = wp.get_device().is_cuda
 
         usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "square_cloth.usd"))
         usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/cloth/cloth"))
@@ -134,7 +152,8 @@ class Example:
             density=0.02,
             tri_ke=1.0e5,
             tri_ka=1.0e5,
-            tri_kd=3.0e-5,
+            tri_kd=2.0e-6,
+            edge_ke=10,
         )
         builder.color()
         self.model = builder.finalize()
@@ -170,6 +189,7 @@ class Example:
         rot_axes = [[1, 0, 0]] * len(right_side) + [[-1, 0, 0]] * len(left_side)
 
         self.rot_point_indices = wp.array(rot_point_indices, dtype=int)
+        self.t = wp.zeros((1,), dtype=float)
         self.rot_centers = wp.zeros(len(rot_point_indices), dtype=wp.vec3)
         self.rot_axes = wp.array(rot_axes, dtype=wp.vec3)
 
@@ -184,41 +204,75 @@ class Example:
                 self.state0.particle_q,
                 self.rot_centers,
                 self.rot_axes,
+                self.t,
+            ],
+            outputs=[
                 self.roots,
                 self.roots_to_ps,
             ],
         )
 
         if stage_path:
-            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=40.0)
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1)
         else:
             self.renderer = None
+        self.cuda_graph = None
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                for _ in range(self.num_substeps):
+                    wp.launch(
+                        kernel=apply_rotation,
+                        dim=self.rot_point_indices.shape[0],
+                        inputs=[
+                            self.rot_point_indices,
+                            self.rot_axes,
+                            self.roots,
+                            self.roots_to_ps,
+                            self.t,
+                            self.rot_angular_velocity,
+                            self.dt,
+                            self.rot_end_time,
+                        ],
+                        outputs=[
+                            self.state0.particle_q,
+                            self.state1.particle_q,
+                        ],
+                    )
+
+                    self.integrator.simulate(self.model, self.state0, self.state1, self.dt, None)
+                    (self.state0, self.state1) = (self.state1, self.state0)
+
+            self.cuda_graph = capture.graph
 
     def step(self):
         with wp.ScopedTimer("step", print=False, dict=self.profiler):
-            for _ in range(self.num_substeps):
-                if self.sim_time < self.rot_end_time:
+            if self.use_cuda_graph:
+                wp.capture_launch(self.cuda_graph)
+            else:
+                for _ in range(self.num_substeps):
                     wp.launch(
                         kernel=apply_rotation,
                         dim=self.rot_point_indices.shape[0],
                         inputs=[
-                            self.sim_time,
-                            self.rot_angular_velocity,
                             self.rot_point_indices,
-                            self.rot_centers,
                             self.rot_axes,
                             self.roots,
                             self.roots_to_ps,
+                            self.t,
+                            self.rot_angular_velocity,
+                            self.dt,
+                            self.rot_end_time,
+                        ],
+                        outputs=[
                             self.state0.particle_q,
                             self.state1.particle_q,
                         ],
                     )
+                    self.integrator.simulate(self.model, self.state0, self.state1, self.dt)
 
-                self.integrator.simulate(self.model, self.state0, self.state1, self.dt)
-
-                (self.state0, self.state1) = (self.state1, self.state0)
+                    (self.state0, self.state1) = (self.state1, self.state0)
 
-                self.sim_time += self.dt
+            self.sim_time += self.dt
 
     def render(self):
         if self.renderer is None:
@@ -241,7 +295,7 @@ if __name__ == "__main__":
         default="example_cloth_self_contact.usd",
         help="Path to the output USD file.",
     )
-    parser.add_argument("--num_frames", type=int, default=1500, help="Total number of frames.")
+    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
 
     args = parser.parse_known_args()[0]
 

warp/examples/tile/example_tile_nbody.py

@@ -26,7 +26,7 @@ import warp as wp
 
 wp.init()
 
-DT = wp.constant(0.01)
+DT = wp.constant(0.016)
 SOFTENING_SQ = wp.constant(0.1**2)  # Softening factor for numerical stability
 TILE_SIZE = wp.constant(64)
 PARTICLE_MASS = wp.constant(1.0)
@@ -75,25 +75,33 @@ def integrate_bodies_tiled(
 
 
 class Example:
-    def __init__(self, headless=False, num_bodies=1024):
+    def __init__(self, headless=False, num_bodies=16384):
         self.num_bodies = num_bodies
 
         rng = np.random.default_rng(42)
 
         # Sample the surface of a sphere
-        r = 10.0 * (num_bodies / 1024) ** (1 / 2)  # Scale factor to maintain a constant density
-        phi = np.arccos(1.0 - 2.0 * rng.uniform(size=self.num_bodies))
+        phi = np.arccos(1.0 - 2.0 * rng.uniform(low=0.0, high=1.0, size=self.num_bodies))
         theta = rng.uniform(low=0.0, high=2.0 * np.pi, size=self.num_bodies)
-        x = r * np.cos(theta) * np.sin(phi)
-        y = r * np.sin(theta) * np.sin(phi)
-        z = r * np.cos(phi)
-
-        self.scale = r
+        x = np.cos(theta) * np.sin(phi)
+        y = np.sin(theta) * np.sin(phi)
+        z = np.cos(phi)
         init_pos_np = np.stack((x, y, z), axis=1)
 
+        scale = (num_bodies / 1024) ** (1 / 2)  # Scale factor to maintain a constant density
+        inner = 0.9625 * scale
+        outer = 1.54 * scale
+        radii = inner + (outer - inner) * rng.uniform(size=(self.num_bodies, 1))
+        init_pos_np = init_pos_np * radii
+
+        axis = np.array([0.0, 0.0, 1.0])
+        v_scale = scale * 3.08
+        init_vel_np = v_scale * np.cross(init_pos_np, axis)
+
+        self.graph_scale = np.max(radii) * 5.0
         self.pos_array_0 = wp.array(init_pos_np, dtype=wp.vec3)
         self.pos_array_1 = wp.empty_like(self.pos_array_0)
-        self.vel_array = wp.zeros(self.num_bodies, dtype=wp.vec3)
+        self.vel_array = wp.array(init_vel_np, dtype=wp.vec3)
 
         if headless:
             self.scatter_plot = None
@@ -108,13 +116,16 @@ class Example:
         ax = self.fig.add_subplot(111, projection="3d")
 
         # Scatter plot of initial positions
+        point_size = 0.05 * self.graph_scale
         init_pos_np = self.pos_array_0.numpy()
-        scatter_plot = ax.scatter(init_pos_np[:, 0], init_pos_np[:, 1], init_pos_np[:, 2], c="#76b900", alpha=0.5)
+        scatter_plot = ax.scatter(
+            init_pos_np[:, 0], init_pos_np[:, 1], init_pos_np[:, 2], s=point_size, c="#76b900", alpha=0.5
+        )
 
         # Set axis limits
-        ax.set_xlim(-self.scale, self.scale)
-        ax.set_ylim(-self.scale, self.scale)
-        ax.set_zlim(-self.scale, self.scale)
+        ax.set_xlim(-self.graph_scale, self.graph_scale)
+        ax.set_ylim(-self.graph_scale, self.graph_scale)
+        ax.set_zlim(-self.graph_scale, self.graph_scale)
 
         return scatter_plot
 
@@ -149,7 +160,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
     parser.add_argument("--num_frames", type=int, default=1000, help="Total number of frames.")
-    parser.add_argument("-N", help="Number of bodies. Should be a multiple of 64.", type=int, default=1024)
+    parser.add_argument("-N", help="Number of bodies. Should be a multiple of 64.", type=int, default=16384)
     parser.add_argument(
         "--headless",
         action="store_true",

warp/native/clang/clang.cpp

@@ -413,7 +413,7 @@ WP_API int load_obj(const char* object_file, const char* module_name)
             SYMBOL(sinhf), SYMBOL_T(sinh, double(*)(double)),
             SYMBOL(coshf), SYMBOL_T(cosh, double(*)(double)),
             SYMBOL(tanhf), SYMBOL_T(tanh, double(*)(double)),
-            SYMBOL(fmaf),
+            SYMBOL(fmaf), SYMBOL_T(fma, double(*)(double, double, double)),
             SYMBOL(memcpy), SYMBOL(memset), SYMBOL(memmove),
             SYMBOL(_wp_assert),
             SYMBOL(_wp_isfinite),

warp/native/crt.h

@@ -303,6 +303,7 @@ double cosh(double);
 float tanhf(float);
 double tanh(double);
 float fmaf(float, float, float);
+double fma(double, double, double);
 
 // stddef.h
 #if defined(_WIN32)

warp/native/mat.h

@@ -681,6 +681,20 @@ inline CUDA_CALLABLE vec_t<Cols,Type> mul(const vec_t<Rows,Type>& b, const mat_t
     return r;
 }
 
+template<typename T>
+inline CUDA_CALLABLE T muladd(T a, T b, T c) {
+    return c + a*b;
+}
+template<>
+inline CUDA_CALLABLE float muladd(float a, float b, float c) {
+    return fmaf(a, b, c);
+}
+template<>
+inline CUDA_CALLABLE double muladd(double a, double b, double c) {
+    return fma(a, b, c);
+}
+
+
 template<unsigned Rows, unsigned Cols, unsigned ColsOut, typename Type>
 inline CUDA_CALLABLE mat_t<Rows,ColsOut,Type> mul(const mat_t<Rows,Cols,Type>& a, const mat_t<Cols,ColsOut,Type>& b)
 {
@@ -693,8 +707,7 @@ inline CUDA_CALLABLE mat_t<Rows,ColsOut,Type> mul(const mat_t<Rows,Cols,Type>& a
 
             for (unsigned k=0; k < Cols; ++k)
             {
-                //t.data[i][j] += a.data[i][k]*b.data[k][j];
-                sum = fmaf(a.data[i][k], b.data[k][j], sum);
+                sum = muladd<Type>(a.data[i][k], b.data[k][j], sum);
             }
 
             t.data[i][j] = sum;
@@ -713,7 +726,7 @@ inline CUDA_CALLABLE Type ddot(const mat_t<Rows,Cols,Type>& a, const mat_t<Rows,
     {
         for (unsigned j=0; j < Cols; ++j)
         {
-            r += a.data[i][j] * b.data[i][j];
+            r = muladd<Type>(a.data[i][j], b.data[i][j], r);
         }
     }
     return r;

warp/native/tile.h

@@ -1079,9 +1079,12 @@ struct tile_shared_t
             constexpr bool contiguous_src = Layout::Stride::dim(lastdim) == 1;
             const bool contiguous_dest = dest.data.strides[lastdim] == sizeof(T);
             const int elements = (dest.data.shape[lastdim] - dest.offset[lastdim]);
-            const bool aligned = (elements*sizeof(T))%sizeof(float4) == 0;
+            const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
            
-            if (contiguous_dest && contiguous_src && aligned)
+            float4* dest128 = (float4*)&dest.data.data[dest.index_from_coord(tile_coord(0,0))];
+            const bool aligned_dst = (uint64_t)(dest128)%sizeof(float4) == 0;
+
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_dst)
             {                    
                 constexpr int M = Layout::Shape::dim(0);
                 constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);            
@@ -1089,10 +1092,9 @@ struct tile_shared_t
                 // alias of shared tile with 128bit type
                 using SrcLayout = tile_layout_strided_t<tile_shape_t<M, N>>;
                 tile_shared_t<float4, SrcLayout> src128((float4*)data.ptr);
-                float4* dest128 = (float4*)&dest.data.data[dest.index_from_coord(tile_coord(0,0))];
 
                 assert(((uint64_t)(data.ptr))%sizeof(float4) == 0);
-                assert(((uint64_t)(ptr))%sizeof(float4) == 0);
+                assert(((uint64_t)(dest128))%sizeof(float4) == 0);
 
                 const int stride_i = dest.data.strides[0]/sizeof(float4);
                 const int stride_j = 1;
@@ -1164,17 +1166,19 @@ struct tile_shared_t
             constexpr bool contiguous_dest = Layout::Stride::dim(lastdim) == 1;
             const bool contiguous_src = src.data.strides[lastdim] == sizeof(T);
             const int elements = (src.data.shape[lastdim] - src.offset[lastdim]);
-            const bool aligned = (elements*sizeof(T))%sizeof(float4) == 0;
+            const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
+
+            float4* src128 = (float4*)&src.data.data[src.index_from_coord(tile_coord(0,0))];
+            const bool aligned_src = (uint64_t)(src128)%sizeof(float4) == 0;
 
-            if (contiguous_dest && contiguous_src && aligned)
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_src)
             {            
                 constexpr int M = Layout::Shape::dim(0);
                 constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
 
                 // alias of shared tile with 128bit type
                 using DestLayout = tile_layout_strided_t<tile_shape_t<M, N>>;
-                tile_shared_t<float4, DestLayout> dest128((float4*)data.ptr);
-                float4* src128 = (float4*)&src.data.data[src.index_from_coord(tile_coord(0,0))];
+                tile_shared_t<float4, DestLayout> dest128((float4*)data.ptr);                
 
                 assert(((uint64_t)(dest128.data.ptr))%sizeof(float4) == 0);
                 assert(((uint64_t)(src128))%sizeof(float4) == 0);

warp/render/render_opengl.py

@@ -1083,7 +1083,7 @@ class OpenGLRenderer:
         self.app = pyglet.app
 
         # making window current opengl rendering context
-        self.window.switch_to()
+        self._switch_context()
 
         self.screen_width, self.screen_height = self.window.get_framebuffer_size()
 
@@ -1435,7 +1435,7 @@ class OpenGLRenderer:
     def clear(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if not self.headless:
             self.app.event_loop.dispatch_event("on_exit")
@@ -1631,7 +1631,7 @@ class OpenGLRenderer:
     def _setup_framebuffer(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if self._frame_texture is None:
             self._frame_texture = gl.GLuint()
@@ -1801,7 +1801,7 @@ class OpenGLRenderer:
     def update_model_matrix(self, model_matrix: Optional[Mat44] = None):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if model_matrix is None:
             self._model_matrix = self.compute_model_matrix(self._camera_axis, self._scaling)
@@ -1898,7 +1898,7 @@ class OpenGLRenderer:
     def _draw(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if not self.headless:
             # catch key hold events
@@ -1999,7 +1999,7 @@ Instances: {len(self._instances)}"""
     def _draw_grid(self, is_tiled=False):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if not is_tiled:
             gl.glUseProgram(self._grid_shader.id)
@@ -2014,7 +2014,7 @@ Instances: {len(self._instances)}"""
     def _draw_sky(self, is_tiled=False):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if not is_tiled:
             gl.glUseProgram(self._sky_shader.id)
@@ -2030,7 +2030,7 @@ Instances: {len(self._instances)}"""
     def _render_scene(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         start_instance_idx = 0
 
@@ -2055,7 +2055,7 @@ Instances: {len(self._instances)}"""
     def _render_scene_tiled(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         for i, viewport in enumerate(self._tile_viewports):
             projection_matrix_ptr = arr_pointer(self._tile_projection_matrices[i])
@@ -2205,7 +2205,7 @@ Instances: {len(self._instances)}"""
     def register_shape(self, geo_hash, vertices, indices, color1=None, color2=None):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         shape = len(self._shapes)
         if color1 is None:
@@ -2256,7 +2256,7 @@ Instances: {len(self._instances)}"""
     def deregister_shape(self, shape):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if shape not in self._shape_gl_buffers:
             return
@@ -2317,7 +2317,7 @@ Instances: {len(self._instances)}"""
     def update_instance_colors(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         colors1, colors2 = [], []
         all_instances = list(self._instances.values())
@@ -2340,7 +2340,7 @@ Instances: {len(self._instances)}"""
     def allocate_shape_instances(self):
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         self._add_shape_instances = False
         self._wp_instance_transforms = wp.array(
@@ -2444,7 +2444,7 @@ Instances: {len(self._instances)}"""
         """
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         if name in self._instances:
             i, body, shape, tf, scale, old_color1, old_color2, v = self._instances[name]
@@ -2558,7 +2558,7 @@ Instances: {len(self._instances)}"""
         """
         gl = OpenGLRenderer.gl
 
-        self.window.switch_to()
+        self._switch_context()
 
         channels = 3 if mode == "rgb" else 1
 
@@ -3497,6 +3497,14 @@ Instances: {len(self._instances)}"""
         # fmt: on
         return np.array(vertices, dtype=np.float32), np.array(indices, dtype=np.uint32)
 
+    def _switch_context(self):
+        try:
+            self.window.switch_to()
+        except AttributeError:
+            # The window could be in the process of being closed, in which case
+            # its corresponding context might have been destroyed and set to `None`.
+            pass
+
 
 if __name__ == "__main__":
     renderer = OpenGLRenderer()

warp/render/render_usd.py

@@ -582,7 +582,12 @@ class UsdRenderer:
         mesh = UsdGeom.Mesh.Get(self.stage, mesh_path)
         if not mesh:
             mesh = UsdGeom.Mesh.Define(self.stage, mesh_path)
-            UsdGeom.Primvar(mesh.GetDisplayColorAttr()).SetInterpolation("vertex")
+            if colors is not None and len(colors) == 3:
+                color_interp = "constant"
+            else:
+                color_interp = "vertex"
+
+            UsdGeom.Primvar(mesh.GetDisplayColorAttr()).SetInterpolation(color_interp)
             _usd_add_xform(mesh)
 
             # force topology update on first frame
@@ -595,7 +600,10 @@ class UsdRenderer:
             mesh.GetFaceVertexIndicesAttr().Set(idxs, self.time)
             mesh.GetFaceVertexCountsAttr().Set([3] * len(idxs), self.time)
 
-        if colors:
+        if colors is not None:
+            if len(colors) == 3:
+                colors = (colors,)
+
             mesh.GetDisplayColorAttr().Set(colors, self.time)
 
         self._shape_constructors[name] = UsdGeom.Mesh