diff --git a/docs/source/api/lab_ov/isaaclab_ov.renderers.rst b/docs/source/api/lab_ov/isaaclab_ov.renderers.rst index 3ad45b7b4e93..a774c1ad1448 100644 --- a/docs/source/api/lab_ov/isaaclab_ov.renderers.rst +++ b/docs/source/api/lab_ov/isaaclab_ov.renderers.rst @@ -15,9 +15,12 @@ The following classes are part of the public :mod:`isaaclab_ov.renderers` API. OVRTXRenderer OVRTXRendererCfg + map_attribute_for_warp_writes .. autoclass:: OVRTXRenderer :show-inheritance: .. autoclass:: OVRTXRendererCfg :show-inheritance: + +.. autofunction:: map_attribute_for_warp_writes diff --git a/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst b/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst new file mode 100644 index 000000000000..d425463225fc --- /dev/null +++ b/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst @@ -0,0 +1,16 @@ +Added +^^^^^ + +* Added :func:`~isaaclab_ov.renderers.map_attribute_for_warp_writes`, a context manager that maps + an OVRTX attribute binding for CUDA writes and unmaps it with the producing Warp stream as the + CUDA sync. Use it instead of ``with binding.map(...)`` for GPU writes: the binding's own context + manager unmaps without a CUDA sync, so OVRTX's commit is not ordered against the fill. + +Fixed +^^^^^ + +* Fixed the OVRTX renderer's GPU transform writes (object and camera ``omni:xform`` mappings) + committing without CUDA synchronization against the Warp kernels that fill the mapped buffers. + The commit is now ordered on the producing Warp stream, as the OVRTX API contract requires; + previously the ordering held only through CUDA legacy default-stream serialization, an + implementation detail the contract does not promise. diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi b/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi index 70526feeb84a..bd16abc54924 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi +++ b/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi @@ -7,8 +7,10 @@ __all__ = [ "OVRTXRenderer", "OVRTXRendererCfg", "Renderer", + "map_attribute_for_warp_writes", ] +from .ovrtx_mapping import map_attribute_for_warp_writes from .ovrtx_renderer import OVRTXRenderer from .ovrtx_renderer import OVRTXRenderer as Renderer from .ovrtx_renderer_cfg import OVRTXRendererCfg diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py new file mode 100644 index 000000000000..a0d272d6abd4 --- /dev/null +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py @@ -0,0 +1,71 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Stream-safe helpers for OVRTX attribute-binding mappings. + +Filling a CUDA-mapped OVRTX attribute buffer from Warp is only correct when the commit at unmap +time is ordered against the Warp stream that produced the data: OVRTX's API contract requires the +mapped data to be ready when the unmap's CUDA sync signals, and an unmap without a sync performs +no synchronization at all. The ``with binding.map(...)`` form cannot carry that sync -- its +``__exit__`` takes no arguments -- so GPU writes through a mapping should use +:func:`map_attribute_for_warp_writes` instead of the binding's own context manager. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +import warp as wp + + +def _cuda_device_id(device: str) -> int: + """CUDA device index parsed from a Warp device string, e.g. ``"cuda:1"`` -> ``1``. + + TODO: A bare ``"cuda"`` parses to ``0`` while Warp enqueues fill work on its *current* CUDA + device, so the mapping and the fill can target different GPUs on multi-GPU processes. The + split predates this helper and is kept here to avoid a behavior change; a follow-up caches + the resolved Warp device on the renderer instead of re-deriving it from strings. + + Args: + device: Warp CUDA device string (``"cuda"`` or ``"cuda:"``). + + Returns: + The parsed CUDA device index, ``0`` when the string carries none. + """ + parts = device.split(":") + return int(parts[1]) if len(parts) > 1 else 0 + + +@contextmanager +def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iterator[wp.array]: + """Map ``binding`` for CUDA writes and yield its buffer as a Warp array; commit after the fill. + + The caller fills the yielded array with Warp work enqueued on ``device``'s current stream (the + default for ``wp.launch``/``wp.copy``). On exit -- error or not -- the mapping is unmapped with + that stream as the CUDA sync, so OVRTX's commit of the mapped data waits for the fill on the + GPU instead of racing it. OVRTX has no discard path (unmap always commits), so a failed fill + still publishes whatever landed in the buffer. + + Args: + binding: OVRTX attribute binding (from ``bind_attribute``) whose buffer is written. + device: Warp CUDA device the fill work runs on (e.g. ``"cuda:0"``). + dtype: Warp dtype the mapped tensor is viewed as (e.g. ``wp.mat44d``). + + Yields: + The mapped buffer as a zero-copy Warp array, valid only inside the ``with`` block. + """ + # Deferred so importing this module (e.g. through the package's lazy exports) cannot initialize + # ovrtx without the guarded environment ``ovrtx_renderer`` establishes (OVRTX_SKIP_USD_CHECK, + # actionable install error). Any real ``binding`` was created through that path, so ovrtx is + # already imported by the time this runs. + from ovrtx import Device # noqa: PLC0415 + + attr_mapping = binding.map(device=Device.CUDA, device_id=_cuda_device_id(device)) + try: + yield wp.from_dlpack(attr_mapping.tensor, dtype=dtype) + finally: + attr_mapping.unmap(stream=wp.get_stream(device).cuda_stream) diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py index 5d7d380bbe34..fefcc840f0d6 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py @@ -87,6 +87,7 @@ decode_stable_id_map, decode_stable_id_semantic_id_map, ) +from .ovrtx_mapping import map_attribute_for_warp_writes from .ovrtx_renderer_cfg import OVRTXRendererCfg from .ovrtx_renderer_kernels import ( create_camera_transforms_kernel, @@ -341,12 +342,6 @@ def supported_output_types(self) -> dict[RenderBufferKind, RenderBufferSpec]: RenderBufferKind.MOTION_VECTORS: RenderBufferSpec(2, wp.float32), } - @property - def _device_id(self) -> int: - """CUDA device index extracted from ``self._device`` for OVRTX ``binding.map()`` calls.""" - parts = self._device.split(":") - return int(parts[1]) if len(parts) > 1 else 0 - def __init__(self, cfg: OVRTXRendererCfg): self.cfg = cfg self._device = "cuda:0" # default; overridden by create_render_data(spec) @@ -852,8 +847,7 @@ def _update_transforms_legacy(self) -> None: if body_q is None: return - with self._object_xform_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping: - ovrtx_transforms = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d) + with map_attribute_for_warp_writes(self._object_xform_binding, self._device, wp.mat44d) as ovrtx_transforms: wp.launch( kernel=sync_newton_transforms_kernel, dim=len(self._object_newton_indices), @@ -955,9 +949,8 @@ def _update_camera_legacy( device=self._device, ) if self._camera_xform_binding is not None: - with self._camera_xform_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping: - wp_transforms_view = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d) - wp.copy(wp_transforms_view, camera_transforms) + with map_attribute_for_warp_writes(self._camera_xform_binding, self._device, wp.mat44d) as transforms_view: + wp.copy(transforms_view, camera_transforms) def read_output( self, @@ -1260,7 +1253,7 @@ def _prepare_ppisp_hdr_source( # FIXME: OVRTX render var mapping can select a different CUDA device # than the camera/output buffers on MGPU systems. Keep this PPISP-only # bridge until render var mapping can be constrained like transform - # bindings, which use ``device_id=self._device_id``. + # bindings, whose maps pin ``device_id`` (see ``ovrtx_mapping``). return wp.clone(tiled_data, device=output_device) def _process_render_frame(self, render_data: OVRTXRenderData, frame, output_buffers: dict) -> None: diff --git a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py index 1cb9de1386f0..c24815628c8a 100644 --- a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py +++ b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py @@ -30,7 +30,10 @@ ] if not _MISSING_MODULES: - from isaaclab_ov.renderers import OVRTXRendererCfg # noqa: E402 + from isaaclab_ov.renderers import ( + OVRTXRendererCfg, # noqa: E402 + ovrtx_mapping, # noqa: E402 + ) from isaaclab_ov.renderers import ovrtx_renderer as ovrtx_renderer_module # noqa: E402 from isaaclab_ov.renderers.ovrtx_renderer import ( # noqa: E402 _DISABLE_LINUX_CUDA_CPU_SYNC_ENV, @@ -476,6 +479,77 @@ def test_ovrtx_map_render_var_orders_the_read_against_render_completion(monkeypa assert render_var.ordering == [expected] +class _RecordingMappedBinding: + """Stand-in for an OVRTX attribute binding that records how its mapping is committed.""" + + def __init__(self): + self.map_calls: list[dict] = [] + self.unmap_calls: list[dict] = [] + + def map(self, *, device, device_id): + self.map_calls.append({"device": device, "device_id": device_id}) + binding = self + + class _Mapping: + tensor = object() + + def unmap(self, *, event=None, stream=None): + binding.unmap_calls.append({"event": event, "stream": stream}) + + return _Mapping() + + +def _patch_warp_device(monkeypatch, *, ordinal: int, cuda_stream: int) -> None: + """Fake the current Warp stream; ``ordinal`` documents the device the test pretends to run on.""" + monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=cuda_stream)) + + +@pytest.mark.parametrize(("device", "expected"), [("cuda:1", 1), ("cuda", 0)]) +def test_cuda_device_id_parses_the_device_string(device, expected): + """The mapping device index is parsed from the string; a bare ``"cuda"`` parses to 0. + + The bare-``"cuda"`` case intentionally preserves pre-existing behavior even though Warp + resolves it to its current CUDA device -- see the TODO on ``_cuda_device_id``. + """ + assert ovrtx_mapping._cuda_device_id(device) == expected + + +def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatch): + """The unmap names the Warp stream that produced the data, so the commit cannot race the fill. + + An unmap without a CUDA sync performs no synchronization at all, so the assertion is on the + unmap's ``stream`` argument, not merely on the unmap happening. + """ + sentinel = object() + binding = _RecordingMappedBinding() + _patch_warp_device(monkeypatch, ordinal=1, cuda_stream=99) + monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: sentinel) + + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda:1", wp.mat44d) as array: + assert array is sentinel + + assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 1}] + assert binding.unmap_calls == [{"event": None, "stream": 99}] + + +def test_map_attribute_for_warp_writes_unmaps_when_the_fill_raises(monkeypatch): + """A failed fill must still release the mapping exactly once, with the same stream ordering. + + Skipping the unmap would leak the mapping to OVRTX's ``__del__`` safety net, which commits + fire-and-forget without any CUDA sync. + """ + binding = _RecordingMappedBinding() + _patch_warp_device(monkeypatch, ordinal=0, cuda_stream=7) + monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: object()) + + with pytest.raises(ValueError, match="fill failed"): + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda:0", wp.mat44d): + raise ValueError("fill failed") + + assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 0}] + assert binding.unmap_calls == [{"event": None, "stream": 7}] + + def test_ovrtx_cleanup_releases_only_the_given_render_data(): """``cleanup`` releases the render data's own buffers and leaves the renderer usable.