WTF

shaders/voxel.slang

@@ -3,6 +3,12 @@ struct PushConstants
     float4x4 view_proj;
     float3   cam_pos;
     uint32_t frame_timestamp;
+    uint32_t width;
+    uint32_t height;
+
+    uint32_t chunk_width;
+    uint32_t chunk_height;
+    uint32_t chunk_alt;
 }
 
 public struct VertexOutput
@@ -180,7 +186,14 @@ float3 floor_scale(float3 position, uint32_t scale_exp)
     return asfloat(asuint(position) & mask);
 }
 
-float4 ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, float dist_offset, out float3 hit_pos)
+struct HitInformation
+{
+    bool hit;
+    float3 hit_pos;
+    float4 color;
+}
+
+HitInformation ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, float dist_offset)
 {
     float fov_deg = 100. / 1920.;
     float fov_rad = (float.getPi() * fov_deg) / 180.;
@@ -189,14 +202,18 @@ float4 ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, floa
     let st_pointer = structure_table_pointer[root_id];
     if(!st_pointer.subdivided())
     {
-        discard;
+        var hit: HitInformation;
+        hit.hit = false;
+        return hit;
     }
 
     if(!st_pointer.pointer_valid())
     {
         // Record request
         structure_table_request_buffer[root_id].add(1);
-        discard;
+        var hit: HitInformation;
+        hit.hit = false;
+        return hit;
     }
 
     ray_origin += float3(1.);
@@ -259,8 +276,13 @@ float4 ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, floa
 
         if(color_pool[current_node_index].colors[child_index].byte_a != 0)
         {
-            hit_pos = pos - float3(1.);
-            return color_pool[current_node_index].colors[child_index].float_color; 
+            //hit_pos = pos - float3(1.);
+            //return color_pool[current_node_index].colors[child_index].float_color; 
+            var hit: HitInformation;
+            hit.hit = true;
+            hit.hit_pos = pos - float3(1.);
+            hit.color = color_pool[current_node_index].colors[child_index].float_color;
+            return hit;
         }
 
         //uint64_t occupancy = get_child_mask(structure_pool[current_node_index].occupancy_low, structure_pool[current_node_index].occupancy_high);
@@ -301,7 +323,7 @@ float4 ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, floa
         int32_t common_depth = (1 + (22 - firstbithigh(diff)) / 2) * 2;
         if(common_depth <= 0)
         {
-            discard;
+            break;
         }
 
         scale_exp = 23 - common_depth;
@@ -309,8 +331,9 @@ float4 ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, floa
     }
 
 
-    hit_pos = pos - float3(1.);
-    return float4(1., 0., 1., 1.);
+    var hit: HitInformation;
+    hit.hit = false;
+    return hit;
 }
 
 struct FragmentOutput
@@ -327,19 +350,96 @@ FragmentOutput fragment(VertexOutput vertex_out)
     let intersection_t = box_intersect(vertex_out.cam_position, ray_direction, vertex_out.chunk_position, vertex_out.chunk_position + float3(1.));
     let local_ray_origin = max(intersection_t.x, 0.) * ray_direction + vertex_out.cam_position - vertex_out.chunk_position;
     // Figure out intersection
-    var hit_pos : float3;
-    let color = ray_march(ray_direction, local_ray_origin, vertex_out.structure_id, max(0., intersection_t.x), hit_pos);
-    let world_hit_pos = hit_pos + vertex_out.chunk_position;
+    let hit = ray_march(ray_direction, local_ray_origin, vertex_out.structure_id, max(0., intersection_t.x));
+    let world_hit_pos = hit.hit_pos + vertex_out.chunk_position;
     let clip = mul(constants.view_proj, float4(world_hit_pos, 1.));
     let depth = clip.z / clip.w;
 
     var frag_out : FragmentOutput;
     frag_out.depth = depth;
-    frag_out.color = color;
+    frag_out.color = hit.color;
 
     return frag_out;
 }
 
