spirv: Implement SpvCapabilitySubgroupShuffleINTEL from SPV_INTEL_subgroups

Reviewed-by: Jason Ekstrand <jason@jlekstrand.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7448>

3 files changed, 53 insertions, 0 deletions

diff --git a/src/compiler/shader_info.h b/src/compiler/shader_info.h
index 59d95416877..7f655bb6ca5 100644
--- a/src/compiler/shader_info.h
+++ b/src/compiler/shader_info.h
@@ -93,6 +93,8 @@ struct spirv_supported_capabilities {
    bool amd_image_read_write_lod;
    bool amd_shader_explicit_vertex_parameter;
    bool amd_image_gather_bias_lod;
+
+   bool intel_subgroup_shuffle;
 };
 
 typedef struct shader_info {
diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index f5d7fe4d160..6eb7b8f258d 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -4393,6 +4393,10 @@ vtn_handle_preamble_instruction(struct vtn_builder *b, SpvOp opcode,
          spv_check_supported(float64_atomic_add, cap);
          break;
 
+      case SpvCapabilitySubgroupShuffleINTEL:
+         spv_check_supported(intel_subgroup_shuffle, cap);
+         break;
+
       default:
          vtn_fail("Unhandled capability: %s (%u)",
                   spirv_capability_to_string(cap), cap);
@@ -5327,6 +5331,10 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpGroupFMaxNonUniformAMD:
    case SpvOpGroupUMaxNonUniformAMD:
    case SpvOpGroupSMaxNonUniformAMD:
+   case SpvOpSubgroupShuffleINTEL:
+   case SpvOpSubgroupShuffleDownINTEL:
+   case SpvOpSubgroupShuffleUpINTEL:
+   case SpvOpSubgroupShuffleXorINTEL:
       vtn_handle_subgroup(b, opcode, w, count);
       break;
 
diff --git a/src/compiler/spirv/vtn_subgroup.c b/src/compiler/spirv/vtn_subgroup.c
index 8e4c3f2ba92..b0256611308 100644
--- a/src/compiler/spirv/vtn_subgroup.c
+++ b/src/compiler/spirv/vtn_subgroup.c
@@ -299,6 +299,49 @@ vtn_handle_subgroup(struct vtn_builder *b, SpvOp opcode,
       break;
    }
 
+   case SpvOpSubgroupShuffleINTEL:
+   case SpvOpSubgroupShuffleXorINTEL: {
+      nir_intrinsic_op op = opcode == SpvOpSubgroupShuffleINTEL ?
+         nir_intrinsic_shuffle : nir_intrinsic_shuffle_xor;
+      vtn_push_ssa_value(b, w[2],
+         vtn_build_subgroup_instr(b, op, vtn_ssa_value(b, w[3]),
+                                  vtn_get_nir_ssa(b, w[4]), 0, 0));
+      break;
+   }
+
+   case SpvOpSubgroupShuffleUpINTEL:
+   case SpvOpSubgroupShuffleDownINTEL: {
+      /* TODO: Move this lower on the compiler stack, where we can move the
+       * current/other data to adjacent registers to avoid doing a shuffle
+       * twice.
+       */
+
+      nir_builder *nb = &b->nb;
+      nir_ssa_def *size = nir_load_subgroup_size(nb);
+      nir_ssa_def *delta = vtn_get_nir_ssa(b, w[5]);
+
+      /* Rewrite UP in terms of DOWN.
+       *
+       *   UP(a, b, delta) == DOWN(a, b, size - delta)
+       */
+      if (opcode == SpvOpSubgroupShuffleUpINTEL)
+         delta = nir_isub(nb, size, delta);
+
+      nir_ssa_def *index = nir_iadd(nb, nir_load_subgroup_invocation(nb), delta);
+      struct vtn_ssa_value *current =
+         vtn_build_subgroup_instr(b, nir_intrinsic_shuffle, vtn_ssa_value(b, w[3]),
+                                  index, 0, 0);
+
+      struct vtn_ssa_value *next =
+         vtn_build_subgroup_instr(b, nir_intrinsic_shuffle, vtn_ssa_value(b, w[4]),
+                                  nir_isub(nb, index, size), 0, 0);
+
+      nir_ssa_def *cond = nir_ilt(nb, index, size);
+      vtn_push_nir_ssa(b, w[2], nir_bcsel(nb, cond, current->def, next->def));
+
+      break;
+   }
+
    case SpvOpGroupNonUniformQuadBroadcast:
       vtn_push_ssa_value(b, w[2],
          vtn_build_subgroup_instr(b, nir_intrinsic_quad_broadcast,