Анимация системы частиц при помощи вычислительных шейдеров в Vulkan

Давайте рассмотрим систему из большого числа частиц (n), двигающихся под действием притяжения двух черный дыр (соответствующий пример на OpenGL был рассмотрен ранее в Вычислительные шейдеры в OpenGL). В этом случае нам понадобится как вычислительный шейдер для анимации системы частиц, так и вершинный и фрагментный для их рендеринга. У нас будет два буфера - в одном (posBuf) мы будем хранить координаты частиц, а в другом (velBuf) - их скорости.

Рассмотрим как анимацию и рендеринг этой системы можно реализовать на Vulkan. В отличии от OpenGL нам понадобятся объекты-конвейеры, причем нам будет нужно два конвейера - один графический и один вычислительный. Также нам понадобятся семафоры для синхронизации расчета и рендеринга.

struct Ubo      // for render pipeline
{
    glm::mat4 model;
    glm::mat4 view;
    glm::mat4 proj;
};

struct  ParticleVertex
{
    glm::vec4   pos;
};

template <>
inline GraphicsPipeline&    registerVertexAttrs<ParticleVertex> ( GraphicsPipeline& pipeline ) 
{
    return pipeline
        .addVertexAttr ( 0, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(ParticleVertex, pos) );     // binding, location, format, offset
}

