warp-lang 1.5.1__py3-none-macosx_10_13_universal2.whl → 1.6.1__py3-none-macosx_10_13_universal2.whl

warp/native/warp.h

@@ -67,21 +67,23 @@ extern "C"
     WP_API void memtile_host(void* dest, const void* src, size_t srcsize, size_t n);
     WP_API void memtile_device(void* context, void* dest, const void* src, size_t srcsize, size_t n);
 
-	WP_API uint64_t bvh_create_host(wp::vec3* lowers, wp::vec3* uppers, int num_items);
+	WP_API uint64_t bvh_create_host(wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type);
 	WP_API void bvh_destroy_host(uint64_t id);
     WP_API void bvh_refit_host(uint64_t id);
 
-	WP_API uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items);
+	WP_API uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type);
 	WP_API void bvh_destroy_device(uint64_t id);
     WP_API void bvh_refit_device(uint64_t id);
 
     // create a user-accessible copy of the mesh, it is the 
     // users responsibility to keep-alive the points/tris data for the duration of the mesh lifetime
-	WP_API uint64_t mesh_create_host(wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number);
+	WP_API uint64_t mesh_create_host(wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, 
+        int num_points, int num_tris, int support_winding_number, int constructor_type);
 	WP_API void mesh_destroy_host(uint64_t id);
     WP_API void mesh_refit_host(uint64_t id);
 
-	WP_API uint64_t mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number);
+	WP_API uint64_t mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, 
+        wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type);
 	WP_API void mesh_destroy_device(uint64_t id);
     WP_API void mesh_refit_device(uint64_t id);
 
@@ -159,6 +161,9 @@ extern "C"
     WP_API void radix_sort_pairs_int_host(uint64_t keys, uint64_t values, int n);
     WP_API void radix_sort_pairs_int_device(uint64_t keys, uint64_t values, int n);
 
+    WP_API void radix_sort_pairs_float_host(uint64_t keys, uint64_t values, int n);
+    WP_API void radix_sort_pairs_float_device(uint64_t keys, uint64_t values, int n);
+
     WP_API void runlength_encode_int_host(uint64_t values, uint64_t run_values, uint64_t run_lengths, uint64_t run_count, int n);
     WP_API void runlength_encode_int_device(uint64_t values, uint64_t run_values, uint64_t run_lengths, uint64_t run_count, int n);
 
@@ -266,6 +271,7 @@ extern "C"
     WP_API int cuda_device_get_pci_device_id(int ordinal);
     WP_API int cuda_device_is_uva(int ordinal);
     WP_API int cuda_device_is_mempool_supported(int ordinal);
+    WP_API int cuda_device_is_ipc_supported(int ordinal);
     WP_API int cuda_device_set_mempool_release_threshold(int ordinal, uint64_t threshold);
     WP_API uint64_t cuda_device_get_mempool_release_threshold(int ordinal);
     WP_API void cuda_device_get_memory_info(int ordinal, size_t* free_mem, size_t* total_mem);
@@ -294,6 +300,13 @@ extern "C"
     WP_API int cuda_is_mempool_access_enabled(int target_ordinal, int peer_ordinal);
     WP_API int cuda_set_mempool_access_enabled(int target_ordinal, int peer_ordinal, int enable);
 
+    // inter-process communication
+    WP_API void cuda_ipc_get_mem_handle(void* ptr, char* out_buffer);
+    WP_API void* cuda_ipc_open_mem_handle(void* context, char* handle);
+    WP_API void cuda_ipc_close_mem_handle(void* ptr);
+    WP_API void cuda_ipc_get_event_handle(void* context, void* event, char* out_buffer);
+    WP_API void* cuda_ipc_open_event_handle(void* context, char* handle);
+
     WP_API void* cuda_stream_create(void* context, int priority);
     WP_API void cuda_stream_destroy(void* context, void* stream);
     WP_API void cuda_stream_register(void* context, void* stream);
