external vulkan memory with vulkano (+ a bit of ash because something is wrong with `RawImage::new`), using dmabuf for memory and opaque fd for semaphore

Cargo.toml

[package]
name = "hello-vk"
version = "0.1.0"
edition = "2024"

[dependencies]
ash = "0.38.0"
image = "0.25.9"
vulkano = "0.35.2"
vulkano-shaders = "0.35.0"

main.rs

/// Exporting and importing memory/semaphore with file descriptor. The exported memory have DMA-BUF flag, allowing
/// different Vulkan (logical) devices to read memory from each other without copying data.
///
/// This code only work on Linux at current moment (since everything are tied to file descriptors). Tested on Intel Iris
/// Xe (Tiger Lake) using experimental xe driver.
use std::{
    error::Error,
    fs::File,
    ptr,
    sync::{Arc, mpsc},
    thread,
};

use ash::vk;
use image::{ImageBuffer, Rgba};
use vulkano::{
    VulkanLibrary, VulkanObject,
    buffer::{Buffer, BufferCreateInfo, BufferUsage},
    command_buffer::{
        AutoCommandBufferBuilder, ClearColorImageInfo, CommandBufferUsage, CopyImageToBufferInfo,
        allocator::StandardCommandBufferAllocator,
    },
    device::{Device, DeviceCreateInfo, DeviceExtensions, DeviceFeatures, Queue, QueueCreateInfo, QueueFlags},
    format::Format,
    image::{ImageCreateInfo, ImageTiling, ImageType, ImageUsage, sys::RawImage},
    instance::Instance,
    memory::{
        DeviceMemory, ExternalMemoryHandleType, ExternalMemoryHandleTypes, MemoryAllocateInfo, MemoryImportInfo,
        ResourceMemory,
        allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
    },
    sync::{
        fence::Fence,
        semaphore::{
            ExternalSemaphoreHandleType, ExternalSemaphoreHandleTypes, ImportSemaphoreFdInfo, Semaphore,
            SemaphoreCreateInfo, SemaphoreType,
        },
    },
};

/// Common function for creating Vulkan device and queue
fn factory(library: Arc<VulkanLibrary>) -> Result<(Arc<Device>, Arc<Queue>), Box<dyn Error>> {
    let instance = Instance::new(library, Default::default())?;
    let physical_device = instance.enumerate_physical_devices()?.next().unwrap();
    let queue_family_index = physical_device
        .queue_family_properties()
        .iter()
        .position(|q| q.queue_flags.contains(QueueFlags::COMPUTE))
        .unwrap() as u32;

    let (device, mut queues) = Device::new(
        physical_device.clone(),
        DeviceCreateInfo {
            enabled_features: DeviceFeatures {
                timeline_semaphore: true,
                ..Default::default()
            },
            enabled_extensions: DeviceExtensions {
                // TODO: If you are modifying the code to make it work on Windows, make sure to remove dma_buf/fd and
                // replace them with win32 handles.
                ext_external_memory_dma_buf: true,
                ext_image_drm_format_modifier: true,
                khr_external_semaphore: true,
                khr_external_semaphore_fd: true,
                ..Default::default()
            },
            queue_create_infos: vec![QueueCreateInfo {
                queue_family_index,
                ..Default::default()
            }],
            ..Default::default()
        },
    )?;

    let queue = queues.next().unwrap();
    Ok((device, queue))
}

struct Message {
    memory_file: File,
    semaphore_file: File,
}