+bool3 min_mask(float3 val)
+{
+    let min_val = min(val.x, min(val.y, val.z));
+    return val == min_val;
+}
+
+[[vk::binding(0, 1)]] 
+[[format("rgba32f")]]
+WTexture2D<float4> output_texture;
+
+[shader("compute")]
+[numthreads(16, 16, 1)]
+void ray_march_compute(uint32_t3 location : SV_DispatchThreadID)
+{
+    if(location.x >= constants.width || location.y >= constants.height)
+    {
+        return;
+    }
+
+    uint32_t2 pixel_loc = uint32_t2(location.x, location.y);
+
+    let ndc_loc_x = (float)location.x / (float)constants.width * 2. - 1.;
+    let ndc_loc_y = 1. - (float)location.y / (float)constants.height * 2.;
+
+    var world_loc = mul(constants.view_proj, float4(ndc_loc_x, ndc_loc_y, 1., 1.));
+    world_loc /= world_loc.w;
+    let ray_direction = normalize(world_loc.xyz - constants.cam_pos);
+
+    var inter = box_intersect(constants.cam_pos, ray_direction, float3(0.), float3(constants.chunk_width, constants.chunk_alt, constants.chunk_height));
+    if(inter.y <= inter.x || inter.y <= 0.)
+    {
+        output_texture.Store(pixel_loc, float4(0.));
+        return;
+    }
+    inter.x = max(0., inter.x);
+
+
+    let start_position = inter.x * ray_direction + constants.cam_pos;
+    // FVT
+    int32_t3 current_voxel = clamp(
+        int32_t3(floor(start_position)), 
+        int32_t3(0),
+        int32_t3(constants.chunk_width - 1, constants.chunk_alt - 1, constants.chunk_height - 1)
+    );
+    int32_t3 offset = int32_t3(sign(ray_direction));
+    float3 delta = abs(1. / ray_direction);
+    float3 position = start_position;
+    float3 t = select(ray_direction > 0., current_voxel + int32_t3(1) - start_position, start_position - current_voxel) * delta;
+
+    var color = float4(0.);
+    var hit_pos = float3(0.);
+    for(int32_t i = 0; i < 256; i++)
+    {
+        let min_mask = min_mask(t);    
+        let t_adv = select(min_mask, t, float3(0.));
+        let t_ray = t_adv.x + t_adv.y + t_adv.z;
+
+
+        let chunk_index = current_voxel.y + current_voxel.z * constants.chunk_alt + current_voxel.x * constants.chunk_alt * constants.chunk_height;
+        let hit = ray_march(ray_direction, position - float3(current_voxel), chunk_index, length(position - constants.cam_pos));
+        position = start_position + ray_direction * t_ray;
+        if(hit.hit)
+        {
+            color = hit.color;
+            break;
+        }
+        t += select(min_mask, delta, float3(0.));
+        current_voxel += select(min_mask, offset, int32_t3(0));
+
+        if(any(current_voxel < 0) || any(current_voxel >= int32_t3(constants.chunk_width, constants.chunk_alt, constants.chunk_height)))
+        {
+            color = float4(0.);
+            break;
+        }
+    }
+    output_texture.Store(pixel_loc, float4(color));
+}
+
 float2 box_intersect(float3 origin, float3 ray_direction, float3 box_min, float3 box_max)
 {
     let min_ts = (box_min - origin) / ray_direction;

src/main.rs

@@ -2,6 +2,7 @@
 #![feature(float_algebraic)]
 
 use std::collections::HashMap;
+use std::ops::Div;
 use std::sync::Arc;
 use std::sync::mpsc::sync_channel;
 
@@ -19,8 +20,10 @@ use itertools::Itertools;
 use rayon::iter::IndexedParallelIterator;
 use rayon::iter::IntoParallelRefIterator;
 use rayon::iter::ParallelIterator;
+use wgpu::BindGroupLayout;
 use wgpu::Buffer;
 use wgpu::BufferUsages;
+use wgpu::ComputePipeline;
 use wgpu::Device;
 use wgpu::Extent3d;
 use wgpu::Features;
@@ -28,7 +31,9 @@ use wgpu::InstanceDescriptor;
 use wgpu::InstanceFlags;
 use wgpu::MemoryBudgetThresholds;
 use wgpu::Operations;
+use wgpu::Origin3d;
 use wgpu::RenderPipeline;
+use wgpu::ShaderStages;
 use wgpu::Texture;
 use wgpu::TextureUsages;
 use wgpu::TextureView;
@@ -82,10 +87,14 @@ struct State
     size: winit::dpi::PhysicalSize<u32>,
     surface: wgpu::Surface<'static>,
     depth_buffer: (wgpu::Texture, wgpu::TextureView),
+    target_texture: (wgpu::Texture, wgpu::TextureView),
     surface_format: wgpu::TextureFormat,
+    target_blitter: wgpu::util::TextureBlitter,
     egui_renderer: EguiRenderer,
 
-    pipeline: RenderPipeline,
+    //pipeline: RenderPipeline,
+    pipeline: ComputePipeline,
+    surface_bg_layout: BindGroupLayout,
     voxel_cache: Arc<Mutex<VoxelCache<4>>>,
     cache_interface: Arc<CacheProducerInterface<4>>,
     terrain_generator: Arc<TerrainGenerator<4>>,
@@ -111,6 +120,12 @@ struct Immediates
     view_proj: Mat4,
     cam_pos: Vec3,
     frame_timestamp: u32,
+    width: u32,
+    height: u32,
+
+    chunk_width: u32,
+    chunk_height: u32,
+    chunk_alt: u32,
 }
 
 #[derive(Debug, Clone, Copy, Zeroable, Pod)]
