vk-graph 0.14.1+beta

mod profile_with_puffin;

use {
    ash::vk,
    clap::Parser,
    hassle_rs::compile_hlsl,
    log::info,
    std::{
        path::{Path, PathBuf},
        sync::Arc,
        time::Instant,
    },
    vk_graph::{
        Graph,
        cmd::{LoadOp, StoreOp},
        driver::{
            buffer::Buffer,
            device::Device,
            graphic::{GraphicPipeline, GraphicPipelineInfo},
            image::{Image, ImageInfo},
            shader::{SamplerInfo, Shader},
        },
        pool::hash::HashPool,
    },
    vk_graph_window::Window,
    vk_shader_macros::glsl,
    vk_sync::AccessType,
};

/// Displays a sequence of image samplers.
///
/// Note that manually specifying image samplers is completely optional, valid defaults will be used
/// if they are not specified when creating the shader which uses them. Additionally, you could
/// instead use use name suffixes such as _llr or _nne for linear/linear repeat or nearest/nearest
/// clamp-to-edge.
///
/// You may run this example program with either --hlsl or --separate arguments as follows:
///
/// cargo run --example image_sampler -- --hlsl --separate
///
/// Run with --help for more information.
///
/// See min_max.rs for more advanced image sampler usage.
fn main() -> anyhow::Result<()> {
    pretty_env_logger::init();
    profile_with_puffin::init();

    let args = Args::parse();
    let window = Window::builder().debug(args.debug).build()?;
    let gulf_image = read_image(&window.device, "examples/res/image/gulf.jpg")?;

    // Sampler info contains the full definition of Vulkan sampler settings using a builder struct
    let edge_edge = SamplerInfo::builder()
        .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE);
    let border_edge_black = SamplerInfo::builder()
        .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_BORDER)
        .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .border_color(vk::BorderColor::FLOAT_OPAQUE_BLACK);
    let edge_border_white = SamplerInfo::builder()
        .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_BORDER)
        .border_color(vk::BorderColor::FLOAT_OPAQUE_WHITE);

    // Image samplers are part of the shader pipeline and so we will create three pipelines total
    let pipeline_modes = [
        ("edge_edge", edge_edge),
        ("border_edge_black", border_edge_black),
        ("edge_border_white", edge_border_white),
    ];
    let pipelines = pipeline_modes
        .into_iter()
        .map(|(_, sampler_info)| create_pipeline(&window.device, sampler_info))
        .collect::<Result<Box<_>, _>>()?;
    let pipeline_names = pipeline_modes.map(|(name, _)| name);
    let mut pipeline_index = 0;
    let mut pipeline_time = 0.0;
    let mut prev_frame_at = Instant::now();

    info!("active sampler mode: {}", pipeline_names[pipeline_index]);

    window.run(|frame| {
        let now = Instant::now();

        let dt = now - prev_frame_at;
        prev_frame_at = now;

        // Periodically change the active pipeline index
        pipeline_time += dt.as_secs_f32();
        if pipeline_time > 5.0 {
            pipeline_time = 0.0;
            pipeline_index += 1;
            pipeline_index %= pipelines.len();
            info!("active sampler mode: {}", pipeline_names[pipeline_index]);
        }

        // Draw gulf.jpg using the active pipeline
        let gulf_image = frame.graph.bind_resource(&gulf_image);
        frame
            .graph
            .begin_cmd()
            .debug_name("Draw gulf image to swapchain")
            .bind_pipeline(&pipelines[pipeline_index])
            .shader_resource_access(
                0,
                gulf_image,
                AccessType::AnyShaderReadSampledImageOrUniformTexelBuffer,
            )
            .color_attachment_image(0, frame.swapchain_image, LoadOp::DontCare, StoreOp::Store)
            .record_cmd(|cmd| {
                cmd.draw(3, 1, 0, 0);
            });
    })?;

    Ok(())
}

