agx: Handle layered block image stores

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>

src/asahi/compiler/agx_compile.c

@@ -635,14 +635,33 @@ agx_emit_block_image_store(agx_builder *b, nir_intrinsic_instr *instr)
 {
    unsigned image = nir_src_as_uint(instr->src[0]);
    agx_index offset = agx_src_index(&instr->src[1]);
+   agx_index layer = agx_src_index(&instr->src[2]);
    enum agx_format format = agx_format_for_pipe(nir_intrinsic_format(instr));
-   enum agx_dim dim = agx_tex_dim(nir_intrinsic_image_dim(instr), false);
+
+   bool ms = nir_intrinsic_image_dim(instr) == GLSL_SAMPLER_DIM_MS;
+   bool array = nir_intrinsic_image_array(instr);
+   enum agx_dim dim = agx_tex_dim(nir_intrinsic_image_dim(instr), array);
+
+   /* Modified coordinate descriptor */
+   agx_index coords;
+   if (array) {
+      coords = agx_temp(b->shader, AGX_SIZE_32);
+      agx_emit_collect_to(
+         b, coords, 2,
+         (agx_index[]){
+            ms ? agx_mov_imm(b, 16, 0) : layer,
+            ms ? layer : agx_mov_imm(b, 16, 0xFFFF) /* TODO: Why can't zero? */,
+         });
+   } else {
+      coords = agx_null();
+   }
 
    // XXX: how does this possibly work
    if (format == AGX_FORMAT_F16)
       format = AGX_FORMAT_I16;
 
-   return agx_block_image_store(b, agx_immediate(image), offset, format, dim);
+   return agx_block_image_store(b, agx_immediate(image), offset, coords, format,
+                                dim);
 }
 
 static agx_instr *

src/asahi/compiler/agx_opcodes.py

@@ -384,9 +384,9 @@ op("signal_pix", (0x58, 0xFF, 4, _), dests = 0, imms = [WRITEOUT],
 op("image_write", (0xF1 | (1 << 23) | (9 << 43), 0xFF, 6, 8), dests = 0, srcs = 5, imms
    = [DIM], can_eliminate = False, schedule_class = "store")
 
-# Sources are the image and the offset within shared memory
+# Sources are the image, the offset within shared memory, and the layer.
 # TODO: Do we need the short encoding?
-op("block_image_store", (0xB1, 0xFF, 10, _), dests = 0, srcs = 2,
+op("block_image_store", (0xB1, 0xFF, 10, _), dests = 0, srcs = 3,
    imms = [FORMAT, DIM], can_eliminate = False, schedule_class = "store")
 
 # Barriers

src/asahi/compiler/agx_pack.c

@@ -890,6 +890,10 @@ agx_pack_instr(struct util_dynarray *emission, struct util_dynarray *fixups,
       unsigned T = agx_pack_texture(agx_zero(), I->src[0], &U, &Tt);
       assert(T < 0x100);
 
+      bool Cs = false;
+      bool Ct = I->src[2].discard;
+      unsigned C = I->src[2].value;
+
       agx_index offset = I->src[1];
       assert(offset.type == AGX_INDEX_REGISTER);
       assert(offset.size == AGX_SIZE_16);
@@ -901,12 +905,14 @@ agx_pack_instr(struct util_dynarray *emission, struct util_dynarray *fixups,
       uint32_t word0 = agx_opcodes_info[I->op].encoding.exact |
                        (1 << 15) /* we always set length bit for now */ |
                        ((F & 1) << 8) | ((R & BITFIELD_MASK(6)) << 9) |
+                       ((C & BITFIELD_MASK(6)) << 16) | (Ct ? (1 << 22) : 0) |
                        (unk1 ? (1u << 31) : 0);
 
       uint32_t word1 = (T & BITFIELD_MASK(6)) | (Tt << 2) |
-                       ((I->dim & BITFIELD_MASK(3)) << 8) |
+                       ((I->dim & BITFIELD_MASK(3)) << 8) | (9 << 11) |
+                       (Cs ? (1 << 15) : 0) |
                        ((I->dim & BITFIELD_BIT(3)) ? (1u << 23) : 0) |
-                       ((R >> 6) << 24);
+                       ((R >> 6) << 24) | ((C >> 6) << 26);
 
       uint32_t word2 = (F >> 1) | (unk3 ? (1 << 3) : 0) | ((T >> 6) << 14);