vkd3d-shader: Store UAV counter information in the descriptor scan info.

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

include/vkd3d_shader.h

@@ -613,6 +613,13 @@ enum vkd3d_shader_resource_data_type
     VKD3D_FORCE_32_BIT_ENUM(VKD3D_SHADER_RESOURCE_DATA_TYPE),
 };
 
+enum vkd3d_shader_descriptor_info_flag
+{
+    VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER = 0x00000001,
+
+    VKD3D_FORCE_32_BIT_ENUM(VKD3D_SHADER_DESCRIPTOR_INFO_FLAG),
+};
+
 struct vkd3d_shader_descriptor_info
 {
     enum vkd3d_shader_descriptor_type type;
@@ -620,6 +627,7 @@ struct vkd3d_shader_descriptor_info
     unsigned int register_index;
     enum vkd3d_shader_resource_type resource_type;
     enum vkd3d_shader_resource_data_type resource_data_type;
+    unsigned int flags; /* vkd3d_shader_descriptor_info_flag */
     unsigned int count;
 };
 
@@ -632,7 +640,6 @@ struct vkd3d_shader_scan_info
     unsigned int descriptor_count;
 
     unsigned int uav_read_mask;    /* VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS */
-    unsigned int uav_counter_mask; /* VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS */
     unsigned int sampler_comparison_mode_mask; /* 16 */
     bool use_vocp;
 };

libs/vkd3d-shader/spirv.c

@@ -5167,6 +5167,24 @@ static SpvImageFormat image_format_for_image_read(enum vkd3d_component_type data
     }
 }
 
+static const struct vkd3d_shader_descriptor_info *vkd3d_dxbc_compiler_get_descriptor_info(
+        struct vkd3d_dxbc_compiler *compiler, enum vkd3d_shader_descriptor_type type,
+        unsigned int register_space, unsigned int register_index)
+{
+    const struct vkd3d_shader_scan_info *scan_info = compiler->scan_info;
+    const struct vkd3d_shader_descriptor_info *d;
+    unsigned int i;
+
+    for (i = 0; i < scan_info->descriptor_count; ++i)
+    {
+        d = &scan_info->descriptors[i];
+        if (d->type == type && d->register_space == register_space && d->register_index == register_index)
+            return d;
+    }
+
+    return NULL;
+}
+
 static uint32_t vkd3d_dxbc_compiler_get_image_type_id(struct vkd3d_dxbc_compiler *compiler,
         const struct vkd3d_shader_register *reg, const struct vkd3d_spirv_resource_type *resource_type_info,
         enum vkd3d_component_type data_type, bool raw_structured, uint32_t depth)
