メモ:wgpu-rs で bloom を実装する

これを wgpu-rs でやってみたときのメモ。例によって GitHub の master 版なのでコードはたぶんすぐに動かなくなる。

Bloomとは

要は光ってるところをぼかすやつです。↑のリンク先がわかりやすいけどあえて書くとこんな感じ。

f:id:yutannihilation:20200809201315p:plain:w400

ということで、3つの描画のステージそれぞれに別々の render pipeline が必要になる。

  1. 明るいとこだけ抽出する
  2. ぼかす
  3. 合成する

明るいところだけ抽出する

render pipeline

この render pipeline はシンプル。出力先が2つなのでcolor_statesが2つになるのを忘れずに。

let extract_render_pipeline_layout =
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: Borrowed(&[]),
        push_constant_ranges: Borrowed(&[]),
    });

let extract_render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
    ...
    color_states: Borrowed(&[
        wgpu::ColorStateDescriptor {
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            color_blend: wgpu::BlendDescriptor::REPLACE,
            alpha_blend: wgpu::BlendDescriptor::REPLACE,
            write_mask: wgpu::ColorWrite::ALL,
        },
        wgpu::ColorStateDescriptor {
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            color_blend: wgpu::BlendDescriptor::REPLACE,
            alpha_blend: wgpu::BlendDescriptor::REPLACE,
            write_mask: wgpu::ColorWrite::ALL,
        }
    ]),
})

Shader

vertex shader は任意なので省略。fragment shader は、通常の出力先の色(ここでは f_color)を計算したあと、 輝度を計算してそれが1.0を超えている部分を2つ目の出力先に。ただ、これは lighting があるときの話で、手元でやったのは 2 次元だったので適切な閾値がわからなかった。とりあえず 0.2 とかにしてる。

#version 450

layout(location = 0) out vec4 f_color;
layout(location = 1) out vec4 b_color;

void main() {
    ...
    f_color = ...;

    float brightness = dot(f_color.rgb, vec3(0.2126, 0.7152, 0.0722));

    // TODO: needs to adjust to proper brightness
    if (brightness > 1.0)
        b_color = vec4(f_color.rgb, 1.0);
    else
        b_color = vec4(0.0, 0.0, 0.0, 1.0);
}

render pass

こんな感じになる。 staging_texture が通常の出力、 blur_texture_views がこのあと gaussina blur にかける部分の出力。 multisample_texture は MSAA 用のやつ。

// Texture views
let staging_texture_view = self.staging_texture.create_default_view();
let blur_texture_views = [
    self.blur_textures[0].create_default_view(),
    self.blur_textures[1].create_default_view(),
];
let multisample_texture_view = [
    self.multisample_texture[0].create_default_view(),
    self.multisample_texture[1].create_default_view(),
];

// Extract ------------------------------------------------------------------------------------------------------------------

{
    let mut extract_render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
        color_attachments: Borrowed(&[
            wgpu::RenderPassColorAttachmentDescriptor {
                attachment: &multisample_texture_view[0],
                resolve_target: Some(&staging_texture_view),
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: true,
                },
            },
            wgpu::RenderPassColorAttachmentDescriptor {
                attachment: &multisample_texture_view[1],
                resolve_target: Some(&blur_texture_views[0]),
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: true,
                },
            },
        ]),
        depth_stencil_attachment: None,
    });

    extract_render_pass.set_pipeline(&self.extract_render_pipeline);
    extract_render_pass.set_index_buffer(index_buffer.slice(..));
    extract_render_pass.set_vertex_buffer(0, vertex_buffer.slice(..));

    extract_render_pass.draw_indexed(0..(self.geometry.indices.len() as u32), 0, 0..1);
}

ぼかす

gaussian blurは、1回だけだとボケ具合が弱いので、2つ texture を用意して(仮に T1T2 とする)、T1 をぼかした結果を T2 に、その T2 をぼかした結果を T1 に、という具合に繰り返しかけていくことになる。 このとき、縦方向のぼかしと横方向のぼかしは分解可能なので、2方向に一気にぼかすのではなく縦横を切り替えながらぼかしていったほうが計算量が少なくて済む(n^2n * 2 か、という差)。two-pass gaussian blur というらしい。

