wined3d: Implement adapter_vk_dispatch_compute().

Signed-off-by: Henri Verbeet <hverbeet@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/wined3d/adapter_vk.c

@@ -1509,7 +1509,45 @@ static void adapter_vk_draw_primitive(struct wined3d_device *device,
 static void adapter_vk_dispatch_compute(struct wined3d_device *device,
         const struct wined3d_state *state, const struct wined3d_dispatch_parameters *parameters)
 {
-    FIXME("device %p, state %p, parameters %p.\n", device, state, parameters);
+    struct wined3d_buffer_vk *indirect_vk = NULL;
+    const struct wined3d_vk_info *vk_info;
+    struct wined3d_context_vk *context_vk;
+    VkCommandBuffer vk_command_buffer;
+
+    TRACE("device %p, state %p, parameters %p.\n", device, state, parameters);
+
+    context_vk = wined3d_context_vk(context_acquire(device, NULL, 0));
+    vk_info = context_vk->vk_info;
+
+    if (parameters->indirect)
+        indirect_vk = wined3d_buffer_vk(parameters->u.indirect.buffer);
+
+    if (!(vk_command_buffer = wined3d_context_vk_apply_compute_state(context_vk, state, indirect_vk)))
+    {
+        ERR("Failed to apply compute state.\n");
+        context_release(&context_vk->c);
+        return;
+    }
+
+    if (parameters->indirect)
+    {
+        struct wined3d_bo_vk *bo = &indirect_vk->bo;
+
+        wined3d_context_vk_reference_bo(context_vk, bo);
+        VK_CALL(vkCmdDispatchIndirect(vk_command_buffer, bo->vk_buffer,
+                bo->buffer_offset + parameters->u.indirect.offset));
+    }
+    else
+    {
+        const struct wined3d_direct_dispatch_parameters *direct = &parameters->u.direct;
+
+        VK_CALL(vkCmdDispatch(vk_command_buffer, direct->group_count_x, direct->group_count_y, direct->group_count_z));
+    }
+
+    VK_CALL(vkCmdPipelineBarrier(vk_command_buffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
+            VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT, 0, 0, NULL, 0, NULL, 0, NULL));
+
+    context_release(&context_vk->c);
 }
 
 void adapter_vk_clear_uav(struct wined3d_context *context,

dlls/wined3d/context_vk.c

@@ -58,6 +58,28 @@ VkCompareOp vk_compare_op_from_wined3d(enum wined3d_cmp_func op)
     }
 }
 