/// Main function for producer thread. The producer thread will:
///
/// 1. Create exportable memory and exportable semaphore, then export both of them as file descriptor
/// 1. Send file descriptors to consumer thread
/// 1. Bind image to exportable memory
/// 1. Execute clear command then signal semaphore
fn producer_thread(library: Arc<VulkanLibrary>, sender: mpsc::Sender<Message>) -> Result<(), Box<dyn Error>> {
    let (device, queue) = factory(library)?;
    let command_buffer_allocator = Arc::new(StandardCommandBufferAllocator::new(device.clone(), Default::default()));
    let exported_device_memory = Arc::new(DeviceMemory::allocate(
        device.clone(),
        MemoryAllocateInfo {
            allocation_size: 256 * 1024 * 1024,
            memory_type_index: 0, // Device-only
            export_handle_types: ExternalMemoryHandleTypes::DMA_BUF,
            ..Default::default()
        },
    )?);

    let ash_device = unsafe { ash::Device::load(&device.instance().fns().v1_0, device.handle()) };
    let image = Arc::new(unsafe {
        let ash_image = ash_device.create_image(
            &vk::ImageCreateInfo {
                usage: vk::ImageUsageFlags::TRANSFER_DST,
                image_type: vk::ImageType::TYPE_2D,
                format: vk::Format::R8G8B8A8_UNORM,
                extent: vk::Extent3D {
                    width: 1024,
                    height: 1024,
                    depth: 1,
                    ..Default::default()
                },
                samples: vk::SampleCountFlags::TYPE_1,
                array_layers: 1,
                mip_levels: 1,
                tiling: vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT,
                sharing_mode: vk::SharingMode::EXCLUSIVE,
                queue_family_index_count: 0,
                p_queue_family_indices: ptr::null(),
                initial_layout: vk::ImageLayout::UNDEFINED,
                ..Default::default()
            }
            .push_next(&mut vk::ExternalMemoryImageCreateInfo {
                handle_types: vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT,
                ..Default::default()
            })
            .push_next(&mut vk::ImageDrmFormatModifierListCreateInfoEXT {
                drm_format_modifier_count: 1,
                p_drm_format_modifiers: [0u64].as_ptr(),
                ..Default::default()
            }),
            None,
        )?;
        RawImage::from_handle(
            device.clone(),
            ash_image,
            ImageCreateInfo {
                usage: ImageUsage::TRANSFER_DST,
                image_type: ImageType::Dim2d,
                format: Format::R8G8B8A8_UNORM,
                extent: [1024, 1024, 1],
                external_memory_handle_types: ExternalMemoryHandleTypes::DMA_BUF,
                tiling: ImageTiling::DrmFormatModifier,
                drm_format_modifiers: vec![0],
                ..Default::default()
            },
        )?
        .bind_memory([ResourceMemory::from_device_memory_unchecked(
            exported_device_memory.clone(),
            0,
            1024 * 1024 * 4,
        )])
        .map_err(|e| e.0)?
    });

    let semaphore = Semaphore::new(
        device.clone(),
        SemaphoreCreateInfo {
            semaphore_type: SemaphoreType::Timeline,
            export_handle_types: ExternalSemaphoreHandleTypes::OPAQUE_FD,
            initial_value: 0,
            ..Default::default()
        },
    )?;

    sender.send(Message {
        memory_file: exported_device_memory.export_fd(ExternalMemoryHandleType::DmaBuf)?,
        semaphore_file: unsafe { semaphore.export_fd(ExternalSemaphoreHandleType::OpaqueFd)? },
    })?;

    let mut builder = AutoCommandBufferBuilder::primary(
        command_buffer_allocator.clone(),
        queue.queue_family_index(),
        CommandBufferUsage::OneTimeSubmit,
    )?;
    builder.clear_color_image(ClearColorImageInfo {
        clear_value: [1.0, 0.0, 0.8, 1.0].into(),
        ..ClearColorImageInfo::image(image.clone())
    })?;
    let command_buffer = builder.build()?;

    unsafe {
        let fence = Fence::new(device.clone(), Default::default())?;
        let ash_semaphore = semaphore.handle();
        let ash_command_buffer = command_buffer.handle();

        ash_device.queue_submit(
            queue.handle(),
            &[vk::SubmitInfo {
                wait_semaphore_count: 0,
                p_wait_semaphores: ptr::null(),
                signal_semaphore_count: 1,
                p_signal_semaphores: &ash_semaphore,
                command_buffer_count: 1,
                p_command_buffers: &ash_command_buffer,
                p_wait_dst_stage_mask: ptr::null(),
                ..Default::default()
            }
            .push_next(&mut vk::TimelineSemaphoreSubmitInfo {
                wait_semaphore_value_count: 0,
                p_wait_semaphore_values: ptr::null(),
                signal_semaphore_value_count: 1,
                p_signal_semaphore_values: [1].as_ptr(),
                ..Default::default()
            })],
            fence.handle(),
        )?;

        fence.wait(None)?;
    }

    println!("Producer thread finished");
    Ok(())
}