f:id:yutannihilation:20200809204946p:plain:w400

uniform buffer

gaussian blur の方向を指定するために、 horizontal という bool を uniform buffer で渡す。のでこんなやつを定義する。

#[repr(C)]
#[derive(Debug, Copy, Clone)]
struct BlurUniforms {
    horizontal: bool,
}

impl BlurUniforms {
    fn new() -> Self {
        Self { horizontal: true }
    }

    fn flip(&mut self) {
        self.horizontal = !self.horizontal;
    }
}

unsafe impl bytemuck::Pod for BlurUniforms {}
unsafe impl bytemuck::Zeroable for BlurUniforms {}

render pipeline

ここは bind group が 2 つ必要になる

  • ぼかす元の texture
  • gaussian blur の方向を指定するためのパラメータの uniform buffer
let blur_bind_group_layout =
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        entries: Borrowed(&[
            wgpu::BindGroupLayoutEntry::new(
                0,
                wgpu::ShaderStage::FRAGMENT,
                wgpu::BindingType::SampledTexture {
                    multisampled: true,
                    dimension: wgpu::TextureViewDimension::D2,
                    component_type: wgpu::TextureComponentType::Float,
                },
            ),
            wgpu::BindGroupLayoutEntry::new(
                1,
                wgpu::ShaderStage::FRAGMENT,
                wgpu::BindingType::Sampler { comparison: false },
            ),
        ]),
        label: None,
    });

let blur_uniform_bind_group_layout =
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        entries: Borrowed(&[wgpu::BindGroupLayoutEntry::new(
            0,
            wgpu::ShaderStage::FRAGMENT,
            wgpu::BindingType::UniformBuffer {
                dynamic: false,
                min_binding_size: wgpu::BufferSize::new(std::mem::size_of::<bool>() as _),
            },
        )]),
        label: None,
    });

let blur_render_pipeline_layout =
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: Borrowed(&[
            &blur_bind_group_layout,
            &blur_uniform_bind_group_layout,
        ]),
        push_constant_ranges: Borrowed(&[]),
    });

let blur_render_pipeline = 
    ...

shader

vertex shader はこう。v_tex_coordsa_position を texture の座標([0, 1]が範囲)に直したもの。

#version 450

layout(location = 0) in vec2 a_position;

layout(location = 0) out vec2 v_tex_coords;

void main() {
    v_tex_coords = 0.5 + vec2(0.5, -0.5) * a_position;
    gl_Position = vec4(a_position, 0.0, 1.0);
}

fragment shader が実際の gaussian blur をやる部分。 ここではLearnOpenGLにあった例と同じく上下4ピクセルをつかってぼかしている。 おそらくこれも uniform buffer で外部から与えられると思うけど、計算の仕方がわからなかったのでいったん保留。

#version 450

layout(location = 0) in vec2 v_tex_coords;
layout(location = 0) out vec4 f_color;

layout(set = 0, binding = 0) uniform texture2D t_staging;
layout(set = 0, binding = 1) uniform sampler s_staging;

float weight[5] = float[] (0.227027, 0.1945946, 0.1216216, 0.054054, 0.016216);

layout(set = 1, binding = 0)
uniform Uniforms {
    bool horizontal;
};

void main() {
    vec2 tex_offset = 1.0 / textureSize(sampler2D(t_staging, s_staging), 0); // gets size of single texel
    f_color = texture(sampler2D(t_staging, s_staging), v_tex_coords) * weight[0]; // current fragment's contribution
    if (horizontal) {
        for(int i = 1; i < 5; ++i) {
            f_color += texture(sampler2D(t_staging, s_staging), v_tex_coords + vec2(tex_offset.x * i, 0.0)) * weight[i];
            f_color += texture(sampler2D(t_staging, s_staging), v_tex_coords - vec2(tex_offset.x * i, 0.0)) * weight[i];
        }
    } else {
        for(int i = 1; i < 5; ++i) {
            f_color += texture(sampler2D(t_staging, s_staging), v_tex_coords + vec2(0.0, tex_offset.y * i)) * weight[i];
            f_color += texture(sampler2D(t_staging, s_staging), v_tex_coords - vec2(0.0, tex_offset.y * i)) * weight[i];
        }
    }
}

