zoey/raindrop
1
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

move to GLSL (rip rust-gpu 😭)

Browse source
This commit is contained in:
zack 2025-01-12 16:25:52 -05:00
parent c1df7bf365
commit d4623ab21f
No known key found for this signature in database
GPG key ID: EE8A2B709E2401D1
14 changed files with 457 additions and 486 deletions
Download patch file Download diff file Expand all files Collapse all files

14
crates/shaders-shared/src/lib.rs
View file

@ -2,13 +2,23 @@
use spirv_std::glam::{Mat4, Vec3, Vec4};
#[repr(C)]
#[derive(Clone, Copy, Debug)]
#[repr(C, align(16))]
pub struct Material {
pub base_color: Vec4,
pub metallic_factor: f32,
pub roughness_factor: f32,
pub _padding: [f32; 2],
}
#[repr(C, align(16))]
#[derive(Clone)]
pub struct UniformBufferObject {
pub model: Mat4,
pub view: Mat4,
pub proj: Mat4,
pub model_color: Vec3,
// pub camera_pos: Vec3,
// pub material: Material,
}
#[repr(C)]

140
crates/shaders/src/lib.rs
View file

@ -1,119 +1,69 @@
#![cfg_attr(target_arch = "spirv", no_std)]
use shaders_shared::UniformBufferObject;
use spirv_std::num_traits::Float;
use spirv_std::{
glam::{Mat3, UVec2, Vec2, Vec3, Vec4, Vec4Swizzles},
image::Image2d,
spirv, Sampler,
};
pub const PI: f32 = core::f32::consts::PI;
fn fresnel_schlick(cos_theta: f32, f0: Vec3) -> Vec3 {
f0 + (Vec3::ONE - f0) * (1.0 - cos_theta).powf(5.0)
}
fn distribution_ggx(n: Vec3, h: Vec3, roughness: f32) -> f32 {
let a = roughness * roughness;
let a2 = a * a;
let n_dot_h = n.dot(h).max(0.0);
let n_dot_h2 = n_dot_h * n_dot_h;
let nom = a2;
let denom = (n_dot_h2 * (a2 - 1.0) + 1.0);
let denom = PI * denom * denom;
nom / denom.max(0.001)
}
fn geometry_schlick_ggx(n_dot_v: f32, roughness: f32) -> f32 {
let r = roughness + 1.0;
let k = (r * r) / 8.0;
let nom = n_dot_v;
let denom = n_dot_v * (1.0 - k) + k;
nom / denom
}
fn geometry_smith(n: Vec3, v: Vec3, l: Vec3, roughness: f32) -> f32 {
let n_dot_v = n.dot(v).max(0.0);
let n_dot_l = n.dot(l).max(0.0);
let ggx2 = geometry_schlick_ggx(n_dot_v, roughness);
let ggx1 = geometry_schlick_ggx(n_dot_l, roughness);
ggx1 * ggx2
}
#[spirv(vertex)]
pub fn main_vs(
// Vertex inputs
in_pos: Vec3,
#[spirv(position)] in_pos: Vec3,
in_normal: Vec3,
in_tex_coord: Vec2,
// Uniform buffer
#[spirv(uniform, descriptor_set = 0, binding = 0)] ubo: &UniformBufferObject,
// Vertex outputs
out_world_position: &mut Vec3,
out_world_pos: &mut Vec3,
out_world_normal: &mut Vec3,
out_tex_coord: &mut Vec2,
#[spirv(position)] gl_position: &mut Vec4,
#[spirv(position)] out_pos: &mut Vec4,
) {
// Transform position to world space
let pos = ubo.model * Vec4::from((in_pos, 1.0));
*out_world_position = (pos / pos.w).xyz();
*out_world_pos = pos.truncate();
// Transform normal to world space
let normal_matrix = Mat3::from_mat4(ubo.model).inverse().transpose();
*out_world_normal = normal_matrix * in_normal;
*out_world_normal = (normal_matrix * in_normal).normalize();
// Calculate clip space position
*gl_position = ubo.proj * ubo.view * pos;
*out_pos = ubo.proj * ubo.view * pos;
*out_tex_coord = in_tex_coord;
}
pub fn get_uv_u(pix: UVec2, tex_size: Vec4) -> spirv_std::glam::Vec2 {
(pix.as_vec2() + Vec2::splat(0.5)) * tex_size.zw()
}
fn ray_triangle_intersect(
ray_origin: Vec3,
ray_direction: Vec3,
v0: Vec3,
v1: Vec3,
v2: Vec3,
) -> f32 {
let epsilon = 0.000001;
let edge1 = v1 - v0;
let edge2 = v2 - v0;
let h = ray_direction.cross(edge2);
let a = edge1.dot(h);
if a > -epsilon && a < epsilon {
return -1.0;
}
let f = 1.0 / a;
let s = ray_origin - v0;
let u = f * s.dot(h);
if !(0.0..=1.0).contains(&u) {
return -1.0;
}
let q = s.cross(edge1);
let v = f * ray_direction.dot(q);
if v < 0.0 || u + v > 1.0 {
return -1.0;
}
let t = f * edge2.dot(q);
if t > epsilon {
return t;
}
-1.0
}
fn material_color(
frag_world_position: Vec3,
frag_normal: Vec3,
light_pos: Vec3,
object_color: Vec3,
) -> Vec3 {
let l = (light_pos - frag_world_position).normalize();
let n = frag_normal.normalize();
let lambertian = f32::max(n.dot(l), 0.0);
object_color * lambertian
}
#[spirv(fragment)]
pub fn main_fs(
frag_world_position: Vec3,
frag_world_normal: Vec3,
frag_tex_coord: Vec2,
#[spirv(descriptor_set = 0, binding = 1)] texture: &Image2d,
#[spirv(descriptor_set = 0, binding = 2)] sampler: &Sampler,
out_color: &mut Vec4,
) {
// Convert fragment coordinate to UV coordinate
// Sample the texture
let base_color = texture.sample(*sampler, frag_tex_coord);
// let light_pos = Vec3::new(2.0, 2.0, -2.0);
//
// // Calculate light direction
// let l = (light_pos - frag_world_position).normalize();
// let n = frag_world_normal.normalize();
//
// // Calculate lambertian lighting
// let lambertian = f32::max(n.dot(l), 0.0);
//
// // Ambient lighting
// let ambient = Vec3::splat(0.2);
//
// // Combine texture color with model color
// let color_from_texture = Vec3::new(base_color.x, base_color.y, base_color.z);
// let combined_color = color_from_texture;
//
// // Final color calculation with gamma correction
// let color = (combined_color * lambertian + ambient).powf(2.2);
*out_color = Vec4::new(base_color.x, base_color.y, base_color.z, base_color.w);
}

