From f44afe766f35942222ec30b2f1128c1126b45974 Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Mon, 19 Dec 2022 23:27:30 -0500
Subject: [PATCH] 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           | 28 +++++++++++++++++-----
 src/asahi/compiler/agx_compiler.h          | 10 ++++++++
 src/asahi/compiler/agx_print.c             |  4 ++++
 src/asahi/compiler/agx_register_allocate.c |  2 +-
 4 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index a26d7630e8b..d76f387d3e4 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/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);
 }
 
 /*
diff --git a/src/asahi/compiler/agx_compiler.h b/src/asahi/compiler/agx_compiler.h
index 6843b5b193f..f5ec42b47eb 100644
--- a/src/asahi/compiler/agx_compiler.h
+++ b/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)
diff --git a/src/asahi/compiler/agx_print.c b/src/asahi/compiler/agx_print.c
index c6f26b67d1d..b9b5240f6e6 100644
--- a/src/asahi/compiler/agx_print.c
+++ b/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;
diff --git a/src/asahi/compiler/agx_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c
index 2a4b2ad37eb..8fac07caac9 100644
--- a/src/asahi/compiler/agx_register_allocate.c
+++ b/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);