Index tiled image buffer in 3D instead of flattening to 1D

Addresses review feedback on #7500.

render() flattened each annotator's tiled buffer into a single 1D
Warp array, whose length is num_envs * tile pixels * channels. That
product can exceed the maximum size of a single Warp array
dimension, so large environment counts / camera resolutions failed
with:
  ValueError: Array shapes must not exceed the maximum representable
  value of a signed 32-bit integer, got 2621440000 in dimension 0

Refactor reshape_tiled_image to index the tiled buffer as a 3D array
(num_tiles_y * image_height, num_tiles_x * image_width, num_channels)
and update the overloads to ndim=3, so each dimension is bounded
individually rather than by their product. Single-channel annotators
return a 2D buffer, so the renderer adds the trailing channel
dimension before launching.

This removes the need for the preflight size check added earlier in
this PR, which is dropped along with the hardcoded Warp dimension
limit: per review, that limit is a Warp implementation detail rather
than a contract, and the check also assumed 4 elements per pixel for
every camera, wrongly rejecting valid single-channel configurations.

Verified against Warp 1.9.1 on GPU: the refactored kernel is
numerically identical to the previous implementation across RGBA,
depth, normals and ragged-grid cases; a (25600, 25600, 4) buffer --
the exact geometry from the reported failure -- still raises on
.flatten() but launches successfully through the 3D path.

Author: shauryadNv

source/isaaclab/changelog.d/shauryad-bug-fixes.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* Changed the ``reshape_tiled_image`` Warp kernel to index the tiled image buffer as a 3D array
+  of shape (num_tiles_y * image_height, num_tiles_x * image_width, num_channels) instead of a
+  flattened 1D array. This keeps every array dimension within Warp's per-dimension size limit, so
+  large environment counts and camera resolutions no longer overflow a single flattened dimension.

source/isaaclab/isaaclab/utils/warp/kernels.py

@@ -344,8 +344,14 @@ def reshape_tiled_image(
     is assumed to be tiled in the x and y directions. The output image is a batch of images with the
     specified height, width, and number of channels.
 
+    The tiled buffer is indexed as a 3D array rather than flattened to 1D so that the number of
+    cameras and the camera resolution are bounded per dimension instead of by their product. A
+    flattened view of a large tiled buffer can exceed the maximum size of a single Warp array
+    dimension, see https://nvidia.github.io/warp/stable/user_guide/limitations.html#arrays.
+
     Args:
-        tiled_image_buffer: The input image buffer. Shape is (height * width * num_channels * num_cameras,).
+        tiled_image_buffer: The input image buffer. Shape is
+            (num_tiles_y * image_height, num_tiles_x * image_width, num_channels).
         batched_image: The output image. Shape is (num_cameras, height, width, num_channels).
         image_width: The width of the image.
         image_height: The height of the image.
@@ -358,32 +364,29 @@ def reshape_tiled_image(
     # resolve the tile indices
     tile_x_id = camera_id % num_tiles_x
     tile_y_id = camera_id // num_tiles_x
-    # compute the start index of the pixel in the tiled image buffer
-    pixel_start = (
-        num_channels * num_tiles_x * image_width * (image_height * tile_y_id + height_id)
-        + num_channels * tile_x_id * image_width
-        + num_channels * width_id
-    )
+    # resolve the pixel position within the tiled image buffer
+    row = image_height * tile_y_id + height_id
+    col = image_width * tile_x_id + width_id
 
     # copy the pixel values into the batched image
     for i in range(num_channels):
-        batched_image[camera_id, height_id, width_id, i] = batched_image.dtype(tiled_image_buffer[pixel_start + i])
+        batched_image[camera_id, height_id, width_id, i] = batched_image.dtype(tiled_image_buffer[row, col, i])
 
 
 # uint32 -> int32 conversion is required for non-colored segmentation annotators
 wp.overload(
     reshape_tiled_image,
-    {"tiled_image_buffer": wp.array(dtype=wp.uint32), "batched_image": wp.array(dtype=wp.uint32, ndim=4)},
+    {"tiled_image_buffer": wp.array(dtype=wp.uint32, ndim=3), "batched_image": wp.array(dtype=wp.uint32, ndim=4)},
 )
 # uint8 is used for 4 channel annotators
 wp.overload(
     reshape_tiled_image,
-    {"tiled_image_buffer": wp.array(dtype=wp.uint8), "batched_image": wp.array(dtype=wp.uint8, ndim=4)},
+    {"tiled_image_buffer": wp.array(dtype=wp.uint8, ndim=3), "batched_image": wp.array(dtype=wp.uint8, ndim=4)},
 )
 # float32 is used for single channel annotators
 wp.overload(
     reshape_tiled_image,
-    {"tiled_image_buffer": wp.array(dtype=wp.float32), "batched_image": wp.array(dtype=wp.float32, ndim=4)},
+    {"tiled_image_buffer": wp.array(dtype=wp.float32, ndim=3), "batched_image": wp.array(dtype=wp.float32, ndim=4)},
 )
 
 ##

source/isaaclab_physx/changelog.d/shauryad-bug-fixes.rst

@@ -4,7 +4,10 @@ Fixed
 * Fixed ``IsaacRtxRenderer.render()`` crashing with ``RuntimeError: Invalid indexing in slice``
   when an annotator's channel buffer has not warmed up yet (e.g. right after attach, at env
   creation) for the ``motion_vectors``, ``normals``, simple-shading, and RGB HDR data types.
-* Added a pre-flight check in ``IsaacRtxRenderer.create_render_data()`` that raises a clear,
-  actionable error when the requested ``num_envs``/camera resolution would overflow Warp's
-  signed-32-bit-representable array shape limit, instead of an opaque ``ValueError`` raised deep
-  inside ``render()``.
+
+Changed
+^^^^^^^
+
+* Changed ``IsaacRtxRenderer.render()`` to pass the tiled annotator buffer to
+  ``reshape_tiled_image`` as a 3D array instead of flattening it to 1D. Large environment counts
+  and camera resolutions no longer overflow the maximum size of a single Warp array dimension.

source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py

@@ -83,32 +83,6 @@ def _raise_missing_ppisp_error(exc: ModuleNotFoundError) -> NoReturn:
 }
 SIMPLE_SHADING_MODE_SETTING = "/rtx/minimal/mode"
 