4
crates/vk-rs/Cargo.toml
View file

@ -26,6 +26,8 @@ gltf = { version = "1.4.1", features = ["import"] }
image = "0.25.5"
rayon = "1.10.0"
bytemuck.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
[build-dependencies]
spirv-builder.workspace = true
shaderc = "0.8"

82
crates/vk-rs/build.rs
View file

@ -1,33 +1,71 @@
use shaderc::{Compiler, ShaderKind};
use std::{
fs::{self, File},
io::{Read, Write},
io::Write,
path::{Path, PathBuf},
};
use spirv_builder::{MetadataPrintout, SpirvBuilder};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Tell Cargo to rerun this script if the shaders crate or its contents change
println!("cargo:rerun-if-changed=../shaders/src");
println!("cargo:rerun-if-changed=../shaders/Cargo.toml");
// Tell Cargo to rerun if shaders directory changes
println!("cargo:rerun-if-changed=../../shaders");
SpirvBuilder::new("../shaders/", "spirv-unknown-vulkan1.2")
.print_metadata(MetadataPrintout::None)
.multimodule(true)
.build()?
.module
.unwrap_multi()
.iter()
.for_each(|(name, path)| {
let mut data = vec![];
File::open(path).unwrap().read_to_end(&mut data).unwrap();
let shader_dir = Path::new("../../shaders");
let cache_dir = Path::new("../../shader-cache");
fs::create_dir_all("./shaders/").unwrap();
// Create shader cache directory if it doesn't exist
fs::create_dir_all(cache_dir)?;
File::create(format!("./shaders/{name}.spv"))
.unwrap()
.write_all(&data)
.unwrap();
});
let compiler = Compiler::new().expect("Failed to create shader compiler");
// Compile all .vert and .frag files
for entry in fs::read_dir(shader_dir)? {
let entry = entry?;
let path = entry.path();
if let Some(extension) = path.extension() {
let kind = match extension.to_str() {
Some("vert") => ShaderKind::Vertex,
Some("frag") => ShaderKind::Fragment,
_ => continue,
};
let source = fs::read_to_string(&path)?;
let file_name = path.file_name().unwrap().to_str().unwrap();
// Create output path
let spirv_path = cache_dir.join(format!("{}.spv", file_name));
// Check if we need to recompile
if should_compile(&path, &spirv_path) {
println!("Compiling shader: {}", file_name);
let compiled =
compiler.compile_into_spirv(&source, kind, file_name, "main", None)?;
let mut file = File::create(&spirv_path)?;
file.write_all(compiled.as_binary_u8())?;
}
}
}
Ok(())
}
fn should_compile(source_path: &Path, output_path: &PathBuf) -> bool {
// If output doesn't exist, we need to compile
if !output_path.exists() {
return true;
}
// Get modification times
let source_modified = fs::metadata(source_path)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
let output_modified = fs::metadata(output_path)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
// Compile if source is newer than output
source_modified > output_modified
}

