agx: make vec widths explicit in IR

no functional changes.

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

src/asahi/compiler/agx_compile.c

@@ -215,7 +215,7 @@ agx_emit_collect_to(agx_builder *b, agx_index dst, unsigned nr_srcs,
 static agx_index
 agx_emit_collect(agx_builder *b, unsigned nr_srcs, agx_index *srcs)
 {
-   agx_index dst = agx_temp(b->shader, srcs[0].size);
+   agx_index dst = agx_vec_temp(b->shader, srcs[0].size, nr_srcs);
    agx_emit_collect_to(b, dst, nr_srcs, srcs);
    return dst;
 }
@@ -667,6 +667,7 @@ agx_emit_block_image_store(agx_builder *b, nir_intrinsic_instr *instr)
     * logically to ensure alignment.
     */
    offset = agx_vec2(b, offset, agx_undef(AGX_SIZE_16));
+   offset.channels_m1--;
    offset.size = AGX_SIZE_32;
 
    /* Modified coordinate descriptor */
@@ -944,6 +945,7 @@ agx_emit_image_load(agx_builder *b, agx_index dst, nir_intrinsic_instr *intr)
       assert(ms_index.size == AGX_SIZE_16);
       agx_index vec = agx_vec2(b, ms_index, layer);
       vec.size = AGX_SIZE_32;
+      vec.channels_m1 = 1 - 1;
       coord[coord_comps++] = vec;
    } else if (is_ms) {
       agx_index tmp = agx_temp(b->shader, AGX_SIZE_32);
@@ -960,7 +962,7 @@ agx_emit_image_load(agx_builder *b, agx_index dst, nir_intrinsic_instr *intr)
    }
 
    agx_index coords = agx_emit_collect(b, coord_comps, coord);
