nir: Add quad operations and lowering

Reviewed-by: Lionel Landwerlin <[email protected]>
Reviewed-by: Iago Toral Quiroga <[email protected]>

3 files changed, 44 insertions, 0 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 4480fff58a1..9522a4f4ff6 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2545,6 +2545,7 @@ typedef struct nir_lower_subgroups_options {
    bool lower_vote_trivial:1;
    bool lower_subgroup_masks:1;
    bool lower_shuffle:1;
+   bool lower_quad:1;
 } nir_lower_subgroups_options;
 
 bool nir_lower_subgroups(nir_shader *shader,
diff --git a/src/compiler/nir/nir_intrinsics.h b/src/compiler/nir/nir_intrinsics.h
index 8c155a609cb..4381541c1fe 100644
--- a/src/compiler/nir/nir_intrinsics.h
+++ b/src/compiler/nir/nir_intrinsics.h
@@ -167,6 +167,16 @@ INTRINSIC(shuffle_up, 2, ARR(0, 1), true, 0, 0,
 INTRINSIC(shuffle_down, 2, ARR(0, 1), true, 0, 0,
           0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
 
+/** Quad operations from SPIR-V. */
+INTRINSIC(quad_broadcast, 2, ARR(0, 1), true, 0, 0,
+          0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+INTRINSIC(quad_swap_horizontal, 1, ARR(0), true, 0, 0,
+          0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+INTRINSIC(quad_swap_vertical, 1, ARR(0), true, 0, 0,
+          0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+INTRINSIC(quad_swap_diagonal, 1, ARR(0), true, 0, 0,
+          0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+
 /**
  * Basic Geometry Shader intrinsics.
  *
diff --git a/src/compiler/nir/nir_lower_subgroups.c b/src/compiler/nir/nir_lower_subgroups.c
index 03c5a1c7b36..1ada6354af3 100644
--- a/src/compiler/nir/nir_lower_subgroups.c
+++ b/src/compiler/nir/nir_lower_subgroups.c
@@ -156,6 +156,29 @@ lower_shuffle(nir_builder *b, nir_intrinsic_instr *intrin,
       assert(intrin->src[1].is_ssa);
       index = nir_iadd(b, index, intrin->src[1].ssa);
       break;
+   case nir_intrinsic_quad_broadcast:
+      assert(intrin->src[1].is_ssa);
+      index = nir_ior(b, nir_iand(b, index, nir_imm_int(b, ~0x3)),
+                         intrin->src[1].ssa);
+      break;
+   case nir_intrinsic_quad_swap_horizontal:
+      /* For Quad operations, subgroups are divided into quads where
+       * (invocation % 4) is the index to a square arranged as follows:
+       *
+       *    +---+---+
+       *    | 0 | 1 |
+       *    +---+---+
+       *    | 2 | 3 |
+       *    +---+---+
+       */
+      index = nir_ixor(b, index, nir_imm_int(b, 0x1));
+      break;
+   case nir_intrinsic_quad_swap_vertical:
+      index = nir_ixor(b, index, nir_imm_int(b, 0x2));
+      break;
+   case nir_intrinsic_quad_swap_diagonal:
+      index = nir_ixor(b, index, nir_imm_int(b, 0x3));
+      break;
    default:
       unreachable("Invalid intrinsic");
    }
@@ -339,6 +362,16 @@ lower_subgroups_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
          return lower_subgroup_op_to_scalar(b, intrin);
       break;
 
+   case nir_intrinsic_quad_broadcast:
+   case nir_intrinsic_quad_swap_horizontal:
+   case nir_intrinsic_quad_swap_vertical:
+   case nir_intrinsic_quad_swap_diagonal:
+      if (options->lower_quad)
+         return lower_shuffle(b, intrin, options->lower_to_scalar);
+      else if (options->lower_to_scalar && intrin->num_components > 1)
+         return lower_subgroup_op_to_scalar(b, intrin);
+      break;
+
    default:
       break;
    }