render pass

Vertex は単純に texture を画面いっぱいに描けばいいので、三角形 2 つで四角形をつくる。

BlurVertex

#[repr(C)]
#[derive(Clone, Copy)]
struct BlurVertex {
    position: [f32; 2],
}

unsafe impl bytemuck::Pod for BlurVertex {}
unsafe impl bytemuck::Zeroable for BlurVertex {}

//   4-1
//   |/|
//   2-3
//
#[rustfmt::skip]
const VERTICES: &[BlurVertex] = &[
    BlurVertex { position: [ 1.0,  1.0], },
    BlurVertex { position: [-1.0, -1.0], },
    BlurVertex { position: [ 1.0, -1.0], },

    BlurVertex { position: [ 1.0,  1.0], },
    BlurVertex { position: [-1.0, -1.0], },
    BlurVertex { position: [-1.0,  1.0], },
];

こういう感じで blur_count 回繰り返す。最後の結果は、blur_textures[blur_count % 2] に入るので次のステップではこれを使う。

for i in 0..blur_count {
    let src_texture = &self.blur_textures[i % 2];
    let dst_texture = &self.blur_textures[(i + 1) % 2];

    // 単純なので省略
    let bind_group = ...;
    let blur_uniform_buffer = ...;
    let blur_uniform_bind_group = ...;

    // Flip the orientation between horizontally and vertically
    self.blur_uniform.flip();

    let resolve_target = dst_texture.create_default_view();
    {
        let mut blur_render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            color_attachments: Borrowed(&[wgpu::RenderPassColorAttachmentDescriptor {
                attachment: &multisample_texture_view[0],
                resolve_target: Some(&resolve_target),
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: true,
                },
            }]),
            depth_stencil_attachment: None,
        });

        blur_render_pass.set_pipeline(&self.blur_render_pipeline);
        blur_render_pass.set_bind_group(0, &bind_group, &[]);
        blur_render_pass.set_bind_group(1, &blur_uniform_bind_group, &[]);
        blur_render_pass.set_vertex_buffer(0, self.square_vertex.slice(..));

        blur_render_pass.draw(0..VERTICES.len() as u32, 0..1);
    }
}

合成する

uniform buffer

exposure (tone mapping)と gamma(gamma correction) をパラメータとして渡す。詳細は LearnOpenGL の HDR の章 に。

#[repr(C)]
#[derive(Debug, Copy, Clone)]
struct BlendUniforms {
    exposure: f32,
    gamma: f32,
}

impl BlendUniforms {
    fn new(exposure: f32, gamma: f32) -> Self {
        Self { exposure, gamma }
    }
}

unsafe impl bytemuck::Pod for BlendUniforms {}
unsafe impl bytemuck::Zeroable for BlendUniforms {}

render pipeline

ここも bind group が 2 つ必要になる。

  • オリジナルの texture とぼかされた texture
  • パラメータの uniform buffer
let blend_bind_group_layout =
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        entries: Borrowed(&[
            wgpu::BindGroupLayoutEntry::new(
                0,
                wgpu::ShaderStage::FRAGMENT,
                wgpu::BindingType::SampledTexture {
                    multisampled: true,
                    dimension: wgpu::TextureViewDimension::D2,
                    component_type: wgpu::TextureComponentType::Float,
                },
            ),
            wgpu::BindGroupLayoutEntry::new(
                1,
                wgpu::ShaderStage::FRAGMENT,
                wgpu::BindingType::SampledTexture {
                    multisampled: true,
                    dimension: wgpu::TextureViewDimension::D2,
                    component_type: wgpu::TextureComponentType::Float,
                },
            ),
            wgpu::BindGroupLayoutEntry::new(
                2,
                wgpu::ShaderStage::FRAGMENT,
                wgpu::BindingType::Sampler { comparison: false },
            ),
        ]),
        label: None,
    });