@@ -166,19 +181,19 @@ impl State
         let mut voxel_cache = VoxelCache::<4>::new(100_000, device.clone(), queue.clone());
         let cache_interface = CacheProducerInterface::new(1024, &device);
 
-        //let terrain_generator = TerrainGenerator::<4>::new(5, "vxls_height.tif", 0.2, "img.jpg");
+        let terrain_generator = TerrainGenerator::<4>::new(5, "vxls_height.tif", 0.2, "img.jpg");
         // let terrain_generator = TerrainGenerator::<4>::new(
         //     5,
         //     "./pointe_percee/height.tif",
         //     0.2,
         //     "./pointe_percee/ortho.jpg",
         // );
-        let terrain_generator = TerrainGenerator::<4>::new(
-            5,
-            "/home/albin/Documents/vxls_maps/lapiz/height.tif",
-            0.2,
-            "/home/albin/Documents/vxls_maps/lapiz/ortho.jpg",
-        );
+        // let terrain_generator = TerrainGenerator::<4>::new(
+        //     5,
+        //     "/home/albin/Documents/vxls_maps/lapiz/height.tif",
+        //     0.2,
+        //     "/home/albin/Documents/vxls_maps/lapiz/ortho.jpg",
+        // );
 
         let mut chunk_pos_map = HashMap::new();
         let chunk_instances = (0..terrain_generator.chunk_width)
@@ -202,7 +217,7 @@ impl State
         let instance_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
             label: Some("Instance buffer"),
             contents: bytemuck::cast_slice(chunk_instances.as_slice()),
-            usage: BufferUsages::COPY_DST | BufferUsages::VERTEX,
+            usage: BufferUsages::COPY_DST | BufferUsages::VERTEX | BufferUsages::STORAGE,
         });
 
         // let shader_module = unsafe {
@@ -229,13 +244,41 @@ impl State
         //let shader_module = device.create_shader_module(include_wgsl!("../shaders/voxel.wgsl"));
         let shader_module = device.create_shader_module(include_spirv!("../shaders/voxel.spv"));
 
+        let surface_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("surface_bg_layout"),
+                entries: &[wgpu::BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: ShaderStages::COMPUTE,
+                    ty: wgpu::BindingType::StorageTexture {
+                        access: wgpu::StorageTextureAccess::WriteOnly,
+                        format: wgpu::TextureFormat::Rgba32Float,
+                        view_dimension: wgpu::TextureViewDimension::D2,
+                    },
+                    count: None,
+                }],
+            });
+
         let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
             label: Some("Voxel pipeline layout"),
 
-            bind_group_layouts: &[Some(&voxel_cache.bind_group_layout())],
+            bind_group_layouts: &[
+                Some(&voxel_cache.bind_group_layout()),
+                Some(&surface_bind_group_layout),
+            ],
             immediate_size: size_of::<Immediates>() as u32,
         });
 
+        let chunk_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label: Some("Compute render"),
+            layout: Some(&pipeline_layout),
+            module: &shader_module,
+            entry_point: Some("ray_march_compute"),
+            compilation_options: wgpu::PipelineCompilationOptions::default(),
+            cache: None,
+        });
+
+        /*
         let chunk_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
             label: Some("Render pipeline"),
             layout: Some(&pipeline_layout),
@@ -290,6 +333,7 @@ impl State
             multiview_mask: None,
             cache: None,
         });
+        */
 
         let state = State {
             instance,
@@ -299,8 +343,13 @@ impl State
             surface_format,
             egui_renderer,
             depth_buffer: Self::create_depth_buffer(&device, size.width, size.height),
+            target_texture: Self::create_target_texture(&device, size.width, size.height),
+            target_blitter: wgpu::util::TextureBlitterBuilder::new(&device, surface_format)
+                .sample_type(wgpu::FilterMode::Nearest)
+                .build(),
             queue,
             device,
+            surface_bg_layout: surface_bind_group_layout,
             usage_vec: Arc::new(Mutex::new(vec![])),
             pipeline: chunk_pipeline,
             voxel_cache: Arc::new(Mutex::new(voxel_cache)),
@@ -320,6 +369,27 @@ impl State
         state
     }
 