@@ -317,9 +330,10 @@ extern "C"
     WP_API bool cuda_graph_launch(void* graph, void* stream);
     WP_API bool cuda_graph_destroy(void* context, void* graph);
 
-    WP_API size_t cuda_compile_program(const char* cuda_src, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes);
+    WP_API size_t cuda_compile_program(const char* cuda_src, const char* program_name, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, bool fuse_fp, bool lineinfo, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes, int* ltoir_input_types);
     WP_API bool cuda_compile_fft(const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int size, int elements_per_thread, int direction, int precision, int* shared_memory_size);
     WP_API bool cuda_compile_dot(const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int M, int N, int K, int precision_A, int precision_B, int precision_C, int type, int arrangement_A, int arrangement_B, int arrangement_C, int num_threads);
+    WP_API bool cuda_compile_solver(const char* fatbin_output_path, const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int M, int N, int function, int precision, int fill_mode, int num_threads);
 
     WP_API void* cuda_load_module(void* context, const char* ptx);
     WP_API void cuda_unload_module(void* context, void* module);

warp/optim/linear.py

@@ -84,10 +84,7 @@ def aslinearoperator(A: _Matrix) -> LinearOperator:
         sparse.bsr_mv(A, x, z, alpha, beta)
 
     def dense_mv(x, y, z, alpha, beta):
-        x = x.reshape((x.shape[0], 1))
-        y = y.reshape((y.shape[0], 1))
-        z = z.reshape((y.shape[0], 1))
-        wp.matmul(A, x, y, z, alpha, beta)
+        wp.launch(_dense_mv_kernel, dim=A.shape[1], device=A.device, inputs=[A, x, y, z, alpha, beta])
 
     def diag_mv(x, y, z, alpha, beta):
         scalar_type = wp.types.type_scalar_type(A.dtype)