/// Main function for consumer thread. The consumer thread will:
///
/// 1. Receive file descriptors from producer thread
/// 1. Import external memory and semaphore from file descriptors
/// 1. Bind image to imported memory
/// 1. Create buffer in order to read from image (because you can't read image directly)
/// 1. Wait for semaphore and execute copy image to buffer command
/// 1. Write content of buffer to image.png
fn consumer_thread(library: Arc<VulkanLibrary>, receiver: mpsc::Receiver<Message>) -> Result<(), Box<dyn Error>> {
    let (device, queue) = factory(library)?;
    let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
    let command_buffer_allocator = Arc::new(StandardCommandBufferAllocator::new(device.clone(), Default::default()));

    let Message {
        memory_file,
        semaphore_file,
    } = receiver.iter().next().unwrap();

    let imported_device_memory = Arc::new(unsafe {
        DeviceMemory::import(
            device.clone(),
            MemoryAllocateInfo {
                allocation_size: 256 * 1024 * 1024,
                memory_type_index: 0,
                ..Default::default()
            },
            MemoryImportInfo::Fd {
                handle_type: ExternalMemoryHandleType::DmaBuf,
                file: memory_file,
            },
        )?
    });

    let ash_device = unsafe { ash::Device::load(&device.instance().fns().v1_0, device.handle()) };
    let image = Arc::new(unsafe {
        let ash_image = ash_device.create_image(
            &vk::ImageCreateInfo {
                usage: vk::ImageUsageFlags::TRANSFER_SRC,
                image_type: vk::ImageType::TYPE_2D,
                format: vk::Format::R8G8B8A8_UNORM,
                extent: vk::Extent3D {
                    width: 1024,
                    height: 1024,
                    depth: 1,
                    ..Default::default()
                },
                samples: vk::SampleCountFlags::TYPE_1,
                array_layers: 1,
                mip_levels: 1,
                tiling: vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT,
                sharing_mode: vk::SharingMode::EXCLUSIVE,
                queue_family_index_count: 0,
                p_queue_family_indices: ptr::null(),
                initial_layout: vk::ImageLayout::UNDEFINED,
                ..Default::default()
            }
            .push_next(&mut vk::ExternalMemoryImageCreateInfo {
                handle_types: vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT,
                ..Default::default()
            })
            .push_next(&mut vk::ImageDrmFormatModifierListCreateInfoEXT {
                drm_format_modifier_count: 1,
                p_drm_format_modifiers: [0u64].as_ptr(),
                ..Default::default()
            }),
            None,
        )?;
        RawImage::from_handle(
            device.clone(),
            ash_image,
            ImageCreateInfo {
                usage: ImageUsage::TRANSFER_SRC,
                image_type: ImageType::Dim2d,
                format: Format::R8G8B8A8_UNORM,
                extent: [1024, 1024, 1],
                external_memory_handle_types: ExternalMemoryHandleTypes::DMA_BUF,
                tiling: ImageTiling::DrmFormatModifier,
                drm_format_modifiers: vec![0],
                ..Default::default()
            },
        )?
        .bind_memory([ResourceMemory::from_device_memory_unchecked(
            imported_device_memory.clone(),
            0,
            1024 * 1024 * 4,
        )])
        .map_err(|e| e.0)?
    });

    let buffer = Buffer::from_iter(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::TRANSFER_DST,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::PREFER_HOST | MemoryTypeFilter::HOST_RANDOM_ACCESS,
            ..Default::default()
        },
        (0..1024 * 1024 * 4).map(|_| 0u8),
    )?;

    let semaphore = Semaphore::new(
        device.clone(),
        SemaphoreCreateInfo {
            semaphore_type: SemaphoreType::Timeline,
            initial_value: 0,
            ..Default::default()
        },
    )?;

    unsafe {
        semaphore.import_fd(ImportSemaphoreFdInfo {
            file: Some(semaphore_file),
            ..ImportSemaphoreFdInfo::handle_type(ExternalSemaphoreHandleType::OpaqueFd)
        })?;
    }

    let mut builder = AutoCommandBufferBuilder::primary(
        command_buffer_allocator.clone(),
        queue.queue_family_index(),
        CommandBufferUsage::OneTimeSubmit,
    )?;
    builder.copy_image_to_buffer(CopyImageToBufferInfo::image_buffer(image, buffer.clone()))?;
    let command_buffer = builder.build()?;

    unsafe {
        let fence = Fence::new(device.clone(), Default::default())?;
        let ash_semaphore = semaphore.handle();
        let ash_command_buffer = command_buffer.handle();

        ash_device.queue_submit(
            queue.handle(),
            &[vk::SubmitInfo {
                wait_semaphore_count: 1,
                p_wait_semaphores: &ash_semaphore,
                signal_semaphore_count: 0,
                p_signal_semaphores: ptr::null(),
                command_buffer_count: 1,
                p_command_buffers: &ash_command_buffer,
                p_wait_dst_stage_mask: [vk::PipelineStageFlags::ALL_COMMANDS].as_ptr(),
                ..Default::default()
            }
            .push_next(&mut vk::TimelineSemaphoreSubmitInfo {
                wait_semaphore_value_count: 1,
                p_wait_semaphore_values: [1].as_ptr(),
                signal_semaphore_value_count: 0,
                p_signal_semaphore_values: ptr::null(),
                ..Default::default()
            })],
            fence.handle(),
        )?;

        fence.wait(None)?;
    }

    let content = buffer.read()?;
    ImageBuffer::<Rgba<u8>, _>::from_raw(1024, 1024, &content[..])
        .unwrap()
        .save("image.png")?;
    println!("Consumer thread wrote image data to image.png");
    println!("Consumer thread finished");
    Ok(())
}

fn main() -> Result<(), Box<dyn Error>> {
    let library = VulkanLibrary::new()?;
    let producer_thread_library = library.clone();
    let consumer_thread_library = library.clone();
    let (sender, receiver) = mpsc::channel();

    let producer_thread = thread::spawn(move || {
        println!("Producer thread started");
        producer_thread(producer_thread_library, sender).unwrap();
    });

    let consumer_thread = thread::spawn(move || {
        println!("Consumer thread started");
        consumer_thread(consumer_thread_library, receiver).unwrap();
    });

    producer_thread.join().unwrap();
    consumer_thread.join().unwrap();
    println!("End of main thread");
    Ok(())
}