-   agx_index tmp = agx_temp(b->shader, dst.size);
+   agx_index tmp = agx_vec_temp(b->shader, dst.size, 4);
 
    agx_instr *I = agx_image_load_to(
       b, tmp, coords, lod, bindless, texture, agx_txf_sampler(b->shader),
@@ -1816,7 +1818,7 @@ agx_emit_tex(agx_builder *b, nir_tex_instr *instr)
          agx_index index2 = agx_src_index(&instr->src[y_idx].src);
 
          /* We explicitly don't cache about the split cache for this */
-         lod = agx_temp(b->shader, AGX_SIZE_32);
+         lod = agx_vec_temp(b->shader, AGX_SIZE_32, 2 * n);
          agx_instr *I = agx_collect_to(b, lod, 2 * n);
 
          for (unsigned i = 0; i < n; ++i) {
@@ -1848,7 +1850,7 @@ agx_emit_tex(agx_builder *b, nir_tex_instr *instr)
    else if (!agx_is_null(compare))
       compare_offset = compare;
 
-   agx_index tmp = agx_temp(b->shader, dst.size);
+   agx_index tmp = agx_vec_temp(b->shader, dst.size, 4);
    agx_instr *I = agx_texture_sample_to(
       b, tmp, coords, lod, bindless, texture, sampler, compare_offset,
       agx_tex_dim(instr->sampler_dim, instr->is_array),

src/asahi/compiler/agx_compiler.h

@@ -56,7 +56,7 @@ agx_size_align_16(enum agx_size size)
 typedef struct {
    /* Sufficient for as many SSA values as we need. Immediates and uniforms fit
     * in 16-bits */
-   unsigned value : 22;
+   unsigned value : 19;
 
    /* Indicates that this source kills the referenced value (because it is the
     * last use in a block and the source is not live after the block). Set by
@@ -71,20 +71,34 @@ typedef struct {
    bool abs : 1;
    bool neg : 1;
 
+   unsigned channels_m1     : 3;
    enum agx_size size       : 2;
    enum agx_index_type type : 3;
 } agx_index;
 
+static inline unsigned
+agx_channels(agx_index idx)
+{
+   return idx.channels_m1 + 1;
+}
+
 static inline agx_index
-agx_get_index(unsigned value, enum agx_size size)
+agx_get_vec_index(unsigned value, enum agx_size size, unsigned channels)
 {
    return (agx_index){
       .value = value,
+      .channels_m1 = channels - 1,
       .size = size,
       .type = AGX_INDEX_NORMAL,
    };
 }
 
+static inline agx_index
+agx_get_index(unsigned value, enum agx_size size)
+{
+   return agx_get_vec_index(value, size, 1);
+}
+
 static inline agx_index
 agx_immediate(uint32_t imm)
 {
@@ -117,6 +131,19 @@ agx_register(uint32_t imm, enum agx_size size)
    };
 }
 
+static inline agx_index
+agx_register_like(uint32_t imm, agx_index like)
+{
+   assert(imm < AGX_NUM_REGS);
+
+   return (agx_index){
+      .value = imm,
+      .channels_m1 = like.channels_m1,
+      .size = like.size,
+      .type = AGX_INDEX_REGISTER,
+   };
+}
+
 static inline agx_index
 agx_undef(enum agx_size size)
 {
@@ -442,6 +469,12 @@ agx_remove_instruction(agx_instr *ins)
    list_del(&ins->link);
 }
 
+static inline agx_index
+agx_vec_temp(agx_context *ctx, enum agx_size size, unsigned channels)
+{
+   return agx_get_vec_index(ctx->alloc++, size, channels);
+}
+
 static inline agx_index
 agx_temp(agx_context *ctx, enum agx_size size)
 {
@@ -468,7 +501,8 @@ agx_size_for_bits(unsigned bits)
 static inline agx_index
 agx_def_index(nir_def *ssa)
 {
-   return agx_get_index(ssa->index, agx_size_for_bits(ssa->bit_size));
+   return agx_get_vec_index(ssa->index, agx_size_for_bits(ssa->bit_size),
+                            ssa->num_components);
 }
 
 static inline agx_index
@@ -480,7 +514,8 @@ agx_src_index(nir_src *src)
 static inline agx_index
 agx_vec_for_def(agx_context *ctx, nir_def *def)
 {
-   return agx_temp(ctx, agx_size_for_bits(def->bit_size));
+   return agx_vec_temp(ctx, agx_size_for_bits(def->bit_size),
+                       def->num_components);
 }
 
 static inline agx_index

src/asahi/compiler/agx_register_allocate.c

@@ -758,7 +758,7 @@ agx_set_sources(struct ra_ctx *rctx, agx_instr *I)
       assert(BITSET_TEST(rctx->visited, I->src[s].value) && "no phis");
 
       unsigned v = rctx->ssa_to_reg[I->src[s].value];
-      agx_replace_src(I, s, agx_register(v, I->src[s].size));
+      agx_replace_src(I, s, agx_register_like(v, I->src[s]));
    }
 }
 
@@ -768,7 +768,7 @@ agx_set_dests(struct ra_ctx *rctx, agx_instr *I)
    agx_foreach_ssa_dest(I, s) {
       unsigned v = rctx->ssa_to_reg[I->dest[s].value];
       I->dest[s] =
-         agx_replace_index(I->dest[s], agx_register(v, I->dest[s].size));
+         agx_replace_index(I->dest[s], agx_register_like(v, I->dest[s]));
    }
 }
 
@@ -1243,6 +1243,7 @@ agx_ra(agx_context *ctx)
 
             agx_index src = ins->src[0];
             src.size = ins->dest[i].size;
+            src.channels_m1 = 0;
             src.value += (i * width);
 
             copies[n++] = (struct agx_copy){

src/asahi/compiler/agx_validate.c

@@ -123,45 +123,45 @@ agx_validate_defs(agx_instr *I, BITSET_WORD *defs)
    return true;
 }
 
-/*
- * Type check the dimensionality of sources and destinations. This occurs in two
- * passes, first to gather all destination sizes, second to validate all source
- * sizes. Depends on SSA form.
- */
+/* Type check the dimensionality of sources and destinations. */
 static bool
 agx_validate_width(agx_context *ctx)
 {
    bool succ = true;
-   uint8_t *width = calloc(ctx->alloc, sizeof(uint8_t));
 
    agx_foreach_instr_global(ctx, I) {
       agx_foreach_dest(I, d) {
-         if (I->dest[d].type != AGX_INDEX_NORMAL)
+         unsigned exp = agx_write_registers(I, d);
+         unsigned act =
+            agx_channels(I->dest[d]) * agx_size_align_16(I->dest[d].size);
+
+         if (exp != act) {
+            succ = false;
+            fprintf(stderr, "destination %u, expected width %u, got width %u\n",
+                    d, exp, act);
+            agx_print_instr(I, stderr);
+            fprintf(stderr, "\n");
+         }
+      }
+
+      agx_foreach_src(I, s) {
+         if (I->src[s].type == AGX_INDEX_NULL)
             continue;
 
-         unsigned v = I->dest[d].value;
-         assert(width[v] == 0 && "broken SSA");
+         unsigned exp = agx_read_registers(I, s);
+         unsigned act =
+            agx_channels(I->src[s]) * agx_size_align_16(I->src[s].size);
 
-         width[v] = agx_write_registers(I, d);
-      }
-   }
-
-   agx_foreach_instr_global(ctx, I) {
-      agx_foreach_ssa_src(I, s) {
-         unsigned v = I->src[s].value;
-         unsigned n = agx_read_registers(I, s);
-
-         if (width[v] != n) {
+         if (exp != act) {
             succ = false;
             fprintf(stderr, "source %u, expected width %u, got width %u\n", s,
-                    n, width[v]);
+                    exp, act);
             agx_print_instr(I, stderr);
             fprintf(stderr, "\n");
          }
       }
    }
 
-   free(width);
    return succ;
 }