Skip to content

Commit

Permalink
Use only one descriptor set in Vulkan.
Browse files Browse the repository at this point in the history 
This reflects how Metal works, and should allow using low-power GPUs.
  • Loading branch information
@zlogic
zlogic committed Apr 6, 2024
1 parent 22b9d57 commit 90e08a3
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 54 deletions.
4 changes: 2 additions & 2 deletions src/correlation/gpu/shaders/cross_check_filter.comp.glsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,11 @@ layout(std430, push_constant) uniform readonly Parameters
float extend_range;
float min_range;
};
layout(std430, set = 1, binding = 0) buffer Img1
layout(std430, set = 0, binding = 0) buffer Img1
{
ivec2 img1[];
};
layout(std430, set = 1, binding = 1) buffer readonly Img2
layout(std430, set = 0, binding = 1) buffer readonly Img2
{
ivec2 img2[];
};
Expand Down
Binary file modified src/correlation/gpu/shaders/cross_check_filter.spv
View file
Binary file not shown.
119 changes: 67 additions & 52 deletions src/correlation/gpu/vulkan.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ pub struct Device {
device: ash::Device,
memory_properties: vk::PhysicalDeviceMemoryProperties,
buffers: Option<DeviceBuffers>,
direction: CorrelationDirection,
max_buffer_size: usize,
descriptor_sets: DescriptorSets,
pipelines: HashMap<ShaderModuleType, ShaderPipeline>,
Expand Down Expand Up @@ -229,6 +230,7 @@ impl Device {
instance.destroy_instance(None);
err
};
let direction = CorrelationDirection::Forward;
// Init control struct - queues, fences, command buffer.
let control =
unsafe { Device::create_control(&device, compute_queue_index).map_err(cleanup_err)? };
Expand All @@ -238,6 +240,7 @@ impl Device {
device,
memory_properties,
buffers: None,
direction,
max_buffer_size,
descriptor_sets,
pipelines,
Expand Down Expand Up @@ -709,7 +712,7 @@ impl Device {
.ty(vk::DescriptorType::STORAGE_BUFFER)
.descriptor_count(6)];
let descriptor_pool_info = vk::DescriptorPoolCreateInfo::default()
.max_sets(2)
.max_sets(1)
.pool_sizes(&descriptor_pool_size);
let descriptor_pool = device.create_descriptor_pool(&descriptor_pool_info, None)?;
let cleanup_err = |err| {
Expand Down Expand Up @@ -836,6 +839,67 @@ impl Device {
Ok(result)
}

fn set_buffer_layout(&mut self, shader: &ShaderModuleType) -> Result<(), GpuError> {
let direction = self.direction;
let descriptor_sets = &self.descriptor_sets;
let buffers = &self.buffers()?;
let create_buffer_infos = |buffers: &[Buffer]| {
buffers
.iter()
.map(|buf| {
vk::DescriptorBufferInfo::default()
.buffer(buf.buffer)
.offset(0)
.range(vk::WHOLE_SIZE)
})
.collect::<Vec<_>>()
};
let (buffer_internal_img1, buffer_internal_img2, buffer_out, buffer_out_reverse) =
match direction {
CorrelationDirection::Forward => (
buffers.buffer_internal_img1,
buffers.buffer_internal_img2,
buffers.buffer_out,
buffers.buffer_out_reverse,
),
CorrelationDirection::Reverse => (
buffers.buffer_internal_img2,
buffers.buffer_internal_img1,
buffers.buffer_out_reverse,
buffers.buffer_out,
),
};
let (buffer_list, descriptor_set) = if matches!(shader, ShaderModuleType::CrossCheckFilter)
{
(
vec![buffer_out, buffer_out_reverse],
descriptor_sets.descriptor_sets[0],
)
} else {
(
vec![
buffers.buffer_img,
buffer_internal_img1,
buffer_internal_img2,
buffers.buffer_internal_int,
buffer_out,
buffers.buffer_out_corr,
],
descriptor_sets.descriptor_sets[1],
)
};
let buffer_infos = create_buffer_infos(buffer_list.as_slice());
let write_descriptor = vk::WriteDescriptorSet::default()
.dst_set(descriptor_set)
.dst_binding(0)
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
.buffer_info(buffer_infos.as_slice());
unsafe {
self.device.update_descriptor_sets(&[write_descriptor], &[]);
}
Ok(())
}

unsafe fn create_control(
device: &ash::Device,
queue_family_index: u32,
Expand Down Expand Up @@ -876,57 +940,7 @@ impl Device {

impl super::Device for Device {
fn set_buffer_direction(&mut self, direction: &CorrelationDirection) -> Result<(), GpuError> {
let descriptor_sets = &self.descriptor_sets;
let buffers = &self.buffers()?;
let create_buffer_infos = |buffers: &[Buffer]| {
buffers
.iter()
.map(|buf| {
vk::DescriptorBufferInfo::default()
.buffer(buf.buffer)
.offset(0)
.range(vk::WHOLE_SIZE)
})
.collect::<Vec<_>>()
};
let (buffer_internal_img1, buffer_internal_img2, buffer_out, buffer_out_reverse) =
match direction {
CorrelationDirection::Forward => (
buffers.buffer_internal_img1,
buffers.buffer_internal_img2,
buffers.buffer_out,
buffers.buffer_out_reverse,
),
CorrelationDirection::Reverse => (
buffers.buffer_internal_img2,
buffers.buffer_internal_img1,
buffers.buffer_out_reverse,
buffers.buffer_out,
),
};
let regular_buffer_infos = create_buffer_infos(&[
buffers.buffer_img,
buffer_internal_img1,
buffer_internal_img2,
buffers.buffer_internal_int,
buffer_out,
buffers.buffer_out_corr,
]);
let regular_write_descriptor = vk::WriteDescriptorSet::default()
.dst_set(descriptor_sets.descriptor_sets[0])
.dst_binding(0)
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
.buffer_info(regular_buffer_infos.as_slice());
let cross_check_buffer_infos = create_buffer_infos(&[buffer_out, buffer_out_reverse]);
let cross_check_write_descriptor = vk::WriteDescriptorSet::default()
.dst_set(descriptor_sets.descriptor_sets[1])
.dst_binding(0)
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
.buffer_info(cross_check_buffer_infos.as_slice());
let write_descriptors = [regular_write_descriptor, cross_check_write_descriptor];
unsafe {
self.device.update_descriptor_sets(&write_descriptors, &[]);
}
self.direction = direction.to_owned();
Ok(())
}

Expand All @@ -952,6 +966,7 @@ impl super::Device for Device {
vk::PipelineBindPoint::COMPUTE,
pipeline_config.pipeline,
);
self.set_buffer_layout(&shader_type)?;
// It's way easier to map all descriptor sets identically, instead of ensuring that every
// kernel gets to use set = 0.
// The cross correlation kernel will need to switch to descriptor set = 1.
Expand Down

0 comments on commit 90e08a3

Please sign in to comment.