src/compiler/nir/nir_opt_vectorize.c

   1 /*
   2  * Copyright © 2015 Connor Abbott
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the "Software"),
   6  * to deal in the Software without restriction, including without limitation
   7  * the rights to use, copy, modify, merge, publish, distribute, sublicense,
   8  * and/or sell copies of the Software, and to permit persons to whom the
   9  * Software is furnished to do so, subject to the following conditions:
  10  *
  11  * The above copyright notice and this permission notice (including the next
  12  * paragraph) shall be included in all copies or substantial portions of the
  13  * Software.
  14  *
  15  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
  18  * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  20  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
  21  * IN THE SOFTWARE.
  22  *
  23  */
  24
  25 #include "nir.h"
  26 #include "nir_vla.h"
  27 #include "nir_builder.h"
  28 #include "util/u_dynarray.h"
  29
  30 #define HASH(hash, data) XXH32(&data, sizeof(data), hash)
  31
  32 static uint32_t
  33 hash_src(uint32_t hash, const nir_src *src)
  34 {
  35    assert(src->is_ssa);
  36    void *hash_data = nir_src_is_const(*src) ? NULL : src->ssa;
  37
  38    return HASH(hash, hash_data);
  39 }
  40
  41 static uint32_t
  42 hash_alu_src(uint32_t hash, const nir_alu_src *src)
  43 {
  44    assert(!src->abs && !src->negate);
  45
  46    /* intentionally don't hash swizzle */
  47
  48    return hash_src(hash, &src->src);
  49 }
  50
  51 static uint32_t
  52 hash_alu(uint32_t hash, const nir_alu_instr *instr)
  53 {
  54    hash = HASH(hash, instr->op);
  55
  56    hash = HASH(hash, instr->dest.dest.ssa.bit_size);
  57
  58    for (unsigned i = 0; i < nir_op_infos[instr->op].num_inputs; i++)
  59       hash = hash_alu_src(hash, &instr->src[i]);
  60
  61    return hash;
  62 }
  63
  64 static uint32_t
  65 hash_instr(const nir_instr *instr)
  66 {
  67    uint32_t hash = 0;
  68
  69    switch (instr->type) {
  70    case nir_instr_type_alu:
  71       return hash_alu(hash, nir_instr_as_alu(instr));
  72    default:
  73       unreachable("bad instruction type");
  74    }
  75 }
  76
  77 static bool
  78 srcs_equal(const nir_src *src1, const nir_src *src2)
  79 {
  80    assert(src1->is_ssa);
  81    assert(src2->is_ssa);
  82
  83    return src1->ssa == src2->ssa ||
  84       nir_src_is_const(*src1) == nir_src_is_const(*src2);
  85 }
  86
  87 static bool
  88 alu_srcs_equal(const nir_alu_src *src1, const nir_alu_src *src2)
  89 {
  90    assert(!src1->abs);
  91    assert(!src1->negate);
  92    assert(!src2->abs);
  93    assert(!src2->negate);
  94
  95    return srcs_equal(&src1->src, &src2->src);
  96 }
  97
  98 static bool
  99 instrs_equal(const nir_instr *instr1, const nir_instr *instr2)
 100 {
 101    switch (instr1->type) {
 102    case nir_instr_type_alu: {
 103       nir_alu_instr *alu1 = nir_instr_as_alu(instr1);
 104       nir_alu_instr *alu2 = nir_instr_as_alu(instr2);
 105
 106       if (alu1->op != alu2->op)
 107          return false;
 108
 109       if (alu1->dest.dest.ssa.bit_size != alu2->dest.dest.ssa.bit_size)
 110          return false;
 111
 112       for (unsigned i = 0; i < nir_op_infos[alu1->op].num_inputs; i++) {
 113          if (!alu_srcs_equal(&alu1->src[i], &alu2->src[i]))
 114             return false;
 115       }
 116
 117       return true;
 118    }
 119
 120    default:
 121       unreachable("bad instruction type");
 122    }
 123 }
 124
 125 static bool
 126 instr_can_rewrite(nir_instr *instr)
 127 {
 128    switch (instr->type) {
 129    case nir_instr_type_alu: {
 130       nir_alu_instr *alu = nir_instr_as_alu(instr);
 131
 132       /* Don't try and vectorize mov's. Either they'll be handled by copy
 133        * prop, or they're actually necessary and trying to vectorize them
 134        * would result in fighting with copy prop.
 135        */
 136       if (alu->op == nir_op_mov)
 137          return false;
 138
 139       if (nir_op_infos[alu->op].output_size != 0)
 140          return false;
 141
 142       for (unsigned i = 0; i < nir_op_infos[alu->op].num_inputs; i++) {
 143          if (nir_op_infos[alu->op].input_sizes[i] != 0)
 144             return false;
 145       }
 146
 147       return true;
 148    }
 149
 150    /* TODO support phi nodes */
 151    default:
 152       break;
 153    }
 154
 155    return false;
 156 }
 157
 158 /*
 159  * Tries to combine two instructions whose sources are different components of
 160  * the same instructions into one vectorized instruction. Note that instr1
 161  * should dominate instr2.
 162  */
 163
 164 static nir_instr *
 165 instr_try_combine(struct nir_shader *nir, nir_instr *instr1, nir_instr *instr2,
 166                   nir_opt_vectorize_cb filter, void *data)
 167 {
 168    assert(instr1->type == nir_instr_type_alu);
 169    assert(instr2->type == nir_instr_type_alu);
 170    nir_alu_instr *alu1 = nir_instr_as_alu(instr1);
 171    nir_alu_instr *alu2 = nir_instr_as_alu(instr2);
 172
 173    assert(alu1->dest.dest.ssa.bit_size == alu2->dest.dest.ssa.bit_size);
 174    unsigned alu1_components = alu1->dest.dest.ssa.num_components;
 175    unsigned alu2_components = alu2->dest.dest.ssa.num_components;
 176    unsigned total_components = alu1_components + alu2_components;
 177
 178    if (total_components > 4)
 179       return NULL;
 180
 181    if (nir->options->vectorize_vec2_16bit &&
 182        (total_components > 2 || alu1->dest.dest.ssa.bit_size != 16))
 183       return NULL;
 184
 185    if (filter && !filter(&alu1->instr, &alu2->instr, data))
 186       return NULL;
 187
 188    nir_builder b;
 189    nir_builder_init(&b, nir_cf_node_get_function(&instr1->block->cf_node));
 190    b.cursor = nir_after_instr(instr1);
 191
 192    nir_alu_instr *new_alu = nir_alu_instr_create(b.shader, alu1->op);
 193    nir_ssa_dest_init(&new_alu->instr, &new_alu->dest.dest,
 194                      total_components, alu1->dest.dest.ssa.bit_size, NULL);
 195    new_alu->dest.write_mask = (1 << total_components) - 1;
 196
 197    for (unsigned i = 0; i < nir_op_infos[alu1->op].num_inputs; i++) {
 198       /* handle constant merging case */
 199       if (alu1->src[i].src.ssa != alu2->src[i].src.ssa) {
 200          nir_const_value *c1 = nir_src_as_const_value(alu1->src[i].src);
 201          nir_const_value *c2 = nir_src_as_const_value(alu2->src[i].src);
 202          assert(c1 && c2);
 203          nir_const_value value[4];
 204          unsigned bit_size = alu1->src[i].src.ssa->bit_size;
 205
 206          for (unsigned j = 0; j < total_components; j++) {
 207             value[j].u64 = j < alu1_components ?
 208                               c1[alu1->src[i].swizzle[j]].u64 :
 209                               c2[alu2->src[i].swizzle[j - alu1_components]].u64;
 210          }
 211          nir_ssa_def *def = nir_build_imm(&b, total_components, bit_size, value);
 212
 213          new_alu->src[i].src = nir_src_for_ssa(def);
 214          for (unsigned j = 0; j < total_components; j++)
 215             new_alu->src[i].swizzle[j] = j;
 216          continue;
 217       }
 218
 219       new_alu->src[i].src = alu1->src[i].src;
 220
 221       for (unsigned j = 0; j < alu1_components; j++)
 222          new_alu->src[i].swizzle[j] = alu1->src[i].swizzle[j];
 223
 224       for (unsigned j = 0; j < alu2_components; j++) {
 225          new_alu->src[i].swizzle[j + alu1_components] =
 226             alu2->src[i].swizzle[j];
 227       }
 228    }
 229
 230    nir_builder_instr_insert(&b, &new_alu->instr);
 231
 232    unsigned swiz[4] = {0, 1, 2, 3};
 233    nir_ssa_def *new_alu1 = nir_swizzle(&b, &new_alu->dest.dest.ssa, swiz,
 234                                        alu1_components);
 235
 236    for (unsigned i = 0; i < alu2_components; i++)
 237       swiz[i] += alu1_components;
 238    nir_ssa_def *new_alu2 = nir_swizzle(&b, &new_alu->dest.dest.ssa, swiz,
 239                                        alu2_components);
 240
 241    nir_foreach_use_safe(src, &alu1->dest.dest.ssa) {
 242       if (src->parent_instr->type == nir_instr_type_alu) {
 243          /* For ALU instructions, rewrite the source directly to avoid a
 244           * round-trip through copy propagation.
 245           */
 246
 247          nir_instr_rewrite_src(src->parent_instr, src,
 248                                nir_src_for_ssa(&new_alu->dest.dest.ssa));
 249       } else {
 250          nir_instr_rewrite_src(src->parent_instr, src,
 251                                nir_src_for_ssa(new_alu1));
 252       }
 253    }
 254
 255    nir_foreach_if_use_safe(src, &alu1->dest.dest.ssa) {
 256       nir_if_rewrite_condition(src->parent_if, nir_src_for_ssa(new_alu1));
 257    }
 258
 259    assert(list_is_empty(&alu1->dest.dest.ssa.uses));
 260    assert(list_is_empty(&alu1->dest.dest.ssa.if_uses));
 261
 262    nir_foreach_use_safe(src, &alu2->dest.dest.ssa) {
 263       if (src->parent_instr->type == nir_instr_type_alu) {
 264          /* For ALU instructions, rewrite the source directly to avoid a
 265           * round-trip through copy propagation.
 266           */
 267
 268          nir_alu_instr *use = nir_instr_as_alu(src->parent_instr);
 269
 270          unsigned src_index = 5;
 271          for (unsigned i = 0; i < nir_op_infos[use->op].num_inputs; i++) {
 272             if (&use->src[i].src == src) {
 273                src_index = i;
 274                break;
 275             }
 276          }
 277          assert(src_index != 5);
 278
 279          nir_instr_rewrite_src(src->parent_instr, src,
 280                                nir_src_for_ssa(&new_alu->dest.dest.ssa));
 281
 282          for (unsigned i = 0;
 283               i < nir_ssa_alu_instr_src_components(use, src_index); i++) {
 284             use->src[src_index].swizzle[i] += alu1_components;
 285          }
 286       } else {
 287          nir_instr_rewrite_src(src->parent_instr, src,
 288                                nir_src_for_ssa(new_alu2));
 289       }
 290    }
 291
 292    nir_foreach_if_use_safe(src, &alu2->dest.dest.ssa) {
 293       nir_if_rewrite_condition(src->parent_if, nir_src_for_ssa(new_alu2));
 294    }
 295
 296    assert(list_is_empty(&alu2->dest.dest.ssa.uses));
 297    assert(list_is_empty(&alu2->dest.dest.ssa.if_uses));
 298
 299    nir_instr_remove(instr1);
 300    nir_instr_remove(instr2);
 301
 302    return &new_alu->instr;
 303 }
 304
 305 /*
 306  * Use an array to represent a stack of instructions that are equivalent.
 307  *
 308  * We push and pop instructions off the stack in dominance order. The first
 309  * element dominates the second element which dominates the third, etc. When
 310  * trying to add to the stack, first we try and combine the instruction with
 311  * each of the instructions on the stack and, if successful, replace the
 312  * instruction on the stack with the newly-combined instruction.
 313  */
 314
 315 static struct util_dynarray *
 316 vec_instr_stack_create(void *mem_ctx)
 317 {
 318    struct util_dynarray *stack = ralloc(mem_ctx, struct util_dynarray);
 319    util_dynarray_init(stack, mem_ctx);
 320    return stack;
 321 }
 322
 323 /* returns true if we were able to successfully replace the instruction */
 324
 325 static bool
 326 vec_instr_stack_push(struct nir_shader *nir, struct util_dynarray *stack,
 327                      nir_instr *instr,
 328                      nir_opt_vectorize_cb filter, void *data)
 329 {
 330    /* Walk the stack from child to parent to make live ranges shorter by
 331     * matching the closest thing we can
 332     */
 333    util_dynarray_foreach_reverse(stack, nir_instr *, stack_instr) {
 334       nir_instr *new_instr = instr_try_combine(nir, *stack_instr, instr,
 335                                                filter, data);
 336       if (new_instr) {
 337          *stack_instr = new_instr;
 338          return true;
 339       }
 340    }
 341
 342    util_dynarray_append(stack, nir_instr *, instr);
 343    return false;
 344 }
 345
 346 static void
 347 vec_instr_stack_pop(struct util_dynarray *stack, nir_instr *instr)
 348 {
 349    ASSERTED nir_instr *last = util_dynarray_pop(stack, nir_instr *);
 350    assert(last == instr);
 351 }
 352
 353 static bool
 354 cmp_func(const void *data1, const void *data2)
 355 {
 356    const struct util_dynarray *arr1 = data1;
 357    const struct util_dynarray *arr2 = data2;
 358
 359    const nir_instr *instr1 = *(nir_instr **)util_dynarray_begin(arr1);
 360    const nir_instr *instr2 = *(nir_instr **)util_dynarray_begin(arr2);
 361
 362    return instrs_equal(instr1, instr2);
 363 }
 364
 365 static uint32_t
 366 hash_stack(const void *data)
 367 {
 368    const struct util_dynarray *stack = data;
 369    const nir_instr *first = *(nir_instr **)util_dynarray_begin(stack);
 370    return hash_instr(first);
 371 }
 372
 373 static struct set *
 374 vec_instr_set_create(void)
 375 {
 376    return _mesa_set_create(NULL, hash_stack, cmp_func);
 377 }
 378
 379 static void
 380 vec_instr_set_destroy(struct set *instr_set)
 381 {
 382    _mesa_set_destroy(instr_set, NULL);
 383 }
 384
 385 static bool
 386 vec_instr_set_add_or_rewrite(struct nir_shader *nir, struct set *instr_set,
 387                              nir_instr *instr,
 388                              nir_opt_vectorize_cb filter, void *data)
 389 {
 390    if (!instr_can_rewrite(instr))
 391       return false;
 392
 393    struct util_dynarray *new_stack = vec_instr_stack_create(instr_set);
 394    vec_instr_stack_push(nir, new_stack, instr, filter, data);
 395
 396    struct set_entry *entry = _mesa_set_search(instr_set, new_stack);
 397
 398    if (entry) {
 399       ralloc_free(new_stack);
 400       struct util_dynarray *stack = (struct util_dynarray *) entry->key;
 401       return vec_instr_stack_push(nir, stack, instr, filter, data);
 402    }
 403
 404    _mesa_set_add(instr_set, new_stack);
 405    return false;
 406 }
 407
 408 static void
 409 vec_instr_set_remove(struct nir_shader *nir, struct set *instr_set,
 410                      nir_instr *instr, nir_opt_vectorize_cb filter, void *data)
 411 {
 412    if (!instr_can_rewrite(instr))
 413       return;
 414
 415    /*
 416     * It's pretty unfortunate that we have to do this, but it's a side effect
 417     * of the hash set interfaces. The hash set assumes that we're only
 418     * interested in storing one equivalent element at a time, and if we try to
 419     * insert a duplicate element it will remove the original. We could hack up
 420     * the comparison function to "know" which input is an instruction we
 421     * passed in and which is an array that's part of the entry, but that
 422     * wouldn't work because we need to pass an array to _mesa_set_add() in
 423     * vec_instr_add_or_rewrite() above, and _mesa_set_add() will call our
 424     * comparison function as well.
 425     */
 426    struct util_dynarray *temp = vec_instr_stack_create(instr_set);
 427    vec_instr_stack_push(nir, temp, instr, filter, data);
 428    struct set_entry *entry = _mesa_set_search(instr_set, temp);
 429    ralloc_free(temp);
 430
 431    if (entry) {
 432       struct util_dynarray *stack = (struct util_dynarray *) entry->key;
 433
 434       if (util_dynarray_num_elements(stack, nir_instr *) > 1)
 435          vec_instr_stack_pop(stack, instr);
 436       else
 437          _mesa_set_remove(instr_set, entry);
 438    }
 439 }
 440
 441 static bool
 442 vectorize_block(struct nir_shader *nir, nir_block *block,
 443                 struct set *instr_set,
 444                 nir_opt_vectorize_cb filter, void *data)
 445 {
 446    bool progress = false;
 447
 448    nir_foreach_instr_safe(instr, block) {
 449       if (vec_instr_set_add_or_rewrite(nir, instr_set, instr, filter, data))
 450          progress = true;
 451    }
 452
 453    for (unsigned i = 0; i < block->num_dom_children; i++) {
 454       nir_block *child = block->dom_children[i];
 455       progress |= vectorize_block(nir, child, instr_set, filter, data);
 456    }
 457
 458    nir_foreach_instr_reverse(instr, block)
 459       vec_instr_set_remove(nir, instr_set, instr, filter, data);
 460
 461    return progress;
 462 }
 463
 464 static bool
 465 nir_opt_vectorize_impl(struct nir_shader *nir, nir_function_impl *impl,
 466                        nir_opt_vectorize_cb filter, void *data)
 467 {
 468    struct set *instr_set = vec_instr_set_create();
 469
 470    nir_metadata_require(impl, nir_metadata_dominance);
 471
 472    bool progress = vectorize_block(nir, nir_start_block(impl), instr_set,
 473                                    filter, data);
 474
 475    if (progress)
 476       nir_metadata_preserve(impl, nir_metadata_block_index |
 477                                   nir_metadata_dominance);
 478
 479    vec_instr_set_destroy(instr_set);
 480    return progress;
 481 }
 482
 483 bool
 484 nir_opt_vectorize(nir_shader *shader, nir_opt_vectorize_cb filter,
 485                   void *data)
 486 {
 487    bool progress = false;
 488
 489    nir_foreach_function(function, shader) {
 490       if (function->impl)
 491          progress |= nir_opt_vectorize_impl(shader, function->impl, filter, data);
 492    }
 493
 494    return progress;
 495 }