fn create_pipeline(
    device: &Device,
    sampler_info: impl Into<SamplerInfo>,
) -> anyhow::Result<GraphicPipeline> {
    let args = Args::parse();

    let mut frag_shader = match (args.hlsl, args.separate) {
        (true, true) => {
            // HLSL separate image sampler
            Shader::new_fragment(
                compile_hlsl(
                    "fragment.hlsl",
                    r#"
                    struct FullscreenVertexOutput
                    {
                        float4 position : SV_Position;
                        [[vk::location(0)]] float2 uv : TEXCOORD0;
                    };

                    [[vk::binding(0, 0)]] Texture2D<float4> screenTexture : register(t0);
                    [[vk::binding(1, 0)]] SamplerState textureSampler : register(s1);

                    float4 main(FullscreenVertexOutput input)
                        : SV_Target
                    {
                        return screenTexture.Sample(textureSampler, input.uv);
                    }
                    "#,
                    "main",
                    "ps_5_0",
                    &["-spirv"],
                    &[],
                )?
                .as_slice(),
            )
        }
        (true, false) => {
            // HLSL combined image sampler: include_glsl uses shaderc which does not support this,
            // so we are using hassle_rs which uses dxc. You must follow the
            // instructions listed here to use hassle_rs:
            // See: https://github.com/Traverse-Research/hassle-rs
            // See: https://github.com/microsoft/DirectXShaderCompiler/wiki/Vulkan-combined-image-sampler-type
            // See: https://github.com/google/shaderc/issues/1310
            Shader::new_fragment(
                compile_hlsl(
                    "fragment.hlsl",
                    r#"
                    struct FullscreenVertexOutput
                    {
                        float4 position : SV_Position;
                        [[vk::location(0)]] float2 uv : TEXCOORD0;
                    };

                    [[vk::combinedImageSampler]][[vk::binding(0, 0)]]
                    Texture2D<float4> screenTexture : register(t0);
                    [[vk::combinedImageSampler]][[vk::binding(0, 0)]]
                    SamplerState textureSampler : register(s0);

                    float4 main(FullscreenVertexOutput input)
                        : SV_Target
                    {
                        return screenTexture.Sample(textureSampler, input.uv);
                    }
                    "#,
                    "main",
                    "ps_5_0",
                    &["-spirv"],
                    &[],
                )?
                .as_slice(),
            )
        }
        (false, true) => {
            // GLSL separate image sampler
            Shader::new_fragment(
                glsl!(
                    r#"
                    #version 460 core
                    #pragma shader_stage(fragment)

                    layout(binding = 0) uniform texture2D image;
                    layout(binding = 1) uniform sampler image_sampler;
                    layout(location = 0) in vec2 vk_TexCoord;
                    layout(location = 0) out vec4 vk_Color;

                    void main() {
                        vk_Color = texture(sampler2D(image, image_sampler), vk_TexCoord);
                    }
                    "#
                )
                .as_slice(),
            )
        }
        (false, false) => {
            // GLSL combined image sampler
            Shader::new_fragment(
                glsl!(
                    r#"
                    #version 460 core
                    #pragma shader_stage(fragment)

                    layout(binding = 0) uniform sampler2D image;
                    layout(location = 0) in vec2 vk_TexCoord;
                    layout(location = 0) out vec4 vk_Color;

                    void main() {
                        vk_Color = texture(image, vk_TexCoord);
                    }
                    "#
                )
                .as_slice(),
            )
        }
    };

    // Use the builder pattern to specify an image sampler at the combined binding index (0) or
    // separate binding index (1).
    let sampler_binding = args.separate as u32;
    frag_shader = frag_shader.image_sampler(sampler_binding, sampler_info);

    Ok(GraphicPipeline::create(
        device,
        GraphicPipelineInfo::default(),
        [
            Shader::new_vertex(
                glsl!(
                    r#"
                    #version 460 core
                    #pragma shader_stage(vertex)

                    const vec2[3] VERTICES = {
                        vec2(-1, -1),
                        vec2(-1,  3),
                        vec2( 3, -1),
                    };

                    layout(location = 0) out vec2 vk_TexCoord;

                    void main() {
                        gl_Position = vec4(VERTICES[gl_VertexIndex], 0, 1);
                        vk_TexCoord = 0.75 * gl_Position.xy + vec2(0.5);
                    }
                    "#
                )
                .as_slice(),
            ),
            frag_shader,
        ],
    )?)
}

fn read_image(device: &Device, path: impl AsRef<Path>) -> anyhow::Result<Arc<Image>> {
    // For another way to loading images, see vk_graph_fx::ImageLoader
    let gulf_jpg = image::open(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path))?;
    let image = Arc::new(Image::create(
        device,
        ImageInfo::image_2d(
            gulf_jpg.width(),
            gulf_jpg.height(),
            vk::Format::R8G8B8A8_UNORM,
            vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST,
        ),
    )?);

    {
        let mut graph = Graph::default();
        let image = graph.bind_resource(&image);
        let image_buf = graph.bind_resource(Buffer::create_from_slice(
            device,
            vk::BufferUsageFlags::TRANSFER_SRC,
            gulf_jpg.into_rgba8().into_vec().as_slice(),
        )?);
        graph.copy_buffer_to_image(image_buf, image);
        graph
            .into_submission()
            .queue_submit(&mut HashPool::new(device), 0, 0)?;

        // Note: There is no need to call wait_until_executed() here
    }

    Ok(image)
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Enable Vulkan SDK validation layers
    #[arg(long)]
    debug: bool,

    /// Use HLSL fragment shaders instead of the default (GLSL)
    #[arg(long)]
    hlsl: bool,

    /// Use separate image sampler objects instead of the default (combined image sampler objects)
    #[arg(long)]
    separate: bool,
}