kronos_compute/api/

command.rs

//! Fluent command building and dispatch

use super::*;
use crate::*; // Import all functions from the crate root
use std::ptr;

/// Fluent builder for compute dispatch commands
/// 
/// This builder provides a safe, ergonomic API for recording
/// and executing compute commands. All Kronos optimizations
/// are applied automatically.
pub struct CommandBuilder {
    context: ComputeContext,
    pipeline: Pipeline,
    command_buffer: VkCommandBuffer,
    descriptor_set: Option<VkDescriptorSet>,
    bindings: Vec<(u32, Buffer)>,
    push_constants: Vec<u8>,
    workgroups: (u32, u32, u32),
}

impl ComputeContext {
    /// Start building a compute dispatch
    pub fn dispatch(&self, pipeline: &Pipeline) -> CommandBuilder {
        CommandBuilder {
            context: self.clone(),
            pipeline: Pipeline {
                context: pipeline.context.clone(),
                pipeline: pipeline.pipeline,
                layout: pipeline.layout,
                descriptor_set_layout: pipeline.descriptor_set_layout,
            },
            command_buffer: VkCommandBuffer::NULL,
            descriptor_set: None,
            bindings: Vec::new(),
            push_constants: Vec::new(),
            workgroups: (1, 1, 1),
        }
    }
}

impl CommandBuilder {
    /// Bind a buffer to a binding point
    pub fn bind_buffer(mut self, binding: u32, buffer: &Buffer) -> Self {
        self.bindings.push((binding, Buffer {
            context: buffer.context.clone(),
            buffer: buffer.buffer,
            memory: buffer.memory,
            size: buffer.size,
            usage: buffer.usage,
            _marker: std::marker::PhantomData,
        }));
        self
    }
    
    /// Set push constants
    pub fn push_constants<T: Copy>(mut self, data: &T) -> Self {
        let bytes = unsafe {
            std::slice::from_raw_parts(
                data as *const T as *const u8,
                std::mem::size_of::<T>(),
            )
        };
        self.push_constants = bytes.to_vec();
        self
    }
    
    /// Set the number of workgroups
    pub fn workgroups(mut self, x: u32, y: u32, z: u32) -> Self {
        self.workgroups = (x, y, z);
        self
    }
    