BIN
crates/vk-rs/shaders/main_fs.spv
View file

Binary file not shown.

BIN
crates/vk-rs/shaders/main_vs.spv
View file

Binary file not shown.

17
crates/vk-rs/src/main.rs
View file

@ -276,21 +276,21 @@ impl MyAppCreator {
p_callback_data: *const vk::DebugUtilsMessengerCallbackDataEXT,
_p_user_data: *mut std::ffi::c_void,
) -> vk::Bool32 {
let severity = match message_severity {
vk::DebugUtilsMessageSeverityFlagsEXT::VERBOSE => "[VERBOSE]",
vk::DebugUtilsMessageSeverityFlagsEXT::WARNING => "[WARNING]",
vk::DebugUtilsMessageSeverityFlagsEXT::ERROR => "[ERROR]",
vk::DebugUtilsMessageSeverityFlagsEXT::INFO => "[INFO]",
_ => panic!("[UNKNOWN]"),
};
let types = match message_types {
vk::DebugUtilsMessageTypeFlagsEXT::GENERAL => "[GENERAL]",
vk::DebugUtilsMessageTypeFlagsEXT::PERFORMANCE => "[PERFORMANCE]",
vk::DebugUtilsMessageTypeFlagsEXT::VALIDATION => "[VALIDATION]",
_ => panic!("[UNKNOWN]"),
};
let message = std::ffi::CStr::from_ptr((*p_callback_data).p_message);
println!("[DEBUG]{}{}{:?}", severity, types, message);
match message_severity {
vk::DebugUtilsMessageSeverityFlagsEXT::VERBOSE => tracing::trace!("{types}{message:?}"),
vk::DebugUtilsMessageSeverityFlagsEXT::WARNING => tracing::warn!("{types}{message:?}"),
vk::DebugUtilsMessageSeverityFlagsEXT::ERROR => tracing::error!("{types}{message:?}"),
vk::DebugUtilsMessageSeverityFlagsEXT::INFO => tracing::info!("{types}{message:?}"),
_ => tracing::debug!("{types}{message:?}"),
};
vk::FALSE
}
@ -680,6 +680,7 @@ impl AppCreator<Arc<Mutex<Allocator>>> for MyAppCreator {
}
fn main() -> std::process::ExitCode {
tracing_subscriber::fmt().pretty().init();
puffin::set_scopes_on(true);
egui_ash::run(

209
crates/vk-rs/src/renderer.rs
View file

@ -11,6 +11,7 @@ use spirv_std::glam::{Mat4, Vec3, Vec4};
use std::{
ffi::CString,
mem::ManuallyDrop,
panic,
path::Path,
sync::{Arc, Mutex},
time::Instant,
@ -33,6 +34,91 @@ macro_rules! include_spirv {
}};
}
pub struct DefaultTextures {
white: Texture, // Default albedo (white)
metallic_roughness: Texture, // Default metallic-roughness (black metallic, 0.5 roughness)
normal: Texture, // Default normal map (flat normal)
sampler: vk::Sampler, // Common sampler for all textures
}
pub struct Material {
base_color: Vec4,
metallic_factor: f32,
roughness_factor: f32,
base_color_texture: Option<Arc<Texture>>,
// metallic_roughness_texture: Option<Arc<Texture>>,
// normal_texture: Option<Arc<Texture>>,
}
impl Material {
fn from_gltf(
material: &gltf::Material,
device: &Device,
allocator: Arc<Mutex<Allocator>>,
command_pool: vk::CommandPool,
queue: vk::Queue,
path: &Path,
texture_cache: &mut TextureCache,
buffers: &[gltf::buffer::Data],
) -> Self {
let pbr = material.pbr_metallic_roughness();
let base_color = pbr.base_color_factor();
let metallic_factor = pbr.metallic_factor();
let roughness_factor = pbr.roughness_factor();
let base_color_texture = pbr.base_color_texture().and_then(|tex| {
let key = format!("{:?}", tex.texture().source().source());
texture_cache.get_or_load_texture(key, || {
load_texture_from_gltf(
device,
allocator.clone(),
command_pool,
queue,
&tex.texture(),
buffers,
path,
)
})
});
// let metallic_roughness_texture = pbr.metallic_roughness_texture().and_then(|tex| {
// load_texture_from_gltf(
// device,
// allocator.clone(),
// command_pool,
// queue,
// &tex.texture(),
// buffers,
// path,
// )
// .map(Arc::new)
// });
//
// let normal_texture = material.normal_texture().and_then(|tex| {
// load_texture_from_gltf(
// device,
// allocator.clone(),
// command_pool,
// queue,
// &tex.texture(),
// buffers,
// path,
// )
// .map(Arc::new)
// });
Self {
base_color: Vec4::from(base_color),
metallic_factor,
roughness_factor,
base_color_texture,
// metallic_roughness_texture,
// normal_texture,
}
}
}
#[repr(C)]
#[derive(Debug, Clone)]
struct Vertex {
@ -94,7 +180,7 @@ impl Texture {
let buffer_size = data.len() as u64;
println!("DATA LEN: {}", buffer_size);
tracing::info!("DATA LEN: {}", buffer_size);
let staging_buffer = unsafe {
device
.create_buffer(
@ -348,12 +434,92 @@ impl Texture {
}
}
impl DefaultTextures {
fn new(
device: &Device,
allocator: Arc<Mutex<Allocator>>,
command_pool: vk::CommandPool,
queue: vk::Queue,
) -> Self {
// Create a 1x1 white texture for default albedo
let white_data = vec![255u8, 255, 255, 255];
let white = Texture::new(
device,
allocator.clone(),
command_pool,
queue,
1,
1,
&white_data,
);
// Create a 1x1 default metallic-roughness texture
// R: unused, G: roughness (0.5), B: metallic (0.0)
let metallic_roughness_data = vec![0u8, 128, 0, 255];
let metallic_roughness = Texture::new(
device,
allocator.clone(),
command_pool,
queue,
1,
1,
&metallic_roughness_data,
);
// Create a 1x1 default normal map (pointing up)
let normal_data = vec![128u8, 128, 255, 255];
let normal = Texture::new(device, allocator, command_pool, queue, 1, 1, &normal_data);
// Create a common sampler
let sampler_create_info = vk::SamplerCreateInfo::builder()
.mag_filter(vk::Filter::LINEAR)
.min_filter(vk::Filter::LINEAR)
.address_mode_u(vk::SamplerAddressMode::REPEAT)
.address_mode_v(vk::SamplerAddressMode::REPEAT)
.address_mode_w(vk::SamplerAddressMode::REPEAT)
.anisotropy_enable(true)
.max_anisotropy(16.0)
.border_color(vk::BorderColor::INT_OPAQUE_BLACK)
.unnormalized_coordinates(false)
.compare_enable(false)
.compare_op(vk::CompareOp::ALWAYS)
.mipmap_mode(vk::SamplerMipmapMode::LINEAR)
.mip_lod_bias(0.0)
.min_lod(0.0)
.max_lod(0.0);
let sampler = unsafe {
device
.create_sampler(&sampler_create_info, None)
.expect("Failed to create sampler")
};
Self {
white,
metallic_roughness,
normal,
sampler,
}
}
fn destroy(&mut self, device: &Device, allocator: &mut Allocator) {
unsafe {
self.white.destroy(device, allocator);
self.metallic_roughness.destroy(device, allocator);
self.normal.destroy(device, allocator);
device.destroy_sampler(self.sampler, None);
}
}
}
pub struct Mesh {
vertex_buffer: Buffer,
vertex_buffer_allocation: Option<Allocation>,
vertex_count: u32,
transform: Mat4,
texture: Option<Arc<Texture>>,
material: Material,
default_textures: DefaultTextures,
descriptor_sets: Vec<vk::DescriptorSet>,
}
@ -373,6 +539,7 @@ impl Model {
.expect("Failed to free memory");
}
self.texture_cache.cleanup(device, allocator);
mesh.default_textures.destroy(device, allocator);
}
}
}
@ -390,15 +557,16 @@ fn load_texture_from_gltf(
let img_data =
gltf::image::Data::from_source(texture.source().source(), Some(path), buffers).ok()?;
println!(
tracing::info!(
"Original texture dimensions: {}x{}",
img_data.width, img_data.height
img_data.width,
img_data.height
);
// Convert to RGB/RGBA
let pixels_rgba = if img_data.pixels.len() == (img_data.width * img_data.height * 3) as usize {
// Image is RGB, convert to RGBA
println!("Converting RGB to RGBA");
tracing::info!("Converting RGB to RGBA");
let mut rgba_data = Vec::with_capacity((img_data.width * img_data.height * 4) as usize);
for chunk in img_data.pixels.chunks_exact(3) {
rgba_data.extend_from_slice(&[chunk[0], chunk[1], chunk[2], 255]);
@ -482,6 +650,17 @@ fn process_node(
if let Some(mesh) = node.mesh() {
for primitive in mesh.primitives() {
let material = Material::from_gltf(
&primitive.material(),
device,
allocator.clone(),
command_pool,
queue,
path,
texture_cache,
buffers,
);
let texture = primitive
.material()
.pbr_metallic_roughness()
@ -545,6 +724,9 @@ fn process_node(
let (vertex_buffer, vertex_buffer_allocation) =
create_vertex_buffer(device, allocator.clone(), command_pool, queue, &vertices);
let default_textures =
DefaultTextures::new(device, allocator.clone(), command_pool, queue);
meshes.push(Mesh {
vertex_buffer,
vertex_buffer_allocation: Some(vertex_buffer_allocation),
@ -552,6 +734,8 @@ fn process_node(
// Store identity matrix, it is not used here anymore
transform: Mat4::IDENTITY,
texture,
material,
default_textures,
descriptor_sets: Vec::new(),
});
}
@ -1199,7 +1383,7 @@ impl RendererInner {
render_pass: vk::RenderPass,
) -> (vk::Pipeline, vk::PipelineLayout) {
let vertex_shader_module = {
let spirv = include_spirv!("../shaders/main_vs.spv");
let spirv = include_spirv!("../../../shader-cache/main.vert.spv");
let shader_module_create_info = vk::ShaderModuleCreateInfo::builder().code(&spirv);
unsafe {
device
@ -1208,7 +1392,7 @@ impl RendererInner {
}
};
let fragment_shader_module = {
let spirv = include_spirv!("../shaders/main_fs.spv");
let spirv = include_spirv!("../../../shader-cache/main.frag.spv");
let shader_module_create_info = vk::ShaderModuleCreateInfo::builder().code(&spirv);
unsafe {
device
@ -1216,8 +1400,8 @@ impl RendererInner {
.expect("Failed to create shader module")
}
};
let main_function_name_fs = CString::new("main_fs").unwrap();
let main_function_name_vs = CString::new("main_vs").unwrap();
let main_function_name_fs = CString::new("main").unwrap();
let main_function_name_vs = CString::new("main").unwrap();
let pipeline_shader_stages = [
vk::PipelineShaderStageCreateInfo::builder()
.stage(vk::ShaderStageFlags::VERTEX)
@ -1558,7 +1742,7 @@ impl RendererInner {
);
let t = Instant::now();
println!("loading model!");
tracing::info!("loading model!");
let mut model = Model::load(
&device,
allocator.clone(),
@ -1566,7 +1750,7 @@ impl RendererInner {
queue,
"./sponza/NewSponza_Main_glTF_003.gltf",
);
println!(
tracing::info!(
"loaded {} meshes in {:.2} seconds!",
model.meshes.len(),
t.elapsed().as_secs_f32()
@ -1839,7 +2023,10 @@ impl RendererInner {
0.1,
1000.0,
),
model_color: self.model_color,
// camera_pos: self.camera_position,
// metallic_factor: mesh.material.metallic_factor,
// roughness_factor: mesh.material.roughness_factor,
// base_color: mesh.material.base_color,
};
let ptr = self.uniform_buffer_allocations[self.current_frame]