+    fn create_target_texture(device: &Device, width: u32, height: u32) -> (Texture, TextureView)
+    {
+        let texture = device.create_texture(&wgpu::TextureDescriptor {
+            label: Some("Target texture"),
+            size: Extent3d {
+                width,
+                height,
+                depth_or_array_layers: 1,
+            },
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: wgpu::TextureDimension::D2,
+            format: wgpu::TextureFormat::Rgba32Float,
+            usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
+            view_formats: &[],
+        });
+
+        let texture_view = texture.create_view(&wgpu::wgt::TextureViewDescriptor::default());
+        (texture, texture_view)
+    }
+
     fn get_window(&self) -> &Window
     {
         &self.window
@@ -395,6 +465,8 @@ impl State
         self.configure_surface();
         self.depth_buffer =
             Self::create_depth_buffer(&self.device, new_size.width, new_size.height);
+        self.target_texture =
+            Self::create_target_texture(&self.device, new_size.width, new_size.height);
     }
 
     fn render(&mut self)
@@ -442,6 +514,15 @@ impl State
                 ..Default::default()
             });
 
+        let surface_bg = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("surface_bg"),
+            layout: &self.surface_bg_layout,
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::TextureView(&self.target_texture.1),
+            }],
+        });
+
         // ~~ Ray-marching timestamp query setup ~~
         let timestamp_query = self.device.create_query_set(&wgpu::QuerySetDescriptor {
             label: Some("timestamp_query_set"),
@@ -458,6 +539,48 @@ impl State
 
         // ~~ Main render pass ~~
         let mut encoder = self.device.create_command_encoder(&Default::default());
+        {
+            let mut renderpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+                label: Some("compute_render_pass"),
+                timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
+                    query_set: &timestamp_query,
+                    beginning_of_pass_write_index: Some(0),
+                    end_of_pass_write_index: Some(1),
+                }),
+            });
+
+            renderpass.set_pipeline(&self.pipeline);
+            renderpass.set_bind_group(0, Some(&self.voxel_cache.lock().bind_group()), &[]);
+            renderpass.set_bind_group(1, Some(&surface_bg), &[]);
+            let imm = [Immediates {
+                view_proj: self.camera.view_proj().inverse(),
+                cam_pos: self.camera.position,
+                frame_timestamp: self.voxel_cache.lock().current_timestamp(),
+                width: self.size.width,
+                height: self.size.height,
+
+                chunk_width: self.terrain_generator.chunk_width as u32,
+                chunk_height: self.terrain_generator.chunk_height as u32,
+                chunk_alt: self.terrain_generator.chunk_alt as u32,
+            }];
+            renderpass.set_immediates(0, unsafe { as_raw_bytes(&imm) });
+            renderpass.dispatch_workgroups(
+                self.size.width.div_ceil(16),
+                self.size.height.div_ceil(16),
+                1,
+            );
+            drop(renderpass);
+
+            encoder.resolve_query_set(&timestamp_query, 0..2, &timestamp_buffer, 0);
+        }
+
+        self.target_blitter.copy(
+            &self.device,
+            &mut encoder,
+            &self.target_texture.1,
+            &texture_view,
+        );
+        /*
         {
             let mut renderpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                 label: None,
@@ -508,6 +631,7 @@ impl State
 
             encoder.resolve_query_set(&timestamp_query, 0..2, &timestamp_buffer, 0);
         }
+        */
 
         // ~~ EGUI Render pass ~~
         {

src/voxel_cache.rs

@@ -95,7 +95,7 @@ where
                     // Structure pool
                     wgpu::BindGroupLayoutEntry {
                         binding: 0,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -106,7 +106,7 @@ where
                     // Color pool
                     wgpu::BindGroupLayoutEntry {
                         binding: 1,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -117,7 +117,7 @@ where
                     // Location pool
                     wgpu::BindGroupLayoutEntry {
                         binding: 2,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -128,7 +128,7 @@ where
                     // Request buffer
                     wgpu::BindGroupLayoutEntry {
                         binding: 3,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -139,7 +139,7 @@ where
                     // Usage buffer
                     wgpu::BindGroupLayoutEntry {
                         binding: 4,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -152,7 +152,7 @@ where
                     // Structure table pointers
                     wgpu::BindGroupLayoutEntry {
                         binding: 5,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,
@@ -164,7 +164,7 @@ where
                     // Structure table request buffer
                     wgpu::BindGroupLayoutEntry {
                         binding: 6,
-                        visibility: ShaderStages::FRAGMENT,
+                        visibility: ShaderStages::COMPUTE,
                         ty: wgpu::BindingType::Buffer {
                             ty: wgpu::BufferBindingType::Storage { read_only: false },
                             has_dynamic_offset: false,