Fixed builtin(primitive_index) for vulkan backend (#2716)

Co-authored-by: Koen Willemsen <koen.willemsen@materialise.be>

wgpu-hal/src/vulkan/adapter.rs

@@ -1237,6 +1237,10 @@ impl super::Adapter {
                 capabilities.push(spv::Capability::MultiView);
             }
 
+            if features.contains(wgt::Features::SHADER_PRIMITIVE_INDEX) {
+                capabilities.push(spv::Capability::Geometry);
+            }
+
             if features.intersects(
                 wgt::Features::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING
                     | wgt::Features::UNIFORM_BUFFER_AND_STORAGE_TEXTURE_ARRAY_NON_UNIFORM_INDEXING,

wgpu/examples/shadow/main.rs

@@ -4,7 +4,7 @@ use std::{borrow::Cow, f32::consts, iter, mem, num::NonZeroU32, ops::Range, rc::
 mod framework;
 
 use bytemuck::{Pod, Zeroable};
-use wgpu::util::DeviceExt;
+use wgpu::util::{align_to, DeviceExt};
 
 #[repr(C)]
 #[derive(Clone, Copy, Pod, Zeroable)]
@@ -291,12 +291,7 @@ impl framework::Example for Example {
         let uniform_alignment = {
             let alignment =
                 device.limits().min_uniform_buffer_offset_alignment as wgpu::BufferAddress;
-            let rem = entity_uniform_size % alignment;
-            if rem != 0 {
-                entity_uniform_size + alignment - rem
-            } else {
-                entity_uniform_size
-            }
+            align_to(entity_uniform_size, alignment)
         };
         // Note: dynamic uniform offsets also have to be aligned to `Limits::min_uniform_buffer_offset_alignment`.
         let entity_uniform_buf = device.create_buffer(&wgpu::BufferDescriptor {

wgpu/src/util/belt.rs

@@ -1,6 +1,6 @@
 use crate::{
-    Buffer, BufferAddress, BufferDescriptor, BufferSize, BufferUsages, BufferViewMut,
-    CommandEncoder, Device, MapMode,
+    util::align_to, Buffer, BufferAddress, BufferDescriptor, BufferSize, BufferUsages,
+    BufferViewMut, CommandEncoder, Device, MapMode,
 };
 use std::fmt;
 use std::sync::{mpsc, Arc};
@@ -95,11 +95,7 @@ impl StagingBelt {
 
         encoder.copy_buffer_to_buffer(&chunk.buffer, chunk.offset, target, offset, size.get());
         let old_offset = chunk.offset;
-        chunk.offset += size.get();
-        let remainder = chunk.offset % crate::MAP_ALIGNMENT;
-        if remainder != 0 {
-            chunk.offset += crate::MAP_ALIGNMENT - remainder;
-        }
+        chunk.offset = align_to(chunk.offset + size.get(), crate::MAP_ALIGNMENT);
 
         self.active_chunks.push(chunk);
         self.active_chunks

wgpu/src/util/mod.rs

@@ -6,6 +6,7 @@ mod encoder;
 mod indirect;
 mod init;
 
+use std::ops::{Add, Rem, Sub};
 use std::sync::Arc;
 use std::{
     borrow::Cow,
@@ -123,3 +124,31 @@ impl std::ops::Deref for DownloadBuffer {
         super::BufferMappedRangeSlice::slice(&self.1)
     }
 }
+
+///
+/// Aligns a `value` to an `alignment`.
+///
+/// Returns the first number greater than or equal to `value` that is also a
+/// multiple of `alignment`. If `value` is already a multiple of `alignment`,
+/// `value` will be returned.
+///
+/// # Examples
+///
+/// ```
+/// # use wgpu::util::align_to;
+/// assert_eq!(align_to(253, 16), 256);
+/// assert_eq!(align_to(256, 16), 256);
+/// assert_eq!(align_to(0, 16), 0);
+/// ```
+///
+pub fn align_to<T>(value: T, alignment: T) -> T
+where
+    T: Add<Output = T> + Copy + Default + PartialEq<T> + Rem<Output = T> + Sub<Output = T>,
+{
+    let remainder = value % alignment;
+    if remainder == T::default() {
+        value
+    } else {
+        value + alignment - remainder
+    }
+}

wgpu/tests/clear_texture.rs

@@ -1,4 +1,5 @@
 use crate::common::{initialize_test, TestParameters, TestingContext};
+use wgpu::util::align_to;
 
 static TEXTURE_FORMATS_UNCOMPRESSED: &[wgpu::TextureFormat] = &[
     wgpu::TextureFormat::R8Unorm,
@@ -239,10 +240,6 @@ fn single_texture_clear_test(
     // TODO: Read back and check zeroness?
 }
 
-fn round_up(value: u32, alignment: u32) -> u32 {
-    ((value + alignment - 1) / alignment) * alignment
-}
-
 fn clear_texture_tests(
     ctx: &TestingContext,
     formats: &[wgpu::TextureFormat],
@@ -251,8 +248,8 @@ fn clear_texture_tests(
 ) {
     for &format in formats {
         let desc = format.describe();
-        let rounded_width = round_up(64, desc.block_dimensions.0 as u32);
-        let rounded_height = round_up(64, desc.block_dimensions.1 as u32);
+        let rounded_width = align_to(64, desc.block_dimensions.0 as u32);
+        let rounded_height = align_to(64, desc.block_dimensions.1 as u32);
 
         // 1D texture
         if supports_1d {

wgpu/tests/root.rs

@@ -6,5 +6,6 @@ mod device;
 mod example_wgsl;
 mod instance;
 mod poll;
+mod shader_primitive_index;
 mod vertex_indices;
 mod zero_init_texture_after_discard;

wgpu/tests/shader_primitive_index/mod.rs

@@ -0,0 +1,247 @@
+use crate::common::{initialize_test, TestParameters, TestingContext};
+use std::num::NonZeroU32;
+use wgpu::util::{align_to, DeviceExt};
+
+//
+// These tests render two triangles to a 2x2 render target. The first triangle
+// in the vertex buffer covers the bottom-left pixel, the second triangle
+// covers the top-right pixel.
+// XY layout of the render target, with two triangles:
+//
+//     (-1,1)   (0,1)   (1,1)
+//        +-------+-------+
+//        |       |   o   |
+//        |       |  / \  |
+//        |       | /   \ |
+//        |       |o-----o|
+// (-1,0) +-------+-------+ (1,0)
+//        |   o   |       |
+//        |  / \  |       |
+//        | /   \ |       |
+//        |o-----o|       |
+//        +-------+-------+
+//     (-1,-1)  (0,-1)  (1,-1)
+//
+//
+// The fragment shader outputs color based on builtin(primitive_index):
+//
+//         if ((index % 2u) == 0u) {
+//             return vec4<f32>(1.0, 0.0, 0.0, 1.0);
+//         } else {
+//             return vec4<f32>(0.0, 0.0, 1.0, 1.0);
+//         }
+//
+// draw() renders directly from the vertex buffer: the first (bottom-left)
+// triangle is colored red, the other one (top-right) will be blue.
+// draw_indexed() draws the triangles in the opposite order, using index
+// buffer [3, 4, 5, 0, 1, 2]. This also swaps the resulting pixel colors.
+//
+#[test]
+fn draw() {
+    //
+    //   +-----+-----+
+    //   |white|blue |
+    //   +-----+-----+
+    //   | red |white|
+    //   +-----+-----+
+    //
+    let expected = [
+        255, 255, 255, 255, 0, 0, 255, 255, 255, 0, 0, 255, 255, 255, 255, 255,
+    ];
+    initialize_test(
+        TestParameters::default()
+            .test_features_limits()
+            .features(wgpu::Features::SHADER_PRIMITIVE_INDEX),
+        |ctx| {
+            pulling_common(ctx, &expected, |rpass| {
+                rpass.draw(0..6, 0..1);
+            })
+        },
+    );
+}
+
+#[test]
+fn draw_indexed() {
+    //
+    //   +-----+-----+
+    //   |white| red |
+    //   +-----+-----+
+    //   |blue |white|
+    //   +-----+-----+
+    //
+    let expected = [
+        255, 255, 255, 255, 255, 0, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,
+    ];
+    initialize_test(
+        TestParameters::default()
+            .test_features_limits()
+            .features(wgpu::Features::SHADER_PRIMITIVE_INDEX),
+        |ctx| {
+            pulling_common(ctx, &expected, |rpass| {
+                rpass.draw_indexed(0..6, 0, 0..1);
+            })
+        },
+    );
+}
+
+fn pulling_common(
+    ctx: TestingContext,
+    expected: &[u8],
+    function: impl FnOnce(&mut wgpu::RenderPass<'_>),
+) {
+    let shader = ctx
+        .device
+        .create_shader_module(&wgpu::include_wgsl!("primitive_index.wgsl"));
+
+    let two_triangles_xy: [f32; 12] = [
+        -1.0, -1.0, 0.0, -1.0, -0.5, 0.0, // left triangle, negative x, negative y
+        0.0, 0.0, 1.0, 0.0, 0.5, 1.0, // right triangle, positive x, positive y
+    ];
+    let vertex_buffer = ctx
+        .device
+        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: None,
+            contents: bytemuck::cast_slice(&two_triangles_xy),
+            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
+        });
+
+    let indices = [3u32, 4, 5, 0, 1, 2]; // index buffer flips triangle order
+    let index_buffer = ctx
+        .device
+        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: None,
+            contents: bytemuck::cast_slice(&indices),
+            usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
+        });
+
+    let pipeline = ctx
+        .device
+        .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: None,
+            layout: None,
+            vertex: wgpu::VertexState {
+                buffers: &[wgpu::VertexBufferLayout {
+                    array_stride: 8,
+                    step_mode: wgpu::VertexStepMode::Vertex,
+                    attributes: &[wgpu::VertexAttribute {
+                        format: wgpu::VertexFormat::Float32x2,
+                        offset: 0,
+                        shader_location: 0,
+                    }],
+                }],
+                entry_point: "vs_main",
+                module: &shader,
+            },
+            primitive: wgpu::PrimitiveState::default(),
+            depth_stencil: None,
+            multisample: wgpu::MultisampleState::default(),
+            fragment: Some(wgpu::FragmentState {
+                entry_point: "fs_main",
+                module: &shader,
+                targets: &[wgpu::ColorTargetState {
+                    format: wgpu::TextureFormat::Rgba8Unorm,
+                    blend: None,
+                    write_mask: wgpu::ColorWrites::ALL,
+                }],
+            }),
+            multiview: None,
+        });
+
+    let width = 2;
+    let height = 2;
+    let texture_size = wgpu::Extent3d {
+        width,
+        height,
+        depth_or_array_layers: 1,
+    };
+    let color_texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
+        label: None,
+        size: texture_size,
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        format: wgpu::TextureFormat::Rgba8Unorm,
+        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
+    });
+    let color_view = color_texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+    let mut encoder = ctx
+        .device
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+
+    let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+        color_attachments: &[wgpu::RenderPassColorAttachment {
+            ops: wgpu::Operations {
+                load: wgpu::LoadOp::Clear(wgpu::Color::WHITE),
+                store: true,
+            },
+            resolve_target: None,
+            view: &color_view,
+        }],
+        depth_stencil_attachment: None,
+        label: None,
+    });
+
+    rpass.set_pipeline(&pipeline);
+    rpass.set_index_buffer(index_buffer.slice(..), wgpu::IndexFormat::Uint32);
+    rpass.set_vertex_buffer(0, vertex_buffer.slice(..));
+    function(&mut rpass);
+
+    drop(rpass);
+
+    ctx.queue.submit(Some(encoder.finish()));
+
+    let data = capture_rgba_u8_texture(ctx, color_texture, texture_size);
+
+    assert_eq!(data, expected);
+}
+
+fn capture_rgba_u8_texture(
+    ctx: TestingContext,
+    color_texture: wgpu::Texture,
+    texture_size: wgpu::Extent3d,
+) -> Vec<u8> {
+    let bytes_per_row = align_to(4 * texture_size.width, wgpu::COPY_BYTES_PER_ROW_ALIGNMENT);
+    let buffer_size = bytes_per_row * texture_size.height;
+    let output_buffer = ctx
+        .device
+        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: None,
+            contents: &vec![0u8; buffer_size as usize],
+            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+        });
+
+    let mut encoder = ctx
+        .device
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+
+    encoder.copy_texture_to_buffer(
+        wgpu::ImageCopyTexture {
+            texture: &color_texture,
+            mip_level: 0,
+            origin: wgpu::Origin3d::ZERO,
+            aspect: wgpu::TextureAspect::All,
+        },
+        wgpu::ImageCopyBuffer {
+            buffer: &output_buffer,
+            layout: wgpu::ImageDataLayout {
+                offset: 0,
+                bytes_per_row: NonZeroU32::new(bytes_per_row),
+                rows_per_image: None,
+            },
+        },
+        texture_size,
+    );
+
+    ctx.queue.submit(Some(encoder.finish()));
+    let slice = output_buffer.slice(..);
+    slice.map_async(wgpu::MapMode::Read, |_| ());
+    ctx.device.poll(wgpu::Maintain::Wait);
+    let data: Vec<u8> = bytemuck::cast_slice(&*slice.get_mapped_range()).to_vec();
+    // Chunk rows from output buffer, take actual pixel
+    // bytes from each row and flatten into a vector.
+    data.chunks_exact(bytes_per_row as usize)
+        .flat_map(|x| x.iter().take(4 * texture_size.width as usize))
+        .copied()
+        .collect()
+}

wgpu/tests/shader_primitive_index/primitive_index.wgsl

@@ -0,0 +1,13 @@
+@vertex
+fn vs_main(@location(0) xy: vec2<f32>) -> @builtin(position) vec4<f32> {
+    return vec4<f32>(xy, 0.0, 1.0);
+}
+
+@fragment
+fn fs_main(@builtin(primitive_index) index: u32) -> @location(0) vec4<f32> {
+    if ((index % 2u) == 0u) {
+        return vec4<f32>(1.0, 0.0, 0.0, 1.0);
+    } else {
+        return vec4<f32>(0.0, 0.0, 1.0, 1.0);
+    }
+}