+VkShaderStageFlagBits vk_shader_stage_from_wined3d(enum wined3d_shader_type shader_type)
+{
+    switch (shader_type)
+    {
+        case WINED3D_SHADER_TYPE_VERTEX:
+            return VK_SHADER_STAGE_VERTEX_BIT;
+        case WINED3D_SHADER_TYPE_HULL:
+            return VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT;
+        case WINED3D_SHADER_TYPE_DOMAIN:
+            return VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT;
+        case WINED3D_SHADER_TYPE_GEOMETRY:
+            return VK_SHADER_STAGE_GEOMETRY_BIT;
+        case WINED3D_SHADER_TYPE_PIXEL:
+            return VK_SHADER_STAGE_FRAGMENT_BIT;
+        case WINED3D_SHADER_TYPE_COMPUTE:
+            return VK_SHADER_STAGE_COMPUTE_BIT;
+        default:
+            ERR("Unhandled shader type %s.\n", debug_shader_type(shader_type));
+            return 0;
+    }
+}
+
 void *wined3d_allocator_chunk_vk_map(struct wined3d_allocator_chunk_vk *chunk_vk,
         struct wined3d_context_vk *context_vk)
 {
@@ -350,6 +372,31 @@ void wined3d_context_vk_destroy_framebuffer(struct wined3d_context_vk *context_v
     o->command_buffer_id = command_buffer_id;
 }
 
+static void wined3d_context_vk_destroy_descriptor_pool(struct wined3d_context_vk *context_vk,
+        VkDescriptorPool vk_descriptor_pool, uint64_t command_buffer_id)
+{
+    struct wined3d_device_vk *device_vk = wined3d_device_vk(context_vk->c.device);
+    const struct wined3d_vk_info *vk_info = context_vk->vk_info;
+    struct wined3d_retired_object_vk *o;
+
+    if (context_vk->completed_command_buffer_id > command_buffer_id)
+    {
+        VK_CALL(vkDestroyDescriptorPool(device_vk->vk_device, vk_descriptor_pool, NULL));
+        TRACE("Destroyed descriptor pool 0x%s.\n", wine_dbgstr_longlong(vk_descriptor_pool));
+        return;
+    }
+
+    if (!(o = wined3d_context_vk_get_retired_object_vk(context_vk)))
+    {
+        ERR("Leaking descriptor pool 0x%s.\n", wine_dbgstr_longlong(vk_descriptor_pool));
+        return;
+    }
+
+    o->type = WINED3D_RETIRED_DESCRIPTOR_POOL_VK;
+    o->u.vk_descriptor_pool = vk_descriptor_pool;
+    o->command_buffer_id = command_buffer_id;
+}
+
 void wined3d_context_vk_destroy_memory(struct wined3d_context_vk *context_vk,
         VkDeviceMemory vk_memory, uint64_t command_buffer_id)
 {
@@ -647,6 +694,11 @@ static void wined3d_context_vk_cleanup_resources(struct wined3d_context_vk *cont
                 TRACE("Destroyed framebuffer 0x%s.\n", wine_dbgstr_longlong(o->u.vk_framebuffer));
                 break;
 
+            case WINED3D_RETIRED_DESCRIPTOR_POOL_VK:
+                VK_CALL(vkDestroyDescriptorPool(device_vk->vk_device, o->u.vk_descriptor_pool, NULL));
+                TRACE("Destroyed descriptor pool 0x%s.\n", wine_dbgstr_longlong(o->u.vk_descriptor_pool));
+                break;
+
             case WINED3D_RETIRED_MEMORY_VK:
                 VK_CALL(vkFreeMemory(device_vk->vk_device, o->u.vk_memory, NULL));
                 TRACE("Freed memory 0x%s.\n", wine_dbgstr_longlong(o->u.vk_memory));
@@ -930,6 +982,15 @@ void wined3d_context_vk_cleanup(struct wined3d_context_vk *context_vk)
 
     wined3d_context_vk_wait_command_buffer(context_vk, buffer->id - 1);
     context_vk->completed_command_buffer_id = buffer->id;
+
+    if (context_vk->vk_pipeline_layout)
+    {
+        VK_CALL(vkDestroyPipelineLayout(device_vk->vk_device, context_vk->vk_pipeline_layout, NULL));
+        VK_CALL(vkDestroyDescriptorSetLayout(device_vk->vk_device, context_vk->vk_set_layout, NULL));
+    }
+    heap_free(context_vk->compute.bindings.bindings);
+    if (context_vk->vk_descriptor_pool)
+        VK_CALL(vkDestroyDescriptorPool(device_vk->vk_device, context_vk->vk_descriptor_pool, NULL));
     wined3d_context_vk_cleanup_resources(context_vk);
     wine_rb_destroy(&context_vk->bo_slab_available, wined3d_context_vk_destroy_bo_slab, context_vk);
     heap_free(context_vk->submitted.buffers);
@@ -1012,6 +1073,10 @@ void wined3d_context_vk_submit_command_buffer(struct wined3d_context_vk *context
     TRACE("Submitting command buffer %p with id 0x%s.\n",
             buffer->vk_command_buffer, wine_dbgstr_longlong(buffer->id));
 
+    context_vk->update_compute_pipeline = 1;
+    context_vk->c.update_compute_shader_resource_bindings = 1;
+    context_vk->c.update_compute_unordered_access_view_bindings = 1;
+
     VK_CALL(vkEndCommandBuffer(buffer->vk_command_buffer));
 
     fence_desc.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
@@ -1122,6 +1187,173 @@ static int wined3d_bo_slab_vk_compare(const void *key, const struct wine_rb_entr
     return k->size - slab->bo.size;
 }
 
+static VkResult wined3d_context_vk_create_descriptor_pool(struct wined3d_device_vk *device_vk,
+        const struct wined3d_vk_info *vk_info, VkDescriptorPool *vk_pool)
+{
+    struct VkDescriptorPoolCreateInfo pool_desc;
+    VkResult vr;
+
+    static const VkDescriptorPoolSize pool_sizes[] =
+    {
+        {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1024},
+        {VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 1024},
+        {VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 1024},
+        {VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 1024},
+        {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1024},
+        {VK_DESCRIPTOR_TYPE_SAMPLER, 1024},
+    };
+
+    pool_desc.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
+    pool_desc.pNext = NULL;
+    pool_desc.flags = 0;
+    pool_desc.maxSets = 512;
+    pool_desc.poolSizeCount = ARRAY_SIZE(pool_sizes);
+    pool_desc.pPoolSizes = pool_sizes;
+
+    if ((vr = VK_CALL(vkCreateDescriptorPool(device_vk->vk_device, &pool_desc, NULL, vk_pool))) < 0)
+        ERR("Failed to create descriptor pool, vr %s.\n", wined3d_debug_vkresult(vr));
+
+    return vr;
+}
+
+static VkResult wined3d_context_vk_create_descriptor_set(struct wined3d_context_vk *context_vk,
+        VkDescriptorSetLayout vk_set_layout, VkDescriptorSet *vk_descriptor_set)
+{
+    struct wined3d_device_vk *device_vk = wined3d_device_vk(context_vk->c.device);
+    const struct wined3d_vk_info *vk_info = context_vk->vk_info;
+    struct VkDescriptorSetAllocateInfo set_desc;
+    VkResult vr;
+
+    if (!context_vk->vk_descriptor_pool && (vr = wined3d_context_vk_create_descriptor_pool(device_vk,
+            vk_info, &context_vk->vk_descriptor_pool)))
+    {
+        WARN("Failed to create descriptor pool, vr %s.\n", wined3d_debug_vkresult(vr));
+        return vr;
+    }
+
+    set_desc.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
+    set_desc.pNext = NULL;
+    set_desc.descriptorPool = context_vk->vk_descriptor_pool;
+    set_desc.descriptorSetCount = 1;
+    set_desc.pSetLayouts = &vk_set_layout;
+    if ((vr = VK_CALL(vkAllocateDescriptorSets(device_vk->vk_device, &set_desc, vk_descriptor_set))) >= 0)
+        return vr;
+
+    if (vr == VK_ERROR_FRAGMENTED_POOL || vr == VK_ERROR_OUT_OF_POOL_MEMORY)
+    {
+        wined3d_context_vk_destroy_descriptor_pool(context_vk,
+                context_vk->vk_descriptor_pool, context_vk->current_command_buffer.id);
+        context_vk->vk_descriptor_pool = VK_NULL_HANDLE;
+        if ((vr = wined3d_context_vk_create_descriptor_pool(device_vk, vk_info, &context_vk->vk_descriptor_pool)))
+        {
+            WARN("Failed to create descriptor pool, vr %s.\n", wined3d_debug_vkresult(vr));
+            return vr;
+        }
+
+        set_desc.descriptorPool = context_vk->vk_descriptor_pool;
+        if ((vr = VK_CALL(vkAllocateDescriptorSets(device_vk->vk_device, &set_desc, vk_descriptor_set))) >= 0)
+            return vr;
+    }
+
+    WARN("Failed to allocate descriptor set, vr %s.\n", wined3d_debug_vkresult(vr));
+
+    return vr;
+}
+
+static bool wined3d_context_vk_update_descriptors(struct wined3d_context_vk *context_vk,
+        VkCommandBuffer vk_command_buffer)
+{
+    const struct wined3d_vk_info *vk_info = context_vk->vk_info;
+    const struct wined3d_shader_resource_binding *binding;
+    struct wined3d_shader_resource_bindings *bindings;
+    VkDescriptorSet vk_descriptor_set;
+    VkResult vr;
+
+    bindings = &context_vk->compute.bindings;
+    if ((vr = wined3d_context_vk_create_descriptor_set(context_vk,
+            context_vk->vk_set_layout, &vk_descriptor_set)))
+    {
+        WARN("Failed to create descriptor set, vr %s.\n", wined3d_debug_vkresult(vr));
+        return false;
+    }
+
+    if (bindings->count)
+    {
+        binding = bindings->bindings;
+        FIXME("Unhandled descriptor type %#x.\n", binding->shader_descriptor_type);
+        return false;
+    }
+
+    VK_CALL(vkCmdBindDescriptorSets(vk_command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE,
+            context_vk->vk_pipeline_layout, 0, 1, &vk_descriptor_set, 0, NULL));
+
+    return true;
+}
+
+VkCommandBuffer wined3d_context_vk_apply_compute_state(struct wined3d_context_vk *context_vk,
+        const struct wined3d_state *state, struct wined3d_buffer_vk *indirect_vk)
+{
+    struct wined3d_device_vk *device_vk = wined3d_device_vk(context_vk->c.device);
+    const struct wined3d_vk_info *vk_info = context_vk->vk_info;
+    VkCommandBuffer vk_command_buffer;
+
+    if (wined3d_context_is_compute_state_dirty(&context_vk->c, STATE_COMPUTE_SHADER))
+        context_vk->c.shader_update_mask |= 1u << WINED3D_SHADER_TYPE_COMPUTE;
+
+    if (context_vk->c.shader_update_mask & (1u << WINED3D_SHADER_TYPE_COMPUTE))
+    {
+        device_vk->d.shader_backend->shader_select_compute(device_vk->d.shader_priv, &context_vk->c, state);
+        if (!context_vk->compute.vk_pipeline)
+        {
+            ERR("No compute pipeline set.\n");
+            return VK_NULL_HANDLE;
+        }
+        context_vk->c.shader_update_mask &= ~(1u << WINED3D_SHADER_TYPE_COMPUTE);
+        context_vk->c.update_compute_shader_resource_bindings = 1;
+        context_vk->c.update_compute_unordered_access_view_bindings = 1;
+        context_vk->update_compute_pipeline = 1;
+    }
+
+    if (indirect_vk)
+        wined3d_buffer_load_location(&indirect_vk->b, &context_vk->c, WINED3D_LOCATION_BUFFER);
+
+    if (!(vk_command_buffer = wined3d_context_vk_get_command_buffer(context_vk)))
+    {
+        ERR("Failed to get command buffer.\n");
+        return VK_NULL_HANDLE;
+    }
+
+    if (context_vk->update_compute_pipeline)
+    {
+        VK_CALL(vkCmdBindPipeline(vk_command_buffer,
+                VK_PIPELINE_BIND_POINT_COMPUTE, context_vk->compute.vk_pipeline));
+        context_vk->update_compute_pipeline = 0;
+    }
+
+    if (wined3d_context_is_compute_state_dirty(&context_vk->c, STATE_COMPUTE_CONSTANT_BUFFER)
+            || wined3d_context_is_compute_state_dirty(&context_vk->c, STATE_COMPUTE_SHADER_RESOURCE_BINDING))
+        context_vk->c.update_compute_shader_resource_bindings = 1;
+    if (wined3d_context_is_compute_state_dirty(&context_vk->c, STATE_COMPUTE_UNORDERED_ACCESS_VIEW_BINDING))
+        context_vk->c.update_compute_unordered_access_view_bindings = 1;
+
+    if (context_vk->c.update_compute_shader_resource_bindings
+            || context_vk->c.update_compute_unordered_access_view_bindings)
+    {
+        if (!wined3d_context_vk_update_descriptors(context_vk, vk_command_buffer))
+        {
+            ERR("Failed to update shader descriptors.\n");
+            return VK_NULL_HANDLE;
+        }
+
+        context_vk->c.update_compute_shader_resource_bindings = 0;
+        context_vk->c.update_compute_unordered_access_view_bindings = 0;
+    }
+
+    memset(context_vk->c.dirty_compute_states, 0, sizeof(context_vk->c.dirty_compute_states));
+
+    return vk_command_buffer;
+}
+
 HRESULT wined3d_context_vk_init(struct wined3d_context_vk *context_vk, struct wined3d_swapchain *swapchain)
 {
     VkCommandPoolCreateInfo command_pool_info;

dlls/wined3d/shader.c

@@ -1123,6 +1123,7 @@ static HRESULT shader_get_registers_used(struct wined3d_shader *shader, DWORD co
                     }
                     reg_maps->resource_info[reg_idx].type = semantic->resource_type;
                     reg_maps->resource_info[reg_idx].data_type = semantic->resource_data_type;
+                    wined3d_bitmap_set(reg_maps->resource_map, reg_idx);
                     break;
 
                 case WINED3DSPR_UAV:
@@ -1146,9 +1147,14 @@ static HRESULT shader_get_registers_used(struct wined3d_shader *shader, DWORD co
         {
             struct wined3d_shader_register *reg = &ins.declaration.src.reg;
             if (reg->idx[0].offset >= WINED3D_MAX_CBS)
+            {
                 ERR("Invalid CB index %u.\n", reg->idx[0].offset);
+            }
             else
+            {
                 reg_maps->cb_sizes[reg->idx[0].offset] = reg->idx[1].offset;
+                wined3d_bitmap_set(&reg_maps->cb_map, reg->idx[0].offset);
+            }
         }
         else if (ins.handler_idx == WINED3DSIH_DCL_GLOBAL_FLAGS)
         {
@@ -1266,6 +1272,7 @@ static HRESULT shader_get_registers_used(struct wined3d_shader *shader, DWORD co
             reg_maps->resource_info[reg_idx].type = WINED3D_SHADER_RESOURCE_BUFFER;
             reg_maps->resource_info[reg_idx].data_type = WINED3D_DATA_UINT;
             reg_maps->resource_info[reg_idx].flags = WINED3D_VIEW_BUFFER_RAW;
+            wined3d_bitmap_set(reg_maps->resource_map, reg_idx);
         }
         else if (ins.handler_idx == WINED3DSIH_DCL_RESOURCE_STRUCTURED)
         {
@@ -1279,6 +1286,7 @@ static HRESULT shader_get_registers_used(struct wined3d_shader *shader, DWORD co
             reg_maps->resource_info[reg_idx].data_type = WINED3D_DATA_UINT;
             reg_maps->resource_info[reg_idx].flags = 0;
             reg_maps->resource_info[reg_idx].stride = ins.declaration.structured_resource.byte_stride / 4;
+            wined3d_bitmap_set(reg_maps->resource_map, reg_idx);
         }
         else if (ins.handler_idx == WINED3DSIH_DCL_SAMPLER)
         {
@@ -1605,6 +1613,7 @@ static HRESULT shader_get_registers_used(struct wined3d_shader *shader, DWORD co
                     reg_maps->resource_info[reg_idx].type = WINED3D_SHADER_RESOURCE_TEXTURE_2D;
                     reg_maps->resource_info[reg_idx].data_type = WINED3D_DATA_FLOAT;
                     shader_record_sample(reg_maps, reg_idx, reg_idx, reg_idx);
+                    wined3d_bitmap_set(reg_maps->resource_map, reg_idx);
 
                     /* texbem is only valid with < 1.4 pixel shaders */
                     if (ins.handler_idx == WINED3DSIH_TEXBEM

dlls/wined3d/shader_spirv.c

@@ -25,17 +25,291 @@ WINE_DEFAULT_DEBUG_CHANNEL(d3d_shader);
 
 static const struct wined3d_shader_backend_ops spirv_shader_backend_vk;
 
+struct shader_spirv_resource_bindings
+{
+    VkDescriptorSetLayoutBinding *vk_bindings;
+    SIZE_T vk_bindings_size, vk_binding_count;
+};
+
 struct shader_spirv_priv
 {
     const struct wined3d_vertex_pipe_ops *vertex_pipe;
     const struct wined3d_fragment_pipe_ops *fragment_pipe;
     bool ffp_proj_control;
+
+    struct shader_spirv_resource_bindings bindings;
+};
+
+struct shader_spirv_compute_program_vk
+{
+    VkShaderModule vk_module;
+    VkPipeline vk_pipeline;
 };
 
 static void shader_spirv_handle_instruction(const struct wined3d_shader_instruction *ins)
 {
 }
 
+static VkShaderModule shader_spirv_compile(struct wined3d_context_vk *context_vk,
+        struct wined3d_shader *shader, const struct shader_spirv_resource_bindings *bindings)
+{
+    FIXME("Not implemented.\n");
+
+    return VK_NULL_HANDLE;
+}
+
+static struct shader_spirv_compute_program_vk *shader_spirv_find_compute_program_vk(struct shader_spirv_priv *priv,
+        struct wined3d_context_vk *context_vk, struct wined3d_shader *shader,
+        const struct shader_spirv_resource_bindings *bindings)
+{
+    struct wined3d_device_vk *device_vk = wined3d_device_vk(context_vk->c.device);
+    const struct wined3d_vk_info *vk_info = context_vk->vk_info;
+    struct shader_spirv_compute_program_vk *program;
+    VkDescriptorSetLayoutCreateInfo set_layout_desc;
+    VkPipelineLayoutCreateInfo pipeline_layout_desc;
+    VkComputePipelineCreateInfo pipeline_info;
+    VkResult vr;
+
+    if ((program = shader->backend_data))
+        return program;
+
+    if (!(program = heap_alloc(sizeof(*program))))
+        return NULL;
+
+    if (!(program->vk_module = shader_spirv_compile(context_vk, shader, bindings)))
+    {
+        heap_free(program);
+        return NULL;
+    }
+
+    if (!context_vk->vk_pipeline_layout)
+    {
+        set_layout_desc.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
+        set_layout_desc.pNext = NULL;
+        set_layout_desc.flags = 0;
+        set_layout_desc.bindingCount = 0;
+        set_layout_desc.pBindings = NULL;
+        if ((vr = VK_CALL(vkCreateDescriptorSetLayout(device_vk->vk_device,
+                &set_layout_desc, NULL, &context_vk->vk_set_layout))) < 0)
+        {
+            WARN("Failed to create Vulkan descriptor set layout, vr %s.\n", wined3d_debug_vkresult(vr));
+            VK_CALL(vkDestroyShaderModule(device_vk->vk_device, program->vk_module, NULL));
+            heap_free(program);
+            return NULL;
+        }
+
+        pipeline_layout_desc.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
+        pipeline_layout_desc.pNext = NULL;
+        pipeline_layout_desc.flags = 0;
+        pipeline_layout_desc.setLayoutCount = 1;
+        pipeline_layout_desc.pSetLayouts = &context_vk->vk_set_layout;
+        pipeline_layout_desc.pushConstantRangeCount = 0;
+        pipeline_layout_desc.pPushConstantRanges = NULL;
+
+        if ((vr = VK_CALL(vkCreatePipelineLayout(device_vk->vk_device,
+                &pipeline_layout_desc, NULL, &context_vk->vk_pipeline_layout))) < 0)
+        {
+            WARN("Failed to create Vulkan pipeline layout, vr %s.\n", wined3d_debug_vkresult(vr));
+            VK_CALL(vkDestroyDescriptorSetLayout(device_vk->vk_device, context_vk->vk_set_layout, NULL));
+            VK_CALL(vkDestroyShaderModule(device_vk->vk_device, program->vk_module, NULL));
+            heap_free(program);
+            return NULL;
+        }
+    }
+
+    pipeline_info.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
+    pipeline_info.pNext = NULL;
+    pipeline_info.flags = 0;
+    pipeline_info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
+    pipeline_info.stage.pNext = NULL;
+    pipeline_info.stage.flags = 0;
+    pipeline_info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
+    pipeline_info.stage.pName = "main";
+    pipeline_info.stage.pSpecializationInfo = NULL;
+    pipeline_info.stage.module = program->vk_module;
+    pipeline_info.layout = context_vk->vk_pipeline_layout;
+    pipeline_info.basePipelineHandle = VK_NULL_HANDLE;
+    pipeline_info.basePipelineIndex = -1;
+    if ((vr = VK_CALL(vkCreateComputePipelines(device_vk->vk_device,
+            VK_NULL_HANDLE, 1, &pipeline_info, NULL, &program->vk_pipeline))) < 0)
+    {
+        ERR("Failed to create Vulkan compute pipeline, vr %s.\n", wined3d_debug_vkresult(vr));
+        VK_CALL(vkDestroyShaderModule(device_vk->vk_device, program->vk_module, NULL));
+        heap_free(program);
+        return NULL;
+    }
+
+    shader->backend_data = program;
+
+    return program;
+}
+
+static void shader_spirv_resource_bindings_cleanup(struct shader_spirv_resource_bindings *bindings)
+{
+    heap_free(bindings->vk_bindings);
+}
+
+static bool shader_spirv_resource_bindings_add_binding(struct shader_spirv_resource_bindings *bindings,
+        VkDescriptorType vk_type, VkShaderStageFlagBits vk_stage, size_t *binding_idx)
+{
+    SIZE_T binding_count = bindings->vk_binding_count;
+    VkDescriptorSetLayoutBinding *binding;
+
+    if (!wined3d_array_reserve((void **)&bindings->vk_bindings, &bindings->vk_bindings_size,
+            binding_count + 1, sizeof(*bindings->vk_bindings)))
+        return false;
+
+    *binding_idx = binding_count;
+    binding = &bindings->vk_bindings[binding_count];
+    binding->binding = binding_count;
+    binding->descriptorType = vk_type;
+    binding->descriptorCount = 1;
+    binding->stageFlags = vk_stage;
+    binding->pImmutableSamplers = NULL;
+    ++bindings->vk_binding_count;
+
+    return true;
+}
+
+static bool wined3d_shader_resource_bindings_add_binding(struct wined3d_shader_resource_bindings *bindings,
+        enum wined3d_shader_type shader_type, enum wined3d_shader_descriptor_type shader_descriptor_type,
+        size_t resource_idx, enum wined3d_shader_resource_type resource_type,
+        enum wined3d_data_type resource_data_type, size_t binding_idx)
+{
+    struct wined3d_shader_resource_binding *binding;
+    SIZE_T binding_count = bindings->count;
+
+    if (!wined3d_array_reserve((void **)&bindings->bindings, &bindings->size,
+            binding_count + 1, sizeof(*bindings->bindings)))
+        return false;
+
+    binding = &bindings->bindings[binding_count];
+    binding->shader_type = shader_type;
+    binding->shader_descriptor_type = shader_descriptor_type;
+    binding->resource_idx = resource_idx;
+    binding->resource_type = resource_type;
+    binding->resource_data_type = resource_data_type;
+    binding->binding_idx = binding_idx;
+    ++bindings->count;
+
+    return true;
+}
+
+static bool shader_spirv_resource_bindings_init(struct shader_spirv_resource_bindings *bindings,
+        struct wined3d_shader_resource_bindings *wined3d_bindings,
+        const struct wined3d_state *state, uint32_t shader_mask)
+{
+    const struct wined3d_shader_resource_info *resource_info;
+    const struct wined3d_shader_reg_maps *reg_maps;
+    enum wined3d_shader_type shader_type;
+    VkDescriptorType vk_descriptor_type;
+    size_t binding_idx, register_idx;
+    VkShaderStageFlagBits vk_stage;
+    struct wined3d_shader *shader;
+    unsigned int i;
+    uint32_t map;
+
+    bindings->vk_binding_count = 0;
+    wined3d_bindings->count = 0;
+
+    for (shader_type = 0; shader_type < WINED3D_SHADER_TYPE_COUNT; ++shader_type)
+    {
+        if (!(shader_mask & (1u << shader_type)) || !(shader = state->shader[shader_type]))
+            continue;
+
+        reg_maps = &shader->reg_maps;
+        vk_stage = vk_shader_stage_from_wined3d(shader_type);
+
+        map = reg_maps->cb_map;
+        while (map)
+        {
+            register_idx = wined3d_bit_scan(&map);
+
+            if (!shader_spirv_resource_bindings_add_binding(bindings,
+                    VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, vk_stage, &binding_idx))
+                return false;
+            if (!wined3d_shader_resource_bindings_add_binding(wined3d_bindings,
+                    shader_type, WINED3D_SHADER_DESCRIPTOR_TYPE_CBV, register_idx,
+                    WINED3D_SHADER_RESOURCE_BUFFER, WINED3D_DATA_UINT, binding_idx))
+                return false;
+        }
+
+        for (i = 0; i < ARRAY_SIZE(reg_maps->resource_map); ++i)
+        {
+            map = reg_maps->resource_map[i];
+            while (map)
+            {
+                register_idx = (i << 5) + wined3d_bit_scan(&map);
+                resource_info = &reg_maps->resource_info[register_idx];
+                if (resource_info->type == WINED3D_SHADER_RESOURCE_BUFFER)
+                    vk_descriptor_type = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
+                else
+                    vk_descriptor_type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
+
+                if (!shader_spirv_resource_bindings_add_binding(bindings,
+                        vk_descriptor_type, vk_stage, &binding_idx))
+                    return false;
+                if (!wined3d_shader_resource_bindings_add_binding(wined3d_bindings,
+                        shader_type, WINED3D_SHADER_DESCRIPTOR_TYPE_SRV, register_idx,
+                        resource_info->type, resource_info->data_type, binding_idx))
+                    return false;
+            }
+        }
+
+        for (register_idx = 0; register_idx < ARRAY_SIZE(reg_maps->uav_resource_info); ++register_idx)
+        {
+            resource_info = &reg_maps->uav_resource_info[register_idx];
+            if (resource_info->type == WINED3D_SHADER_RESOURCE_NONE)
+                continue;
+            if (resource_info->type == WINED3D_SHADER_RESOURCE_BUFFER)
+                vk_descriptor_type = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
+            else
+                vk_descriptor_type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
+
+            if (!shader_spirv_resource_bindings_add_binding(bindings,
+                    vk_descriptor_type, vk_stage, &binding_idx))
+                return false;
+            if (!wined3d_shader_resource_bindings_add_binding(wined3d_bindings,
+                    shader_type, WINED3D_SHADER_DESCRIPTOR_TYPE_UAV, register_idx,
+                    resource_info->type, resource_info->data_type, binding_idx))
+                return false;
+
+            if (reg_maps->uav_counter_mask & (1u << register_idx))
+            {
+                if (!shader_spirv_resource_bindings_add_binding(bindings,
+                        VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, vk_stage, &binding_idx))
+                    return false;
+                if (!wined3d_shader_resource_bindings_add_binding(wined3d_bindings,
+                        shader_type, WINED3D_SHADER_DESCRIPTOR_TYPE_UAV_COUNTER, register_idx,
+                        WINED3D_SHADER_RESOURCE_BUFFER, WINED3D_DATA_UINT, binding_idx))
+                    return false;
+            }
+        }
+
+        map = 0;
+        for (i = 0; i < reg_maps->sampler_map.count; ++i)
+        {
+            if (reg_maps->sampler_map.entries[i].sampler_idx != WINED3D_SAMPLER_DEFAULT)
+                map |= 1u << reg_maps->sampler_map.entries[i].sampler_idx;
+        }
+
+        while (map)
+        {
+            register_idx = wined3d_bit_scan(&map);
+
+            if (!shader_spirv_resource_bindings_add_binding(bindings,
+                    VK_DESCRIPTOR_TYPE_SAMPLER, vk_stage, &binding_idx))
+                return false;
+            if (!wined3d_shader_resource_bindings_add_binding(wined3d_bindings,
+                    shader_type, WINED3D_SHADER_DESCRIPTOR_TYPE_SAMPLER, register_idx,
+                    WINED3D_SHADER_RESOURCE_NONE, WINED3D_DATA_SAMPLER, binding_idx))
+                return false;
+        }
+    }
+
+    return true;
+}
+
 static void shader_spirv_precompile(void *shader_priv, struct wined3d_shader *shader)
 {
     WARN("Not implemented.\n");
@@ -53,16 +327,35 @@ static void shader_spirv_select(void *shader_priv, struct wined3d_context *conte
 static void shader_spirv_select_compute(void *shader_priv,
         struct wined3d_context *context, const struct wined3d_state *state)
 {
-    FIXME("Not implemented.\n");
+    struct wined3d_context_vk *context_vk = wined3d_context_vk(context);
+    struct shader_spirv_compute_program_vk *program;
+    struct shader_spirv_priv *priv = shader_priv;
+    struct wined3d_shader *shader;
+
+    if (!shader_spirv_resource_bindings_init(&priv->bindings,
+            &context_vk->compute.bindings, state, 1u << WINED3D_SHADER_TYPE_COMPUTE))
+        ERR("Failed to initialise shader resource bindings.\n");
+
+    if ((shader = state->shader[WINED3D_SHADER_TYPE_COMPUTE]))
+        program = shader_spirv_find_compute_program_vk(priv, context_vk, shader, &priv->bindings);
+    else
+        program = NULL;
+
+    if (program)
+        context_vk->compute.vk_pipeline = program->vk_pipeline;
+    else
+        context_vk->compute.vk_pipeline = VK_NULL_HANDLE;
 }
 
 static void shader_spirv_disable(void *shader_priv, struct wined3d_context *context)
 {
+    struct wined3d_context_vk *context_vk = wined3d_context_vk(context);
     struct shader_spirv_priv *priv = shader_priv;
 
     priv->vertex_pipe->vp_enable(context, false);
     priv->fragment_pipe->fp_enable(context, false);
 
+    context_vk->compute.vk_pipeline = VK_NULL_HANDLE;
     context->shader_update_mask = (1u << WINED3D_SHADER_TYPE_PIXEL)
             | (1u << WINED3D_SHADER_TYPE_VERTEX)
             | (1u << WINED3D_SHADER_TYPE_GEOMETRY)
@@ -87,8 +380,41 @@ static void shader_spirv_load_constants(void *shader_priv, struct wined3d_contex
     WARN("Not implemented.\n");
 }
 
+static void shader_spirv_invalidate_compute_program(struct wined3d_context_vk *context_vk,
+        const struct shader_spirv_compute_program_vk *program)
+{
+    if (context_vk->compute.vk_pipeline == program->vk_pipeline)
+    {
+        context_vk->c.shader_update_mask |= (1u << WINED3D_SHADER_TYPE_COMPUTE);
+        context_vk->compute.vk_pipeline = VK_NULL_HANDLE;
+    }
+}
+
+static void shader_spirv_invalidate_contexts_compute_program(struct wined3d_device *device,
+    const struct shader_spirv_compute_program_vk *program)
+{
+    unsigned int i;
+
+    for (i = 0; i < device->context_count; ++i)
+    {
+        shader_spirv_invalidate_compute_program(wined3d_context_vk(device->contexts[i]), program);
+    }
+}
+
 static void shader_spirv_destroy(struct wined3d_shader *shader)
 {
+    struct wined3d_device_vk *device_vk = wined3d_device_vk(shader->device);
+    struct wined3d_vk_info *vk_info = &device_vk->vk_info;
+    struct shader_spirv_compute_program_vk *program;
+
+    if (!(program = shader->backend_data))
+        return;
+
+    shader_spirv_invalidate_contexts_compute_program(&device_vk->d, program);
+    VK_CALL(vkDestroyPipeline(device_vk->vk_device, program->vk_pipeline, NULL));
+    VK_CALL(vkDestroyShaderModule(device_vk->vk_device, program->vk_module, NULL));
+    shader->backend_data = NULL;
+    heap_free(program);
 }
 
 static HRESULT shader_spirv_alloc(struct wined3d_device *device,
@@ -120,6 +446,7 @@ static HRESULT shader_spirv_alloc(struct wined3d_device *device,
     priv->fragment_pipe = fragment_pipe;
     fragment_pipe->get_caps(device->adapter, &fragment_caps);
     priv->ffp_proj_control = fragment_caps.wined3d_caps & WINED3D_FRAGMENT_CAP_PROJ_CONTROL;
+    memset(&priv->bindings, 0, sizeof(priv->bindings));
 
     device->vertex_priv = vertex_priv;
     device->fragment_priv = fragment_priv;
@@ -132,6 +459,7 @@ static void shader_spirv_free(struct wined3d_device *device, struct wined3d_cont
 {
     struct shader_spirv_priv *priv = device->shader_priv;
 
+    shader_spirv_resource_bindings_cleanup(&priv->bindings);
     priv->fragment_pipe->free_private(device, context);
     priv->vertex_pipe->vp_free(device, context);
     heap_free(priv);

dlls/wined3d/wined3d_private.h

@@ -84,6 +84,7 @@
 
 struct wined3d_fragment_pipe_ops;
 struct wined3d_adapter;
+struct wined3d_buffer_vk;
 struct wined3d_context;
 struct wined3d_gl_info;
 struct wined3d_state;
@@ -304,6 +305,7 @@ GLenum wined3d_gl_compare_func(enum wined3d_cmp_func f) DECLSPEC_HIDDEN;
 VkAccessFlags vk_access_mask_from_bind_flags(uint32_t bind_flags) DECLSPEC_HIDDEN;
 VkCompareOp vk_compare_op_from_wined3d(enum wined3d_cmp_func op) DECLSPEC_HIDDEN;
 VkImageViewType vk_image_view_type_from_wined3d(enum wined3d_resource_type type, uint32_t flags) DECLSPEC_HIDDEN;
+VkShaderStageFlagBits vk_shader_stage_from_wined3d(enum wined3d_shader_type shader_type) DECLSPEC_HIDDEN;
 
 static inline enum wined3d_cmp_func wined3d_sanitize_cmp_func(enum wined3d_cmp_func func)
 {
@@ -1034,8 +1036,10 @@ struct wined3d_shader_reg_maps
     WORD local_int_consts;                          /* WINED3D_MAX_CONSTS_I, 16 */
     WORD local_bool_consts;                         /* WINED3D_MAX_CONSTS_B, 16 */
     UINT cb_sizes[WINED3D_MAX_CBS];
+    uint32_t cb_map;                                /* WINED3D_MAX_CBS, 15 */
 
     struct wined3d_shader_resource_info resource_info[MAX_SHADER_RESOURCE_VIEWS];
+    uint32_t resource_map[WINED3D_BITMAP_SIZE(MAX_SHADER_RESOURCE_VIEWS)];
     struct wined3d_shader_sampler_map sampler_map;
     DWORD sampler_comparison_mode;
     BYTE bumpmat;                                   /* WINED3D_MAX_TEXTURES, 8 */
@@ -2230,6 +2234,7 @@ enum wined3d_retired_object_type_vk
 {
     WINED3D_RETIRED_FREE_VK,
     WINED3D_RETIRED_FRAMEBUFFER_VK,
+    WINED3D_RETIRED_DESCRIPTOR_POOL_VK,
     WINED3D_RETIRED_MEMORY_VK,
     WINED3D_RETIRED_ALLOCATOR_BLOCK_VK,
     WINED3D_RETIRED_BO_SLAB_SLICE_VK,
@@ -2247,6 +2252,7 @@ struct wined3d_retired_object_vk
     {
         struct wined3d_retired_object_vk *next;
         VkFramebuffer vk_framebuffer;
+        VkDescriptorPool vk_descriptor_pool;
         VkDeviceMemory vk_memory;
         struct wined3d_allocator_block *block;
         struct
@@ -2293,12 +2299,46 @@ struct wined3d_render_pass_vk
     VkRenderPass vk_render_pass;
 };
 
+enum wined3d_shader_descriptor_type
+{
+    WINED3D_SHADER_DESCRIPTOR_TYPE_CBV,
+    WINED3D_SHADER_DESCRIPTOR_TYPE_SRV,
+    WINED3D_SHADER_DESCRIPTOR_TYPE_UAV,
+    WINED3D_SHADER_DESCRIPTOR_TYPE_UAV_COUNTER,
+    WINED3D_SHADER_DESCRIPTOR_TYPE_SAMPLER,
+};
+
+struct wined3d_shader_resource_binding
+{
+    enum wined3d_shader_type shader_type;
+    enum wined3d_shader_descriptor_type shader_descriptor_type;
+    size_t resource_idx;
+    enum wined3d_shader_resource_type resource_type;
+    enum wined3d_data_type resource_data_type;
+    size_t binding_idx;
+};
+
+struct wined3d_shader_resource_bindings
+{
+    struct wined3d_shader_resource_binding *bindings;
+    SIZE_T size, count;
+};
+
 struct wined3d_context_vk
 {
     struct wined3d_context c;
 
     const struct wined3d_vk_info *vk_info;
 
+    uint32_t update_compute_pipeline : 1;
+    uint32_t padding : 31;
+
+    struct
+    {
+        VkPipeline vk_pipeline;
+        struct wined3d_shader_resource_bindings bindings;
+    } compute;
+
     VkCommandPool vk_command_pool;
     struct wined3d_command_buffer_vk current_command_buffer;
     uint64_t completed_command_buffer_id;
@@ -2310,6 +2350,10 @@ struct wined3d_context_vk
         SIZE_T buffer_count;
     } submitted;
 
+    VkDescriptorPool vk_descriptor_pool;
+    VkPipelineLayout vk_pipeline_layout;
+    VkDescriptorSetLayout vk_set_layout;
+
     struct wined3d_retired_objects_vk retired;
     struct wine_rb_tree render_passes;
     struct wine_rb_tree bo_slab_available;
@@ -2324,6 +2368,8 @@ struct wined3d_allocator_block *wined3d_context_vk_allocate_memory(struct wined3
         unsigned int memory_type, VkDeviceSize size, VkDeviceMemory *vk_memory) DECLSPEC_HIDDEN;
 VkDeviceMemory wined3d_context_vk_allocate_vram_chunk_memory(struct wined3d_context_vk *context_vk,
         unsigned int pool, size_t size) DECLSPEC_HIDDEN;
+VkCommandBuffer wined3d_context_vk_apply_compute_state(struct wined3d_context_vk *context_vk,
+        const struct wined3d_state *state, struct wined3d_buffer_vk *indirect_vk) DECLSPEC_HIDDEN;
 void wined3d_context_vk_cleanup(struct wined3d_context_vk *context_vk) DECLSPEC_HIDDEN;
 BOOL wined3d_context_vk_create_bo(struct wined3d_context_vk *context_vk, VkDeviceSize size,
         VkBufferUsageFlags usage, VkMemoryPropertyFlags memory_type, struct wined3d_bo_vk *bo) DECLSPEC_HIDDEN;
@@ -3621,13 +3667,6 @@ static inline struct wined3d_device_vk *wined3d_device_vk(struct wined3d_device
     return CONTAINING_RECORD(device, struct wined3d_device_vk, d);
 }
 
-static inline BOOL isStateDirty(const struct wined3d_context *context, unsigned int state_id)
-{
-    unsigned int idx = state_id / (sizeof(*context->dirty_graphics_states) * CHAR_BIT);
-    unsigned int shift = state_id & ((sizeof(*context->dirty_graphics_states) * CHAR_BIT) - 1);
-    return context->dirty_graphics_states[idx] & (1u << shift);
-}
-
 static inline float wined3d_alpha_ref(const struct wined3d_state *state)
 {
     return (state->render_states[WINED3D_RS_ALPHAREF] & 0xff) / 255.0f;
@@ -5634,6 +5673,11 @@ static inline void wined3d_viewport_get_z_range(const struct wined3d_viewport *v
     *max_z = max(vp->max_z, vp->min_z + 0.001f);
 }
 
+static inline BOOL wined3d_bitmap_set(uint32_t *map, unsigned int idx)
+{
+    return map[idx >> 5] |= (1u << (idx & 0x1f));
+}
+
 static inline BOOL wined3d_bitmap_is_set(const uint32_t *map, unsigned int idx)
 {
     return map[idx >> 5] & (1u << (idx & 0x1f));
@@ -5693,6 +5737,21 @@ static inline BOOL wined3d_bitmap_get_range(const DWORD *bitmap, unsigned int bi
     return TRUE;
 }
 
+static inline bool wined3d_context_is_graphics_state_dirty(const struct wined3d_context *context, unsigned int state_id)
+{
+    return wined3d_bitmap_is_set(context->dirty_graphics_states, state_id);
+}
+
+static inline bool wined3d_context_is_compute_state_dirty(const struct wined3d_context *context, unsigned int state_id)
+{
+    return wined3d_bitmap_is_set(context->dirty_compute_states, state_id - STATE_COMPUTE_OFFSET);
+}
+
+static inline bool isStateDirty(const struct wined3d_context *context, unsigned int state_id)
+{
+    return wined3d_context_is_graphics_state_dirty(context, state_id);
+}
+
 static inline VkImageAspectFlags vk_aspect_mask_from_format(const struct wined3d_format *format)
 {
     VkImageAspectFlags mask = 0;