xref: /aosp_15_r20/external/mesa3d/src/asahi/compiler/agx_opt_promote_constants.c (revision 6104692788411f58d303aa86923a9ff6ecaded22)
1 /*
2  * Copyright 2023 Valve Corporation
3  * SPDX-License-Identifier: MIT
4  */
5 
6 #include <stdlib.h>
7 #include "util/bitset.h"
8 #include "util/hash_table.h"
9 #include "util/ralloc.h"
10 #include "agx_compiler.h"
11 #include "agx_opcodes.h"
12 
13 /*
14  * Information about a constant, indexed by its 64-bit value. This describes the
15  * value, not the move that generated it. If there are multiple moves in the
16  * shader with the same immediate value, they resolve to the same constant.
17  */
18 struct constant_info {
19    uint64_t value;
20 
21    /* Number of uses of the constant that could be promoted */
22    unsigned nr_promotable_uses;
23 
24    /* If we push, the uniform used */
25    uint16_t uniform;
26 
27    /* Alignment in 16-bit units needed for the constant */
28    uint8_t align_16;
29 
30    /* True if the constant was promoted to a uniform */
31    bool promoted;
32 };
33 
34 /*
35  * Choosing constants to promote is similar to the 0-1 knapsack problem. We use
36  * a well-known heuristic: sort by benefit divided by size. We approximate
37  * benefit by use count.
38  */
39 static int
constant_priority(const struct constant_info * const info)40 constant_priority(const struct constant_info *const info)
41 {
42    int size = info->align_16;
43    assert(size == 1 || size == 2 || size == 4);
44    int inverse_size = (size == 1) ? 4 : (size == 2) ? 2 : 1;
45 
46    return info->nr_promotable_uses * inverse_size;
47 }
48 
49 static int
priority_compare(const void * A_,const void * B_)50 priority_compare(const void *A_, const void *B_)
51 {
52    const struct constant_info *const *A = A_;
53    const struct constant_info *const *B = B_;
54 
55    /* This is backwards from qsort's documentation, because we want descending
56     * order and qsort returns ascending.
57     */
58    return constant_priority(*B) - constant_priority(*A);
59 }
60 
61 static void
record_use(void * memctx,struct hash_table_u64 * constants,uint64_t imm,enum agx_size size)62 record_use(void *memctx, struct hash_table_u64 *constants, uint64_t imm,
63            enum agx_size size)
64 {
65    struct constant_info *info = _mesa_hash_table_u64_search(constants, imm);
66 
67    if (!info) {
68       info = rzalloc(memctx, struct constant_info);
69       info->value = imm;
70       _mesa_hash_table_u64_insert(constants, imm, info);
71    }
72 
73    info->nr_promotable_uses++;
74    info->align_16 = MAX2(info->align_16, agx_size_align_16(size));
75 }
76 
77 static void
pass(agx_context * ctx,void * memctx)78 pass(agx_context *ctx, void *memctx)
79 {
80    /* Map from SSA indices to struct constant_info */
81    struct hash_table_u64 *constants = _mesa_hash_table_u64_create(memctx);
82 
83    /* Map from SSA indices to immediate values */
84    uint64_t *values = rzalloc_array(memctx, uint64_t, ctx->alloc);
85 
86    /* Set of SSA indices that map to immediate values */
87    BITSET_WORD *is_immediate =
88       rzalloc_array(memctx, BITSET_WORD, BITSET_WORDS(ctx->alloc));
89 
90    /* Gather constant definitions and use */
91    agx_foreach_instr_global(ctx, I) {
92       if (I->op == AGX_OPCODE_MOV_IMM) {
93          assert(I->dest[0].type == AGX_INDEX_NORMAL);
94          BITSET_SET(is_immediate, I->dest[0].value);
95          values[I->dest[0].value] = I->imm;
96       } else {
97          agx_foreach_ssa_src(I, s) {
98             if (BITSET_TEST(is_immediate, I->src[s].value) &&
99                 agx_instr_accepts_uniform(I->op, s, ctx->out->push_count,
100                                           I->src[s].size)) {
101 
102                record_use(memctx, constants, values[I->src[s].value],
103                           I->src[s].size);
104             }
105          }
106       }
107    }
108 
109    /* Early exit if there were no constants */
110    unsigned nr_nodes = _mesa_hash_table_u64_num_entries(constants);
111    if (nr_nodes == 0)
112       return;
113 
114    /* Collect nodes that are promotable */
115    struct constant_info **flat =
116       rzalloc_array(memctx, struct constant_info *, nr_nodes);
117 
118    unsigned flat_count = 0;
119    hash_table_u64_foreach(constants, entry) {
120       flat[flat_count++] = entry.data;
121    }
122 
123    /* Select constants. Even when we can promote everything, sorting keeps hot
124     * constants in lower uniforms, required by some instructions.
125     */
126    qsort(flat, flat_count, sizeof(*flat), priority_compare);
127 
128    ctx->out->immediate_base_uniform = ctx->out->push_count;
129 
130    /* Promote as many constants as we can */
131    for (unsigned i = 0; i < flat_count; ++i) {
132       struct constant_info *info = flat[i];
133       assert(info->nr_promotable_uses > 0);
134 
135       /* Try to assign a uniform */
136       unsigned uniform = ALIGN_POT(ctx->out->push_count, info->align_16);
137       unsigned new_count = uniform + info->align_16;
138       if (new_count > AGX_NUM_UNIFORMS)
139          break;
140 
141       info->uniform = uniform;
142       info->promoted = true;
143       ctx->out->push_count = new_count;
144 
145       unsigned size_B = info->align_16 * 2;
146       memcpy(&ctx->out->immediates[uniform - ctx->out->immediate_base_uniform],
147              &info->value, size_B);
148 
149       ctx->out->immediate_size_16 =
150          new_count - ctx->out->immediate_base_uniform;
151    }
152 
153    /* Promote in the IR */
154    agx_foreach_instr_global(ctx, I) {
155       agx_foreach_ssa_src(I, s) {
156          if (!BITSET_TEST(is_immediate, I->src[s].value))
157             continue;
158 
159          struct constant_info *info =
160             _mesa_hash_table_u64_search(constants, values[I->src[s].value]);
161 
162          if (info && info->promoted &&
163              agx_instr_accepts_uniform(I->op, s, info->uniform,
164                                        I->src[s].size)) {
165 
166             agx_replace_src(I, s, agx_uniform(info->uniform, I->src[s].size));
167          }
168       }
169    }
170 }
171 
172 void
agx_opt_promote_constants(agx_context * ctx)173 agx_opt_promote_constants(agx_context *ctx)
174 {
175    /* We do not promote constants in preambles since it's pointless and wastes
176     * uniform slots.
177     */
178    if (ctx->is_preamble)
179       return;
180 
181    void *memctx = ralloc_context(NULL);
182    pass(ctx, memctx);
183    ralloc_free(memctx);
184 }
185