@@ -803,6 +800,27 @@ def _run_solver_loop(
     return cur_iter, err, atol
 
 
+@wp.func
+def _calc_mv_product(i: wp.int32, A: wp.array2d(dtype=Any), x: wp.array1d(dtype=Any)):
+    sum = A.dtype(0)
+    for j in range(A.shape[1]):
+        sum += A[i, j] * x[j]
+    return sum
+
+
+@wp.kernel
+def _dense_mv_kernel(
+    A: wp.array2d(dtype=Any),
+    x: wp.array1d(dtype=Any),
+    y: wp.array1d(dtype=Any),
+    z: wp.array1d(dtype=Any),
+    alpha: Any,
+    beta: Any,
+):
+    i = wp.tid()
+    z[i] = z.dtype(beta) * y[i] + z.dtype(alpha) * _calc_mv_product(i, A, x)
+
+
 @wp.kernel
 def _diag_mv_kernel(
     A: wp.array(dtype=Any),

warp/render/render_opengl.py

@@ -671,6 +671,15 @@ class ShapeInstancer:
         [3D point, 3D normal, UV texture coordinates]
     """
 
+    gl = None  # Class-level variable to hold the imported module
+
+    @classmethod
+    def initialize_gl(cls):
+        if cls.gl is None:  # Only import if not already imported
+            from pyglet import gl
+
+            cls.gl = gl
+
     def __new__(cls, *args, **kwargs):
         instance = super(ShapeInstancer, cls).__new__(cls)
         instance.instance_transform_gl_buffer = None
@@ -690,8 +699,10 @@ class ShapeInstancer:
         self.scalings = None
         self._instance_transform_cuda_buffer = None
 
+        ShapeInstancer.initialize_gl()
+
     def __del__(self):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         if self.instance_transform_gl_buffer is not None:
             try:
@@ -709,7 +720,7 @@ class ShapeInstancer:
                 pass
 
     def register_shape(self, vertices, indices, color1=(1.0, 1.0, 1.0), color2=(0.0, 0.0, 0.0)):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         if color1 is not None and color2 is None:
             color2 = np.clip(np.array(color1) + 0.25, 0.0, 1.0)
@@ -750,7 +761,7 @@ class ShapeInstancer:
         self.face_count = len(indices)
 
     def update_colors(self, colors1, colors2):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         if colors1 is None:
             colors1 = np.tile(self.color1, (self.num_instances, 1))
@@ -789,7 +800,7 @@ class ShapeInstancer:
         gl.glVertexAttribDivisor(8, 1)
 
     def allocate_instances(self, positions, rotations=None, colors1=None, colors2=None, scalings=None):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         gl.glBindVertexArray(self.vao)
 
@@ -836,6 +847,7 @@ class ShapeInstancer:
                 vbo_transforms,
             ],
             device=self.device,
+            record_tape=False,
         )
 
         vbo_transforms = vbo_transforms.numpy()
@@ -864,7 +876,7 @@ class ShapeInstancer:
         gl.glBindVertexArray(0)
 
     def update_instances(self, transforms: wp.array = None, scalings: wp.array = None, colors1=None, colors2=None):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         if transforms is not None:
             if transforms.device.is_cuda:
@@ -897,6 +909,7 @@ class ShapeInstancer:
                     vbo_transforms,
                 ],
                 device=self.device,
+                record_tape=False,
             )
 
             self._instance_transform_cuda_buffer.unmap()
@@ -905,7 +918,7 @@ class ShapeInstancer:
             self.update_colors(colors1, colors2)
 
     def render(self):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         gl.glUseProgram(self.shape_shader.id)
 
@@ -915,7 +928,7 @@ class ShapeInstancer:
 
     # scope exposes VBO transforms to be set directly by a warp kernel
     def __enter__(self):
-        from pyglet import gl
+        gl = ShapeInstancer.gl
 
         gl.glBindVertexArray(self.vao)
         self.vbo_transforms = self._instance_transform_cuda_buffer.map(dtype=wp.mat44, shape=(self.num_instances,))
@@ -941,6 +954,15 @@ class OpenGLRenderer:
     # number of segments to use for rendering spheres, capsules, cones and cylinders
     default_num_segments = 32
 
+    gl = None  # Class-level variable to hold the imported module
+
+    @classmethod
+    def initialize_gl(cls):
+        if cls.gl is None:  # Only import if not already imported
+            from pyglet import gl
+
+            cls.gl = gl
+
     def __init__(
         self,
         title="Warp",
@@ -968,6 +990,7 @@ class OpenGLRenderer:
         enable_backface_culling=True,
         enable_mouse_interaction=True,
         enable_keyboard_interaction=True,
+        device=None,
     ):
         """
         Args:
@@ -997,6 +1020,7 @@ class OpenGLRenderer:
             enable_backface_culling (bool): Whether to enable backface culling.
             enable_mouse_interaction (bool): Whether to enable mouse interaction.
             enable_keyboard_interaction (bool): Whether to enable keyboard interaction.
+            device (Devicelike): Where to store the internal data.
 
         Note:
 
@@ -1021,9 +1045,11 @@ class OpenGLRenderer:
             # disable error checking for performance
             pyglet.options["debug_gl"] = False
 
-            from pyglet import gl
             from pyglet.graphics.shader import Shader, ShaderProgram
             from pyglet.math import Vec3 as PyVec3
+
+            OpenGLRenderer.initialize_gl()
+            gl = OpenGLRenderer.gl
         except ImportError as e:
             raise Exception("OpenGLRenderer requires pyglet (version >= 2.0) to be installed.") from e
 
@@ -1040,7 +1066,11 @@ class OpenGLRenderer:
         self.render_depth = render_depth
         self.enable_backface_culling = enable_backface_culling
 
-        self._device = wp.get_preferred_device()
+        if device is None:
+            self._device = wp.get_preferred_device()
+        else:
+            self._device = wp.get_device(device)
+
         self._title = title
 
         self.window = pyglet.window.Window(
@@ -1052,9 +1082,8 @@ class OpenGLRenderer:
             self.headless = headless
         self.app = pyglet.app
 
-        if not headless:
-            # making window current opengl rendering context
-            self.window.switch_to()
+        # making window current opengl rendering context
+        self._switch_context()
 
         self.screen_width, self.screen_height = self.window.get_framebuffer_size()
 
@@ -1379,7 +1408,6 @@ class OpenGLRenderer:
 
             Window._enable_event_queue = False
 
-            self.window.switch_to()
             self.window.dispatch_pending_events()
 
             platform_event_loop = self.app.platform_event_loop
@@ -1405,7 +1433,9 @@ class OpenGLRenderer:
         return self.app.event_loop.has_exit
 
     def clear(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if not self.headless:
             self.app.event_loop.dispatch_event("on_exit")
@@ -1525,9 +1555,9 @@ class OpenGLRenderer:
             if rescale_window:
                 self.window.set_size(self._tile_width * self._tile_ncols, self._tile_height * self._tile_nrows)
         else:
-            assert (
-                len(tile_positions) == n and len(tile_sizes) == n
-            ), "Number of tiles does not match number of instances."
+            assert len(tile_positions) == n and len(tile_sizes) == n, (
+                "Number of tiles does not match number of instances."
+            )
             self._tile_ncols = None
             self._tile_nrows = None
             self._tile_width = None
@@ -1599,7 +1629,9 @@ class OpenGLRenderer:
         self._tile_viewports[tile_id] = (x, y, w, h)
 
     def _setup_framebuffer(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if self._frame_texture is None:
             self._frame_texture = gl.GLuint()
@@ -1767,7 +1799,9 @@ class OpenGLRenderer:
         return np.array((scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 1), dtype=np.float32)
 
     def update_model_matrix(self, model_matrix: Optional[Mat44] = None):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if model_matrix is None:
             self._model_matrix = self.compute_model_matrix(self._camera_axis, self._scaling)
@@ -1862,7 +1896,9 @@ class OpenGLRenderer:
                 self._draw()
 
     def _draw(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if not self.headless:
             # catch key hold events
@@ -1961,7 +1997,9 @@ Instances: {len(self._instances)}"""
             cb()
 
     def _draw_grid(self, is_tiled=False):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if not is_tiled:
             gl.glUseProgram(self._grid_shader.id)
@@ -1974,7 +2012,9 @@ Instances: {len(self._instances)}"""
         gl.glBindVertexArray(0)
 
     def _draw_sky(self, is_tiled=False):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if not is_tiled:
             gl.glUseProgram(self._sky_shader.id)
@@ -1988,7 +2028,9 @@ Instances: {len(self._instances)}"""
         gl.glBindVertexArray(0)
 
     def _render_scene(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         start_instance_idx = 0
 
@@ -2011,7 +2053,9 @@ Instances: {len(self._instances)}"""
         gl.glBindVertexArray(0)
 
     def _render_scene_tiled(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         for i, viewport in enumerate(self._tile_viewports):
             projection_matrix_ptr = arr_pointer(self._tile_projection_matrices[i])
@@ -2066,6 +2110,7 @@ Instances: {len(self._instances)}"""
             return
 
         import pyglet
+        from pyglet.math import Vec3 as PyVec3
 
         if buttons & pyglet.window.mouse.LEFT:
             sensitivity = 0.1
@@ -2077,10 +2122,12 @@ Instances: {len(self._instances)}"""
 
             self._pitch = max(min(self._pitch, 89.0), -89.0)
 
-            self._camera_front.x = np.cos(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch))
-            self._camera_front.y = np.sin(np.deg2rad(self._pitch))
-            self._camera_front.z = np.sin(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch))
-            self._camera_front = self._camera_front.normalize()
+            self._camera_front = PyVec3(
+                np.cos(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch)),
+                np.sin(np.deg2rad(self._pitch)),
+                np.sin(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch)),
+            ).normalize()
+
             self.update_view_matrix()
 
     def _scroll_callback(self, x, y, scroll_x, scroll_y):
@@ -2156,7 +2203,9 @@ Instances: {len(self._instances)}"""
         self._setup_framebuffer()
 
     def register_shape(self, geo_hash, vertices, indices, color1=None, color2=None):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         shape = len(self._shapes)
         if color1 is None:
@@ -2205,7 +2254,9 @@ Instances: {len(self._instances)}"""
         return shape
 
     def deregister_shape(self, shape):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if shape not in self._shape_gl_buffers:
             return
@@ -2264,7 +2315,9 @@ Instances: {len(self._instances)}"""
         del self._instances[name]
 
     def update_instance_colors(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         colors1, colors2 = [], []
         all_instances = list(self._instances.values())
@@ -2285,7 +2338,9 @@ Instances: {len(self._instances)}"""
         gl.glBufferData(gl.GL_ARRAY_BUFFER, colors2.nbytes, colors2.ctypes.data, gl.GL_STATIC_DRAW)
 
     def allocate_shape_instances(self):
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         self._add_shape_instances = False
         self._wp_instance_transforms = wp.array(
@@ -2387,7 +2442,9 @@ Instances: {len(self._instances)}"""
             color2: The second color of the checker pattern
             visible: Whether the shape is visible
         """
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         if name in self._instances:
             i, body, shape, tf, scale, old_color1, old_color2, v = self._instances[name]
@@ -2452,6 +2509,7 @@ Instances: {len(self._instances)}"""
                 vbo_transforms,
             ],
             device=self._device,
+            record_tape=False,
         )
 
         self._instance_transform_cuda_buffer.unmap()
@@ -2498,29 +2556,30 @@ Instances: {len(self._instances)}"""
         Returns:
             bool: Whether the pixels were successfully read.
         """
-        from pyglet import gl
+        gl = OpenGLRenderer.gl
+
+        self._switch_context()
 
         channels = 3 if mode == "rgb" else 1
 
         if split_up_tiles:
-            assert (
-                self._tile_width is not None and self._tile_height is not None
-            ), "Tile width and height are not set, tiles must all have the same size"
-            assert all(
-                vp[2] == self._tile_width for vp in self._tile_viewports
-            ), "Tile widths do not all equal global tile_width, use `get_tile_pixels` instead to retrieve pixels for a single tile"
-            assert all(
-                vp[3] == self._tile_height for vp in self._tile_viewports
-            ), "Tile heights do not all equal global tile_height, use `get_tile_pixels` instead to retrieve pixels for a single tile"
-            assert (
-                target_image.shape
-                == (
-                    self.num_tiles,
-                    self._tile_height,
-                    self._tile_width,
-                    channels,
-                )
-            ), f"Shape of `target_image` array does not match {self.num_tiles} x {self._tile_height} x {self._tile_width} x {channels}"
+            assert self._tile_width is not None and self._tile_height is not None, (
+                "Tile width and height are not set, tiles must all have the same size"
+            )
+            assert all(vp[2] == self._tile_width for vp in self._tile_viewports), (
+                "Tile widths do not all equal global tile_width, use `get_tile_pixels` instead to retrieve pixels for a single tile"
+            )
+            assert all(vp[3] == self._tile_height for vp in self._tile_viewports), (
+                "Tile heights do not all equal global tile_height, use `get_tile_pixels` instead to retrieve pixels for a single tile"
+            )
+            assert target_image.shape == (
+                self.num_tiles,
+                self._tile_height,
+                self._tile_width,
+                channels,
+            ), (
+                f"Shape of `target_image` array does not match {self.num_tiles} x {self._tile_height} x {self._tile_width} x {channels}"
+            )
         else:
             assert target_image.shape == (
                 self.screen_height,
@@ -2784,7 +2843,7 @@ Instances: {len(self._instances)}"""
             up_axis: The axis of the capsule that points up (0: x, 1: y, 2: z)
             color: The color of the capsule
         """
-        geo_hash = hash(("capsule", radius, half_height))
+        geo_hash = hash(("capsule", radius, half_height, up_axis))
         if geo_hash in self._shape_geo_hash:
             shape = self._shape_geo_hash[geo_hash]
             if self.update_shape_instance(name, pos, rot):
@@ -2819,7 +2878,7 @@ Instances: {len(self._instances)}"""
             up_axis: The axis of the cylinder that points up (0: x, 1: y, 2: z)
             color: The color of the capsule
         """
-        geo_hash = hash(("cylinder", radius, half_height))
+        geo_hash = hash(("cylinder", radius, half_height, up_axis))
         if geo_hash in self._shape_geo_hash:
             shape = self._shape_geo_hash[geo_hash]
             if self.update_shape_instance(name, pos, rot):
@@ -2854,7 +2913,7 @@ Instances: {len(self._instances)}"""
             up_axis: The axis of the cone that points up (0: x, 1: y, 2: z)
             color: The color of the cone
         """
-        geo_hash = hash(("cone", radius, half_height))
+        geo_hash = hash(("cone", radius, half_height, up_axis))
         if geo_hash in self._shape_geo_hash:
             shape = self._shape_geo_hash[geo_hash]
             if self.update_shape_instance(name, pos, rot):
@@ -2933,7 +2992,7 @@ Instances: {len(self._instances)}"""
         indices = np.array(indices, dtype=np.int32).reshape((-1, 3))
         idx_count = len(indices)
 
-        geo_hash = hash((indices.tobytes(),))
+        geo_hash = hash((points.tobytes(), indices.tobytes()))
 
         if name in self._instances:
             # We've already registered this mesh instance and its associated shape.
@@ -2955,6 +3014,12 @@ Instances: {len(self._instances)}"""
             if shape is not None:
                 # Update the shape's point positions.
                 self.update_shape_vertices(shape, points, scale)
+
+                if not is_template and name not in self._instances:
+                    # Create a new instance.
+                    body = self._resolve_body_id(parent_body)
+                    self.add_shape_instance(name, shape, body, pos, rot, color1=colors)
+
                 return shape
 
         # No existing shape for the given mesh was found, or its topology may have changed,
@@ -3032,16 +3097,16 @@ Instances: {len(self._instances)}"""
             name: A name for the USD prim on the stage
             up_axis: The axis of the arrow that points up (0: x, 1: y, 2: z)
         """
-        geo_hash = hash(("arrow", base_radius, base_height, cap_radius, cap_height))
+        geo_hash = hash(("arrow", base_radius, base_height, cap_radius, cap_height, up_axis))
         if geo_hash in self._shape_geo_hash:
             shape = self._shape_geo_hash[geo_hash]
-            if self.update_shape_instance(name, pos, rot):
+            if self.update_shape_instance(name, pos, rot, color1=color, color2=color):
                 return shape
         else:
             vertices, indices = self._create_arrow_mesh(
                 base_radius, base_height, cap_radius, cap_height, up_axis=up_axis
             )
-            shape = self.register_shape(geo_hash, vertices, indices)
+            shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
         if not is_template:
             body = self._resolve_body_id(parent_body)
             self.add_shape_instance(name, shape, body, pos, rot, color1=color, color2=color)
@@ -3432,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

warp/sim/__init__.py

@@ -50,4 +50,9 @@ from .model import (
     ModelShapeMaterials,
     State,
 )
-from .utils import load_mesh, quat_from_euler, quat_to_euler, velocity_at_point
+from .utils import (
+    load_mesh,
+    quat_from_euler,
+    quat_to_euler,
+    velocity_at_point,
+)