Skip to content

Commit

Permalink
Code-golf the shader down some more to work around the crash again
Browse files Browse the repository at this point in the history
  • Loading branch information
@jleibs
jleibs committed Feb 6, 2024
1 parent f780e7e commit 22bb6ab
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 95 deletions.
50 changes: 25 additions & 25 deletions crates/re_renderer/shader/decodings.wgsl
View file
Original file line number Diff line number Diff line change
@@ -1,41 +1,41 @@
#import <./types.wgsl>


/// Loads an RGBA texel from a texture holding an NV12 encoded image at the given screen space coordinates.
fn decode_nv12(texture: texture_2d<u32>, coords: vec2i) -> vec4f {
/// Loads an RGBA texel from a texture holding an NV12 or YUY2 encoded image at the given screen space coordinates.
fn decode_nv12_or_yuy2(sample_type: u32, texture: texture_2d<u32>, coords: vec2i) -> vec4f {
let texture_dim = vec2f(textureDimensions(texture).xy);
let uv_offset = u32(floor(texture_dim.y / 1.5));
let uv_row = u32(coords.y / 2);
var uv_col = u32(coords.x / 2) * 2u;
var y: f32;
var u: f32;
var v: f32;

let y = f32(textureLoad(texture, vec2u(coords), 0).r);
let u = f32(textureLoad(texture, vec2u(u32(uv_col), uv_offset + uv_row), 0).r);
let v = f32(textureLoad(texture, vec2u((u32(uv_col) + 1u), uv_offset + uv_row), 0).r);
if sample_type == SAMPLE_TYPE_NV12 {
let uv_offset = u32(floor(texture_dim.y / 1.5));
let uv_row = u32(coords.y / 2);
var uv_col = u32(coords.x / 2) * 2u;

let rgb = set_color_standard(vec3f(y, u, v));

return vec4f(rgb, 1.0);
}

/// Loads an RGBA texel from a texture holding an YUY2 encoded image at the given screen space coordinates.
fn decode_yuy2(texture: texture_2d<u32>, coords: vec2i) -> vec4f {
// texture is 2 * width * height
// every 4 bytes is 2 pixels
let uv_row = u32(coords.y);
// multiply by 2 because the width is multiplied by 2
let y_col = u32(coords.x) * 2u;
let y = f32(textureLoad(texture, vec2u(y_col, uv_row), 0).r);
y = f32(textureLoad(texture, vec2u(coords), 0).r);
u = f32(textureLoad(texture, vec2u(u32(uv_col), uv_offset + uv_row), 0).r);
v = f32(textureLoad(texture, vec2u((u32(uv_col) + 1u), uv_offset + uv_row), 0).r);
} else if sample_type == SAMPLE_TYPE_YUY2 {
// texture is 2 * width * height
// every 4 bytes is 2 pixels
let uv_row = u32(coords.y);
// multiply by 2 because the width is multiplied by 2
let y_col = u32(coords.x) * 2u;
y = f32(textureLoad(texture, vec2u(y_col, uv_row), 0).r);

// at odd pixels we're in the second half of the yuyu block, offset back by 2
let uv_col = y_col - u32(coords.x % 2) * 2u;
let u = f32(textureLoad(texture, vec2u(uv_col + 1u, uv_row), 0).r);
let v = f32(textureLoad(texture, vec2u(uv_col + 3u, uv_row), 0).r);
// at odd pixels we're in the second half of the yuyu block, offset back by 2
let uv_col = y_col - u32(coords.x % 2) * 2u;
u = f32(textureLoad(texture, vec2u(uv_col + 1u, uv_row), 0).r);
v = f32(textureLoad(texture, vec2u(uv_col + 3u, uv_row), 0).r);
}

let rgb = set_color_standard(vec3f(y, u, v));

return vec4f(rgb, 1.0);
}


/// Sets the color standard for the given YUV color.
///
/// This conversion mirrors the function in `crates/re_types/src/datatypes/tensor_data_ext.rs`
Expand Down
128 changes: 58 additions & 70 deletions crates/re_renderer/shader/rectangle_fs.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,20 +51,28 @@ fn clamp_to_edge_nearest_neighbor(coord: vec2f, texture_dimension: vec2f) -> vec
return vec2i(clamp(floor(coord), vec2f(0.0), texture_dimension - vec2f(1.0)));
}

