src/compiler/spirv/vtn_subgroup.c

   1 /*
   2  * Copyright © 2016 Intel Corporation
   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 #include "vtn_private.h"
  25
  26 static void
  27 vtn_build_subgroup_instr(struct vtn_builder *b,
  28                          nir_intrinsic_op nir_op,
  29                          struct vtn_ssa_value *dst,
  30                          struct vtn_ssa_value *src0,
  31                          nir_ssa_def *index,
  32                          unsigned const_idx0,
  33                          unsigned const_idx1)
  34 {
  35    /* Some of the subgroup operations take an index.  SPIR-V allows this to be
  36     * any integer type.  To make things simpler for drivers, we only support
  37     * 32-bit indices.
  38     */
  39    if (index && index->bit_size != 32)
  40       index = nir_u2u32(&b->nb, index);
  41
  42    vtn_assert(dst->type == src0->type);
  43    if (!glsl_type_is_vector_or_scalar(dst->type)) {
  44       for (unsigned i = 0; i < glsl_get_length(dst->type); i++) {
  45          vtn_build_subgroup_instr(b, nir_op, dst->elems[i],
  46                                   src0->elems[i], index,
  47                                   const_idx0, const_idx1);
  48       }
  49       return;
  50    }
  51
  52    nir_intrinsic_instr *intrin =
  53       nir_intrinsic_instr_create(b->nb.shader, nir_op);
  54    nir_ssa_dest_init_for_type(&intrin->instr, &intrin->dest,
  55                               dst->type, NULL);
  56    intrin->num_components = intrin->dest.ssa.num_components;
  57
  58    intrin->src[0] = nir_src_for_ssa(src0->def);
  59    if (index)
  60       intrin->src[1] = nir_src_for_ssa(index);
  61
  62    intrin->const_index[0] = const_idx0;
  63    intrin->const_index[1] = const_idx1;
  64
  65    nir_builder_instr_insert(&b->nb, &intrin->instr);
  66
  67    dst->def = &intrin->dest.ssa;
  68 }
  69
  70 void
  71 vtn_handle_subgroup(struct vtn_builder *b, SpvOp opcode,
  72                     const uint32_t *w, unsigned count)
  73 {
  74    struct vtn_value *val = vtn_push_value(b, w[2], vtn_value_type_ssa);
  75
  76    val->ssa = vtn_create_ssa_value(b, val->type->type);
  77
  78    switch (opcode) {
  79    case SpvOpGroupNonUniformElect: {
  80       vtn_fail_if(val->type->type != glsl_bool_type(),
  81                   "OpGroupNonUniformElect must return a Bool");
  82       nir_intrinsic_instr *elect =
  83          nir_intrinsic_instr_create(b->nb.shader, nir_intrinsic_elect);
  84       nir_ssa_dest_init_for_type(&elect->instr, &elect->dest,
  85                                  val->type->type, NULL);
  86       nir_builder_instr_insert(&b->nb, &elect->instr);
  87       val->ssa->def = &elect->dest.ssa;
  88       break;
  89    }
  90
  91    case SpvOpGroupNonUniformBallot: {
  92       vtn_fail_if(val->type->type != glsl_vector_type(GLSL_TYPE_UINT, 4),
  93                   "OpGroupNonUniformBallot must return a uvec4");
  94       nir_intrinsic_instr *ballot =
  95          nir_intrinsic_instr_create(b->nb.shader, nir_intrinsic_ballot);
  96       ballot->src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[4])->def);
  97       nir_ssa_dest_init(&ballot->instr, &ballot->dest, 4, 32, NULL);
  98       ballot->num_components = 4;
  99       nir_builder_instr_insert(&b->nb, &ballot->instr);
 100       val->ssa->def = &ballot->dest.ssa;
 101       break;
 102    }
 103
 104    case SpvOpGroupNonUniformInverseBallot: {
 105       /* This one is just a BallotBitfieldExtract with subgroup invocation.
 106        * We could add a NIR intrinsic but it's easier to just lower it on the
 107        * spot.
 108        */
 109       nir_intrinsic_instr *intrin =
 110          nir_intrinsic_instr_create(b->nb.shader,
 111                                     nir_intrinsic_ballot_bitfield_extract);
 112
 113       intrin->src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[4])->def);
 114       intrin->src[1] = nir_src_for_ssa(nir_load_subgroup_invocation(&b->nb));
 115
 116       nir_ssa_dest_init_for_type(&intrin->instr, &intrin->dest,
 117                                  val->type->type, NULL);
 118       nir_builder_instr_insert(&b->nb, &intrin->instr);
 119
 120       val->ssa->def = &intrin->dest.ssa;
 121       break;
 122    }
 123
 124    case SpvOpGroupNonUniformBallotBitExtract:
 125    case SpvOpGroupNonUniformBallotBitCount:
 126    case SpvOpGroupNonUniformBallotFindLSB:
 127    case SpvOpGroupNonUniformBallotFindMSB: {
 128       nir_ssa_def *src0, *src1 = NULL;
 129       nir_intrinsic_op op;
 130       switch (opcode) {
 131       case SpvOpGroupNonUniformBallotBitExtract:
 132          op = nir_intrinsic_ballot_bitfield_extract;
 133          src0 = vtn_ssa_value(b, w[4])->def;
 134          src1 = vtn_ssa_value(b, w[5])->def;
 135          break;
 136       case SpvOpGroupNonUniformBallotBitCount:
 137          switch ((SpvGroupOperation)w[4]) {
 138          case SpvGroupOperationReduce:
 139             op = nir_intrinsic_ballot_bit_count_reduce;
 140             break;
 141          case SpvGroupOperationInclusiveScan:
 142             op = nir_intrinsic_ballot_bit_count_inclusive;
 143             break;
 144          case SpvGroupOperationExclusiveScan:
 145             op = nir_intrinsic_ballot_bit_count_exclusive;
 146             break;
 147          default:
 148             unreachable("Invalid group operation");
 149          }
 150          src0 = vtn_ssa_value(b, w[5])->def;
 151          break;
 152       case SpvOpGroupNonUniformBallotFindLSB:
 153          op = nir_intrinsic_ballot_find_lsb;
 154          src0 = vtn_ssa_value(b, w[4])->def;
 155          break;
 156       case SpvOpGroupNonUniformBallotFindMSB:
 157          op = nir_intrinsic_ballot_find_msb;
 158          src0 = vtn_ssa_value(b, w[4])->def;
 159          break;
 160       default:
 161          unreachable("Unhandled opcode");
 162       }
 163
 164       nir_intrinsic_instr *intrin =
 165          nir_intrinsic_instr_create(b->nb.shader, op);
 166
 167       intrin->src[0] = nir_src_for_ssa(src0);
 168       if (src1)
 169          intrin->src[1] = nir_src_for_ssa(src1);
 170
 171       nir_ssa_dest_init_for_type(&intrin->instr, &intrin->dest,
 172                                  val->type->type, NULL);
 173       nir_builder_instr_insert(&b->nb, &intrin->instr);
 174
 175       val->ssa->def = &intrin->dest.ssa;
 176       break;
 177    }
 178
 179    case SpvOpGroupNonUniformBroadcastFirst:
 180       vtn_build_subgroup_instr(b, nir_intrinsic_read_first_invocation,
 181                                val->ssa, vtn_ssa_value(b, w[4]), NULL, 0, 0);
 182       break;
 183
 184    case SpvOpGroupNonUniformBroadcast:
 185       vtn_build_subgroup_instr(b, nir_intrinsic_read_invocation,
 186                                val->ssa, vtn_ssa_value(b, w[4]),
 187                                vtn_ssa_value(b, w[5])->def, 0, 0);
 188       break;
 189
 190    case SpvOpGroupNonUniformAll:
 191    case SpvOpGroupNonUniformAny:
 192    case SpvOpGroupNonUniformAllEqual:
 193    case SpvOpSubgroupAllKHR:
 194    case SpvOpSubgroupAnyKHR:
 195    case SpvOpSubgroupAllEqualKHR: {
 196       vtn_fail_if(val->type->type != glsl_bool_type(),
 197                   "OpGroupNonUniform(All|Any|AllEqual) must return a bool");
 198       nir_intrinsic_op op;
 199       switch (opcode) {
 200       case SpvOpGroupNonUniformAll:
 201       case SpvOpSubgroupAllKHR:
 202          op = nir_intrinsic_vote_all;
 203          break;
 204       case SpvOpGroupNonUniformAny:
 205       case SpvOpSubgroupAnyKHR:
 206          op = nir_intrinsic_vote_any;
 207          break;
 208       case SpvOpGroupNonUniformAllEqual:
 209       case SpvOpSubgroupAllEqualKHR: {
 210          switch (glsl_get_base_type(val->type->type)) {
 211          case GLSL_TYPE_FLOAT:
 212          case GLSL_TYPE_DOUBLE:
 213             op = nir_intrinsic_vote_feq;
 214             break;
 215          case GLSL_TYPE_UINT:
 216          case GLSL_TYPE_INT:
 217          case GLSL_TYPE_UINT64:
 218          case GLSL_TYPE_INT64:
 219          case GLSL_TYPE_BOOL:
 220             op = nir_intrinsic_vote_ieq;
 221             break;
 222          default:
 223             unreachable("Unhandled type");
 224          }
 225          break;
 226       }
 227       default:
 228          unreachable("Unhandled opcode");
 229       }
 230
 231       nir_ssa_def *src0;
 232       if (opcode == SpvOpGroupNonUniformAll ||
 233           opcode == SpvOpGroupNonUniformAny ||
 234           opcode == SpvOpGroupNonUniformAllEqual) {
 235          src0 = vtn_ssa_value(b, w[4])->def;
 236       } else {
 237          src0 = vtn_ssa_value(b, w[3])->def;
 238       }
 239       nir_intrinsic_instr *intrin =
 240          nir_intrinsic_instr_create(b->nb.shader, op);
 241       intrin->num_components = src0->num_components;
 242       intrin->src[0] = nir_src_for_ssa(src0);
 243       nir_ssa_dest_init_for_type(&intrin->instr, &intrin->dest,
 244                                  val->type->type, NULL);
 245       nir_builder_instr_insert(&b->nb, &intrin->instr);
 246
 247       val->ssa->def = &intrin->dest.ssa;
 248       break;
 249    }
 250
 251    case SpvOpGroupNonUniformShuffle:
 252    case SpvOpGroupNonUniformShuffleXor:
 253    case SpvOpGroupNonUniformShuffleUp:
 254    case SpvOpGroupNonUniformShuffleDown: {
 255       nir_intrinsic_op op;
 256       switch (opcode) {
 257       case SpvOpGroupNonUniformShuffle:
 258          op = nir_intrinsic_shuffle;
 259          break;
 260       case SpvOpGroupNonUniformShuffleXor:
 261          op = nir_intrinsic_shuffle_xor;
 262          break;
 263       case SpvOpGroupNonUniformShuffleUp:
 264          op = nir_intrinsic_shuffle_up;
 265          break;
 266       case SpvOpGroupNonUniformShuffleDown:
 267          op = nir_intrinsic_shuffle_down;
 268          break;
 269       default:
 270          unreachable("Invalid opcode");
 271       }
 272       vtn_build_subgroup_instr(b, op, val->ssa, vtn_ssa_value(b, w[4]),
 273                                vtn_ssa_value(b, w[5])->def, 0, 0);
 274       break;
 275    }
 276
 277    case SpvOpGroupNonUniformQuadBroadcast:
 278       vtn_build_subgroup_instr(b, nir_intrinsic_quad_broadcast,
 279                                val->ssa, vtn_ssa_value(b, w[4]),
 280                                vtn_ssa_value(b, w[5])->def, 0, 0);
 281       break;
 282
 283    case SpvOpGroupNonUniformQuadSwap: {
 284       unsigned direction = vtn_constant_uint(b, w[5]);
 285       nir_intrinsic_op op;
 286       switch (direction) {
 287       case 0:
 288          op = nir_intrinsic_quad_swap_horizontal;
 289          break;
 290       case 1:
 291          op = nir_intrinsic_quad_swap_vertical;
 292          break;
 293       case 2:
 294          op = nir_intrinsic_quad_swap_diagonal;
 295          break;
 296       default:
 297          vtn_fail("Invalid constant value in OpGroupNonUniformQuadSwap");
 298       }
 299       vtn_build_subgroup_instr(b, op, val->ssa, vtn_ssa_value(b, w[4]),
 300                                NULL, 0, 0);
 301       break;
 302    }
 303
 304    case SpvOpGroupNonUniformIAdd:
 305    case SpvOpGroupNonUniformFAdd:
 306    case SpvOpGroupNonUniformIMul:
 307    case SpvOpGroupNonUniformFMul:
 308    case SpvOpGroupNonUniformSMin:
 309    case SpvOpGroupNonUniformUMin:
 310    case SpvOpGroupNonUniformFMin:
 311    case SpvOpGroupNonUniformSMax:
 312    case SpvOpGroupNonUniformUMax:
 313    case SpvOpGroupNonUniformFMax:
 314    case SpvOpGroupNonUniformBitwiseAnd:
 315    case SpvOpGroupNonUniformBitwiseOr:
 316    case SpvOpGroupNonUniformBitwiseXor:
 317    case SpvOpGroupNonUniformLogicalAnd:
 318    case SpvOpGroupNonUniformLogicalOr:
 319    case SpvOpGroupNonUniformLogicalXor: {
 320       nir_op reduction_op;
 321       switch (opcode) {
 322       case SpvOpGroupNonUniformIAdd:
 323          reduction_op = nir_op_iadd;
 324          break;
 325       case SpvOpGroupNonUniformFAdd:
 326          reduction_op = nir_op_fadd;
 327          break;
 328       case SpvOpGroupNonUniformIMul:
 329          reduction_op = nir_op_imul;
 330          break;
 331       case SpvOpGroupNonUniformFMul:
 332          reduction_op = nir_op_fmul;
 333          break;
 334       case SpvOpGroupNonUniformSMin:
 335          reduction_op = nir_op_imin;
 336          break;
 337       case SpvOpGroupNonUniformUMin:
 338          reduction_op = nir_op_umin;
 339          break;
 340       case SpvOpGroupNonUniformFMin:
 341          reduction_op = nir_op_fmin;
 342          break;
 343       case SpvOpGroupNonUniformSMax:
 344          reduction_op = nir_op_imax;
 345          break;
 346       case SpvOpGroupNonUniformUMax:
 347          reduction_op = nir_op_umax;
 348          break;
 349       case SpvOpGroupNonUniformFMax:
 350          reduction_op = nir_op_fmax;
 351          break;
 352       case SpvOpGroupNonUniformBitwiseAnd:
 353       case SpvOpGroupNonUniformLogicalAnd:
 354          reduction_op = nir_op_iand;
 355          break;
 356       case SpvOpGroupNonUniformBitwiseOr:
 357       case SpvOpGroupNonUniformLogicalOr:
 358          reduction_op = nir_op_ior;
 359          break;
 360       case SpvOpGroupNonUniformBitwiseXor:
 361       case SpvOpGroupNonUniformLogicalXor:
 362          reduction_op = nir_op_ixor;
 363          break;
 364       default:
 365          unreachable("Invalid reduction operation");
 366       }
 367
 368       nir_intrinsic_op op;
 369       unsigned cluster_size = 0;
 370       switch ((SpvGroupOperation)w[4]) {
 371       case SpvGroupOperationReduce:
 372          op = nir_intrinsic_reduce;
 373          break;
 374       case SpvGroupOperationInclusiveScan:
 375          op = nir_intrinsic_inclusive_scan;
 376          break;
 377       case SpvGroupOperationExclusiveScan:
 378          op = nir_intrinsic_exclusive_scan;
 379          break;
 380       case SpvGroupOperationClusteredReduce:
 381          op = nir_intrinsic_reduce;
 382          assert(count == 7);
 383          cluster_size = vtn_constant_uint(b, w[6]);
 384          break;
 385       default:
 386          unreachable("Invalid group operation");
 387       }
 388
 389       vtn_build_subgroup_instr(b, op, val->ssa, vtn_ssa_value(b, w[5]),
 390                                NULL, reduction_op, cluster_size);
 391       break;
 392    }
 393
 394    default:
 395       unreachable("Invalid SPIR-V opcode");
 396    }
 397 }