-# Warp arrays require every dimension to fit in a signed 32-bit int (see
-# ``warp.types.array.__getitem__``/``check_array_shape``). ``render()`` flattens the whole
-# tiled buffer (all environments' camera tiles concatenated) into one Warp array per data
-# type, so this bounds how large ``num_envs * tile pixels`` can get before that call raises.
-_WARP_MAX_ARRAY_DIM = 2_147_483_647
-
-
-def _validate_tiled_buffer_size(spec: CameraRenderSpec) -> None:
-    """Raise a clear error if this camera's tiled buffer would overflow Warp's array shape limit.
-
-    Worst case is 4 elements per pixel: RGBA/HDR annotators return 4 channels directly, and
-    colorized segmentation annotators reinterpret a uint32 id as 4 uint8 channels.
-    """
-    num_tile_cols = math.ceil(math.sqrt(spec.view_count))
-    num_tile_rows = math.ceil(spec.view_count / num_tile_cols)
-    tiled_pixels = (num_tile_cols * spec.cfg.width) * (num_tile_rows * spec.cfg.height)
-    max_elements_per_pixel = 4
-    if tiled_pixels * max_elements_per_pixel > _WARP_MAX_ARRAY_DIM:
-        raise ValueError(
-            f"Camera '{spec.cfg.prim_path}' would allocate a tiled render buffer of up to"
-            f" {tiled_pixels * max_elements_per_pixel} elements ({spec.view_count} environments tiled into a"
-            f" {num_tile_cols}x{num_tile_rows} grid at {spec.cfg.width}x{spec.cfg.height} resolution), which"
-            f" exceeds Warp's signed-32-bit-representable array shape limit ({_WARP_MAX_ARRAY_DIM})."
-            " Reduce the number of environments (--num_envs) or the camera resolution for this task."
-        )
-
 
 def _camera_semantic_filter_predicate(semantic_filter: str | list[str]) -> str:
     """Build the instance-mapping semantics predicate from :attr:`isaaclab.sensors.camera.CameraCfg.semantic_filter`.
@@ -276,10 +250,6 @@ def prepare_stage(self, stage: Usd.Stage, num_envs: int) -> None:
     def create_render_data(self, spec: CameraRenderSpec) -> IsaacRtxRenderData:
         """Create render product and annotators for the tiled camera.
         See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.create_render_data`."""
-        # Fail fast with an actionable error instead of a cryptic Warp ``ValueError`` raised deep
-        # inside render()'s ``.flatten()`` call once the annotators are already attached.
-        _validate_tiled_buffer_size(spec)
-
         import omni.replicator.core as rep
         from omni.syntheticdata import SyntheticData
         from pxr import UsdGeom
@@ -617,11 +587,20 @@ def tiling_grid_shape():
                 buf_wp = render_data._hdr_scratch_wp
             else:
                 buf_wp = output_data[data_type].warp
+
+            # ``reshape_tiled_image`` indexes the tiled buffer as (rows, cols, channels). Single-channel
+            # annotators (e.g. depth, non-colorized segmentation) hand back a 2D buffer, so add the
+            # trailing channel dimension. Passing the buffer in 3D rather than flattening it keeps every
+            # dimension well inside Warp's per-dimension array size limit, so large environment counts and
+            # camera resolutions no longer overflow a single flattened dimension.
+            if tiled_data_buffer.ndim == 2:
+                tiled_data_buffer = tiled_data_buffer.reshape((*tiled_data_buffer.shape, 1))
+
             wp.launch(
                 kernel=reshape_tiled_image,
                 dim=(view_count, cfg.height, cfg.width),
                 inputs=[
-                    tiled_data_buffer.flatten(),
+                    tiled_data_buffer,
                     buf_wp,
                     *list(buf_wp.shape[1:]),
                     num_tiles_x,