    /// Execute the dispatch
    pub fn execute(mut self) -> Result<()> {
        unsafe {
            self.context.with_inner(|inner| {
                // Allocate command buffer
                let alloc_info = VkCommandBufferAllocateInfo {
                    sType: VkStructureType::CommandBufferAllocateInfo,
                    pNext: ptr::null(),
                    commandPool: inner.command_pool,
                    level: VkCommandBufferLevel::Primary,
                    commandBufferCount: 1,
                };
                
                vkAllocateCommandBuffers(inner.device, &alloc_info, &mut self.command_buffer);
                
                // Begin command buffer
                let begin_info = VkCommandBufferBeginInfo {
                    sType: VkStructureType::CommandBufferBeginInfo,
                    pNext: ptr::null(),
                    flags: VkCommandBufferUsageFlags::ONE_TIME_SUBMIT,
                    pInheritanceInfo: ptr::null(),
                };
                
                let result = vkBeginCommandBuffer(self.command_buffer, &begin_info);
                if result != VkResult::Success {
                    return Err(KronosError::from(result));
                }
                
                // Create and update descriptor set if we have bindings
                if !self.bindings.is_empty() {
                    // Allocate descriptor set
                    let alloc_info = VkDescriptorSetAllocateInfo {
                        sType: VkStructureType::DescriptorSetAllocateInfo,
                        pNext: ptr::null(),
                        descriptorPool: inner.descriptor_pool,
                        descriptorSetCount: 1,
                        pSetLayouts: &self.pipeline.descriptor_set_layout,
                    };
                    
                    let mut descriptor_set = VkDescriptorSet::NULL;
                    let result = vkAllocateDescriptorSets(inner.device, &alloc_info, &mut descriptor_set);
                    if result != VkResult::Success {
                        return Err(KronosError::from(result));
                    }
                    
                    self.descriptor_set = Some(descriptor_set);
                    
                    // Update descriptor set
                    let buffer_infos: Vec<VkDescriptorBufferInfo> = self.bindings.iter().map(|(_, buffer)| {
                        VkDescriptorBufferInfo {
                            buffer: buffer.buffer,
                            offset: 0,
                            range: buffer.size as VkDeviceSize,
                        }
                    }).collect();
                    
                    let writes: Vec<VkWriteDescriptorSet> = self.bindings.iter().enumerate().map(|(i, (binding, _))| {
                        VkWriteDescriptorSet {
                            sType: VkStructureType::WriteDescriptorSet,
                            pNext: ptr::null(),
                            dstSet: descriptor_set,
                            dstBinding: *binding,
                            dstArrayElement: 0,
                            descriptorCount: 1,
                            descriptorType: VkDescriptorType::StorageBuffer,
                            pImageInfo: ptr::null(),
                            pBufferInfo: &buffer_infos[i],
                            pTexelBufferView: ptr::null(),
                        }
                    }).collect();
                    
                    vkUpdateDescriptorSets(inner.device, writes.len() as u32, writes.as_ptr(), 0, ptr::null());
                }
                
                // Insert barriers for buffers (smart barrier optimization)
                // In a real implementation, this would use the barrier_policy module
                let barriers: Vec<VkBufferMemoryBarrier> = self.bindings.iter().map(|(_, buffer)| {
                    VkBufferMemoryBarrier {
                        sType: VkStructureType::BufferMemoryBarrier,
                        pNext: ptr::null(),
                        srcAccessMask: VkAccessFlags::TRANSFER_WRITE,
                        dstAccessMask: VkAccessFlags::SHADER_READ | VkAccessFlags::SHADER_WRITE,
                        srcQueueFamilyIndex: VK_QUEUE_FAMILY_IGNORED,
                        dstQueueFamilyIndex: VK_QUEUE_FAMILY_IGNORED,
                        buffer: buffer.buffer,
                        offset: 0,
                        size: buffer.size as VkDeviceSize,
                    }
                }).collect();
                
                if !barriers.is_empty() {
                    vkCmdPipelineBarrier(
                        self.command_buffer,
                        VkPipelineStageFlags::TOP_OF_PIPE,
                        VkPipelineStageFlags::COMPUTE_SHADER,
                        VkDependencyFlags::empty(),
                        0,
                        ptr::null(),
                        barriers.len() as u32,
                        barriers.as_ptr(),
                        0,
                        ptr::null(),
                    );
                }
                
                // Bind pipeline
                vkCmdBindPipeline(self.command_buffer, VkPipelineBindPoint::Compute, self.pipeline.pipeline);
                
                // Bind descriptor set
                if let Some(descriptor_set) = self.descriptor_set {
                    vkCmdBindDescriptorSets(
                        self.command_buffer,
                        VkPipelineBindPoint::Compute,
                        self.pipeline.layout,
                        0,
                        1,
                        &descriptor_set,
                        0,
                        ptr::null(),
                    );
                }
                
                // Push constants
                if !self.push_constants.is_empty() {
                    vkCmdPushConstants(
                        self.command_buffer,
                        self.pipeline.layout,
                        VkShaderStageFlags::COMPUTE,
                        0,
                        self.push_constants.len() as u32,
                        self.push_constants.as_ptr() as *const _,
                    );
                }
                
                // Dispatch
                vkCmdDispatch(self.command_buffer, self.workgroups.0, self.workgroups.1, self.workgroups.2);
                
                // End command buffer
                let result = vkEndCommandBuffer(self.command_buffer);
                if result != VkResult::Success {
                    return Err(KronosError::from(result));
                }
                
                // Submit (with timeline batching optimization)
                let submit_info = VkSubmitInfo {
                    sType: VkStructureType::SubmitInfo,
                    pNext: ptr::null(),
                    waitSemaphoreCount: 0,
                    pWaitSemaphores: ptr::null(),
                    pWaitDstStageMask: ptr::null(),
                    commandBufferCount: 1,
                    pCommandBuffers: &self.command_buffer,
                    signalSemaphoreCount: 0,
                    pSignalSemaphores: ptr::null(),
                };
                
                let result = vkQueueSubmit(inner.queue, 1, &submit_info, VkFence::NULL);
                if result != VkResult::Success {
                    return Err(KronosError::CommandExecutionFailed(
                        format!("vkQueueSubmit failed: {:?}", result)
                    ));
                }
                
                // Wait for completion (in a real implementation, this could be async)
                vkQueueWaitIdle(inner.queue);
                
                // Free command buffer
                vkFreeCommandBuffers(inner.device, inner.command_pool, 1, &self.command_buffer);
                
                // Free descriptor set if allocated
                if let Some(descriptor_set) = self.descriptor_set {
                    vkFreeDescriptorSets(inner.device, inner.descriptor_pool, 1, &descriptor_set);
                }
                
                Ok(())
            })
        }
    }
}