agx: Use texture write mask

We do need to use undefs instead of zeroes in this internal collect. While this
vector gets copypropped out, it'd cause us to fail compilation if noopt is on.
Fix that.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/20446>

src/asahi/compiler/agx_compile.c

@@ -1275,20 +1275,36 @@ agx_emit_tex(agx_builder *b, nir_tex_instr *instr)
    else if (!agx_is_null(compare))
       compare_offset = compare;
 
-   unsigned channels = nir_dest_num_components(instr->dest);
+   unsigned nr_channels = nir_dest_num_components(instr->dest);
+   nir_component_mask_t mask = nir_ssa_def_components_read(&instr->dest.ssa);
+
+   agx_index tmp = agx_temp(b->shader, dst.size);
 
    agx_instr *I = agx_texture_sample_to(
-      b, dst, coords, lod, texture, sampler, compare_offset,
+      b, tmp, coords, lod, texture, sampler, compare_offset,
       agx_tex_dim(instr->sampler_dim, instr->is_array),
-      agx_lod_mode_for_nir(instr->op),
-      BITFIELD_MASK(channels), /* TODO: wrmask */
-      0, !agx_is_null(packed_offset), !agx_is_null(compare));
+      agx_lod_mode_for_nir(instr->op), mask, 0, !agx_is_null(packed_offset),
+      !agx_is_null(compare));
 
    if (txf)
       I->op = AGX_OPCODE_TEXTURE_LOAD;
 
    agx_wait(b, 0);
-   agx_emit_cached_split(b, dst, channels);
+
+   agx_index packed_channels[4] = {agx_null()};
+   agx_index unpacked_channels[4] = {agx_null()};
+
+   /* Hardware writes the masked components contiguously, expand out for NIR */
+   agx_emit_split(b, packed_channels, tmp, 4 /* XXX: why not nr_channels */);
+
+   for (unsigned i = 0; i < nr_channels; ++i) {
+      unpacked_channels[i] =
+         (mask & BITFIELD_BIT(i))
+            ? packed_channels[util_bitcount(mask & BITFIELD_MASK(i))]
+            : agx_undef(tmp.size);
+   }
+
+   agx_emit_collect_to(b, dst, nr_channels, unpacked_channels);
 }
 
 /*

src/asahi/compiler/agx_compiler.h

@@ -64,6 +64,7 @@ enum agx_index_type {
    AGX_INDEX_IMMEDIATE = 2,
    AGX_INDEX_UNIFORM = 3,
    AGX_INDEX_REGISTER = 4,
+   AGX_INDEX_UNDEF = 5,
 };
 
 enum agx_size { AGX_SIZE_16 = 0, AGX_SIZE_32 = 1, AGX_SIZE_64 = 2 };
@@ -148,6 +149,15 @@ agx_register(uint32_t imm, enum agx_size size)
    };
 }
 
+static inline agx_index
+agx_undef(enum agx_size size)
+{
+   return (agx_index){
+      .size = size,
+      .type = AGX_INDEX_UNDEF,
+   };
+}
+
 /* Also in half-words */
 static inline agx_index
 agx_uniform(uint32_t imm, enum agx_size size)

src/asahi/compiler/agx_print.c

@@ -75,6 +75,10 @@ agx_print_index(agx_index index, bool is_float, FILE *fp)
 
       break;
 
+   case AGX_INDEX_UNDEF:
+      fprintf(fp, "undef");
+      break;
+
    case AGX_INDEX_UNIFORM:
       agx_print_sized('u', index.value, index.size, fp);
       break;

src/asahi/compiler/agx_register_allocate.c

@@ -459,7 +459,7 @@ agx_ra(agx_context *ctx)
 
          /* Move the sources */
          agx_foreach_src(ins, i) {
-            if (agx_is_null(ins->src[i]))
+            if (agx_is_null(ins->src[i]) || ins->src[i].type == AGX_INDEX_UNDEF)
                continue;
             assert(ins->src[i].size == ins->src[0].size);