Starting coarse rasterizer

Working down the pipeline.

WIP

piet-gpu/bin/cli.rs

@@ -41,7 +41,7 @@ fn main() -> Result<(), Error> {
 
         let fence = device.create_fence(false)?;
         let mut cmd_buf = device.create_cmd_buf()?;
-        let query_pool = device.create_query_pool(3)?;
+        let query_pool = device.create_query_pool(4)?;
 
         let mut ctx = PietGpuRenderContext::new();
         render_scene(&mut ctx);
@@ -61,11 +61,16 @@ fn main() -> Result<(), Error> {
         let ts = device.reap_query_pool(&query_pool).unwrap();
         println!("Element kernel time: {:.3}ms", ts[0] * 1e3);
         println!("Binning kernel time: {:.3}ms", (ts[1] - ts[0]) * 1e3);
+        println!("Coarse kernel time: {:.3}ms", (ts[2] - ts[1]) * 1e3);
 
         /*
         let mut data: Vec<u32> = Default::default();
         device.read_buffer(&renderer.bin_buf, &mut data).unwrap();
         piet_gpu::dump_k1_data(&data);
+
+        let mut data: Vec<u32> = Default::default();
+        device.read_buffer(&renderer.ptcl_buf, &mut data).unwrap();
+        piet_gpu::dump_k1_data(&data);
         */
 
         let mut img_data: Vec<u8> = Default::default();

piet-gpu/shader/binning.comp

@@ -3,20 +3,7 @@
 #version 450
 #extension GL_GOOGLE_include_directive : enable
 
-#define N_ROWS 4
+#include "setup.h"
-#define WG_SIZE 32
-#define LG_WG_SIZE 5
-#define TILE_SIZE (WG_SIZE * N_ROWS)
-
-// TODO: move these to setup file
-#define N_TILE_X 16
-#define N_TILE_Y 16
-#define N_TILE (N_TILE_X * N_TILE_Y)
-#define N_SLICE (N_TILE / 32)
-#define N_WG 16 // Number of workgroups, should be 1 per SM
-
-#define BIN_INITIAL_ALLOC 64
-#define BIN_ALLOC 256
 
 layout(local_size_x = N_TILE, local_size_y = 1) in;
 
@@ -37,7 +24,6 @@ layout(set = 0, binding = 2) buffer BinsBuf {
 
 #include "annotated.h"
 #include "bins.h"
-#include "setup.h"
 
 // scale factors useful for converting coordinates to bins
 #define SX (1.0 / float(N_TILE_X * TILE_WIDTH_PX))
@@ -135,9 +121,7 @@ void main() {
                 uint next_chunk = chunk_ref.offset + BinChunk_size + chunk_n * 4;
                 if (next_chunk + BinChunk_size + min(24, element_count * 4) > wr_limit) {
                     uint alloc_amount = max(BIN_ALLOC, BinChunk_size + element_count * 4);
-                    if (alloc_amount - BIN_ALLOC < 64) {
+                    // could try to reduce fragmentation if BIN_ALLOC is only a bit above needed
-                        alloc_amount = BIN_ALLOC;
-                    }
                     next_chunk = atomicAdd(alloc, alloc_amount);
                     wr_limit = next_chunk + alloc_amount;
                 }
@@ -149,9 +133,6 @@ void main() {
                 chunk_end = wr_limit;
                 chunk_n = (wr_limit - instance_ref.offset) / 4;
                 uint alloc_amount = max(BIN_ALLOC, BinChunk_size + (element_count - chunk_n) * 4);
-                if (alloc_amount - BIN_ALLOC < 64) {
-                    alloc_amount = BIN_ALLOC;
-                }
                 chunk_new_start = atomicAdd(alloc, alloc_amount);
                 wr_limit = chunk_new_start + alloc_amount;
                 BinChunk_write(chunk_ref, BinChunk(chunk_n, BinChunkRef(chunk_new_start)));

piet-gpu/shader/build.ninja

@@ -22,4 +22,6 @@ build kernel4.spv: glsl kernel4.comp | ptcl.h segment.h fill_seg.h setup.h
 
 build elements.spv: glsl elements.comp | scene.h state.h annotated.h
 
-build binning.spv: glsl binning.comp | annotated.h setup.h
+build binning.spv: glsl binning.comp | annotated.h bins.h setup.h
+
+build coarse.spv: glsl coarse.comp | annotated.h bins.h ptcl.h setup.h

piet-gpu/shader/coarse.comp

@@ -0,0 +1,129 @@
+// The coarse rasterizer stage of the pipeline.
+
+#version 450
+#extension GL_GOOGLE_include_directive : enable
+
+#include "setup.h"
+
+layout(local_size_x = N_TILE, local_size_y = 1) in;
+
+layout(set = 0, binding = 0) buffer AnnotatedBuf {
+    uint[] annotated;
+};
+
+layout(set = 0, binding = 1) buffer BinsBuf {
+    uint[] bins;
+};
+
+layout(set = 0, binding = 2) buffer AllocBuf {
+    uint alloc;
+};
+
+layout(set = 0, binding = 3) buffer PtclBuf {
+    uint[] ptcl;
+};
+
+#include "annotated.h"
+#include "bins.h"
+#include "ptcl.h"
+
+#define N_RINGBUF 512
+
+shared uint sh_elements[N_RINGBUF];
+shared uint sh_chunk[N_WG];
+shared uint sh_chunk_next[N_WG];
+shared uint sh_chunk_n[N_WG];
+shared uint sh_min_buf;
+// Some of these are kept in shared memory to ease register
+// pressure, but it could go either way.
+shared uint sh_first_el[N_WG];
+shared uint sh_selected_n;
+shared uint sh_elements_ref;
+
+shared uint sh_bitmaps[N_SLICE][N_TILE];
+
+void main() {
+    // Could use either linear or 2d layouts for both dispatch and
+    // invocations within the workgroup. We'll use variables to abstract.
+    uint bin_ix = N_TILE_X * gl_WorkGroupID.y + gl_WorkGroupID.x;
+    uint th_ix = gl_LocalInvocationID.x;
+    uint wr_ix = 0;
+    uint rd_ix = 0;
+    uint first_el;
+    if (th_ix < N_WG) {
+        uint start_chunk = (bin_ix * N_WG + th_ix) * BIN_INITIAL_ALLOC;
+        sh_chunk[th_ix] = start_chunk;
+        BinChunk chunk = BinChunk_read(BinChunkRef(start_chunk));
+        sh_chunk_n[th_ix] = chunk.n;
+        sh_chunk_next[th_ix] = chunk.next.offset;
+        sh_first_el[th_ix] = chunk.n > 0 ?
+            BinInstance_read(BinInstanceRef(start_chunk + BinChunk_size)).element_ix : ~0;
+    }
+    uint probe = 0; // for debugging
+    do {
+        for (uint i = 0; i < N_SLICE; i++) {
+            sh_bitmaps[i][th_ix] = 0;
+        }
+
+        while (wr_ix - rd_ix <= N_TILE) {
+            // Choose segment with least element.
+            uint my_min;
+            if (th_ix < N_WG) {
+                if (th_ix == 0) {
+                    sh_selected_n = 0;
+                    sh_min_buf = ~1;
+                }
+            }
+            barrier();
+            // Tempting to do this with subgroups, but atomic should be good enough.
+            my_min = sh_first_el[th_ix];
+            if (th_ix < N_WG) {
+                atomicMin(sh_min_buf, my_min);
+            }
+            barrier();
+            if (th_ix < N_WG) {
+                if (sh_first_el[th_ix] == sh_min_buf) {
+                    sh_elements_ref = sh_chunk[th_ix] + BinChunk_size;
+                    uint selected_n = sh_chunk_n[th_ix];
+                    sh_selected_n = selected_n;
+                    uint next_chunk = sh_chunk_next[th_ix];
+                    if (next_chunk == 0) {
+                        sh_first_el[th_ix] = ~0;
+                    } else {
+                        sh_chunk[th_ix] = next_chunk;
+                        BinChunk chunk = BinChunk_read(BinChunkRef(next_chunk));
+                        sh_chunk_n[th_ix] = chunk.n;
+                        sh_chunk_next[th_ix] = chunk.next.offset;
+                        sh_first_el[th_ix] = BinInstance_read(
+                            BinInstanceRef(next_chunk + BinChunk_size)).element_ix;
+                    }
+                }
+            }
+            barrier();
+            uint chunk_n = sh_selected_n;
+            if (chunk_n == 0) {
+                // All chunks consumed
+                break;
+            }
+            BinInstanceRef inst_ref = BinInstanceRef(sh_elements_ref);
+            if (th_ix < chunk_n) {
+                uint el = BinInstance_read(BinInstance_index(inst_ref, th_ix)).element_ix;
+                sh_elements[(wr_ix + th_ix) % N_RINGBUF] = el;
+                probe = el;
+            }
+            wr_ix += chunk_n;
+        }
+
+        // We've done the merge and filled the buffer.
+        uint tag = Annotated_Nop;
+        AnnotatedRef ref;
+        if (th_ix + rd_ix < wr_ix) {
+            uint element_ix = (sh_elements[rd_ix] + th_ix) % N_RINGBUF;
+            ref = AnnotatedRef(element_ix * Annotated_size);
+            tag = Annotated_tag(ref);
+            probe = tag;
+        }
+        rd_ix += N_TILE;
+    } while (wr_ix > rd_ix);
+    ptcl[bin_ix * N_TILE + th_ix] = probe;
+}

piet-gpu/shader/setup.h

@@ -40,3 +40,20 @@
 // Maximum number of segments in a SegChunk
 #define SEG_CHUNK_N 32
 #define SEG_CHUNK_ALLOC 512
+
+// Stuff for new algorithm follows; some of the above should get
+// deleted.
+
+// These should probably be renamed and/or reworked. In the binning
+// kernel, they represent the number of bins. Also, the workgroup size
+// of that kernel is equal to the number of bins, but should probably
+// be more flexible (it's 512 in the K&L paper).
+#define N_TILE_X 16
+#define N_TILE_Y 16
+#define N_TILE (N_TILE_X * N_TILE_Y)
+#define N_SLICE (N_TILE / 32)
+// Number of workgroups for binning kernel
+#define N_WG 16
+
+#define BIN_INITIAL_ALLOC 64
+#define BIN_ALLOC 256

piet-gpu/src/lib.rs

@@ -117,6 +117,7 @@ pub struct Renderer<D: Device> {
     pub state_buf: D::Buffer,
     pub anno_buf: D::Buffer,
     pub bin_buf: D::Buffer,
+    pub ptcl_buf: D::Buffer,
 
     el_pipeline: D::Pipeline,
     el_ds: D::DescriptorSet,
@@ -127,6 +128,12 @@ pub struct Renderer<D: Device> {
     bin_alloc_buf_host: D::Buffer,
     bin_alloc_buf_dev: D::Buffer,
 
+    coarse_pipeline: D::Pipeline,
+    coarse_ds: D::DescriptorSet,
+
+    coarse_alloc_buf_host: D::Buffer,
+    coarse_alloc_buf_dev: D::Buffer,
+
     /*
     k1_alloc_buf_host: D::Buffer,
     k1_alloc_buf_dev: D::Buffer,
@@ -172,6 +179,7 @@ impl<D: Device> Renderer<D> {
         let state_buf = device.create_buffer(64 * 1024 * 1024, dev)?;
         let anno_buf = device.create_buffer(64 * 1024 * 1024, dev)?;
         let bin_buf = device.create_buffer(64 * 1024 * 1024, dev)?;
+        let ptcl_buf = device.create_buffer(48 * 1024 * 1024, dev)?;
         let image_dev = device.create_image2d(WIDTH as u32, HEIGHT as u32, dev)?;
 
         let el_code = include_bytes!("../shader/elements.spv");
@@ -202,6 +210,23 @@ impl<D: Device> Renderer<D> {
             &[],
         )?;
 
+        let coarse_alloc_buf_host = device.create_buffer(4, host)?;
+        let coarse_alloc_buf_dev = device.create_buffer(4, dev)?;
+
+        let coarse_alloc_start = 256 * 64 * N_WG;
+        device
+            .write_buffer(&coarse_alloc_buf_host, &[
+                coarse_alloc_start,
+            ])
+            ?;
+        let coarse_code = include_bytes!("../shader/coarse.spv");
+        let coarse_pipeline = device.create_simple_compute_pipeline(coarse_code, 4, 0)?;
+        let coarse_ds = device.create_descriptor_set(
+            &coarse_pipeline,
+            &[&anno_buf, &bin_buf, &coarse_alloc_buf_dev, &ptcl_buf],
+            &[],
+        )?;
+
         /*
         let tilegroup_buf = device.create_buffer(4 * 1024 * 1024, dev)?;
         let ptcl_buf = device.create_buffer(48 * 1024 * 1024, dev)?;
@@ -285,11 +310,16 @@ impl<D: Device> Renderer<D> {
             el_ds,
             bin_pipeline,
             bin_ds,
+            coarse_pipeline,
+            coarse_ds,
             state_buf,
             anno_buf,
             bin_buf,
+            ptcl_buf,
             bin_alloc_buf_host,
             bin_alloc_buf_dev,
+            coarse_alloc_buf_host,
+            coarse_alloc_buf_dev,
             n_elements,
         })
     }
@@ -297,6 +327,7 @@ impl<D: Device> Renderer<D> {
     pub unsafe fn record(&self, cmd_buf: &mut impl CmdBuf<D>, query_pool: &D::QueryPool) {
         cmd_buf.copy_buffer(&self.scene_buf, &self.scene_dev);
         cmd_buf.copy_buffer(&self.bin_alloc_buf_host, &self.bin_alloc_buf_dev);
+        cmd_buf.copy_buffer(&self.coarse_alloc_buf_host, &self.coarse_alloc_buf_dev);
         cmd_buf.memory_barrier();
         cmd_buf.image_barrier(
             &self.image_dev,
@@ -319,6 +350,13 @@ impl<D: Device> Renderer<D> {
         );
         cmd_buf.write_timestamp(&query_pool, 2);
         cmd_buf.memory_barrier();
+        cmd_buf.dispatch(
+            &self.coarse_pipeline,
+            &self.coarse_ds,
+            (WIDTH as u32 / 256, HEIGHT as u32 / 256, 1),
+        );
+        cmd_buf.write_timestamp(&query_pool, 3);
+        cmd_buf.memory_barrier();
         cmd_buf.image_barrier(&self.image_dev, ImageLayout::General, ImageLayout::BlitSrc);
     }
 }