julia_set/
vulkan.rs

1//! GLSL / Vulkan backend for Julia sets.
2
3use std::{slice, sync::Arc};
4
5use anyhow::{Context as _, anyhow};
6use shaderc::{CompilationArtifact, CompileOptions, OptimizationLevel, ShaderKind};
7use vulkano::{
8    VulkanLibrary,
9    buffer::{Buffer, BufferContents, BufferCreateInfo, BufferUsage},
10    command_buffer::{
11        AutoCommandBufferBuilder, CommandBufferUsage,
12        allocator::{StandardCommandBufferAllocator, StandardCommandBufferAllocatorCreateInfo},
13    },
14    descriptor_set::{
15        DescriptorSet, WriteDescriptorSet,
16        allocator::{StandardDescriptorSetAllocator, StandardDescriptorSetAllocatorCreateInfo},
17    },
18    device::{
19        Device, DeviceCreateInfo, DeviceExtensions, Queue, QueueCreateInfo, QueueFlags,
20        physical::{PhysicalDevice, PhysicalDeviceType},
21    },
22    instance::{Instance, InstanceCreateFlags, InstanceCreateInfo},
23    memory::allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
24    pipeline::{
25        ComputePipeline, Pipeline, PipelineBindPoint, PipelineLayout,
26        PipelineShaderStageCreateInfo, compute::ComputePipelineCreateInfo,
27        layout::PipelineDescriptorSetLayoutCreateInfo,
28    },
29    shader::{ShaderModule, ShaderModuleCreateInfo},
30    sync::{self, GpuFuture},
31};
32
33use crate::{Backend, Function, ImageBuffer, Params, Render, compiler::Compiler};
34
35const PROGRAM: &str = include_str!(concat!(env!("OUT_DIR"), "/program.glsl"));
36
37const LOCAL_WORKGROUP_SIZES: [u32; 2] = [16, 16];
38
39fn compile_shader(function: &str) -> shaderc::Result<CompilationArtifact> {
40    let compiler = shaderc::Compiler::new()?;
41    let mut options = CompileOptions::new()?;
42    options.add_macro_definition("COMPUTE", Some(function));
43    options.set_optimization_level(OptimizationLevel::Performance);
44    compiler.compile_into_spirv(
45        PROGRAM,
46        ShaderKind::Compute,
47        "program.glsl",
48        "main",
49        Some(&options),
50    )
51}
52
53#[derive(Debug, Clone, Copy, BufferContents)]
54#[repr(C, packed)]
55struct VulkanParams {
56    view_center: [f32; 2],
57    view_size: [f32; 2],
58    image_size: [u32; 2],
59    inf_distance_sq: f32,
60    max_iterations: u32,
61}
62
63/// Backend based on [Vulkan].
64///
65/// [Vulkan]: https://www.khronos.org/vulkan/
66#[cfg_attr(docsrs, doc(cfg(feature = "vulkan_backend")))]
67#[derive(Debug, Clone, Default)]
68pub struct Vulkan;
69
70impl Backend<&Function> for Vulkan {
71    type Error = anyhow::Error;
72    type Program = VulkanProgram;
73
74    fn create_program(&self, function: &Function) -> Result<Self::Program, Self::Error> {
75        let compiled = Compiler::for_gl().compile(function);
76        VulkanProgram::new(&compiled)
77    }
78}
79
80/// Program produced by the [`Vulkan`] backend.
81#[cfg_attr(docsrs, doc(cfg(feature = "vulkan_backend")))]
82#[derive(Debug)]
83pub struct VulkanProgram {
84    device: Arc<Device>,
85    queue: Arc<Queue>,
86    pipeline: Arc<ComputePipeline>,
87    memory_allocator: Arc<StandardMemoryAllocator>,
88    descriptor_set_allocator: Arc<StandardDescriptorSetAllocator>,
89    command_buffer_allocator: Arc<StandardCommandBufferAllocator>,
90}
91
92impl VulkanProgram {
93    fn new(compiled_function: &str) -> anyhow::Result<Self> {
94        let library = VulkanLibrary::new()?;
95        let create_info = InstanceCreateInfo {
96            flags: InstanceCreateFlags::ENUMERATE_PORTABILITY,
97            ..InstanceCreateInfo::default()
98        };
99        let instance = Instance::new(library, create_info)?;
100
101        let device_extensions = DeviceExtensions {
102            khr_storage_buffer_storage_class: true,
103            ..DeviceExtensions::empty()
104        };
105        let (device, queue_family_index) =
106            Self::select_physical_device(&instance, &device_extensions)?;
107        let (device, mut queues) = Device::new(
108            device,
109            DeviceCreateInfo {
110                enabled_extensions: device_extensions,
111                queue_create_infos: vec![QueueCreateInfo {
112                    queue_family_index,
113                    ..QueueCreateInfo::default()
114                }],
115                ..DeviceCreateInfo::default()
116            },
117        )?;
118        let queue = queues
119            .next()
120            .with_context(|| format!("cannot initialize compute queue on device {device:?}"))?;
121
122        let shader = compile_shader(compiled_function)?;
123        let create_info = ShaderModuleCreateInfo::new(shader.as_binary());
124        let shader = unsafe { ShaderModule::new(device.clone(), create_info)? };
125        let entry_point = shader
126            .entry_point("main")
127            .context("cannot find entry point `main` in Julia set compute shader")?;
128
129        let stage = PipelineShaderStageCreateInfo::new(entry_point);
130        let layout = PipelineLayout::new(
131            device.clone(),
132            PipelineDescriptorSetLayoutCreateInfo::from_stages([&stage])
133                .into_pipeline_layout_create_info(device.clone())
134                .context("cannot create compute pipeline layout")?,
135        )
136        .context("failed validating pipeline layout")?;
137
138        let create_info = ComputePipelineCreateInfo::stage_layout(stage, layout);
139        let pipeline = ComputePipeline::new(device.clone(), None, create_info)?;
140        let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
141        let descriptor_set_allocator = Arc::new(StandardDescriptorSetAllocator::new(
142            device.clone(),
143            StandardDescriptorSetAllocatorCreateInfo::default(),
144        ));
145        let command_buffer_allocator = Arc::new(StandardCommandBufferAllocator::new(
146            device.clone(),
147            StandardCommandBufferAllocatorCreateInfo::default(),
148        ));
149
150        Ok(Self {
151            device,
152            queue,
153            pipeline,
154            memory_allocator,
155            descriptor_set_allocator,
156            command_buffer_allocator,
157        })
158    }
159
160    #[allow(clippy::cast_possible_truncation)]
161    fn select_physical_device(
162        instance: &Arc<Instance>,
163        device_extensions: &DeviceExtensions,
164    ) -> anyhow::Result<(Arc<PhysicalDevice>, u32)> {
165        let devices = instance.enumerate_physical_devices()?;
166        let devices =
167            devices.filter(|device| device.supported_extensions().contains(device_extensions));
168        let devices = devices.filter_map(|device| {
169            device
170                .queue_family_properties()
171                .iter()
172                .position(|q| q.queue_flags.intersects(QueueFlags::COMPUTE))
173                .map(|idx| (device, idx as u32))
174        });
175
176        let device = devices
177            .min_by_key(|(device, _)| Self::device_type_priority(device.properties().device_type));
178        device.ok_or_else(|| anyhow!("failed selecting physical device with compute queue"))
179    }
180
181    fn device_type_priority(ty: PhysicalDeviceType) -> usize {
182        match ty {
183            PhysicalDeviceType::DiscreteGpu => 0,
184            PhysicalDeviceType::IntegratedGpu => 1,
185            PhysicalDeviceType::VirtualGpu => 2,
186            PhysicalDeviceType::Cpu => 3,
187            _ => 4,
188        }
189    }
190}
191
192impl Render for VulkanProgram {
193    type Error = anyhow::Error;
194
195    #[allow(clippy::cast_possible_truncation)]
196    fn render(&self, params: &Params) -> anyhow::Result<ImageBuffer> {
197        // Bind uniforms: the output image buffer and the rendering params.
198        let pixel_count = (params.image_size[0] * params.image_size[1]) as usize;
199        let image_buffer = Buffer::new_slice::<u32>(
200            self.memory_allocator.clone(),
201            BufferCreateInfo {
202                usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_DST,
203                ..BufferCreateInfo::default()
204            },
205            AllocationCreateInfo {
206                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
207                    | MemoryTypeFilter::HOST_RANDOM_ACCESS,
208                ..AllocationCreateInfo::default()
209            },
210            pixel_count as u64,
211        )?;
212
213        let gl_params = VulkanParams {
214            view_center: params.view_center,
215            view_size: [params.view_width(), params.view_height],
216            image_size: params.image_size,
217            inf_distance_sq: params.inf_distance * params.inf_distance,
218            max_iterations: u32::from(params.max_iterations),
219        };
220
221        let layout = self.pipeline.layout();
222        let layout = &layout.set_layouts()[0];
223        let descriptor_set = DescriptorSet::new(
224            self.descriptor_set_allocator.clone(),
225            layout.clone(),
226            [WriteDescriptorSet::buffer(0, image_buffer.clone())],
227            [],
228        )?;
229
230        // Create the commands to render the image and copy it to the buffer.
231        let task_dimensions = [
232            params.image_size[0].div_ceil(LOCAL_WORKGROUP_SIZES[0]),
233            params.image_size[1].div_ceil(LOCAL_WORKGROUP_SIZES[1]),
234            1,
235        ];
236        let layout = self.pipeline.layout();
237        let mut builder = AutoCommandBufferBuilder::primary(
238            self.command_buffer_allocator.clone(),
239            self.queue.queue_family_index(),
240            CommandBufferUsage::OneTimeSubmit,
241        )?;
242        builder
243            .bind_pipeline_compute(self.pipeline.clone())
244            .context("failed binding compute pipeline to command buffer")?
245            .bind_descriptor_sets(
246                PipelineBindPoint::Compute,
247                layout.clone(),
248                0,
249                descriptor_set,
250            )
251            .context("failed binding descriptor sets to command buffer")?
252            .fill_buffer(image_buffer.clone(), 0)?
253            .push_constants(layout.clone(), 0, gl_params)
254            .context("failed pushing constants to command buffer")?;
255        unsafe {
256            builder.dispatch(task_dimensions)?;
257        }
258        let command_buffer = builder.build()?;
259
260        sync::now(self.device.clone())
261            .then_execute(self.queue.clone(), command_buffer)?
262            .then_signal_fence_and_flush()?
263            .wait(None)?;
264
265        // Convert the buffer into an `ImageBuffer`.
266        let buffer_content = image_buffer.read()?;
267        debug_assert!(buffer_content.len() * 4 >= pixel_count);
268        let buffer_content = unsafe {
269            // SAFETY: Buffer length is correct by construction, and `[u8]` doesn't require
270            // any special alignment.
271            slice::from_raw_parts(buffer_content.as_ptr().cast::<u8>(), pixel_count)
272        };
273
274        Ok(ImageBuffer::from_vec(
275            params.image_size[0],
276            params.image_size[1],
277            buffer_content.to_vec(),
278        )
279        .unwrap())
280    }
281}
julia_set/vulkan.rs

julia_set/
vulkan.rs
