ibl.rs (32202B)
1 use rayon::prelude::*; 2 pub struct IblData { 3 pub bindgroup: wgpu::BindGroup, 4 } 5 6 pub fn create_ibl_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout { 7 device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { 8 label: Some("ibl_bgl"), 9 entries: &[ 10 // binding 0: irradiance cubemap 11 wgpu::BindGroupLayoutEntry { 12 binding: 0, 13 visibility: wgpu::ShaderStages::FRAGMENT, 14 ty: wgpu::BindingType::Texture { 15 multisampled: false, 16 view_dimension: wgpu::TextureViewDimension::Cube, 17 sample_type: wgpu::TextureSampleType::Float { filterable: true }, 18 }, 19 count: None, 20 }, 21 // binding 1: sampler 22 wgpu::BindGroupLayoutEntry { 23 binding: 1, 24 visibility: wgpu::ShaderStages::FRAGMENT, 25 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), 26 count: None, 27 }, 28 // binding 2: pre-filtered environment cubemap (with mipmaps) 29 wgpu::BindGroupLayoutEntry { 30 binding: 2, 31 visibility: wgpu::ShaderStages::FRAGMENT, 32 ty: wgpu::BindingType::Texture { 33 multisampled: false, 34 view_dimension: wgpu::TextureViewDimension::Cube, 35 sample_type: wgpu::TextureSampleType::Float { filterable: true }, 36 }, 37 count: None, 38 }, 39 // binding 3: BRDF LUT (2D texture) 40 wgpu::BindGroupLayoutEntry { 41 binding: 3, 42 visibility: wgpu::ShaderStages::FRAGMENT, 43 ty: wgpu::BindingType::Texture { 44 multisampled: false, 45 view_dimension: wgpu::TextureViewDimension::D2, 46 sample_type: wgpu::TextureSampleType::Float { filterable: true }, 47 }, 48 count: None, 49 }, 50 ], 51 }) 52 } 53 54 /// Create IBL data with a procedural gradient cubemap for testing. 55 /// Replace this with a real irradiance map later. 56 #[allow(dead_code)] 57 pub fn create_test_ibl( 58 device: &wgpu::Device, 59 queue: &wgpu::Queue, 60 layout: &wgpu::BindGroupLayout, 61 ) -> IblData { 62 let size = 32u32; // small for testing 63 let irradiance_view = create_gradient_cubemap(device, queue, size); 64 65 // For test IBL, use the same gradient cubemap for prefiltered (not accurate but works) 66 let prefiltered_view = create_test_prefiltered_cubemap(device, queue, 64, 5); 67 68 // Generate BRDF LUT (this is environment-independent) 69 let brdf_lut_view = generate_brdf_lut(device, queue, 256); 70 71 let sampler = device.create_sampler(&wgpu::SamplerDescriptor { 72 label: Some("ibl_sampler"), 73 address_mode_u: wgpu::AddressMode::ClampToEdge, 74 address_mode_v: wgpu::AddressMode::ClampToEdge, 75 address_mode_w: wgpu::AddressMode::ClampToEdge, 76 mag_filter: wgpu::FilterMode::Linear, 77 min_filter: wgpu::FilterMode::Linear, 78 mipmap_filter: wgpu::FilterMode::Linear, 79 ..Default::default() 80 }); 81 82 let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { 83 label: Some("ibl_bg"), 84 layout, 85 entries: &[ 86 wgpu::BindGroupEntry { 87 binding: 0, 88 resource: wgpu::BindingResource::TextureView(&irradiance_view), 89 }, 90 wgpu::BindGroupEntry { 91 binding: 1, 92 resource: wgpu::BindingResource::Sampler(&sampler), 93 }, 94 wgpu::BindGroupEntry { 95 binding: 2, 96 resource: wgpu::BindingResource::TextureView(&prefiltered_view), 97 }, 98 wgpu::BindGroupEntry { 99 binding: 3, 100 resource: wgpu::BindingResource::TextureView(&brdf_lut_view), 101 }, 102 ], 103 }); 104 105 IblData { bindgroup } 106 } 107 108 /// Creates a simple gradient cubemap for testing IBL pipeline. 109 /// Sky-ish blue on top, ground-ish brown on bottom, neutral sides. 110 fn create_gradient_cubemap( 111 device: &wgpu::Device, 112 queue: &wgpu::Queue, 113 size: u32, 114 ) -> wgpu::TextureView { 115 let extent = wgpu::Extent3d { 116 width: size, 117 height: size, 118 depth_or_array_layers: 6, 119 }; 120 121 // Use Rgba16Float for HDR values > 1.0 122 let texture = device.create_texture(&wgpu::TextureDescriptor { 123 label: Some("irradiance_cubemap"), 124 size: extent, 125 mip_level_count: 1, 126 sample_count: 1, 127 dimension: wgpu::TextureDimension::D2, 128 format: wgpu::TextureFormat::Rgba16Float, 129 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, 130 view_formats: &[], 131 }); 132 133 // Face order: +X, -X, +Y, -Y, +Z, -Z 134 // HDR values - will be tonemapped in shader 135 let face_colors: [[f32; 3]; 6] = [ 136 [0.4, 0.38, 0.35], // +X (right) - warm neutral 137 [0.35, 0.38, 0.4], // -X (left) - cool neutral 138 [0.5, 0.6, 0.8], // +Y (up/sky) - blue sky 139 [0.25, 0.2, 0.15], // -Y (down/ground) - brown ground 140 [0.4, 0.4, 0.4], // +Z (front) - neutral 141 [0.38, 0.38, 0.42], // -Z (back) - slightly cool 142 ]; 143 144 let bytes_per_pixel = 8usize; // 4 x f16 = 8 bytes 145 let unpadded_row = size as usize * bytes_per_pixel; 146 let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; 147 let padded_row = unpadded_row.div_ceil(align) * align; 148 149 for (face_idx, color) in face_colors.iter().enumerate() { 150 let mut data = vec![0u8; padded_row * size as usize]; 151 152 for y in 0..size { 153 for x in 0..size { 154 let offset = (y as usize * padded_row) + (x as usize * bytes_per_pixel); 155 let r = half::f16::from_f32(color[0]); 156 let g = half::f16::from_f32(color[1]); 157 let b = half::f16::from_f32(color[2]); 158 let a = half::f16::from_f32(1.0); 159 160 data[offset..offset + 2].copy_from_slice(&r.to_le_bytes()); 161 data[offset + 2..offset + 4].copy_from_slice(&g.to_le_bytes()); 162 data[offset + 4..offset + 6].copy_from_slice(&b.to_le_bytes()); 163 data[offset + 6..offset + 8].copy_from_slice(&a.to_le_bytes()); 164 } 165 } 166 167 queue.write_texture( 168 wgpu::TexelCopyTextureInfo { 169 texture: &texture, 170 mip_level: 0, 171 origin: wgpu::Origin3d { 172 x: 0, 173 y: 0, 174 z: face_idx as u32, 175 }, 176 aspect: wgpu::TextureAspect::All, 177 }, 178 &data, 179 wgpu::TexelCopyBufferLayout { 180 offset: 0, 181 bytes_per_row: Some(padded_row as u32), 182 rows_per_image: Some(size), 183 }, 184 wgpu::Extent3d { 185 width: size, 186 height: size, 187 depth_or_array_layers: 1, 188 }, 189 ); 190 } 191 192 texture.create_view(&wgpu::TextureViewDescriptor { 193 label: Some("irradiance_cubemap_view"), 194 dimension: Some(wgpu::TextureViewDimension::Cube), 195 ..Default::default() 196 }) 197 } 198 199 /// Creates a simple test prefiltered cubemap with mip levels. 200 /// Uses solid colors that get darker with higher mip levels (simulating blur). 201 #[allow(dead_code)] 202 fn create_test_prefiltered_cubemap( 203 device: &wgpu::Device, 204 queue: &wgpu::Queue, 205 face_size: u32, 206 mip_count: u32, 207 ) -> wgpu::TextureView { 208 let texture = device.create_texture(&wgpu::TextureDescriptor { 209 label: Some("test_prefiltered_cubemap"), 210 size: wgpu::Extent3d { 211 width: face_size, 212 height: face_size, 213 depth_or_array_layers: 6, 214 }, 215 mip_level_count: mip_count, 216 sample_count: 1, 217 dimension: wgpu::TextureDimension::D2, 218 format: wgpu::TextureFormat::Rgba16Float, 219 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, 220 view_formats: &[], 221 }); 222 223 // Face colors (same as gradient cubemap) 224 let face_colors: [[f32; 3]; 6] = [ 225 [0.4, 0.38, 0.35], 226 [0.35, 0.38, 0.4], 227 [0.5, 0.6, 0.8], 228 [0.25, 0.2, 0.15], 229 [0.4, 0.4, 0.4], 230 [0.38, 0.38, 0.42], 231 ]; 232 233 for mip in 0..mip_count { 234 let mip_size = face_size >> mip; 235 let bytes_per_pixel = 8usize; 236 let unpadded_row = mip_size as usize * bytes_per_pixel; 237 let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; 238 let padded_row = unpadded_row.div_ceil(align) * align; 239 240 for (face_idx, color) in face_colors.iter().enumerate() { 241 let mut data = vec![0u8; padded_row * mip_size as usize]; 242 243 for y in 0..mip_size { 244 for x in 0..mip_size { 245 let offset = (y as usize * padded_row) + (x as usize * bytes_per_pixel); 246 let r = half::f16::from_f32(color[0]); 247 let g = half::f16::from_f32(color[1]); 248 let b = half::f16::from_f32(color[2]); 249 let a = half::f16::from_f32(1.0); 250 251 data[offset..offset + 2].copy_from_slice(&r.to_le_bytes()); 252 data[offset + 2..offset + 4].copy_from_slice(&g.to_le_bytes()); 253 data[offset + 4..offset + 6].copy_from_slice(&b.to_le_bytes()); 254 data[offset + 6..offset + 8].copy_from_slice(&a.to_le_bytes()); 255 } 256 } 257 258 queue.write_texture( 259 wgpu::TexelCopyTextureInfo { 260 texture: &texture, 261 mip_level: mip, 262 origin: wgpu::Origin3d { 263 x: 0, 264 y: 0, 265 z: face_idx as u32, 266 }, 267 aspect: wgpu::TextureAspect::All, 268 }, 269 &data, 270 wgpu::TexelCopyBufferLayout { 271 offset: 0, 272 bytes_per_row: Some(padded_row as u32), 273 rows_per_image: Some(mip_size), 274 }, 275 wgpu::Extent3d { 276 width: mip_size, 277 height: mip_size, 278 depth_or_array_layers: 1, 279 }, 280 ); 281 } 282 } 283 284 texture.create_view(&wgpu::TextureViewDescriptor { 285 label: Some("test_prefiltered_cubemap_view"), 286 dimension: Some(wgpu::TextureViewDimension::Cube), 287 ..Default::default() 288 }) 289 } 290 291 /// Load an HDR environment map from raw bytes (e.g. from `include_bytes!`). 292 pub fn load_hdr_ibl_from_bytes( 293 device: &wgpu::Device, 294 queue: &wgpu::Queue, 295 layout: &wgpu::BindGroupLayout, 296 bytes: &[u8], 297 ) -> Result<IblData, image::ImageError> { 298 let img = image::load_from_memory(bytes)?.into_rgb32f(); 299 load_hdr_ibl_from_image(device, queue, layout, img) 300 } 301 302 fn load_hdr_ibl_from_image( 303 device: &wgpu::Device, 304 queue: &wgpu::Queue, 305 layout: &wgpu::BindGroupLayout, 306 img: image::Rgb32FImage, 307 ) -> Result<IblData, image::ImageError> { 308 let width = img.width(); 309 let height = img.height(); 310 let pixels: Vec<_> = img.pixels().cloned().collect(); 311 312 // Convolve for diffuse irradiance (CPU-side, relatively slow but correct) 313 let irradiance_view = equirect_to_irradiance_cubemap(device, queue, &pixels, width, height, 32); 314 315 // Generate pre-filtered specular cubemap with mip chain 316 let prefiltered_view = 317 generate_prefiltered_cubemap(device, queue, &pixels, width, height, 128, 5); 318 319 // Generate BRDF integration LUT (environment-independent, could be cached) 320 let brdf_lut_view = generate_brdf_lut(device, queue, 256); 321 322 let sampler = device.create_sampler(&wgpu::SamplerDescriptor { 323 label: Some("ibl_sampler"), 324 address_mode_u: wgpu::AddressMode::ClampToEdge, 325 address_mode_v: wgpu::AddressMode::ClampToEdge, 326 address_mode_w: wgpu::AddressMode::ClampToEdge, 327 mag_filter: wgpu::FilterMode::Linear, 328 min_filter: wgpu::FilterMode::Linear, 329 mipmap_filter: wgpu::FilterMode::Linear, 330 ..Default::default() 331 }); 332 333 let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { 334 label: Some("ibl_bg"), 335 layout, 336 entries: &[ 337 wgpu::BindGroupEntry { 338 binding: 0, 339 resource: wgpu::BindingResource::TextureView(&irradiance_view), 340 }, 341 wgpu::BindGroupEntry { 342 binding: 1, 343 resource: wgpu::BindingResource::Sampler(&sampler), 344 }, 345 wgpu::BindGroupEntry { 346 binding: 2, 347 resource: wgpu::BindingResource::TextureView(&prefiltered_view), 348 }, 349 wgpu::BindGroupEntry { 350 binding: 3, 351 resource: wgpu::BindingResource::TextureView(&brdf_lut_view), 352 }, 353 ], 354 }); 355 356 Ok(IblData { bindgroup }) 357 } 358 359 /// Convert equirectangular panorama to irradiance cubemap (with hemisphere convolution). 360 fn equirect_to_irradiance_cubemap( 361 device: &wgpu::Device, 362 queue: &wgpu::Queue, 363 pixels: &[image::Rgb<f32>], 364 src_width: u32, 365 src_height: u32, 366 face_size: u32, 367 ) -> wgpu::TextureView { 368 let extent = wgpu::Extent3d { 369 width: face_size, 370 height: face_size, 371 depth_or_array_layers: 6, 372 }; 373 374 let texture = device.create_texture(&wgpu::TextureDescriptor { 375 label: Some("irradiance_cubemap"), 376 size: extent, 377 mip_level_count: 1, 378 sample_count: 1, 379 dimension: wgpu::TextureDimension::D2, 380 format: wgpu::TextureFormat::Rgba16Float, 381 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, 382 view_formats: &[], 383 }); 384 385 let bytes_per_pixel = 8usize; 386 let unpadded_row = face_size as usize * bytes_per_pixel; 387 let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; 388 let padded_row = unpadded_row.div_ceil(align) * align; 389 390 for face in 0..6 { 391 // Compute all pixels in parallel 392 let pixel_colors: Vec<[f32; 3]> = (0..face_size * face_size) 393 .into_par_iter() 394 .map(|idx| { 395 let x = idx % face_size; 396 let y = idx / face_size; 397 let u = (x as f32 + 0.5) / face_size as f32 * 2.0 - 1.0; 398 let v = (y as f32 + 0.5) / face_size as f32 * 2.0 - 1.0; 399 400 let dir = face_uv_to_direction(face, u, v); 401 let n = normalize(dir); 402 convolve_irradiance(pixels, src_width, src_height, n) 403 }) 404 .collect(); 405 406 // Write results to buffer 407 let mut data = vec![0u8; padded_row * face_size as usize]; 408 for (idx, color) in pixel_colors.iter().enumerate() { 409 let x = idx as u32 % face_size; 410 let y = idx as u32 / face_size; 411 let offset = (y as usize * padded_row) + (x as usize * bytes_per_pixel); 412 413 let r = half::f16::from_f32(color[0]); 414 let g = half::f16::from_f32(color[1]); 415 let b = half::f16::from_f32(color[2]); 416 let a = half::f16::from_f32(1.0); 417 418 data[offset..offset + 2].copy_from_slice(&r.to_le_bytes()); 419 data[offset + 2..offset + 4].copy_from_slice(&g.to_le_bytes()); 420 data[offset + 4..offset + 6].copy_from_slice(&b.to_le_bytes()); 421 data[offset + 6..offset + 8].copy_from_slice(&a.to_le_bytes()); 422 } 423 424 queue.write_texture( 425 wgpu::TexelCopyTextureInfo { 426 texture: &texture, 427 mip_level: 0, 428 origin: wgpu::Origin3d { 429 x: 0, 430 y: 0, 431 z: face, 432 }, 433 aspect: wgpu::TextureAspect::All, 434 }, 435 &data, 436 wgpu::TexelCopyBufferLayout { 437 offset: 0, 438 bytes_per_row: Some(padded_row as u32), 439 rows_per_image: Some(face_size), 440 }, 441 wgpu::Extent3d { 442 width: face_size, 443 height: face_size, 444 depth_or_array_layers: 1, 445 }, 446 ); 447 } 448 449 texture.create_view(&wgpu::TextureViewDescriptor { 450 label: Some("irradiance_cubemap_view"), 451 dimension: Some(wgpu::TextureViewDimension::Cube), 452 ..Default::default() 453 }) 454 } 455 456 fn normalize(v: [f32; 3]) -> [f32; 3] { 457 let len = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt(); 458 [v[0] / len, v[1] / len, v[2] / len] 459 } 460 461 /// Integrate the environment map over a hemisphere for diffuse irradiance. 462 /// Uses discrete sampling over the hemisphere. 463 fn convolve_irradiance( 464 pixels: &[image::Rgb<f32>], 465 width: u32, 466 height: u32, 467 normal: [f32; 3], 468 ) -> [f32; 3] { 469 let mut irradiance = [0.0f32; 3]; 470 471 // Build tangent frame from normal 472 let up = if normal[1].abs() < 0.999 { 473 [0.0, 1.0, 0.0] 474 } else { 475 [1.0, 0.0, 0.0] 476 }; 477 let tangent = normalize(cross(up, normal)); 478 let bitangent = cross(normal, tangent); 479 480 // Sample hemisphere with uniform spacing 481 let sample_delta = 0.05; // Adjust for quality vs speed 482 let mut n_samples = 0u32; 483 484 let mut phi = 0.0f32; 485 while phi < 2.0 * std::f32::consts::PI { 486 let mut theta = 0.0f32; 487 while theta < 0.5 * std::f32::consts::PI { 488 // Spherical to cartesian (in tangent space) 489 let sin_theta = theta.sin(); 490 let cos_theta = theta.cos(); 491 let sin_phi = phi.sin(); 492 let cos_phi = phi.cos(); 493 494 let tangent_sample = [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta]; 495 496 // Transform to world space 497 let sample_dir = [ 498 tangent_sample[0] * tangent[0] 499 + tangent_sample[1] * bitangent[0] 500 + tangent_sample[2] * normal[0], 501 tangent_sample[0] * tangent[1] 502 + tangent_sample[1] * bitangent[1] 503 + tangent_sample[2] * normal[1], 504 tangent_sample[0] * tangent[2] 505 + tangent_sample[1] * bitangent[2] 506 + tangent_sample[2] * normal[2], 507 ]; 508 509 let color = sample_equirect(pixels, width, height, sample_dir); 510 511 // Weight by cos(theta) * sin(theta) for hemisphere integration 512 let weight = cos_theta * sin_theta; 513 irradiance[0] += color[0] * weight; 514 irradiance[1] += color[1] * weight; 515 irradiance[2] += color[2] * weight; 516 n_samples += 1; 517 518 theta += sample_delta; 519 } 520 phi += sample_delta; 521 } 522 523 // Normalize 524 let scale = std::f32::consts::PI / n_samples as f32; 525 [ 526 irradiance[0] * scale, 527 irradiance[1] * scale, 528 irradiance[2] * scale, 529 ] 530 } 531 532 fn cross(a: [f32; 3], b: [f32; 3]) -> [f32; 3] { 533 [ 534 a[1] * b[2] - a[2] * b[1], 535 a[2] * b[0] - a[0] * b[2], 536 a[0] * b[1] - a[1] * b[0], 537 ] 538 } 539 540 /// Convert face index + UV to 3D direction. 541 /// Face order: +X, -X, +Y, -Y, +Z, -Z 542 fn face_uv_to_direction(face: u32, u: f32, v: f32) -> [f32; 3] { 543 match face { 544 0 => [1.0, -v, -u], // +X 545 1 => [-1.0, -v, u], // -X 546 2 => [u, 1.0, v], // +Y 547 3 => [u, -1.0, -v], // -Y 548 4 => [u, -v, 1.0], // +Z 549 5 => [-u, -v, -1.0], // -Z 550 _ => [0.0, 0.0, 1.0], 551 } 552 } 553 554 /// Sample equirectangular panorama given a 3D direction. 555 fn sample_equirect(pixels: &[image::Rgb<f32>], width: u32, height: u32, dir: [f32; 3]) -> [f32; 3] { 556 let len = (dir[0] * dir[0] + dir[1] * dir[1] + dir[2] * dir[2]).sqrt(); 557 let x = dir[0] / len; 558 let y = dir[1] / len; 559 let z = dir[2] / len; 560 561 // Convert to spherical (theta = azimuth, phi = elevation) 562 let theta = z.atan2(x); // -PI to PI 563 let phi = y.asin(); // -PI/2 to PI/2 564 565 // Convert to UV 566 let u = (theta / std::f32::consts::PI + 1.0) * 0.5; // 0 to 1 567 let v = (-phi / std::f32::consts::FRAC_PI_2 + 1.0) * 0.5; // 0 to 1 568 569 let px = ((u * width as f32) as u32).min(width - 1); 570 let py = ((v * height as f32) as u32).min(height - 1); 571 572 let idx = (py * width + px) as usize; 573 let p = &pixels[idx]; 574 [p.0[0], p.0[1], p.0[2]] 575 } 576 577 // ============================================================================ 578 // Specular IBL: Pre-filtered environment map and BRDF LUT 579 // ============================================================================ 580 581 /// Generate a 2D BRDF integration LUT for split-sum approximation. 582 /// X axis: NdotV (0..1), Y axis: roughness (0..1) 583 /// Output: RG16Float with (scale, bias) for Fresnel term 584 fn generate_brdf_lut(device: &wgpu::Device, queue: &wgpu::Queue, size: u32) -> wgpu::TextureView { 585 let texture = device.create_texture(&wgpu::TextureDescriptor { 586 label: Some("brdf_lut"), 587 size: wgpu::Extent3d { 588 width: size, 589 height: size, 590 depth_or_array_layers: 1, 591 }, 592 mip_level_count: 1, 593 sample_count: 1, 594 dimension: wgpu::TextureDimension::D2, 595 format: wgpu::TextureFormat::Rg16Float, 596 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, 597 view_formats: &[], 598 }); 599 600 let bytes_per_pixel = 4usize; // 2 x f16 = 4 bytes 601 let unpadded_row = size as usize * bytes_per_pixel; 602 let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; 603 let padded_row = unpadded_row.div_ceil(align) * align; 604 605 let sample_count = 1024u32; 606 607 // Compute all BRDF values in parallel 608 let brdf_values: Vec<(f32, f32)> = (0..size * size) 609 .into_par_iter() 610 .map(|idx| { 611 let x = idx % size; 612 let y = idx / size; 613 let ndot_v = (x as f32 + 0.5) / size as f32; 614 let roughness = (y as f32 + 0.5) / size as f32; 615 integrate_brdf(ndot_v.max(0.001), roughness, sample_count) 616 }) 617 .collect(); 618 619 // Write results to buffer 620 let mut data = vec![0u8; padded_row * size as usize]; 621 for (idx, (scale, bias)) in brdf_values.iter().enumerate() { 622 let x = idx as u32 % size; 623 let y = idx as u32 / size; 624 let offset = (y as usize * padded_row) + (x as usize * bytes_per_pixel); 625 626 let r = half::f16::from_f32(*scale); 627 let g = half::f16::from_f32(*bias); 628 629 data[offset..offset + 2].copy_from_slice(&r.to_le_bytes()); 630 data[offset + 2..offset + 4].copy_from_slice(&g.to_le_bytes()); 631 } 632 633 queue.write_texture( 634 wgpu::TexelCopyTextureInfo { 635 texture: &texture, 636 mip_level: 0, 637 origin: wgpu::Origin3d::ZERO, 638 aspect: wgpu::TextureAspect::All, 639 }, 640 &data, 641 wgpu::TexelCopyBufferLayout { 642 offset: 0, 643 bytes_per_row: Some(padded_row as u32), 644 rows_per_image: Some(size), 645 }, 646 wgpu::Extent3d { 647 width: size, 648 height: size, 649 depth_or_array_layers: 1, 650 }, 651 ); 652 653 texture.create_view(&wgpu::TextureViewDescriptor::default()) 654 } 655 656 /// Integrate the BRDF over the hemisphere using importance sampling. 657 /// Returns (scale, bias) for the split-sum: F0 * scale + bias 658 fn integrate_brdf(ndot_v: f32, roughness: f32, sample_count: u32) -> (f32, f32) { 659 // View direction in tangent space (N = [0,0,1]) 660 let v = [ 661 (1.0 - ndot_v * ndot_v).sqrt(), // sin(theta) 662 0.0, 663 ndot_v, // cos(theta) 664 ]; 665 666 let mut a = 0.0f32; 667 let mut b = 0.0f32; 668 669 let alpha = roughness * roughness; 670 671 for i in 0..sample_count { 672 // Hammersley sequence for quasi-random sampling 673 let (xi_x, xi_y) = hammersley(i, sample_count); 674 675 // Importance sample GGX 676 let h = importance_sample_ggx(xi_x, xi_y, alpha); 677 678 // Compute light direction by reflecting view around half vector 679 let v_dot_h = dot(v, h).max(0.0); 680 let l = [ 681 2.0 * v_dot_h * h[0] - v[0], 682 2.0 * v_dot_h * h[1] - v[1], 683 2.0 * v_dot_h * h[2] - v[2], 684 ]; 685 686 let n_dot_l = l[2].max(0.0); // N = [0,0,1] 687 let n_dot_h = h[2].max(0.0); 688 689 if n_dot_l > 0.0 { 690 let g = geometry_smith(ndot_v, n_dot_l, roughness); 691 let g_vis = (g * v_dot_h) / (n_dot_h * ndot_v).max(0.001); 692 let fc = (1.0 - v_dot_h).powf(5.0); 693 694 a += (1.0 - fc) * g_vis; 695 b += fc * g_vis; 696 } 697 } 698 699 let inv_samples = 1.0 / sample_count as f32; 700 (a * inv_samples, b * inv_samples) 701 } 702 703 /// Hammersley quasi-random sequence 704 fn hammersley(i: u32, n: u32) -> (f32, f32) { 705 (i as f32 / n as f32, radical_inverse_vdc(i)) 706 } 707 708 /// Van der Corput radical inverse 709 fn radical_inverse_vdc(mut bits: u32) -> f32 { 710 bits = bits.rotate_right(16); 711 bits = ((bits & 0x55555555) << 1) | ((bits & 0xAAAAAAAA) >> 1); 712 bits = ((bits & 0x33333333) << 2) | ((bits & 0xCCCCCCCC) >> 2); 713 bits = ((bits & 0x0F0F0F0F) << 4) | ((bits & 0xF0F0F0F0) >> 4); 714 bits = ((bits & 0x00FF00FF) << 8) | ((bits & 0xFF00FF00) >> 8); 715 bits as f32 * 2.328_306_4e-10 // 0x100000000 716 } 717 718 /// Importance sample the GGX NDF to get a half-vector in tangent space 719 fn importance_sample_ggx(xi_x: f32, xi_y: f32, alpha: f32) -> [f32; 3] { 720 let a2 = alpha * alpha; 721 722 let phi = 2.0 * std::f32::consts::PI * xi_x; 723 let cos_theta = ((1.0 - xi_y) / (1.0 + (a2 - 1.0) * xi_y)).sqrt(); 724 let sin_theta = (1.0 - cos_theta * cos_theta).sqrt(); 725 726 [sin_theta * phi.cos(), sin_theta * phi.sin(), cos_theta] 727 } 728 729 /// Smith geometry function for GGX 730 fn geometry_smith(n_dot_v: f32, n_dot_l: f32, roughness: f32) -> f32 { 731 let r = roughness + 1.0; 732 let k = (r * r) / 8.0; 733 734 let g1_v = n_dot_v / (n_dot_v * (1.0 - k) + k); 735 let g1_l = n_dot_l / (n_dot_l * (1.0 - k) + k); 736 g1_v * g1_l 737 } 738 739 fn dot(a: [f32; 3], b: [f32; 3]) -> f32 { 740 a[0] * b[0] + a[1] * b[1] + a[2] * b[2] 741 } 742 743 /// Generate pre-filtered environment cubemap with mip levels for different roughness. 744 fn generate_prefiltered_cubemap( 745 device: &wgpu::Device, 746 queue: &wgpu::Queue, 747 pixels: &[image::Rgb<f32>], 748 src_width: u32, 749 src_height: u32, 750 face_size: u32, 751 mip_count: u32, 752 ) -> wgpu::TextureView { 753 let texture = device.create_texture(&wgpu::TextureDescriptor { 754 label: Some("prefiltered_cubemap"), 755 size: wgpu::Extent3d { 756 width: face_size, 757 height: face_size, 758 depth_or_array_layers: 6, 759 }, 760 mip_level_count: mip_count, 761 sample_count: 1, 762 dimension: wgpu::TextureDimension::D2, 763 format: wgpu::TextureFormat::Rgba16Float, 764 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, 765 view_formats: &[], 766 }); 767 768 let sample_count = 512u32; 769 770 for mip in 0..mip_count { 771 let mip_size = face_size >> mip; 772 let roughness = mip as f32 / (mip_count - 1) as f32; 773 774 let bytes_per_pixel = 8usize; // 4 x f16 775 let unpadded_row = mip_size as usize * bytes_per_pixel; 776 let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; 777 let padded_row = unpadded_row.div_ceil(align) * align; 778 779 for face in 0..6u32 { 780 // Compute all pixels in parallel 781 let pixel_colors: Vec<[f32; 3]> = (0..mip_size * mip_size) 782 .into_par_iter() 783 .map(|idx| { 784 let x = idx % mip_size; 785 let y = idx / mip_size; 786 let u = (x as f32 + 0.5) / mip_size as f32 * 2.0 - 1.0; 787 let v = (y as f32 + 0.5) / mip_size as f32 * 2.0 - 1.0; 788 789 let n = normalize(face_uv_to_direction(face, u, v)); 790 prefilter_env_map(pixels, src_width, src_height, n, roughness, sample_count) 791 }) 792 .collect(); 793 794 // Write results to buffer 795 let mut data = vec![0u8; padded_row * mip_size as usize]; 796 for (idx, color) in pixel_colors.iter().enumerate() { 797 let x = idx as u32 % mip_size; 798 let y = idx as u32 / mip_size; 799 let offset = (y as usize * padded_row) + (x as usize * bytes_per_pixel); 800 801 let r = half::f16::from_f32(color[0]); 802 let g = half::f16::from_f32(color[1]); 803 let b = half::f16::from_f32(color[2]); 804 let a = half::f16::from_f32(1.0); 805 806 data[offset..offset + 2].copy_from_slice(&r.to_le_bytes()); 807 data[offset + 2..offset + 4].copy_from_slice(&g.to_le_bytes()); 808 data[offset + 4..offset + 6].copy_from_slice(&b.to_le_bytes()); 809 data[offset + 6..offset + 8].copy_from_slice(&a.to_le_bytes()); 810 } 811 812 queue.write_texture( 813 wgpu::TexelCopyTextureInfo { 814 texture: &texture, 815 mip_level: mip, 816 origin: wgpu::Origin3d { 817 x: 0, 818 y: 0, 819 z: face, 820 }, 821 aspect: wgpu::TextureAspect::All, 822 }, 823 &data, 824 wgpu::TexelCopyBufferLayout { 825 offset: 0, 826 bytes_per_row: Some(padded_row as u32), 827 rows_per_image: Some(mip_size), 828 }, 829 wgpu::Extent3d { 830 width: mip_size, 831 height: mip_size, 832 depth_or_array_layers: 1, 833 }, 834 ); 835 } 836 } 837 838 texture.create_view(&wgpu::TextureViewDescriptor { 839 label: Some("prefiltered_cubemap_view"), 840 dimension: Some(wgpu::TextureViewDimension::Cube), 841 ..Default::default() 842 }) 843 } 844 845 /// Pre-filter the environment map for a given roughness using GGX importance sampling. 846 fn prefilter_env_map( 847 pixels: &[image::Rgb<f32>], 848 src_width: u32, 849 src_height: u32, 850 n: [f32; 3], 851 roughness: f32, 852 sample_count: u32, 853 ) -> [f32; 3] { 854 // For roughness = 0, just sample the environment directly 855 if roughness < 0.001 { 856 return sample_equirect(pixels, src_width, src_height, n); 857 } 858 859 // Use N = V = R assumption for pre-filtering 860 let r = n; 861 let v = r; 862 863 let mut prefilt = [0.0f32; 3]; 864 let mut total_weight = 0.0f32; 865 866 let alpha = roughness * roughness; 867 868 for i in 0..sample_count { 869 let (xi_x, xi_y) = hammersley(i, sample_count); 870 let h = importance_sample_ggx_world(xi_x, xi_y, n, alpha); 871 872 // Reflect V around H to get L 873 let v_dot_h = dot(v, h).max(0.0); 874 let l = [ 875 2.0 * v_dot_h * h[0] - v[0], 876 2.0 * v_dot_h * h[1] - v[1], 877 2.0 * v_dot_h * h[2] - v[2], 878 ]; 879 880 let n_dot_l = dot(n, l); 881 if n_dot_l > 0.0 { 882 let color = sample_equirect(pixels, src_width, src_height, l); 883 prefilt[0] += color[0] * n_dot_l; 884 prefilt[1] += color[1] * n_dot_l; 885 prefilt[2] += color[2] * n_dot_l; 886 total_weight += n_dot_l; 887 } 888 } 889 890 if total_weight > 0.0 { 891 let inv = 1.0 / total_weight; 892 [prefilt[0] * inv, prefilt[1] * inv, prefilt[2] * inv] 893 } else { 894 [0.0, 0.0, 0.0] 895 } 896 } 897 898 /// Importance sample GGX and return half-vector in world space. 899 fn importance_sample_ggx_world(xi_x: f32, xi_y: f32, n: [f32; 3], alpha: f32) -> [f32; 3] { 900 // Sample in tangent space 901 let h_tangent = importance_sample_ggx(xi_x, xi_y, alpha); 902 903 // Build tangent frame 904 let up = if n[1].abs() < 0.999 { 905 [0.0, 1.0, 0.0] 906 } else { 907 [1.0, 0.0, 0.0] 908 }; 909 let tangent = normalize(cross(up, n)); 910 let bitangent = cross(n, tangent); 911 912 // Transform to world space 913 normalize([ 914 h_tangent[0] * tangent[0] + h_tangent[1] * bitangent[0] + h_tangent[2] * n[0], 915 h_tangent[0] * tangent[1] + h_tangent[1] * bitangent[1] + h_tangent[2] * n[1], 916 h_tangent[0] * tangent[2] + h_tangent[1] * bitangent[2] + h_tangent[2] * n[2], 917 ]) 918 }