Support any dtype for depth images (#2602)

Closes #2550
* #2550

### What
Support any `dtype` for depth images. Previously only `uint16` and
`float32` were supported.

This also unifies the textures used for depth cloud projection and for
displaying images ("rectangle" shader). This means less time spent
uploading data to the GPU, half the GPU memory used, and less room for
bugs and inconsistencies.

I also took an unrelated pass at reducing log spam when running with
`RUST_LOG=debug`.

Example of using `float16`:

![image](https://github.com/rerun-io/rerun/assets/1148717/579b61af-6554-4c8c-b894-81b134a34d5c)


### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/2602) (if
applicable)

- [PR Build Summary](https://build.rerun.io/pr/2602)
- [Docs
preview](https://rerun.io/preview/pr%3Aemilk%2Fgeneralize-depth-texture/docs)
- [Examples
preview](https://rerun.io/preview/pr%3Aemilk%2Fgeneralize-depth-texture/examples)

crates/re_analytics/src/sink_native.rs

@@ -85,7 +85,10 @@ impl PostHogSink {
             .collect::<Vec<_>>();
         let batch = PostHogBatch::from_events(&events);
 
-        re_log::debug!("{}", serde_json::to_string_pretty(&batch)?);
+        re_log::trace!(
+            "Sending analytics: {}",
+            serde_json::to_string_pretty(&batch)?
+        );
         self.agent.post(resolved_url).send_json(&batch)?;
         Ok(())
     }

crates/re_arrow_store/src/store_gc.rs

@@ -82,7 +82,7 @@ impl DataStore {
                 let num_bytes_to_drop = initial_num_bytes * p;
                 let target_num_bytes = initial_num_bytes - num_bytes_to_drop;
 
-                re_log::debug!(
+                re_log::trace!(
                     kind = "gc",
                     id = self.gc_id,
                     %target,
@@ -106,7 +106,7 @@ impl DataStore {
         let new_num_bytes =
             (stats_after.temporal.num_bytes + stats_after.metadata_registry.num_bytes) as f64;
 
-        re_log::debug!(
+        re_log::trace!(
             kind = "gc",
             id = self.gc_id,
             %target,

crates/re_data_store/src/entity_tree.rs

@@ -298,10 +298,12 @@ impl EntityTree {
             components: fields,
         } = self;
 
-        for (timeline, histogram) in &mut prefix_times.0 {
+        {
             re_tracing::profile_scope!("prefix_times");
-            if let Some(cutoff_time) = cutoff_times.get(timeline) {
-                histogram.remove(..cutoff_time.as_i64());
+            for (timeline, histogram) in &mut prefix_times.0 {
+                if let Some(cutoff_time) = cutoff_times.get(timeline) {
+                    histogram.remove(..cutoff_time.as_i64());
+                }
             }
         }
         {

crates/re_data_store/src/store_db.rs

@@ -345,7 +345,7 @@ impl StoreDb {
         let (drop_row_ids, stats_diff) = self.entity_db.data_store.gc(
             re_arrow_store::GarbageCollectionTarget::DropAtLeastFraction(fraction_to_purge as _),
         );
-        re_log::debug!(
+        re_log::trace!(
             num_row_ids_dropped = drop_row_ids.len(),
             size_bytes_dropped = re_format::format_bytes(stats_diff.total.num_bytes as _),
             "purged datastore"

crates/re_renderer/shader/depth_cloud.wgsl

@@ -13,8 +13,13 @@
 
 // ---
 
-/// Keep in sync with `DepthCloudInfoUBO` in `depth_cloud.rs`.
-///
+// Keep in sync with `DepthCloudInfoUBO` in `depth_cloud.rs`.
+
+// Which texture to read from?
+const SAMPLE_TYPE_FLOAT = 1u;
+const SAMPLE_TYPE_SINT  = 2u;
+const SAMPLE_TYPE_UINT  = 3u;
+
 /// Same for all draw-phases.
 struct DepthCloudInfo {
     /// The extrinsincs of the camera used for the projection.
@@ -43,6 +48,9 @@ struct DepthCloudInfo {
     /// Configures color mapping mode, see `colormap.wgsl`.
     colormap: u32,
 
+    /// Which texture sample to use
+    sample_type: u32,
+
     /// Changes between the opaque and outline draw-phases.
     radius_boost_in_ui_points: f32,
 };
@@ -51,7 +59,13 @@ struct DepthCloudInfo {
 var<uniform> depth_cloud_info: DepthCloudInfo;
 
 @group(1) @binding(1)
-var depth_texture: texture_2d<f32>;
+var texture_float: texture_2d<f32>;
+
+@group(1) @binding(2)
+var texture_sint: texture_2d<i32>;
+
+@group(1) @binding(3)
+var texture_uint: texture_2d<u32>;
 
 struct VertexOut {
     @builtin(position)
@@ -83,11 +97,24 @@ struct PointData {
 
 // Backprojects the depth texture using the intrinsics passed in the uniform buffer.
 fn compute_point_data(quad_idx: u32) -> PointData {
-    let wh = textureDimensions(depth_texture);
-    let texcoords = UVec2(quad_idx % wh.x, quad_idx / wh.x);
+    var texcoords: UVec2;
+    var texture_value = 0.0;
+    if depth_cloud_info.sample_type == SAMPLE_TYPE_FLOAT {
+        let wh = textureDimensions(texture_float);
+        texcoords = UVec2(quad_idx % wh.x, quad_idx / wh.x);
+        texture_value = textureLoad(texture_float, texcoords, 0).x;
+    } else if depth_cloud_info.sample_type == SAMPLE_TYPE_SINT {
+        let wh = textureDimensions(texture_sint);
+        texcoords = UVec2(quad_idx % wh.x, quad_idx / wh.x);
+        texture_value = f32(textureLoad(texture_sint, texcoords, 0).x);
+    } else {
+        let wh = textureDimensions(texture_uint);
+        texcoords = UVec2(quad_idx % wh.x, quad_idx / wh.x);
+        texture_value = f32(textureLoad(texture_uint, texcoords, 0).x);
+    }
 
     // TODO(cmc): expose knobs to linearize/normalize/flip/cam-to-plane depth.
-    let world_space_depth = depth_cloud_info.world_depth_from_texture_value * textureLoad(depth_texture, texcoords, 0).x;
+    let world_space_depth = depth_cloud_info.world_depth_from_texture_value * texture_value;
 
     var data: PointData;
 

crates/re_renderer/shader/rectangle.wgsl

@@ -62,21 +62,18 @@ struct UniformBuffer {
 var<uniform> rect_info: UniformBuffer;
 
 @group(1) @binding(1)
-var texture_sampler: sampler;
-
-@group(1) @binding(2)
 var texture_float: texture_2d<f32>;
 
-@group(1) @binding(3)
+@group(1) @binding(2)
 var texture_sint: texture_2d<i32>;
 
-@group(1) @binding(4)
+@group(1) @binding(3)
 var texture_uint: texture_2d<u32>;
 
-@group(1) @binding(5)
+@group(1) @binding(4)
 var colormap_texture: texture_2d<f32>;
 
-@group(1) @binding(6)
+@group(1) @binding(5)
 var texture_float_filterable: texture_2d<f32>;
 
 struct VertexOut {

crates/re_renderer/src/renderer/depth_cloud.rs

@@ -11,6 +11,7 @@
 //! The vertex shader backprojects the depth texture using the user-specified intrinsics, and then
 //! behaves pretty much exactly like our point cloud renderer (see [`point_cloud.rs`]).
 
+use itertools::Itertools;
 use smallvec::smallvec;
 
 use crate::{
@@ -36,7 +37,15 @@ use super::{
 mod gpu_data {
     use crate::{wgpu_buffer_types, PickingLayerObjectId};
 
-    /// Keep in sync with mirror in `depth_cloud.wgsl.`
+    use super::DepthCloudDrawDataError;
+
+    // Keep in sync with mirror in `depth_cloud.wgsl.`
+
+    // Which texture to read from?
+    const SAMPLE_TYPE_FLOAT: u32 = 1;
+    const SAMPLE_TYPE_SINT: u32 = 2;
+    const SAMPLE_TYPE_UINT: u32 = 3;
+
     #[repr(C, align(256))]
     #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
     pub struct DepthCloudInfoUBO {
@@ -48,6 +57,7 @@ mod gpu_data {
         pub outline_mask_id: wgpu_buffer_types::UVec2,
         pub picking_layer_object_id: PickingLayerObjectId,
 
+        // ---
         /// Multiplier to get world-space depth from whatever is in the texture.
         pub world_depth_from_texture_depth: f32,
 
@@ -60,42 +70,63 @@ mod gpu_data {
         /// Which colormap should be used.
         pub colormap: u32,
 
+        // ---
+        /// One of `SAMPLE_TYPE_*`.
+        pub sample_type: u32,
+
         /// Changes over different draw-phases.
-        pub radius_boost_in_ui_points: wgpu_buffer_types::F32RowPadded,
+        pub radius_boost_in_ui_points: f32,
 
-        pub end_padding: [wgpu_buffer_types::PaddingRow; 16 - 4 - 3 - 1 - 1 - 1],
+        pub _row_padding: [f32; 2],
+
+        // ---
+        pub _end_padding: [wgpu_buffer_types::PaddingRow; 16 - 4 - 3 - 1 - 1 - 1],
     }
 
     impl DepthCloudInfoUBO {
         pub fn from_depth_cloud(
             radius_boost_in_ui_points: f32,
             depth_cloud: &super::DepthCloud,
-        ) -> Self {
+        ) -> Result<Self, DepthCloudDrawDataError> {
             let super::DepthCloud {
                 world_from_obj,
                 depth_camera_intrinsics,
                 world_depth_from_texture_depth,
                 point_radius_from_world_depth,
                 max_depth_in_world,
                 depth_dimensions: _,
-                depth_texture: _,
+                depth_texture,
                 colormap,
                 outline_mask_id,
                 picking_object_id,
             } = depth_cloud;
 
-            Self {
+            let texture_format = depth_texture.texture.format();
+            let sample_type = match texture_format.sample_type(None) {
+                Some(wgpu::TextureSampleType::Float { .. }) => SAMPLE_TYPE_FLOAT,
+                Some(wgpu::TextureSampleType::Sint) => SAMPLE_TYPE_SINT,
+                Some(wgpu::TextureSampleType::Uint) => SAMPLE_TYPE_UINT,
+                _ => {
+                    return Err(DepthCloudDrawDataError::TextureFormatNotSupported(
+                        texture_format,
+                    ));
+                }
+            };
+
+            Ok(Self {
                 world_from_obj: (*world_from_obj).into(),
                 depth_camera_intrinsics: (*depth_camera_intrinsics).into(),
                 outline_mask_id: outline_mask_id.0.unwrap_or_default().into(),
                 world_depth_from_texture_depth: *world_depth_from_texture_depth,
                 point_radius_from_world_depth: *point_radius_from_world_depth,
                 max_depth_in_world: *max_depth_in_world,
                 colormap: *colormap as u32,
-                radius_boost_in_ui_points: radius_boost_in_ui_points.into(),
+                sample_type,
+                radius_boost_in_ui_points,
                 picking_layer_object_id: *picking_object_id,
-                end_padding: Default::default(),
-            }
+                _row_padding: Default::default(),
+                _end_padding: Default::default(),
+            })
         }
     }
 }
@@ -191,8 +222,8 @@ impl DrawData for DepthCloudDrawData {
 
 #[derive(thiserror::Error, Debug)]
 pub enum DepthCloudDrawDataError {
-    #[error("Depth texture format was {0:?}, only formats with sample type float are supported")]
-    InvalidDepthTextureFormat(wgpu::TextureFormat),
+    #[error("Texture format not supported: {0:?} - use float or integer textures instead.")]
+    TextureFormatNotSupported(wgpu::TextureFormat),
 
     #[error(transparent)]
     ResourceManagerError(#[from] ResourceManagerError),
@@ -227,37 +258,50 @@ impl DepthCloudDrawData {
             });
         }
 
-        let depth_cloud_ubo_binding_outlines = create_and_fill_uniform_buffer_batch(
-            ctx,
-            "depth_cloud_ubos".into(),
-            depth_clouds.iter().map(|dc| {
-                gpu_data::DepthCloudInfoUBO::from_depth_cloud(
-                    *radius_boost_in_ui_points_for_outlines,
-                    dc,
-                )
-            }),
-        );
-        let depth_cloud_ubo_binding_opaque = create_and_fill_uniform_buffer_batch(
-            ctx,
-            "depth_cloud_ubos".into(),
-            depth_clouds
+        let depth_cloud_ubo_binding_outlines = {
+            let radius_boost = *radius_boost_in_ui_points_for_outlines;
+            let ubos: Vec<gpu_data::DepthCloudInfoUBO> = depth_clouds
                 .iter()
-                .map(|dc| gpu_data::DepthCloudInfoUBO::from_depth_cloud(0.0, dc)),
-        );
+                .map(|dc| gpu_data::DepthCloudInfoUBO::from_depth_cloud(radius_boost, dc))
+                .try_collect()?;
+            create_and_fill_uniform_buffer_batch(ctx, "depth_cloud_ubos".into(), ubos.into_iter())
+        };
+
+        let depth_cloud_ubo_binding_opaque = {
+            let ubos: Vec<gpu_data::DepthCloudInfoUBO> = depth_clouds
+                .iter()
+                .map(|dc| gpu_data::DepthCloudInfoUBO::from_depth_cloud(0.0, dc))
+                .try_collect()?;
+            create_and_fill_uniform_buffer_batch(ctx, "depth_cloud_ubos".into(), ubos.into_iter())
+        };
 
         let mut instances = Vec::with_capacity(depth_clouds.len());
         for (depth_cloud, ubo_outlines, ubo_opaque) in itertools::izip!(
             depth_clouds,
             depth_cloud_ubo_binding_outlines,
             depth_cloud_ubo_binding_opaque
         ) {
-            if !matches!(
-                depth_cloud.depth_texture.format().sample_type(None),
-                Some(wgpu::TextureSampleType::Float { filterable: _ })
-            ) {
-                return Err(DepthCloudDrawDataError::InvalidDepthTextureFormat(
-                    depth_cloud.depth_texture.format(),
-                ));
+            // We set up several texture sources, then instruct the shader to read from at most one of them.
+            let mut texture_float = ctx.texture_manager_2d.zeroed_texture_float().handle;
+            let mut texture_sint = ctx.texture_manager_2d.zeroed_texture_sint().handle;
+            let mut texture_uint = ctx.texture_manager_2d.zeroed_texture_uint().handle;
+
+            let texture_format = depth_cloud.depth_texture.format();
+            match texture_format.sample_type(None) {
+                Some(wgpu::TextureSampleType::Float { .. }) => {
+                    texture_float = depth_cloud.depth_texture.handle;
+                }
+                Some(wgpu::TextureSampleType::Sint) => {
+                    texture_sint = depth_cloud.depth_texture.handle;
+                }
+                Some(wgpu::TextureSampleType::Uint) => {
+                    texture_uint = depth_cloud.depth_texture.handle;
+                }
+                _ => {
+                    return Err(DepthCloudDrawDataError::TextureFormatNotSupported(
+                        texture_format,
+                    ));
+                }
             }
 
             let mk_bind_group = |label, ubo: BindGroupEntry| {
@@ -268,7 +312,9 @@ impl DepthCloudDrawData {
                         label,
                         entries: smallvec![
                             ubo,
-                            BindGroupEntry::DefaultTextureView(depth_cloud.depth_texture.handle),
+                            BindGroupEntry::DefaultTextureView(texture_float),
+                            BindGroupEntry::DefaultTextureView(texture_sint),
+                            BindGroupEntry::DefaultTextureView(texture_uint),
                         ],
                         layout: bg_layout,
                     },
@@ -334,6 +380,7 @@ impl Renderer for DepthCloudRenderer {
                         },
                         count: None,
                     },
+                    // float texture:
                     wgpu::BindGroupLayoutEntry {
                         binding: 1,
                         visibility: wgpu::ShaderStages::VERTEX,
@@ -344,6 +391,28 @@ impl Renderer for DepthCloudRenderer {
                         },
                         count: None,
                     },
+                    // sint texture:
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 2,
+                        visibility: wgpu::ShaderStages::VERTEX,
+                        ty: wgpu::BindingType::Texture {
+                            sample_type: wgpu::TextureSampleType::Sint,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            multisampled: false,
+                        },
+                        count: None,
+                    },
+                    // uint texture:
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 3,
+                        visibility: wgpu::ShaderStages::VERTEX,
+                        ty: wgpu::BindingType::Texture {
+                            sample_type: wgpu::TextureSampleType::Uint,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            multisampled: false,
+                        },
+                        count: None,
+                    },
                 ],
             },
         );