fn decode_color_and_filter_bilinear(coord: vec2f, v00: vec4f, v01: vec4f, v10: vec4f, v11: vec4f) -> vec4f {
fn decode_color_and_filter_nearest_or_bilinear(filter_nearest: bool, coord: vec2f, v00: vec4f, v01: vec4f, v10: vec4f, v11: vec4f) -> vec4f {
let c00 = decode_color(v00);
let c01 = decode_color(v01);
let c10 = decode_color(v10);
let c11 = decode_color(v11);
let top = mix(c00, c10, fract(coord.x - 0.5));
let bottom = mix(c01, c11, fract(coord.x - 0.5));
return mix(top, bottom, fract(coord.y - 0.5));
if filter_nearest {
return c00;
} else {
let c01 = decode_color(v01);
let c10 = decode_color(v10);
let c11 = decode_color(v11);
let top = mix(c00, c10, fract(coord.x - 0.5));
let bottom = mix(c01, c11, fract(coord.x - 0.5));
return mix(top, bottom, fract(coord.y - 0.5));
}
}

@fragment
fn fs_main(in: VertexOut) -> @location(0) vec4f {
// Sample the main texture:
var normalized_value: vec4f;
var v00_coord: vec2i;
var v01_coord: vec2i;
var v10_coord: vec2i;
var v11_coord: vec2i;

var texture_dimensions: vec2f;
if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
Expand All @@ -80,73 +88,53 @@ fn fs_main(in: VertexOut) -> @location(0) vec4f {
}

let coord = in.texcoord * texture_dimensions;
let clamped_coord = clamp_to_edge_nearest_neighbor(coord, texture_dimensions);
let v00_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(-0.5, -0.5), texture_dimensions);
let v01_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(-0.5, 0.5), texture_dimensions);
let v10_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(0.5, -0.5), texture_dimensions);
let v11_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(0.5, 0.5), texture_dimensions);

let filter_nearest = (tex_filter(coord) == FILTER_NEAREST);

if filter_nearest {
v00_coord = clamp_to_edge_nearest_neighbor(coord, texture_dimensions);
v01_coord = v00_coord;
v10_coord = v00_coord;
v11_coord = v00_coord;
} else {
v00_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(-0.5, -0.5), texture_dimensions);
v01_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(-0.5, 0.5), texture_dimensions);
v10_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(0.5, -0.5), texture_dimensions);
v11_coord = clamp_to_edge_nearest_neighbor(coord + vec2f(0.5, 0.5), texture_dimensions);
}

