Skip to content

Commit

Permalink
Premultiply the alpha on the GPU (#2190)
Browse files Browse the repository at this point in the history 
* Premultiply the alpha on the GPU

Premultiplying on the GPU saves us A LOT of cpu.

But texture filtering must happen AFTER premultiplying, which means
we need to do texture filtering in software on the shader.

* Add a TODO

* Cleanup

* Docfix
  • Loading branch information
@emilk
emilk authored May 24, 2023
1 parent 043b7a3 commit 1606471
Show file tree
Hide file tree
Showing 8 changed files with 87 additions and 97 deletions.
14 changes: 3 additions & 11 deletions crates/re_renderer/examples/2d.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,7 @@ impl framework::Example for Render2D {
let rerun_logo =
image::load_from_memory(include_bytes!("../../re_ui/data/logo_dark_mode.png")).unwrap();

let mut image_data = rerun_logo.as_rgba8().unwrap().to_vec();

// Premultiply alpha.
for color in image_data.chunks_exact_mut(4) {
color.clone_from_slice(
&ecolor::Color32::from_rgba_unmultiplied(color[0], color[1], color[2], color[3])
.to_array(),
);
}
let image_data = rerun_logo.as_rgba8().unwrap().to_vec();

let rerun_logo_texture = re_ctx
.texture_manager_2d
Expand Down Expand Up @@ -229,7 +221,7 @@ impl framework::Example for Render2D {
top_left_corner_position: glam::vec3(500.0, 120.0, -0.05),
extent_u: self.rerun_logo_texture_width as f32 * image_scale * glam::Vec3::X,
extent_v: self.rerun_logo_texture_height as f32 * image_scale * glam::Vec3::Y,
colormapped_texture: ColormappedTexture::from_unorm_srgba(
colormapped_texture: ColormappedTexture::from_unorm_rgba(
self.rerun_logo_texture.clone(),
),
options: RectangleOptions {
Expand All @@ -247,7 +239,7 @@ impl framework::Example for Render2D {
),
extent_u: self.rerun_logo_texture_width as f32 * image_scale * glam::Vec3::X,
extent_v: self.rerun_logo_texture_height as f32 * image_scale * glam::Vec3::Y,
colormapped_texture: ColormappedTexture::from_unorm_srgba(
colormapped_texture: ColormappedTexture::from_unorm_rgba(
self.rerun_logo_texture.clone(),
),
options: RectangleOptions {
Expand Down
2 changes: 1 addition & 1 deletion crates/re_renderer/examples/depth_cloud.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ impl framework::Example for RenderDepthClouds {
.transform_point3(glam::Vec3::new(1.0, 1.0, 0.0)),
extent_u: world_from_model.transform_vector3(-glam::Vec3::X),
extent_v: world_from_model.transform_vector3(-glam::Vec3::Y),
colormapped_texture: ColormappedTexture::from_unorm_srgba(albedo.texture.clone()),
colormapped_texture: ColormappedTexture::from_unorm_rgba(albedo.texture.clone()),
options: RectangleOptions {
texture_filter_magnification: re_renderer::renderer::TextureFilterMag::Nearest,
texture_filter_minification: re_renderer::renderer::TextureFilterMin::Linear,
Expand Down
10 changes: 6 additions & 4 deletions crates/re_renderer/shader/rectangle.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,9 @@
// Keep in sync with mirror in rectangle.rs

// Which texture to read from?
const SAMPLE_TYPE_FLOAT_FILTER = 1u;
const SAMPLE_TYPE_FLOAT_NOFILTER = 2u;
const SAMPLE_TYPE_SINT_NOFILTER = 3u;
const SAMPLE_TYPE_UINT_NOFILTER = 4u;
const SAMPLE_TYPE_FLOAT = 1u;
const SAMPLE_TYPE_SINT = 2u;
const SAMPLE_TYPE_UINT = 3u;

// How do we do colormapping?
const COLOR_MAPPER_OFF = 1u;
Expand Down Expand Up @@ -51,6 +50,9 @@ struct UniformBuffer {

minification_filter: u32,
magnification_filter: u32,

/// Boolean: decode 0-1 sRGB gamma to linear space before filtering?
decode_srgb: u32,
};

@group(1) @binding(0)
Expand Down
36 changes: 24 additions & 12 deletions crates/re_renderer/shader/rectangle_fs.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,29 +14,40 @@ fn tex_filter(pixel_coord: Vec2) -> u32 {
}
}

fn decode_color(rgba_arg: Vec4) -> Vec4 {
var rgba = rgba_arg;

// Convert to linear space:
if rect_info.decode_srgb != 0u {
rgba = linear_from_srgba(rgba);
}

// Premultiply alpha:
rgba = vec4(rgba.xyz * rgba.a, rgba.a);

return rgba;
}

@fragment
fn fs_main(in: VertexOut) -> @location(0) Vec4 {
// Sample the main texture:
var sampled_value: Vec4;
if rect_info.sample_type == SAMPLE_TYPE_FLOAT_FILTER {
// TODO(emilk): support mipmaps
sampled_value = textureSampleLevel(texture_float_filterable, texture_sampler, in.texcoord, 0.0);
} else if rect_info.sample_type == SAMPLE_TYPE_FLOAT_NOFILTER {
if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
let coord = in.texcoord * Vec2(textureDimensions(texture_float).xy);
if tex_filter(coord) == FILTER_NEAREST {
// nearest
sampled_value = textureLoad(texture_float, IVec2(coord + vec2(0.5)), 0);
sampled_value = decode_color(textureLoad(texture_float, IVec2(coord + vec2(0.5)), 0));
} else {
// bilinear
let v00 = textureLoad(texture_float, IVec2(coord) + IVec2(0, 0), 0);
let v01 = textureLoad(texture_float, IVec2(coord) + IVec2(0, 1), 0);
let v10 = textureLoad(texture_float, IVec2(coord) + IVec2(1, 0), 0);
let v11 = textureLoad(texture_float, IVec2(coord) + IVec2(1, 1), 0);
let v00 = decode_color(textureLoad(texture_float, IVec2(coord) + IVec2(0, 0), 0));
let v01 = decode_color(textureLoad(texture_float, IVec2(coord) + IVec2(0, 1), 0));
let v10 = decode_color(textureLoad(texture_float, IVec2(coord) + IVec2(1, 0), 0));
let v11 = decode_color(textureLoad(texture_float, IVec2(coord) + IVec2(1, 1), 0));
let top = mix(v00, v10, fract(coord.x));
let bottom = mix(v01, v11, fract(coord.x));
sampled_value = mix(top, bottom, fract(coord.y));
}
} else if rect_info.sample_type == SAMPLE_TYPE_SINT_NOFILTER {
} else if rect_info.sample_type == SAMPLE_TYPE_SINT {
let coord = in.texcoord * Vec2(textureDimensions(texture_sint).xy);
if tex_filter(coord) == FILTER_NEAREST {
// nearest
Expand All @@ -51,7 +62,8 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
let bottom = mix(v01, v11, fract(coord.x));
sampled_value = mix(top, bottom, fract(coord.y));
}
} else if rect_info.sample_type == SAMPLE_TYPE_UINT_NOFILTER {
} else if rect_info.sample_type == SAMPLE_TYPE_UINT {
// TODO(emilk): support premultiplying alpha on this path. Requires knowing the alpha range (255, 65535, …).
let coord = in.texcoord * Vec2(textureDimensions(texture_uint).xy);
if tex_filter(coord) == FILTER_NEAREST {
// nearest
Expand All @@ -75,7 +87,7 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
var normalized_value: Vec4 = (sampled_value - range.x) / (range.y - range.x);

// Apply gamma:
normalized_value = vec4(pow(normalized_value.rgb, vec3(rect_info.gamma)), normalized_value.a); // TODO(emilk): handle premultiplied alpha
normalized_value = vec4(pow(normalized_value.rgb, vec3(rect_info.gamma)), normalized_value.a);

// Apply colormap, if any:
var texture_color: Vec4;
Expand Down
76 changes: 35 additions & 41 deletions crates/re_renderer/src/renderer/rectangles.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ use crate::{
draw_phases::{DrawPhase, OutlineMaskProcessor},
include_shader_module,
resource_managers::{GpuTexture2D, ResourceManagerError},
texture_info,
view_builder::ViewBuilder,
wgpu_resources::{
BindGroupDesc, BindGroupEntry, BindGroupLayoutDesc, GpuBindGroup, GpuBindGroupLayoutHandle,
Expand Down Expand Up @@ -54,6 +53,11 @@ pub enum TextureFilterMin {
pub struct ColormappedTexture {
pub texture: GpuTexture2D,

/// Decode 0-1 sRGB gamma values to linear space before filtering?
///
/// Only applies to [`wgpu::TextureFormat::Rgba8Unorm`] and float textures.
pub decode_srgb: bool,

/// Min/max range of the values in the texture.
/// Used to normalize the input values (squash them to the 0-1 range).
pub range: [f32; 2],
Expand Down Expand Up @@ -89,9 +93,13 @@ pub enum ColorMapper {
}

impl ColormappedTexture {
pub fn from_unorm_srgba(texture: GpuTexture2D) -> Self {
/// Assumes a separate/unmultiplied alpha.
pub fn from_unorm_rgba(texture: GpuTexture2D) -> Self {
// If the texture is an sRGB texture, the GPU will decode it for us.
let decode_srgb = !texture.format().is_srgb();
Self {
texture,
decode_srgb,
range: [0.0, 1.0],
gamma: 1.0,
color_mapper: None,
Expand Down Expand Up @@ -166,6 +174,9 @@ pub enum RectangleError {

#[error("Invalid color map texture format: {0:?}")]
UnsupportedColormapTextureFormat(wgpu::TextureFormat),

#[error("decode_srgb set to true, but the texture was already sRGB aware")]
DoubleDecodingSrgbTexture,
}

mod gpu_data {
Expand All @@ -176,10 +187,9 @@ mod gpu_data {
// Keep in sync with mirror in rectangle.wgsl

// Which texture to read from?
const SAMPLE_TYPE_FLOAT_FILTER: u32 = 1;
const SAMPLE_TYPE_FLOAT_NOFILTER: u32 = 2;
const SAMPLE_TYPE_SINT_NOFILTER: u32 = 3;
const SAMPLE_TYPE_UINT_NOFILTER: u32 = 4;
const SAMPLE_TYPE_FLOAT: u32 = 1;
const SAMPLE_TYPE_SINT: u32 = 2;
const SAMPLE_TYPE_UINT: u32 = 3;

// How do we do colormapping?
const COLOR_MAPPER_OFF: u32 = 1;
Expand Down Expand Up @@ -213,16 +223,20 @@ mod gpu_data {
minification_filter: u32,
magnification_filter: u32,

_end_padding: [wgpu_buffer_types::PaddingRow; 16 - 6],
decode_srgb: u32,
_row_padding: [u32; 3],

_end_padding: [wgpu_buffer_types::PaddingRow; 16 - 7],
}

impl UniformBuffer {
pub fn from_textured_rect(
rectangle: &super::TexturedRect,
device_features: wgpu::Features,
) -> Result<Self, RectangleError> {
pub fn from_textured_rect(rectangle: &super::TexturedRect) -> Result<Self, RectangleError> {
let texture_format = rectangle.colormapped_texture.texture.format();

if texture_format.is_srgb() && rectangle.colormapped_texture.decode_srgb {
return Err(RectangleError::DoubleDecodingSrgbTexture);
}

let TexturedRect {
top_left_corner_position,
extent_u,
Expand All @@ -233,6 +247,7 @@ mod gpu_data {

let super::ColormappedTexture {
texture: _,
decode_srgb,
range,
gamma,
color_mapper,
Expand All @@ -247,15 +262,9 @@ mod gpu_data {
} = options;

let sample_type = match texture_format.sample_type(None) {
Some(wgpu::TextureSampleType::Float { .. }) => {
if texture_info::is_float_filterable(texture_format, device_features) {
SAMPLE_TYPE_FLOAT_FILTER
} else {
SAMPLE_TYPE_FLOAT_NOFILTER
}
}
Some(wgpu::TextureSampleType::Sint) => SAMPLE_TYPE_SINT_NOFILTER,
Some(wgpu::TextureSampleType::Uint) => SAMPLE_TYPE_UINT_NOFILTER,
Some(wgpu::TextureSampleType::Float { .. }) => SAMPLE_TYPE_FLOAT,
Some(wgpu::TextureSampleType::Sint) => SAMPLE_TYPE_SINT,
Some(wgpu::TextureSampleType::Uint) => SAMPLE_TYPE_UINT,
_ => {
return Err(RectangleError::DepthTexturesNotSupported);
}
Expand Down Expand Up @@ -312,6 +321,8 @@ mod gpu_data {
gamma: *gamma,
minification_filter,
magnification_filter,
decode_srgb: *decode_srgb as _,
_row_padding: Default::default(),
_end_padding: Default::default(),
})
}
Expand Down Expand Up @@ -357,7 +368,7 @@ impl RectangleDrawData {
// TODO(emilk): continue on error (skipping just that rectangle)?
let uniform_buffers: Vec<_> = rectangles
.iter()
.map(|rect| gpu_data::UniformBuffer::from_textured_rect(rect, ctx.device.features()))
.map(gpu_data::UniformBuffer::from_textured_rect)
.try_collect()?;

let uniform_buffer_bindings = create_and_fill_uniform_buffer_batch(
Expand Down Expand Up @@ -401,18 +412,13 @@ impl RectangleDrawData {
}

// We set up several texture sources, then instruct the shader to read from at most one of them.
let mut texture_float_filterable = ctx.texture_manager_2d.zeroed_texture_float().handle;
let mut texture_float_nofilter = ctx.texture_manager_2d.zeroed_texture_float().handle;
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;

match texture_format.sample_type(None) {
Some(wgpu::TextureSampleType::Float { .. }) => {
if texture_info::is_float_filterable(texture_format, ctx.device.features()) {
texture_float_filterable = texture.handle;
} else {
texture_float_nofilter = texture.handle;
}
texture_float = texture.handle;
}
Some(wgpu::TextureSampleType::Sint) => {
texture_sint = texture.handle;
Expand Down Expand Up @@ -447,11 +453,10 @@ impl RectangleDrawData {
entries: smallvec![
uniform_buffer,
BindGroupEntry::Sampler(sampler),
BindGroupEntry::DefaultTextureView(texture_float_nofilter),
BindGroupEntry::DefaultTextureView(texture_float),
BindGroupEntry::DefaultTextureView(texture_sint),
BindGroupEntry::DefaultTextureView(texture_uint),
BindGroupEntry::DefaultTextureView(colormap_texture),
BindGroupEntry::DefaultTextureView(texture_float_filterable),
],
layout: rectangle_renderer.bind_group_layout,
},
Expand Down Expand Up @@ -553,17 +558,6 @@ impl Renderer for RectangleRenderer {
},
count: None,
},
// float textures with filtering (e.g. Rgba8UnormSrgb):
wgpu::BindGroupLayoutEntry {
binding: 6,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
],
},
);
Expand Down
1 change: 1 addition & 0 deletions crates/re_viewer/src/ui/view_tensor/tensor_slice_to_gpu.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ pub fn colormapped_texture(

Ok(ColormappedTexture {
texture,
decode_srgb: false,
range,
gamma: color_mapping.gamma,
color_mapper: Some(re_renderer::renderer::ColorMapper::Function(
Expand Down
1 change: 1 addition & 0 deletions crates/re_viewer/src/ui/view_tensor/ui.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,7 @@ fn paint_colormap_gradient(

let colormapped_texture = re_renderer::renderer::ColormappedTexture {
texture: horizontal_gradient,
decode_srgb: false,
range: [0.0, 1.0],
gamma: 1.0,
color_mapper: Some(re_renderer::renderer::ColorMapper::Function(colormap)),
Expand Down
Loading

0 comments on commit 1606471

Please sign in to comment.