1use crate::types::error::Result;
2
3use ash::{Entry, Instance};
4use ash::vk;
5
6#[cfg(feature = "nvidia")]
7use std::{ffi::c_void, os::unix::io::RawFd};
8#[cfg(feature = "nvidia")]
9use ash::{ext, khr, Device};
10#[cfg(feature = "nvidia")]
11use crate::types::video_frame::DmaBufPlane;
12
13#[cfg(feature = "nvidia")]
14const DRM_FORMAT_MOD_INVALID: u64 = 0x00ff_ffff_ffff_ffff;
15
16#[cfg(feature = "vaapi")]
17#[derive(Clone, Copy)]
18#[allow(clippy::upper_case_acronyms)]
19pub enum GpuVendor {
20 NVIDIA,
21 AMD,
22 INTEL,
23 UNKNOWN,
24}
25
26#[cfg(feature = "vaapi")]
27impl GpuVendor {
28 fn from_vendor_id(id: u32) -> Self {
29 match id {
30 0x10DE => GpuVendor::NVIDIA,
31 0x1002 => GpuVendor::AMD,
32 0x8086 => GpuVendor::INTEL,
33 _ => {
34 log::error!("Unknown GPU vendor ID: 0x{id:04X}");
35 GpuVendor::UNKNOWN
36 }
37 }
38 }
39}
40
41#[cfg(feature = "vaapi")]
42pub fn detect_gpu_vendor() -> Result<GpuVendor> {
43 let entry = unsafe { Entry::load() }.map_err(|e| format!("Failed to load Vulkan: {e}"))?;
44
45 let app_name = c"waycap-rs";
46 let app_info = vk::ApplicationInfo::default()
47 .application_name(app_name)
48 .api_version(vk::API_VERSION_1_1);
49 let instance_ci = vk::InstanceCreateInfo::default().application_info(&app_info);
50 let instance = unsafe { entry.create_instance(&instance_ci, None) }
51 .map_err(|e| format!("Failed to create Vulkan instance: {e}"))?;
52
53 let physical_devices = unsafe { instance.enumerate_physical_devices() }
54 .map_err(|e| format!("Failed to enumerate physical devices: {e}"))?;
55
56 let vendor = physical_devices
57 .first()
58 .map(|&pd| {
59 let props = unsafe { instance.get_physical_device_properties(pd) };
60 GpuVendor::from_vendor_id(props.vendor_id)
61 })
62 .unwrap_or(GpuVendor::UNKNOWN);
63
64 unsafe { instance.destroy_instance(None) };
65 Ok(vendor)
66}
67
68unsafe impl Send for VulkanContext {}
69unsafe impl Sync for VulkanContext {}
70
71pub struct VulkanContext {
72 _entry: Entry,
73 instance: Instance,
74 physical_device: vk::PhysicalDevice,
75 device: Device,
76 queue: vk::Queue,
77 #[allow(dead_code)]
78 queue_family_index: u32,
79 command_pool: vk::CommandPool,
80
81 external_memory_fd: khr::external_memory_fd::Device,
82
83 persistent_buffer: vk::Buffer,
84 persistent_buffer_memory: vk::DeviceMemory,
85 persistent_buffer_size: u64,
86
87 #[allow(dead_code)]
88 width: u32,
89 #[allow(dead_code)]
90 height: u32,
91}
92
93impl VulkanContext {
94 pub fn new(width: u32, height: u32) -> Result<Self> {
95 let entry = unsafe { Entry::load() }.map_err(|e| format!("Failed to load Vulkan: {e}"))?;
96
97 let app_name = c"waycap-rs";
98 let app_info = vk::ApplicationInfo::default()
99 .application_name(app_name)
100 .api_version(vk::API_VERSION_1_1);
101
102 let instance_ci = vk::InstanceCreateInfo::default().application_info(&app_info);
103 let instance = unsafe { entry.create_instance(&instance_ci, None) }
104 .map_err(|e| format!("Failed to create Vulkan instance: {e}"))?;
105
106 let physical_devices = unsafe { instance.enumerate_physical_devices() }
107 .map_err(|e| format!("Failed to enumerate physical devices: {e}"))?;
108 if physical_devices.is_empty() {
109 return Err("No Vulkan physical devices found".into());
110 }
111
112 let (physical_device, queue_family_index) = physical_devices
113 .iter()
114 .find_map(|&pd| {
115 let qfs = unsafe { instance.get_physical_device_queue_family_properties(pd) };
116 qfs.iter().enumerate().find_map(|(i, qf)| {
117 if qf
118 .queue_flags
119 .contains(vk::QueueFlags::GRAPHICS | vk::QueueFlags::TRANSFER)
120 {
121 Some((pd, i as u32))
122 } else {
123 None
124 }
125 })
126 })
127 .ok_or("No suitable Vulkan queue family found")?;
128
129 let queue_priorities = [1.0_f32];
130 let queue_ci = vk::DeviceQueueCreateInfo::default()
131 .queue_family_index(queue_family_index)
132 .queue_priorities(&queue_priorities);
133
134 let device_exts = [
135 khr::external_memory_fd::NAME.as_ptr(),
136 ext::external_memory_dma_buf::NAME.as_ptr(),
137 ext::image_drm_format_modifier::NAME.as_ptr(),
138 ];
139 let device_ci = vk::DeviceCreateInfo::default()
140 .queue_create_infos(std::slice::from_ref(&queue_ci))
141 .enabled_extension_names(&device_exts);
142 let device = unsafe { instance.create_device(physical_device, &device_ci, None) }
143 .map_err(|e| format!("Failed to create Vulkan device: {e}"))?;
144
145 let queue = unsafe { device.get_device_queue(queue_family_index, 0) };
146
147 let pool_ci = vk::CommandPoolCreateInfo::default()
148 .queue_family_index(queue_family_index)
149 .flags(
150 vk::CommandPoolCreateFlags::TRANSIENT
151 | vk::CommandPoolCreateFlags::RESET_COMMAND_BUFFER,
152 );
153 let command_pool = unsafe { device.create_command_pool(&pool_ci, None) }
154 .map_err(|e| format!("Failed to create command pool: {e}"))?;
155
156 let external_memory_fd = khr::external_memory_fd::Device::new(&instance, &device);
157
158 let mem_props = unsafe { instance.get_physical_device_memory_properties(physical_device) };
159
160 let persistent_buffer_size = (width * height * 4) as u64;
161 let (persistent_buffer, persistent_buffer_memory) =
162 Self::create_persistent_buffer(&device, &mem_props, persistent_buffer_size)?;
163
164 Ok(Self {
165 _entry: entry,
166 instance,
167 physical_device,
168 device,
169 queue,
170 queue_family_index,
171 command_pool,
172 external_memory_fd,
173 persistent_buffer,
174 persistent_buffer_memory,
175 persistent_buffer_size,
176 width,
177 height,
178 })
179 }
180
181 fn create_persistent_buffer(
182 device: &Device,
183 mem_props: &vk::PhysicalDeviceMemoryProperties,
184 size: u64,
185 ) -> Result<(vk::Buffer, vk::DeviceMemory)> {
186 let mut export_info = vk::ExternalMemoryBufferCreateInfo::default()
187 .handle_types(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD);
188
189 let buf_ci = vk::BufferCreateInfo::default()
190 .size(size)
191 .usage(vk::BufferUsageFlags::TRANSFER_DST)
192 .sharing_mode(vk::SharingMode::EXCLUSIVE)
193 .push_next(&mut export_info);
194
195 let buffer = unsafe { device.create_buffer(&buf_ci, None) }
196 .map_err(|e| format!("Failed to create persistent buffer: {e}"))?;
197
198 let mem_reqs = unsafe { device.get_buffer_memory_requirements(buffer) };
199
200 let memory_type_index = find_memory_type(
201 mem_props,
202 mem_reqs.memory_type_bits,
203 vk::MemoryPropertyFlags::DEVICE_LOCAL,
204 )
205 .ok_or("No DEVICE_LOCAL memory type for persistent buffer")?;
206
207 let mut export_alloc = vk::ExportMemoryAllocateInfo::default()
208 .handle_types(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD);
209
210 let alloc_info = vk::MemoryAllocateInfo::default()
211 .allocation_size(mem_reqs.size)
212 .memory_type_index(memory_type_index)
213 .push_next(&mut export_alloc);
214
215 let memory = unsafe { device.allocate_memory(&alloc_info, None) }.map_err(|e| {
216 unsafe { device.destroy_buffer(buffer, None) };
217 format!("Failed to allocate persistent buffer memory: {e}")
218 })?;
219
220 unsafe { device.bind_buffer_memory(buffer, memory, 0) }.map_err(|e| {
221 unsafe {
222 device.free_memory(memory, None);
223 device.destroy_buffer(buffer, None);
224 }
225 format!("Failed to bind persistent buffer memory: {e}")
226 })?;
227
228 Ok((buffer, memory))
229 }
230
231 pub fn export_persistent_memory_fd(&self) -> Result<RawFd> {
232 let get_fd_info = vk::MemoryGetFdInfoKHR::default()
233 .memory(self.persistent_buffer_memory)
234 .handle_type(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD);
235 unsafe { self.external_memory_fd.get_memory_fd(&get_fd_info) }
236 .map_err(|e| format!("Failed to export Vulkan memory FD: {e}").into())
237 }
238
239 pub fn get_persistent_buffer_size(&self) -> u64 {
240 self.persistent_buffer_size
241 }
242
243 pub fn copy_dmabuf_to_persistent_buffer(
244 &self,
245 planes: &[DmaBufPlane],
246 modifier: u64,
247 width: u32,
248 height: u32,
249 ) -> Result<()> {
250 let plane = planes.first().ok_or("No DMA-BUF planes provided")?;
251
252 let dup_fd = unsafe { libc::dup(plane.fd) };
253 if dup_fd < 0 {
254 return Err("Failed to dup DMA-BUF fd".into());
255 }
256
257 let (temp_image, temp_memory) = self
258 .import_dmabuf_as_image(dup_fd, plane, modifier, width, height)
259 .inspect_err(|_| {
260 unsafe { libc::close(dup_fd) };
261 })?;
262
263 let copy_result = self.record_and_submit_copy(temp_image, width, height);
264
265 unsafe {
266 self.device.destroy_image(temp_image, None);
267 self.device.free_memory(temp_memory, None);
268 }
269
270 copy_result
271 }
272
273 fn import_dmabuf_as_image(
274 &self,
275 fd: RawFd,
276 plane: &DmaBufPlane,
277 modifier: u64,
278 width: u32,
279 height: u32,
280 ) -> Result<(vk::Image, vk::DeviceMemory)> {
281 let mem_props = unsafe {
282 self.instance
283 .get_physical_device_memory_properties(self.physical_device)
284 };
285
286 let plane_layout = vk::SubresourceLayout {
287 offset: plane.offset as u64,
288 size: 0,
289 row_pitch: plane.stride as u64,
290 array_pitch: 0,
291 depth_pitch: 0,
292 };
293
294 let effective_modifier = if modifier == DRM_FORMAT_MOD_INVALID {
295 0
296 } else {
297 modifier
298 };
299
300 let mut modifier_info = vk::ImageDrmFormatModifierExplicitCreateInfoEXT::default()
301 .drm_format_modifier(effective_modifier)
302 .plane_layouts(std::slice::from_ref(&plane_layout));
303
304 let mut external_image_info = vk::ExternalMemoryImageCreateInfo::default()
305 .handle_types(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT);
306 external_image_info.p_next = (&raw mut modifier_info).cast::<c_void>();
307
308 let image_ci = vk::ImageCreateInfo::default()
309 .image_type(vk::ImageType::TYPE_2D)
310 .format(vk::Format::B8G8R8A8_UNORM)
311 .extent(vk::Extent3D {
312 width,
313 height,
314 depth: 1,
315 })
316 .mip_levels(1)
317 .array_layers(1)
318 .samples(vk::SampleCountFlags::TYPE_1)
319 .tiling(vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT)
320 .usage(vk::ImageUsageFlags::TRANSFER_SRC)
321 .sharing_mode(vk::SharingMode::EXCLUSIVE)
322 .initial_layout(vk::ImageLayout::UNDEFINED)
323 .push_next(&mut external_image_info);
324
325 let image = unsafe { self.device.create_image(&image_ci, None) }
326 .map_err(|e| format!("Failed to create DMA-BUF image: {e}"))?;
327
328 let mut dedicated_reqs = vk::MemoryDedicatedRequirements::default();
329 let mut mem_reqs2 = vk::MemoryRequirements2 {
330 p_next: (&raw mut dedicated_reqs).cast::<c_void>(),
331 ..Default::default()
332 };
333
334 let image_mem_reqs_info = vk::ImageMemoryRequirementsInfo2::default().image(image);
335 unsafe {
336 self.device
337 .get_image_memory_requirements2(&image_mem_reqs_info, &mut mem_reqs2)
338 };
339
340 let memory_type_index = find_memory_type(
341 &mem_props,
342 mem_reqs2.memory_requirements.memory_type_bits,
343 vk::MemoryPropertyFlags::DEVICE_LOCAL,
344 )
345 .ok_or("No suitable memory type for DMA-BUF import")?;
346
347 let mut import_info = vk::ImportMemoryFdInfoKHR::default()
348 .handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT)
349 .fd(fd);
350
351 let mut dedicated_alloc = vk::MemoryDedicatedAllocateInfo::default().image(image);
352 dedicated_alloc.p_next = (&raw mut import_info).cast::<c_void>();
353
354 let alloc_info = vk::MemoryAllocateInfo::default()
355 .allocation_size(mem_reqs2.memory_requirements.size)
356 .memory_type_index(memory_type_index)
357 .push_next(&mut dedicated_alloc);
358
359 let memory = unsafe { self.device.allocate_memory(&alloc_info, None) }.map_err(|e| {
360 unsafe { self.device.destroy_image(image, None) };
361 format!("Failed to allocate DMA-BUF memory: {e}")
362 })?;
363
364 unsafe { self.device.bind_image_memory(image, memory, 0) }.map_err(|e| {
365 unsafe {
366 self.device.free_memory(memory, None);
367 self.device.destroy_image(image, None);
368 }
369 format!("Failed to bind DMA-BUF image memory: {e}")
370 })?;
371
372 Ok((image, memory))
373 }
374
375 fn record_and_submit_copy(&self, src_image: vk::Image, width: u32, height: u32) -> Result<()> {
376 let cmd_buf = alloc_cmd_buf(&self.device, self.command_pool)?;
377
378 let begin_info = vk::CommandBufferBeginInfo::default()
379 .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT);
380 unsafe { self.device.begin_command_buffer(cmd_buf, &begin_info) }
381 .map_err(|e| format!("begin_command_buffer: {e}"))?;
382
383 let src_barrier = vk::ImageMemoryBarrier::default()
384 .old_layout(vk::ImageLayout::UNDEFINED)
385 .new_layout(vk::ImageLayout::GENERAL)
386 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
387 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
388 .image(src_image)
389 .subresource_range(color_subresource_range())
390 .src_access_mask(vk::AccessFlags::empty())
391 .dst_access_mask(vk::AccessFlags::TRANSFER_READ);
392
393 unsafe {
394 self.device.cmd_pipeline_barrier(
395 cmd_buf,
396 vk::PipelineStageFlags::TOP_OF_PIPE,
397 vk::PipelineStageFlags::TRANSFER,
398 vk::DependencyFlags::empty(),
399 &[],
400 &[],
401 std::slice::from_ref(&src_barrier),
402 );
403 }
404
405 let region = vk::BufferImageCopy {
406 buffer_offset: 0,
407 buffer_row_length: 0,
408 buffer_image_height: 0,
409 image_subresource: vk::ImageSubresourceLayers {
410 aspect_mask: vk::ImageAspectFlags::COLOR,
411 mip_level: 0,
412 base_array_layer: 0,
413 layer_count: 1,
414 },
415 image_offset: vk::Offset3D::default(),
416 image_extent: vk::Extent3D {
417 width,
418 height,
419 depth: 1,
420 },
421 };
422
423 unsafe {
424 self.device.cmd_copy_image_to_buffer(
425 cmd_buf,
426 src_image,
427 vk::ImageLayout::GENERAL,
428 self.persistent_buffer,
429 std::slice::from_ref(®ion),
430 );
431 }
432
433 let buf_barrier = vk::BufferMemoryBarrier::default()
434 .src_access_mask(vk::AccessFlags::TRANSFER_WRITE)
435 .dst_access_mask(vk::AccessFlags::MEMORY_READ)
436 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
437 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
438 .buffer(self.persistent_buffer)
439 .offset(0)
440 .size(vk::WHOLE_SIZE);
441
442 unsafe {
443 self.device.cmd_pipeline_barrier(
444 cmd_buf,
445 vk::PipelineStageFlags::TRANSFER,
446 vk::PipelineStageFlags::BOTTOM_OF_PIPE,
447 vk::DependencyFlags::empty(),
448 &[],
449 std::slice::from_ref(&buf_barrier),
450 &[],
451 );
452 }
453
454 submit_and_wait(&self.device, self.queue, cmd_buf)?;
455 unsafe {
456 self.device
457 .free_command_buffers(self.command_pool, &[cmd_buf])
458 };
459
460 Ok(())
461 }
462}
463
464impl Drop for VulkanContext {
465 fn drop(&mut self) {
466 unsafe {
467 let _ = self.device.device_wait_idle();
468 self.device.destroy_buffer(self.persistent_buffer, None);
469 self.device.free_memory(self.persistent_buffer_memory, None);
470 self.device.destroy_command_pool(self.command_pool, None);
471 self.device.destroy_device(None);
472 self.instance.destroy_instance(None);
473 }
474 }
475}
476
477fn find_memory_type(
478 mem_props: &vk::PhysicalDeviceMemoryProperties,
479 type_filter: u32,
480 flags: vk::MemoryPropertyFlags,
481) -> Option<u32> {
482 (0..mem_props.memory_type_count).find(|&i| {
483 (type_filter & (1 << i)) != 0
484 && mem_props.memory_types[i as usize]
485 .property_flags
486 .contains(flags)
487 })
488}
489
490fn color_subresource_range() -> vk::ImageSubresourceRange {
491 vk::ImageSubresourceRange {
492 aspect_mask: vk::ImageAspectFlags::COLOR,
493 base_mip_level: 0,
494 level_count: 1,
495 base_array_layer: 0,
496 layer_count: 1,
497 }
498}
499
500fn alloc_cmd_buf(device: &Device, pool: vk::CommandPool) -> Result<vk::CommandBuffer> {
501 let alloc_info = vk::CommandBufferAllocateInfo::default()
502 .command_pool(pool)
503 .level(vk::CommandBufferLevel::PRIMARY)
504 .command_buffer_count(1);
505 let bufs = unsafe { device.allocate_command_buffers(&alloc_info) }
506 .map_err(|e| format!("Failed to allocate command buffer: {e}"))?;
507 Ok(bufs[0])
508}
509
510fn submit_and_wait(device: &Device, queue: vk::Queue, cmd_buf: vk::CommandBuffer) -> Result<()> {
511 unsafe { device.end_command_buffer(cmd_buf) }
512 .map_err(|e| format!("end_command_buffer: {e}"))?;
513
514 let fence = unsafe { device.create_fence(&vk::FenceCreateInfo::default(), None) }
515 .map_err(|e| format!("create_fence: {e}"))?;
516
517 let submit = vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&cmd_buf));
518 unsafe {
519 device
520 .queue_submit(queue, std::slice::from_ref(&submit), fence)
521 .map_err(|e| format!("queue_submit: {e}"))?;
522 device
523 .wait_for_fences(std::slice::from_ref(&fence), true, u64::MAX)
524 .map_err(|e| format!("wait_for_fences: {e}"))?;
525 device.destroy_fence(fence, None);
526 }
527 Ok(())
528}