zink: implement nir_texop_txs

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/merge_requests/3275>

3 files changed, 50 insertions, 4 deletions

diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
index b9fb5380b17..35c2fc4c0f3 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
@@ -1366,8 +1366,10 @@ emit_tex(struct ntv_context *ctx, nir_tex_instr *tex)
           tex->op == nir_texop_txb ||
           tex->op == nir_texop_txl ||
           tex->op == nir_texop_txd ||
-          tex->op == nir_texop_txf);
-   assert(nir_alu_type_get_base_type(tex->dest_type) == nir_type_float);
+          tex->op == nir_texop_txf ||
+          tex->op == nir_texop_txs);
+   assert(tex->op == nir_texop_txs ||
+          nir_alu_type_get_base_type(tex->dest_type) == nir_type_float);
    assert(tex->texture_index == tex->sampler_index);
 
    SpvId coord = 0, proj = 0, bias = 0, lod = 0, dref = 0, dx = 0, dy = 0;
@@ -1396,7 +1398,8 @@ emit_tex(struct ntv_context *ctx, nir_tex_instr *tex)
 
       case nir_tex_src_lod:
          assert(nir_src_num_components(tex->src[i].src) == 1);
-         if (tex->op == nir_texop_txf)
+         if (tex->op == nir_texop_txf ||
+             tex->op == nir_texop_txs)
             lod = get_src_int(ctx, &tex->src[i].src);
          else
             lod = get_src_float(ctx, &tex->src[i].src);
@@ -1445,6 +1448,15 @@ emit_tex(struct ntv_context *ctx, nir_tex_instr *tex)
 
    SpvId dest_type = get_dest_type(ctx, &tex->dest, tex->dest_type);
 
+   if (tex->op == nir_texop_txs) {
+      SpvId image = spirv_builder_emit_image(&ctx->builder, image_type, load);
+      SpvId result = spirv_builder_emit_image_query_size(&ctx->builder,
+                                                         dest_type, image,
+                                                         lod);
+      store_dest(ctx, &tex->dest, result, tex->dest_type);
+      return;
+   }
+
    if (proj) {
       SpvId constituents[coord_components + 1];
       if (coord_components == 1)
@@ -1779,8 +1791,10 @@ nir_to_spirv(struct nir_shader *s)
    }
 
    // TODO: only enable when needed
-   if (s->info.stage == MESA_SHADER_FRAGMENT)
+   if (s->info.stage == MESA_SHADER_FRAGMENT) {
       spirv_builder_emit_cap(&ctx.builder, SpvCapabilitySampled1D);
+      spirv_builder_emit_cap(&ctx.builder, SpvCapabilityImageQuery);
+   }
 
    ctx.stage = s->info.stage;
    ctx.GLSL_std_450 = spirv_builder_import(&ctx.builder, "GLSL.std.450");
diff --git a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
index 572b38b291d..bae23c97d4c 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
@@ -608,6 +608,32 @@ spirv_builder_emit_image_fetch(struct spirv_builder *b,
 }
 
 SpvId
+spirv_builder_emit_image_query_size(struct spirv_builder *b,
+                                    SpvId result_type,
+                                    SpvId image,
+                                    SpvId lod)
+{
+   int opcode = SpvOpImageQuerySize;
+   int words = 4;
+   if (lod) {
+      words++;
+      opcode = SpvOpImageQuerySizeLod;
+   }
+
+   SpvId result = spirv_builder_new_id(b);
+   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_emit_word(&b->instructions, opcode | (words << 16));
+   spirv_buffer_emit_word(&b->instructions, result_type);
+   spirv_buffer_emit_word(&b->instructions, result);
+   spirv_buffer_emit_word(&b->instructions, image);
+
+   if (lod)
+      spirv_buffer_emit_word(&b->instructions, lod);
+
+   return result;
+}
+
+SpvId
 spirv_builder_emit_ext_inst(struct spirv_builder *b, SpvId result_type,
                             SpvId set, uint32_t instruction,
                             const SpvId *args, size_t num_args)
diff --git a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
index d7213d85f65..56db0a11cf5 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
+++ b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
@@ -229,6 +229,12 @@ spirv_builder_emit_image_fetch(struct spirv_builder *b,
                                SpvId lod);
 
 SpvId
+spirv_builder_emit_image_query_size(struct spirv_builder *b,
+                                    SpvId result_type,
+                                    SpvId image,
+                                    SpvId lod);
+
+SpvId
 spirv_builder_emit_ext_inst(struct spirv_builder *b, SpvId result_type,
                             SpvId set, uint32_t instruction,
                             const SpvId args[], size_t num_args);