turnip: Add limited support for storage images.

So far this doesn't handle the texture state-based storage image access
loads, and doesn't support descriptor arrays (same as SSBOs).  The texture
side is more tricky, since we have another remapping table to work around.

This is enough to get some of dEQP-VK.image.atomic_operations.* working.

Reviewed-by: Jonathan Marek <[email protected]>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/merge_requests/3433>

6 files changed, 116 insertions, 19 deletions

diff --git a/src/freedreno/vulkan/tu_cmd_buffer.c b/src/freedreno/vulkan/tu_cmd_buffer.c
index c922c34a275..e8098dd72e1 100644
--- a/src/freedreno/vulkan/tu_cmd_buffer.c
+++ b/src/freedreno/vulkan/tu_cmd_buffer.c
@@ -2597,6 +2597,28 @@ write_tex_const(struct tu_cmd_buffer *cmd,
    }
 }
 
+static void
+write_image_ibo(struct tu_cmd_buffer *cmd,
+                uint32_t *dst,
+                struct tu_descriptor_state *descriptors_state,
+                const struct tu_descriptor_map *map,
+                unsigned i, unsigned array_index)
+{
+   assert(descriptors_state->valid & (1 << map->set[i]));
+
+   struct tu_descriptor_set *set = descriptors_state->sets[map->set[i]];
+   assert(map->binding[i] < set->layout->binding_count);
+
+   const struct tu_descriptor_set_binding_layout *layout =
+      &set->layout->binding[map->binding[i]];
+
+   assert(layout->type == VK_DESCRIPTOR_TYPE_STORAGE_IMAGE);
+
+   memcpy(dst, &set->mapped_ptr[layout->offset / 4 +
+                                (array_index * 2 + 1) * A6XX_TEX_CONST_DWORDS],
+          A6XX_TEX_CONST_DWORDS * 4);
+}
+
 static uint64_t
 buffer_ptr(struct tu_descriptor_state *descriptors_state,
            const struct tu_descriptor_map *map,
@@ -2956,7 +2978,7 @@ tu6_emit_ibo(struct tu_cmd_buffer *cmd,
       &pipeline->program.link[type];
    VkResult result;
 
-   unsigned num_desc = link->ssbo_map.num_desc;
+   unsigned num_desc = link->ssbo_map.num_desc + link->image_map.num_desc;
 
    if (num_desc == 0) {
       *entry = (struct tu_cs_entry) {};
@@ -2993,6 +3015,18 @@ tu6_emit_ibo(struct tu_cmd_buffer *cmd,
          ssbo_index++;
       }
    }
+
+   for (unsigned i = 0; i < link->image_map.num; i++) {
+      for (int j = 0; j < link->image_map.array_size[i]; j++) {
+         uint32_t *dst = &ibo_const.map[A6XX_TEX_CONST_DWORDS * ssbo_index];
+
+         write_image_ibo(cmd, dst,
+                         descriptors_state, &link->image_map, i, j);
+
+         ssbo_index++;
+      }
+   }
+
    assert(ssbo_index == num_desc);
 
    struct tu_cs cs;
diff --git a/src/freedreno/vulkan/tu_descriptor_set.c b/src/freedreno/vulkan/tu_descriptor_set.c
index c2edc549438..713a246c9a9 100644
--- a/src/freedreno/vulkan/tu_descriptor_set.c
+++ b/src/freedreno/vulkan/tu_descriptor_set.c
@@ -90,9 +90,11 @@ descriptor_size(enum VkDescriptorType type)
       /* 64bit pointer */
       return 8;
    case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
-   case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
    case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
       return A6XX_TEX_CONST_DWORDS * 4;
+   case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
+      /* We may need the IBO or the TEX representation, or both. */
+      return A6XX_TEX_CONST_DWORDS * 4 * 2;
    case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
       /* texture const + tu_sampler struct (includes border color) */
       return A6XX_TEX_CONST_DWORDS * 4 + sizeof(struct tu_sampler);
@@ -744,16 +746,13 @@ write_image_descriptor(struct tu_device *device,
              const VkDescriptorImageInfo *image_info)
 {
    TU_FROM_HANDLE(tu_image_view, iview, image_info->imageView);
-   uint32_t *descriptor;
 
+   memcpy(dst, iview->descriptor, sizeof(iview->descriptor));
    if (descriptor_type == VK_DESCRIPTOR_TYPE_STORAGE_IMAGE) {
-      descriptor = iview->storage_descriptor;
-   } else {
-      descriptor = iview->descriptor;
+      memcpy(&dst[A6XX_TEX_CONST_DWORDS], iview->storage_descriptor,
+             sizeof(iview->storage_descriptor));
    }
 
-   memcpy(dst, descriptor, sizeof(iview->descriptor));
-
    if (cmd_buffer)
       tu_bo_list_add(&cmd_buffer->bo_list, iview->image->bo, MSM_SUBMIT_BO_READ);
    else
diff --git a/src/freedreno/vulkan/tu_image.c b/src/freedreno/vulkan/tu_image.c
index 1f4c454299d..978b1b18453 100644
--- a/src/freedreno/vulkan/tu_image.c
+++ b/src/freedreno/vulkan/tu_image.c
@@ -269,6 +269,8 @@ tu_image_view_init(struct tu_image_view *iview,
    enum a6xx_tile_mode tile_mode = tu6_get_image_tile_mode(image, iview->base_mip);
    uint32_t width = u_minify(image->extent.width, iview->base_mip);
    uint32_t height = u_minify(image->extent.height, iview->base_mip);
+   uint32_t depth = pCreateInfo->viewType == VK_IMAGE_VIEW_TYPE_3D ?
+      u_minify(image->extent.depth, iview->base_mip) : iview->layer_count;
 
    unsigned fmt_tex = fmt->tex;
    if (iview->aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT &&
@@ -290,7 +292,7 @@ tu_image_view_init(struct tu_image_view *iview,
       A6XX_TEX_CONST_2_TYPE(tu6_tex_type(pCreateInfo->viewType));
    iview->descriptor[3] = A6XX_TEX_CONST_3_ARRAY_PITCH(tu_layer_size(image, iview->base_mip));
    iview->descriptor[4] = base_addr;
-   iview->descriptor[5] = base_addr >> 32;
+   iview->descriptor[5] = (base_addr >> 32) | A6XX_TEX_CONST_5_DEPTH(depth);
 
    if (image->layout.ubwc_size) {
       uint32_t block_width, block_height;
@@ -307,13 +309,36 @@ tu_image_view_init(struct tu_image_view *iview,
          A6XX_TEX_CONST_10_FLAG_BUFFER_LOGH(util_logbase2_ceil(DIV_ROUND_UP(height, block_height)));
    }
 
-   if (pCreateInfo->viewType != VK_IMAGE_VIEW_TYPE_3D) {
-      iview->descriptor[5] |= A6XX_TEX_CONST_5_DEPTH(iview->layer_count);
-   } else {
+   if (pCreateInfo->viewType == VK_IMAGE_VIEW_TYPE_3D) {
       iview->descriptor[3] |=
          A6XX_TEX_CONST_3_MIN_LAYERSZ(image->layout.slices[image->level_count - 1].size0);
-      iview->descriptor[5] |=
-         A6XX_TEX_CONST_5_DEPTH(u_minify(image->extent.depth, iview->base_mip));
+   }
+
+   if (image->usage & VK_IMAGE_USAGE_STORAGE_BIT) {
+      memset(iview->storage_descriptor, 0, sizeof(iview->storage_descriptor));
+
+      iview->storage_descriptor[0] =
+         A6XX_IBO_0_FMT(fmt->tex) |
+         A6XX_IBO_0_TILE_MODE(tile_mode);
+      iview->storage_descriptor[1] =
+         A6XX_IBO_1_WIDTH(width) |
+         A6XX_IBO_1_HEIGHT(height);
+      iview->storage_descriptor[2] =
+         A6XX_IBO_2_PITCH(pitch) |
+         A6XX_IBO_2_TYPE(tu6_tex_type(pCreateInfo->viewType));
+      iview->storage_descriptor[3] = A6XX_IBO_3_ARRAY_PITCH(tu_layer_size(image, iview->base_mip));
+
+      iview->storage_descriptor[4] = base_addr;
+      iview->storage_descriptor[5] = (base_addr >> 32) | A6XX_IBO_5_DEPTH(depth);
+
+      if (image->layout.ubwc_size) {
+         iview->storage_descriptor[3] |= A6XX_IBO_3_FLAG | A6XX_IBO_3_UNK27;
+         iview->storage_descriptor[7] |= ubwc_addr;
+         iview->storage_descriptor[8] |= ubwc_addr >> 32;
+         iview->storage_descriptor[9] = A6XX_IBO_9_FLAG_BUFFER_ARRAY_PITCH(tu_image_ubwc_size(image, iview->base_mip) >> 2);
+         iview->storage_descriptor[10] =
+            A6XX_IBO_10_FLAG_BUFFER_PITCH(tu_image_ubwc_pitch(image, iview->base_mip));
+      }
    }
 }
 
diff --git a/src/freedreno/vulkan/tu_pipeline.c b/src/freedreno/vulkan/tu_pipeline.c
index 1b43264b5bc..7593efb7bbd 100644
--- a/src/freedreno/vulkan/tu_pipeline.c
+++ b/src/freedreno/vulkan/tu_pipeline.c
@@ -361,10 +361,10 @@ tu6_blend_op(VkBlendOp op)
 static unsigned
 tu_shader_nibo(const struct tu_shader *shader)
 {
-   /* In tu_cmd_buffer.c we emit the SSBO's IBOS, but not yet storage image
-    * IBOs.
+   /* Don't use ir3_shader_nibo(), because that would include declared but
+    * unused storage images and SSBOs.
     */
-   return shader->ssbo_map.num_desc;
+   return shader->ssbo_map.num_desc + shader->image_map.num_desc;
 }
 
 static void
@@ -1641,7 +1641,7 @@ tu_pipeline_set_linkage(struct tu_program_descriptor_linkage *link,
    link->sampler_map = shader->sampler_map;
    link->ubo_map = shader->ubo_map;
    link->ssbo_map = shader->ssbo_map;
-   link->image_mapping =  v->image_mapping;
+   link->image_map = shader->image_map;
 }
 
 static void
diff --git a/src/freedreno/vulkan/tu_private.h b/src/freedreno/vulkan/tu_private.h
index b8bab7cde52..4169d015ce1 100644
--- a/src/freedreno/vulkan/tu_private.h
+++ b/src/freedreno/vulkan/tu_private.h
@@ -1058,6 +1058,7 @@ struct tu_shader
    struct tu_descriptor_map sampler_map;
    struct tu_descriptor_map ubo_map;
    struct tu_descriptor_map ssbo_map;
+   struct tu_descriptor_map image_map;
 
    /* This may be true for vertex shaders.  When true, variants[1] is the
     * binning variant and binning_binary is non-NULL.
@@ -1105,7 +1106,7 @@ struct tu_program_descriptor_linkage
    struct tu_descriptor_map sampler_map;
    struct tu_descriptor_map ubo_map;
    struct tu_descriptor_map ssbo_map;
-   struct ir3_ibo_mapping image_mapping;
+   struct tu_descriptor_map image_map;
 };
 
 struct tu_pipeline
diff --git a/src/freedreno/vulkan/tu_shader.c b/src/freedreno/vulkan/tu_shader.c
index 514f3d852f9..6303f222264 100644
--- a/src/freedreno/vulkan/tu_shader.c
+++ b/src/freedreno/vulkan/tu_shader.c
@@ -295,6 +295,24 @@ lower_vulkan_resource_index(nir_builder *b, nir_intrinsic_instr *instr,
    nir_instr_remove(&instr->instr);
 }
 
+static void
+add_image_deref_mapping(nir_intrinsic_instr *instr, struct tu_shader *shader,
+                const struct tu_pipeline_layout *layout)
+{
+   nir_deref_instr *deref = nir_src_as_deref(instr->src[0]);
+   nir_variable *var = nir_deref_instr_get_variable(deref);
+
+   uint32_t set = var->data.descriptor_set;
+   uint32_t binding = var->data.binding;
+   struct tu_descriptor_set_layout *set_layout = layout->set[set].layout;
+   struct tu_descriptor_set_binding_layout *binding_layout =
+      &set_layout->binding[binding];
+
+   var->data.driver_location =
+      map_add(&shader->image_map, set, binding, var->data.index,
+              binding_layout->array_size);
+}
+
 static bool
 lower_intrinsic(nir_builder *b, nir_intrinsic_instr *instr,
                 struct tu_shader *shader,
@@ -316,6 +334,26 @@ lower_intrinsic(nir_builder *b, nir_intrinsic_instr *instr,
       lower_vulkan_resource_index(b, instr, shader, layout);
       return true;
 
+   case nir_intrinsic_image_deref_load:
+   case nir_intrinsic_image_deref_store:
+   case nir_intrinsic_image_deref_atomic_add:
+   case nir_intrinsic_image_deref_atomic_imin:
+   case nir_intrinsic_image_deref_atomic_umin:
+   case nir_intrinsic_image_deref_atomic_imax:
+   case nir_intrinsic_image_deref_atomic_umax:
+   case nir_intrinsic_image_deref_atomic_and:
+   case nir_intrinsic_image_deref_atomic_or:
+   case nir_intrinsic_image_deref_atomic_xor:
+   case nir_intrinsic_image_deref_atomic_exchange:
+   case nir_intrinsic_image_deref_atomic_comp_swap:
+   case nir_intrinsic_image_deref_size:
+   case nir_intrinsic_image_deref_samples:
+   case nir_intrinsic_image_deref_load_param_intel:
+   case nir_intrinsic_image_deref_load_raw_intel:
+   case nir_intrinsic_image_deref_store_raw_intel:
+      add_image_deref_mapping(instr, shader, layout);
+      return true;
+
    default:
       return false;
    }