vulkano_test/crates/engine_vulkan/src/device.rs

use std::sync::Arc;

use bevy_ecs::world::World;
use engine_window::raw_handle::DisplayHandleWrapper;
use vulkano::{
    command_buffer::allocator::StandardCommandBufferAllocator,
    descriptor_set::allocator::StandardDescriptorSetAllocator,
    device::{
        Device, DeviceCreateInfo, DeviceExtensions, Queue,
        physical::{PhysicalDevice, PhysicalDeviceType},
    },
    memory::allocator::StandardMemoryAllocator,
};

use crate::{
    VulkanCommandBufferAllocator, VulkanComputeQueue, VulkanConfig, VulkanDescriptorSetAllocator,
    VulkanDevice, VulkanGraphicsQueue, VulkanInstance, VulkanMemoryAllocator, VulkanTransferQueue,
    queues::{VulkanQueueFamilyIndices, VulkanQueuesQuery, find_queues_family_indices},
};

pub fn create_and_insert_device(world: &mut World, config: &VulkanConfig) {
    let picked_device =
        pick_physical_device(world, &config).expect("Failed to pick physical device");

    let device = picked_device.device;
    let physical_device = device.physical_device();

    log::debug!("Vulkan device created");
    log::debug!(
        "\tPhysical device: {:?} ({:?})",
        physical_device.properties().device_name,
        physical_device.properties().device_type
    );
    log::debug!("\tDevice extensions: {:?}", device.enabled_extensions());
    log::debug!("\tDevice features: {:?}", device.enabled_features());

    world.insert_resource(VulkanDevice(device.clone()));

    log::debug!("\tDevice selected queues:");
    if config.with_graphics_queue {
        world.insert_resource(VulkanGraphicsQueue(
            picked_device
                .graphics_queue
                .expect("Failed to get graphics queue"),
        ));
        log::debug!("\t\t- Graphics queue");
    }

    if config.with_compute_queue {
        world.insert_resource(VulkanComputeQueue(
            picked_device
                .compute_queue
                .expect("Failed to get compute queue"),
        ));
        log::debug!("\t\t- Compute queue");
    }

    if config.with_transfer_queue {
        world.insert_resource(VulkanTransferQueue(
            picked_device
                .transfer_queue
                .expect("Failed to get transfer queue"),
        ));
        log::debug!("\t\t- Transfer queue");
    }

    world.insert_resource(VulkanMemoryAllocator(Arc::new(
        StandardMemoryAllocator::new_default(device.clone()),
    )));
    world.insert_resource(VulkanCommandBufferAllocator(Arc::new(
        StandardCommandBufferAllocator::new(device.clone(), Default::default()),
    )));
    world.insert_resource(VulkanDescriptorSetAllocator(Arc::new(
        StandardDescriptorSetAllocator::new(device.clone(), Default::default()),
    )));
}

struct PickedDevice {
    pub device: Arc<Device>,
    pub graphics_queue: Option<Arc<Queue>>,
    pub compute_queue: Option<Arc<Queue>>,
    pub transfer_queue: Option<Arc<Queue>>,
}

fn pick_physical_device(world: &World, config: &VulkanConfig) -> Option<PickedDevice> {
    let instance = world
        .get_resource::<VulkanInstance>()
        .expect("Failed to get VulkanInstance during vulkan plugin initialization");

    instance
        .0
        .enumerate_physical_devices()
        .expect("Failed to enumerate physical devices")
        .filter_map(|p| check_physical_device_support(world, &p, config).and_then(|r| Some((p, r))))
        .min_by_key(|(p, _)| match p.properties().device_type {
            PhysicalDeviceType::DiscreteGpu => 0,
            PhysicalDeviceType::IntegratedGpu => 1,
            PhysicalDeviceType::VirtualGpu => 2,
            PhysicalDeviceType::Cpu => 3,
            PhysicalDeviceType::Other => 4,
            _ => 5,
        })
        .take()
        .and_then(|(p, (device_extensions, queue_family_indices))| {
            Some(create_device(
                config,
                &p,
                device_extensions,
                queue_family_indices,
            ))
        })
}

fn check_device_extensions_support(
    physical_device: &Arc<PhysicalDevice>,
    config: &VulkanConfig,
) -> Option<DeviceExtensions> {
    let device_extensions = DeviceExtensions {
        khr_swapchain: config.with_window_surface,
        ..config.device_extensions
    };

    if physical_device
        .supported_extensions()
        .contains(&device_extensions)
    {
        log::debug!(
            "\t\t[OK] Device supports required extensions {:?}",
            device_extensions
        );
        Some(device_extensions)
    } else {
        log::debug!(
            "\t\t[FAILED] Device does not support required extensions {:?}",
            device_extensions
        );
        None
    }
}

fn check_physical_device_support(
    world: &World,
    physical_device: &Arc<PhysicalDevice>,
    config: &VulkanConfig,
) -> Option<(DeviceExtensions, VulkanQueueFamilyIndices)> {
    log::debug!("Checking physical device");
    log::debug!("\tProperties");
    log::debug!("\t\tName: {}", physical_device.properties().device_name);
    log::debug!("\t\tAPI version: {}", physical_device.api_version());
    log::debug!(
        "\t\tDevice type: {:?}",
        physical_device.properties().device_type
    );
    log::debug!("\tRequired supports checking report");

    let device_extensions = check_device_extensions_support(physical_device, config)?;

    let display_handle = world
        .get_resource::<DisplayHandleWrapper>()
        .expect("DisplayHandleWrapper must be added before VulkanPlugin");

    let queue_support_query = VulkanQueuesQuery {
        with_surface: Some(&display_handle),
        with_graphics_queue: config.with_graphics_queue,
        with_compute_queue: config.with_compute_queue,
        with_transfer_queue: config.with_transfer_queue,
    };

    let queue_family_indices = find_queues_family_indices(physical_device, &queue_support_query);
    log::debug!("\t\tQueue family indices: {:#?}", queue_family_indices);

    Some((device_extensions, queue_family_indices))
}

fn create_device(
    config: &VulkanConfig,
    physical_device: &Arc<PhysicalDevice>,
    device_extensions: DeviceExtensions,
    queue_family_indices: VulkanQueueFamilyIndices,
) -> PickedDevice {
    let (device, mut queues) = Device::new(
        physical_device.clone(),
        DeviceCreateInfo {
            queue_create_infos: queue_family_indices.into(),
            enabled_extensions: device_extensions,
            enabled_features: config.device_features,
            ..Default::default()
        },
    )
    .expect("Failed to create device");

    let mut graphics_queue = None;
    let mut compute_queue = None;
    let mut transfer_queue = None;

    if config.with_graphics_queue {
        graphics_queue = queues.next();
    }

    if config.with_compute_queue {
        compute_queue = queues.next();
    }

    if config.with_transfer_queue {
        transfer_queue = queues.next();
    }

    PickedDevice {
        device,
        graphics_queue,
        compute_queue,
        transfer_queue,
    }
}