@@ -5295,25 +5313,32 @@ static void vkd3d_dxbc_compiler_emit_resource_declaration(struct vkd3d_dxbc_comp
     if (is_uav && !(scan_info->uav_read_mask & (1u << reg->idx[0].offset)))
         vkd3d_spirv_build_op_decorate(builder, var_id, SpvDecorationNonReadable, NULL, 0);
 
-    if (is_uav && (scan_info->uav_counter_mask & (1u << reg->idx[0].offset)))
+    if (is_uav)
     {
-        assert(structure_stride); /* counters are valid only for structured buffers */
+        const struct vkd3d_shader_descriptor_info *d;
 
-        if (vkd3d_dxbc_compiler_is_opengl_target(compiler))
+        d = vkd3d_dxbc_compiler_get_descriptor_info(compiler,
+                VKD3D_SHADER_DESCRIPTOR_TYPE_UAV, register_space, register_index);
+        if (d->flags & VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER)
         {
-            vkd3d_spirv_enable_capability(builder, SpvCapabilityAtomicStorage);
+            assert(structure_stride); /* counters are valid only for structured buffers */
-            storage_class = SpvStorageClassAtomicCounter;
+
-            counter_type_id = vkd3d_spirv_get_type_id(builder, VKD3D_TYPE_UINT, 1);
+            if (vkd3d_dxbc_compiler_is_opengl_target(compiler))
-            ptr_type_id = vkd3d_spirv_get_op_type_pointer(builder, storage_class, counter_type_id);
+            {
+                vkd3d_spirv_enable_capability(builder, SpvCapabilityAtomicStorage);
+                storage_class = SpvStorageClassAtomicCounter;
+                counter_type_id = vkd3d_spirv_get_type_id(builder, VKD3D_TYPE_UINT, 1);
+                ptr_type_id = vkd3d_spirv_get_op_type_pointer(builder, storage_class, counter_type_id);
+            }
+
+            counter_var_id = vkd3d_spirv_build_op_variable(builder, &builder->global_stream,
+                    ptr_type_id, storage_class, 0);
+
+            vkd3d_dxbc_compiler_emit_descriptor_binding_for_reg(compiler,
+                    counter_var_id, reg, register_space, register_index, resource_type, true);
+
+            vkd3d_spirv_build_op_name(builder, counter_var_id, "u%u_counter", reg->idx[0].offset);
         }
-
-        counter_var_id = vkd3d_spirv_build_op_variable(builder, &builder->global_stream,
-                ptr_type_id, storage_class, 0);
-
-        vkd3d_dxbc_compiler_emit_descriptor_binding_for_reg(compiler,
-                counter_var_id, reg, register_space, register_index, resource_type, true);
-
-        vkd3d_spirv_build_op_name(builder, counter_var_id, "u%u_counter", reg->idx[0].offset);
     }
 
     vkd3d_symbol_make_resource(&resource_symbol, reg);

libs/vkd3d-shader/vkd3d_shader_main.c

@@ -22,7 +22,6 @@
 
 VKD3D_DEBUG_ENV_NAME("VKD3D_SHADER_DEBUG");
 
-STATIC_ASSERT(MEMBER_SIZE(struct vkd3d_shader_scan_info, uav_counter_mask) * CHAR_BIT >= VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS);
 STATIC_ASSERT(MEMBER_SIZE(struct vkd3d_shader_scan_info, uav_read_mask) * CHAR_BIT >= VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS);
 
 static void vkd3d_shader_dump_blob(const char *path, const char *prefix, const void *data, size_t size)
@@ -195,8 +194,30 @@ struct vkd3d_shader_scan_context
 {
     struct vkd3d_shader_scan_info *scan_info;
     size_t descriptors_size;
+
+    struct
+    {
+        unsigned int id;
+        unsigned int descriptor_idx;
+    } *uav_ranges;
+    size_t uav_ranges_size;
+    size_t uav_range_count;
 };
 
+static struct vkd3d_shader_descriptor_info *vkd3d_shader_scan_get_uav_descriptor_info(
+        const struct vkd3d_shader_scan_context *context, unsigned int range_id)
+{
+    unsigned int i;
+
+    for (i = 0; i < context->uav_range_count; ++i)
+    {
+        if (context->uav_ranges[i].id == range_id)
+            return &context->scan_info->descriptors[context->uav_ranges[i].descriptor_idx];
+    }
+
+    return NULL;
+}
+
 static bool vkd3d_shader_instruction_is_uav_read(const struct vkd3d_shader_instruction *instruction)
 {
     enum VKD3D_SHADER_INSTRUCTION_HANDLER handler_idx = instruction->handler_idx;
@@ -221,11 +242,13 @@ static bool vkd3d_shader_instruction_is_uav_counter(const struct vkd3d_shader_in
             || handler_idx == VKD3DSIH_IMM_ATOMIC_CONSUME;
 }
 
-static void vkd3d_shader_scan_record_uav_counter(struct vkd3d_shader_scan_info *scan_info,
+static void vkd3d_shader_scan_record_uav_counter(struct vkd3d_shader_scan_context *context,
         const struct vkd3d_shader_register *reg)
 {
-    assert(reg->idx[0].offset < VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS);
+    struct vkd3d_shader_descriptor_info *d;
-    scan_info->uav_counter_mask |= 1u << reg->idx[0].offset;
+
+    d = vkd3d_shader_scan_get_uav_descriptor_info(context, reg->idx[0].offset);
+    d->flags |= VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER;
 }
 
 static bool vkd3d_shader_scan_add_descriptor(struct vkd3d_shader_scan_context *context,
@@ -248,12 +271,30 @@ static bool vkd3d_shader_scan_add_descriptor(struct vkd3d_shader_scan_context *c
     d->register_index = register_index;
     d->resource_type = resource_type;
     d->resource_data_type = resource_data_type;
+    d->flags = 0;
     d->count = 1;
     ++scan_info->descriptor_count;
 
     return true;
 }
 
+static bool vkd3d_shader_scan_add_uav_range(struct vkd3d_shader_scan_context *context,
+        unsigned int id, unsigned int descriptor_idx)
+{
+    if (!vkd3d_array_reserve((void **)&context->uav_ranges, &context->uav_ranges_size,
+            context->uav_range_count + 1, sizeof(*context->uav_ranges)))
+    {
+        ERR("Failed to allocate UAV range.\n");
+        return false;
+    }
+
+    context->uav_ranges[context->uav_range_count].id = id;
+    context->uav_ranges[context->uav_range_count].descriptor_idx = descriptor_idx;
+    ++context->uav_range_count;
+
+    return true;
+}
+
 static void vkd3d_shader_scan_constant_buffer_declaration(struct vkd3d_shader_scan_context *context,
         const struct vkd3d_shader_instruction *instruction)
 {
@@ -324,6 +365,9 @@ static void vkd3d_shader_scan_resource_declaration(struct vkd3d_shader_scan_cont
     }
     vkd3d_shader_scan_add_descriptor(context, type, semantic->register_space,
             semantic->register_index, semantic->resource_type, resource_data_type);
+    if (type == VKD3D_SHADER_DESCRIPTOR_TYPE_UAV)
+        vkd3d_shader_scan_add_uav_range(context, semantic->reg.reg.idx[0].offset,
+                context->scan_info->descriptor_count - 1);
 }
 
 static void vkd3d_shader_scan_resource_declaration_raw(struct vkd3d_shader_scan_context *context,
@@ -338,6 +382,9 @@ static void vkd3d_shader_scan_resource_declaration_raw(struct vkd3d_shader_scan_
         type = VKD3D_SHADER_DESCRIPTOR_TYPE_SRV;
     vkd3d_shader_scan_add_descriptor(context, type, resource->register_space,
             resource->register_index, VKD3D_SHADER_RESOURCE_BUFFER, VKD3D_SHADER_RESOURCE_DATA_UINT);
+    if (type == VKD3D_SHADER_DESCRIPTOR_TYPE_UAV)
+        vkd3d_shader_scan_add_uav_range(context, resource->dst.reg.idx[0].offset,
+                context->scan_info->descriptor_count - 1);
 }
 
 static void vkd3d_shader_scan_resource_declaration_structured(struct vkd3d_shader_scan_context *context,
@@ -352,6 +399,9 @@ static void vkd3d_shader_scan_resource_declaration_structured(struct vkd3d_shade
         type = VKD3D_SHADER_DESCRIPTOR_TYPE_SRV;
     vkd3d_shader_scan_add_descriptor(context, type, resource->register_space,
             resource->register_index, VKD3D_SHADER_RESOURCE_BUFFER, VKD3D_SHADER_RESOURCE_DATA_UINT);
+    if (type == VKD3D_SHADER_DESCRIPTOR_TYPE_UAV)
+        vkd3d_shader_scan_add_uav_range(context, resource->reg.reg.idx[0].offset,
+                context->scan_info->descriptor_count - 1);
 }
 
 static void vkd3d_shader_scan_instruction(struct vkd3d_shader_scan_context *context,
@@ -401,7 +451,7 @@ static void vkd3d_shader_scan_instruction(struct vkd3d_shader_scan_context *cont
     }
 
     if (vkd3d_shader_instruction_is_uav_counter(instruction))
-        vkd3d_shader_scan_record_uav_counter(context->scan_info, &instruction->src[0].reg);
+        vkd3d_shader_scan_record_uav_counter(context, &instruction->src[0].reg);
 }
 
 int vkd3d_shader_scan_dxbc(const struct vkd3d_shader_code *dxbc,
@@ -435,6 +485,7 @@ int vkd3d_shader_scan_dxbc(const struct vkd3d_shader_code *dxbc,
         if (instruction.handler_idx == VKD3DSIH_INVALID)
         {
             WARN("Encountered unrecognized or invalid instruction.\n");
+            vkd3d_free(context.uav_ranges);
             vkd3d_shader_free_scan_info(scan_info);
             vkd3d_shader_parser_destroy(&parser);
             return VKD3D_ERROR_INVALID_ARGUMENT;
@@ -443,6 +494,7 @@ int vkd3d_shader_scan_dxbc(const struct vkd3d_shader_code *dxbc,
         vkd3d_shader_scan_instruction(&context, &instruction);
     }
 
+    vkd3d_free(context.uav_ranges);
     vkd3d_shader_parser_destroy(&parser);
     return VKD3D_OK;
 }

libs/vkd3d/state.c

@@ -1413,11 +1413,20 @@ static HRESULT d3d12_pipeline_state_init_compute_uav_counters(struct d3d12_pipel
     struct vkd3d_descriptor_set_context context;
     VkDescriptorSetLayoutBinding *binding_desc;
     VkDescriptorSetLayout set_layouts[3];
-    unsigned int uav_counter_count;
+    unsigned int uav_counter_count = 0;
     unsigned int i, j;
     HRESULT hr;
 
-    if (!(uav_counter_count = vkd3d_popcount(shader_info->uav_counter_mask)))
+    for (i = 0; i < shader_info->descriptor_count; ++i)
+    {
+        const struct vkd3d_shader_descriptor_info *d = &shader_info->descriptors[i];
+
+        if (d->type == VKD3D_SHADER_DESCRIPTOR_TYPE_UAV
+                && (d->flags & VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER))
+            ++uav_counter_count;
+    }
+
+    if (!uav_counter_count)
         return S_OK;
 
     if (!(binding_desc = vkd3d_calloc(uav_counter_count, sizeof(*binding_desc))))
@@ -1435,20 +1444,22 @@ static HRESULT d3d12_pipeline_state_init_compute_uav_counters(struct d3d12_pipel
     if (root_signature->vk_set_layout)
         set_layouts[context.set_index++] = root_signature->vk_set_layout;
 
-    for (i = 0, j = 0; i < VKD3D_SHADER_MAX_UNORDERED_ACCESS_VIEWS; ++i)
+    for (i = 0, j = 0; i < shader_info->descriptor_count; ++i)
     {
-        if (!(shader_info->uav_counter_mask & (1u << i)))
+        const struct vkd3d_shader_descriptor_info *d = &shader_info->descriptors[i];
+
+        if (d->type != VKD3D_SHADER_DESCRIPTOR_TYPE_UAV
+                || !(d->flags & VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER))
             continue;
 
-        state->uav_counters[j].register_space = 0;
+        state->uav_counters[j].register_space = d->register_space;
-        state->uav_counters[j].register_index = i;
+        state->uav_counters[j].register_index = d->register_index;
         state->uav_counters[j].shader_visibility = VKD3D_SHADER_VISIBILITY_COMPUTE;
         state->uav_counters[j].binding.set = context.set_index;
         state->uav_counters[j].binding.binding = context.descriptor_binding;
 
-        /* FIXME: For graphics pipeline we have to take the shader visibility
+        /* FIXME: For the graphics pipeline we have to take the shader
-         * into account.
+         * visibility into account. */
-         */
         binding_desc[j].binding = context.descriptor_binding;
         binding_desc[j].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
         binding_desc[j].descriptorCount = 1;
@@ -2266,8 +2277,17 @@ static HRESULT d3d12_pipeline_state_init_graphics(struct d3d12_pipeline_state *s
             hr = hresult_from_vkd3d_result(ret);
             goto fail;
         }
-        if (shader_info.uav_counter_mask)
+        for (j = 0; j < shader_info.descriptor_count; ++j)
-            FIXME("UAV counters not implemented for graphics pipelines.\n");
+        {
+            const struct vkd3d_shader_descriptor_info *d = &shader_info.descriptors[j];
+
+            if (d->type == VKD3D_SHADER_DESCRIPTOR_TYPE_UAV
+                    && (d->flags & VKD3D_SHADER_DESCRIPTOR_INFO_FLAG_UAV_COUNTER))
+            {
+                FIXME("UAV counters not implemented for graphics pipelines.\n");
+                break;
+            }
+        }
         vkd3d_shader_free_scan_info(&shader_info);
 
         target_info = NULL;