let blend_uniform_bind_group_layout =
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        entries: Borrowed(&[wgpu::BindGroupLayoutEntry::new(
            0,
            wgpu::ShaderStage::FRAGMENT,
            wgpu::BindingType::UniformBuffer {
                dynamic: false,
                min_binding_size: wgpu::BufferSize::new(
                    (std::mem::size_of::<f32>() * 2) as _,
                ),
            },
        )]),
        label: None,
    });

let blend_render_pipeline_layout =
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: Borrowed(&[
            &blend_bind_group_layout,
            &blend_uniform_bind_group_layout,
        ]),
        push_constant_ranges: Borrowed(&[]),
    });

let blend_render_pipeline = ...;

shader

vertex shader は、単純に texture を画面いっぱいに描けばいいので、blur のと同じのを使い回す。

fragment shader はこんな感じになる。これも LearnOpenGL の例 ほぼそのまま。

#version 450

layout(location = 0) in vec2 v_tex_coords;
layout(location = 0) out vec4 f_color;     // frame

layout(set = 0, binding = 0) uniform texture2D t_base;   // original texture
layout(set = 0, binding = 1) uniform texture2D t_blur;   // blur texture
layout(set = 0, binding = 2) uniform sampler s_base;

layout(set = 1, binding = 0)
uniform Uniforms {
    float exposure;
    float gamma;
};

void main() {
    vec3 hdr_color = texture(sampler2D(t_base, s_base), v_tex_coords).rgb;      
    vec3 blur_color = texture(sampler2D(t_blur, s_base), v_tex_coords).rgb;

    // additive blending
    hdr_color += blur_color;

    // converting HDR values to LDR values (tone mapping)
    // c.f. https://en.wikipedia.org/wiki/Tone_mapping, https://learnopengl.com/Advanced-Lighting/HDR
    vec3 result = vec3(1.0) - exp(-hdr_color * exposure);

    // Adjust the luminance (gamma correction)
    // c.f. https://en.wikipedia.org/wiki/Gamma_correction
    result = pow(result, vec3(1.0 / gamma));
    f_color = vec4(result, 1.0);
}

render pass

staging_texture_view がオリジナルの描画結果が入ってるやつ、blur_textures[blur_count % 2] がぼかされた結果が入ってるやつ。

let blend_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    layout: &self.blend_bind_group_layout,
    entries: Borrowed(&[
        wgpu::BindGroupEntry {
            binding: 0,
            // a texture that contains the unmodified version
            resource: wgpu::BindingResource::TextureView(&staging_texture_view),
        },
        wgpu::BindGroupEntry {
            binding: 1,
            // a texture that contains the last result of gaussian blur
            resource: wgpu::BindingResource::TextureView(
                &self.blur_textures[blur_count % 2].create_default_view(),
            ),
        },
        wgpu::BindGroupEntry {
            binding: 2,
            resource: wgpu::BindingResource::Sampler(&sampler),
        },
    ]),
    label: None,
});

{
    let mut blend_render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
        color_attachments: Borrowed(&[
            wgpu::RenderPassColorAttachmentDescriptor {
                attachment: &multisample_texture_view[0],
                resolve_target: Some(&frame.output.view),
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: true,
                },
            },
        ]),
        depth_stencil_attachment: None,
    });

    blend_render_pass.set_pipeline(&self.blend_render_pipeline);
    blend_render_pass.set_bind_group(0, &blend_bind_group, &[]);
    blend_render_pass.set_bind_group(1, &blend_uniform_bind_group, &[]);
    blend_render_pass.set_vertex_buffer(0, self.square_vertex.slice(..));

    blend_render_pass.draw(0..VERTICES.len() as u32, 0..1);
}

結果

github.com

youtu.be