class   ParticleWindow : public VulkanWindow
{
    std::vector<CommandBuffer>      commandBuffers;
    GraphicsPipeline                graphicsPipeline;
    ComputePipeline                 computePipeline;
    Renderpass                      renderPass;
    std::vector<Uniform<Ubo>>       uniformBuffers;
    std::vector<DescriptorSet>      descriptorSets;
    Buffer                          posBuffer;
    Buffer                          velBuffer;
    DescriptorSet                   computeDescriptorSet;
    CommandBuffer                   computeCommandBuffer;
    Semaphore                       computeSemaphore, graphicsSemaphore;
    CommandPool                     computeCommandPool;
    std::vector<Fence>              fences;
    size_t                          n;
    size_t                          numParticles;
    float                           t     = 0;              // current time in seconds
    float                           zNear = 0.1f;
    float                           zFar  = 100.0f;

Также нам будет нужен отдельный объект CommandPool, который будет выделять командные буфера для вычислений и будет нужен выделенный из него командный буфер (computeCommand). Еще нам понадобятся два отдельных множества дескрипторов - descriptorSet и computeDescriptorSet.

void    createDescriptorSets ()
{
    descriptorSets.resize ( swapChain.imageCount () );

    for ( uint32_t i = 0; i < swapChain.imageCount (); i++ )
        descriptorSets  [i]
            .setLayout        ( device, descAllocator, graphicsPipeline.getDescLayout () )
            .addUniformBuffer ( 0, uniformBuffers [i], 0, sizeof ( Ubo ) )
            .create           ();
        
    computeDescriptorSet
        .setLayout        ( device, descAllocator, computePipeline.getDescLayout () )
        .addStorageBuffer ( 0, posBuffer )
        .addStorageBuffer ( 1, velBuffer )
        .create           ();
}

Метод createPipelines, как и ранее, отвечает за создание конвейеров, на этот раз сразу двух.

virtual void    createPipelines () override 
{
    VkDeviceSize bufferSize = sizeof ( Ubo );

    uniformBuffers.resize ( swapChain.imageCount() );

    for ( size_t i = 0; i < swapChain.imageCount (); i++ )
        uniformBuffers [i].create ( device, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT );

        // current app code
    renderPass.addAttachment   ( swapChain.getFormat (),                VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR )
        .addAttachment   ( depthTexture.getImage ().getFormat (), VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL )
        .addSubpass      ( 0 )
        .addDepthSubpass ( 1 )
        .create          ( device );
        
    graphicsPipeline
        .setDevice         ( device )
        .setVertexShader   ( "shaders/particles-render.vert.spv" )
        .setFragmentShader ( "shaders/particles-render.frag.spv" )
        .setSize           ( swapChain.getExtent () )
        .addVertexBinding  ( sizeof ( ParticleVertex ) )
        .addVertexAttributes <ParticleVertex> ()
        .addDescLayout     ( 0, DescSetLayout ()
            .add ( 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT ) )
        .setTopology       ( VK_PRIMITIVE_TOPOLOGY_POINT_LIST )
        .create            ( renderPass );
            
    computePipeline
        .setDevice     ( device )
        .setShader     ( "shaders/particles-compute.comp.spv" )
        .addDescriptor ( 0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_SHADER_STAGE_COMPUTE_BIT )
        .addDescriptor ( 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_SHADER_STAGE_COMPUTE_BIT )
        .create        ();
        
            // create before command buffers
    swapChain.createFramebuffers ( renderPass, depthTexture.getImageView () );

    createDescriptorSets       ();
    createCommandBuffers       ( renderPass );
    createComputeCommandBuffer ();
        
    fences.resize ( swapChain.imageCount () );

    for ( auto& f : fences )
        f.create ( device );
}

Сильно изменится метод submit, который сперва должен поместить вычислительный командный буфер в вычислительную очередь, а потом графический командный буфер в графическую.

virtual void    submit ( uint32_t imageIndex ) override 
{
    updateUniformBuffer ( imageIndex );

    SubmitInfo ()
        .wait    ( { {graphicsSemaphore, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT} } )
        .signal  ( { computeSemaphore } )
        .buffers ( { computeCommandBuffer } )
        .submit  ( device.getComputeQueue (), fences [currentImage].getHandle () );

    fences [currentImage].wait  ( UINT64_MAX );
    fences [currentImage].reset ();

    VkFence     currentFence  = swapChain.currentInFlightFence ();

    vkResetFences ( device.getDevice (), 1, &currentFence );

    SubmitInfo ()
        .wait    ( { { swapChain.currentAvailableSemaphore (), VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT }, { computeSemaphore.getHandle (), VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT } } )
        .buffers ( { commandBuffers [imageIndex]  } )
        .signal  ( { swapChain.currentRenderFinishedSemaphore (), graphicsSemaphore.getHandle () } )
        .submit  ( device.getGraphicsQueue (), swapChain.currentInFlightFence () );
}

Сперва мы отдаем на выполнение вычислительный командный буфер, а после него - графический. Каждый буфер отдается в соответствующую очередь (в общем случае это могут быть разные очереди). При этом мы используем семафоры для того, чтобы гарантировать что соответствующая команда начнет работать после выполнения предыдущей.

Также стоит обратить внимание на то, каким образом создаются соответствующие командные буфера.

void    createCommandBuffers ( Renderpass& renderPass )
{
    auto    framebuffers = swapChain.getFramebuffers ();

    commandBuffers = device.allocCommandBuffers ( (uint32_t)framebuffers.size ());

    for ( size_t i = 0; i < commandBuffers.size(); i++ )
    {
        commandBuffers [i]
            .begin             ()
            .beginRenderPass   ( RenderPassInfo ( renderPass ).framebuffer ( framebuffers [i] ).extent ( swapChain.getExtent () ).clearColor ().clearDepthStencil () )
            .pipeline          ( graphicsPipeline )
            .addDescriptorSets ( { descriptorSets[i] } )
            .bindVertexBuffers ( { {posBuffer, 0}, { velBuffer, 1 } } )
            .draw              ( (uint32_t)numParticles, 1, 0, 0 )
            .end               ();
    }
}

void    createComputeCommandBuffer ()
{
    computeCommandBuffer.create ( device );
        
    computeCommandBuffer.begin ();

        // Add memory barrier to ensure that the (graphics) vertex shader has fetched 
        // attributes before compute starts to write to the buffer
    if ( device.getGraphicsQueue () != device.getComputeQueue () )
    {
        transitionBuffer ( posBuffer, 0, VK_ACCESS_SHADER_WRITE_BIT, device.getGraphicsFamilyIndex (), device.getComputeFamilyIndex (), VK_PIPELINE_STAGE_VERTEX_INPUT_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT );
        transitionBuffer ( velBuffer, 0, VK_ACCESS_SHADER_WRITE_BIT, device.getGraphicsFamilyIndex (), device.getComputeFamilyIndex (), VK_PIPELINE_STAGE_VERTEX_INPUT_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT );
    }

    computeCommandBuffer
        .pipeline ( computePipeline )
        .addDescriptorSets ( { computeDescriptorSet } )
        .bindVertexBuffers ( { {posBuffer, 0 } } )
        .dispatch ( (uint32_t) (numParticles + 511) / 512, 1, 1 );

        // Add barrier to ensure that compute shader has finished writing to the buffer
        // Without this the (rendering) vertex shader may display incomplete results (partial data from last frame) 
    if ( device.getGraphicsQueue () != device.getComputeQueue () )
    {
        transitionBuffer ( posBuffer, VK_ACCESS_SHADER_WRITE_BIT, 0, device.getComputeFamilyIndex (), device.getGraphicsFamilyIndex (), VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_VERTEX_INPUT_BIT );
        transitionBuffer ( velBuffer, VK_ACCESS_SHADER_WRITE_BIT, 0, device.getComputeFamilyIndex (), device.getGraphicsFamilyIndex (), VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_VERTEX_INPUT_BIT );
    }

    computeCommandBuffer.end ();
}

Методы updateUniformBuffer и initParticles приводятся ниже.

void updateUniformBuffer ( uint32_t currentImage )
{
    uniformBuffers [currentImage]->model = glm::mat4 ( 1.0f );
    uniformBuffers [currentImage]->view  = controller->getModelView ();
    uniformBuffers [currentImage]->proj  = controller->getProjection ();
}
    
void initParticles ( int num )
{
    n            = num;
    numParticles = n * n * n;

            // init buffers with particle data
    std::vector<glm::vec4>  vb;
    std::vector<glm::vec4>  pb;
    float                   h = 2.0f / (n - 1);

    for ( size_t i = 0; i < n; i++ )
        for ( size_t j = 0; j < n; j++ )
            for ( size_t k = 0; k < n; k++ ) 
            {
                glm::vec4   p ( h * i - 1, h * j - 1, h * k - 1, 1 );

                pb.push_back ( p );
                vb.push_back ( glm::vec4 ( 0 ) );
            }

    posBuffer.create ( device, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, pb, Buffer::hostWrite );
    velBuffer.create ( device, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, vb, Buffer::hostWrite );
}

Исходный код для этого примера можно скачать в https://github.com/steps3d/vulkan-with-classes.