if rect_info.sample_type == SAMPLE_TYPE_FLOAT {
if tex_filter(coord) == FILTER_NEAREST {
// nearest
normalized_value = decode_color(textureLoad(texture_float,
clamp_to_edge_nearest_neighbor(coord, texture_dimensions), 0));
} else {
// bilinear
let v00 = textureLoad(texture_float, v00_coord, 0);
let v01 = textureLoad(texture_float, v01_coord, 0);
let v10 = textureLoad(texture_float, v10_coord, 0);
let v11 = textureLoad(texture_float, v11_coord, 0);
normalized_value = decode_color_and_filter_bilinear(coord, v00, v01, v10, v11);
}
normalized_value = decode_color_and_filter_nearest_or_bilinear(
filter_nearest,
coord,
textureLoad(texture_float, v00_coord, 0),
textureLoad(texture_float, v01_coord, 0),
textureLoad(texture_float, v10_coord, 0),
textureLoad(texture_float, v11_coord, 0));
} else if rect_info.sample_type == SAMPLE_TYPE_SINT {
if tex_filter(coord) == FILTER_NEAREST {
// nearest
normalized_value = decode_color(vec4f(textureLoad(texture_sint, clamped_coord, 0)));
} else {
// bilinear
let v00 = vec4f(textureLoad(texture_sint, v00_coord, 0));
let v01 = vec4f(textureLoad(texture_sint, v01_coord, 0));
let v10 = vec4f(textureLoad(texture_sint, v10_coord, 0));
let v11 = vec4f(textureLoad(texture_sint, v11_coord, 0));
normalized_value = decode_color_and_filter_bilinear(coord, v00, v01, v10, v11);
}
normalized_value = decode_color_and_filter_nearest_or_bilinear(
filter_nearest,
coord,
vec4f(textureLoad(texture_sint, v00_coord, 0)),
vec4f(textureLoad(texture_sint, v01_coord, 0)),
vec4f(textureLoad(texture_sint, v10_coord, 0)),
vec4f(textureLoad(texture_sint, v11_coord, 0)));
} else if rect_info.sample_type == SAMPLE_TYPE_UINT {
if tex_filter(coord) == FILTER_NEAREST {
// nearest
normalized_value = decode_color(vec4f(textureLoad(texture_uint, clamped_coord, 0)));
} else {
// bilinear
let v00 = vec4f(textureLoad(texture_uint, v00_coord, 0));
let v01 = vec4f(textureLoad(texture_uint, v01_coord, 0));
let v10 = vec4f(textureLoad(texture_uint, v10_coord, 0));
let v11 = vec4f(textureLoad(texture_uint, v11_coord, 0));
normalized_value = decode_color_and_filter_bilinear(coord, v00, v01, v10, v11);
}
} else if rect_info.sample_type == SAMPLE_TYPE_NV12 {
if tex_filter(coord) == FILTER_NEAREST {
// nearest
normalized_value = decode_color(vec4f(decode_nv12(texture_uint, clamped_coord)));
} else {
// bilinear
let v00 = decode_nv12(texture_uint, v00_coord);
let v01 = decode_nv12(texture_uint, v01_coord);
let v10 = decode_nv12(texture_uint, v10_coord);
let v11 = decode_nv12(texture_uint, v11_coord);
normalized_value = decode_color_and_filter_bilinear(coord, v00, v01, v10, v11);
}
} else if rect_info.sample_type == SAMPLE_TYPE_YUY2 {
if tex_filter(coord) == FILTER_NEAREST {
// nearest
normalized_value = decode_color(vec4f(decode_yuy2(texture_uint, clamped_coord)));
} else {
// bilinear
let v00 = decode_yuy2(texture_uint, v00_coord);
let v01 = decode_yuy2(texture_uint, v01_coord);
let v10 = decode_yuy2(texture_uint, v10_coord);
let v11 = decode_yuy2(texture_uint, v11_coord);
normalized_value = decode_color_and_filter_bilinear(coord, v00, v01, v10, v11);
}
normalized_value = decode_color_and_filter_nearest_or_bilinear(
filter_nearest,
coord,
vec4f(textureLoad(texture_uint, v00_coord, 0)),
vec4f(textureLoad(texture_uint, v01_coord, 0)),
vec4f(textureLoad(texture_uint, v10_coord, 0)),
vec4f(textureLoad(texture_uint, v11_coord, 0)));
} else if rect_info.sample_type == SAMPLE_TYPE_NV12 || rect_info.sample_type == SAMPLE_TYPE_YUY2{
normalized_value = decode_color_and_filter_nearest_or_bilinear(
filter_nearest,
coord,
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v00_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v01_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v10_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v11_coord));
} else {
return ERROR_RGBA; // unknown sample type
}
Expand Down

0 comments on commit 22bb6ab

Please sign in to comment.