nir: Add subgroup arithmetic reduction intrinsics

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

4 files changed, 34 insertions, 0 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 9522a4f4ff6..65c5343de4d 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1081,6 +1081,16 @@ typedef enum {
     */
    NIR_INTRINSIC_INTERP_MODE = 9,
 
+   /**
+    * A binary nir_op to use when performing a reduction or scan operation
+    */
+   NIR_INTRINSIC_REDUCTION_OP = 10,
+
+   /**
+    * Cluster size for reduction operations
+    */
+   NIR_INTRINSIC_CLUSTER_SIZE = 11,
+
    NIR_INTRINSIC_NUM_INDEX_FLAGS,
 
 } nir_intrinsic_index_flag;
@@ -1149,6 +1159,8 @@ INTRINSIC_IDX_ACCESSORS(desc_set, DESC_SET, unsigned)
 INTRINSIC_IDX_ACCESSORS(binding, BINDING, unsigned)
 INTRINSIC_IDX_ACCESSORS(component, COMPONENT, unsigned)
 INTRINSIC_IDX_ACCESSORS(interp_mode, INTERP_MODE, unsigned)
+INTRINSIC_IDX_ACCESSORS(reduction_op, REDUCTION_OP, unsigned)
+INTRINSIC_IDX_ACCESSORS(cluster_size, CLUSTER_SIZE, unsigned)
 
 /**
  * \group texture information
diff --git a/src/compiler/nir/nir_intrinsics.h b/src/compiler/nir/nir_intrinsics.h
index 4381541c1fe..7b737559d5a 100644
--- a/src/compiler/nir/nir_intrinsics.h
+++ b/src/compiler/nir/nir_intrinsics.h
@@ -177,6 +177,13 @@ INTRINSIC(quad_swap_vertical, 1, ARR(0), true, 0, 0,
 INTRINSIC(quad_swap_diagonal, 1, ARR(0), true, 0, 0,
           0, xx, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
 
+INTRINSIC(reduce, 1, ARR(0), true, 0, 0,
+          2, REDUCTION_OP, CLUSTER_SIZE, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+INTRINSIC(inclusive_scan, 1, ARR(0), true, 0, 0,
+          1, REDUCTION_OP, xx, xx, NIR_INTRINSIC_CAN_ELIMINATE)
+INTRINSIC(exclusive_scan, 1, ARR(0), true, 0, 0,
+          1, REDUCTION_OP, 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 1ada6354af3..f18ad00c370 100644
--- a/src/compiler/nir/nir_lower_subgroups.c
+++ b/src/compiler/nir/nir_lower_subgroups.c
@@ -104,6 +104,9 @@ lower_subgroup_op_to_scalar(nir_builder *b, nir_intrinsic_instr *intrin)
          nir_src_copy(&chan_intrin->src[1], &intrin->src[1], chan_intrin);
       }
 
+      chan_intrin->const_index[0] = intrin->const_index[0];
+      chan_intrin->const_index[1] = intrin->const_index[1];
+
       nir_builder_instr_insert(b, &chan_intrin->instr);
 
       reads[i] = &chan_intrin->dest.ssa;
@@ -372,6 +375,13 @@ lower_subgroups_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
          return lower_subgroup_op_to_scalar(b, intrin);
       break;
 
+   case nir_intrinsic_reduce:
+   case nir_intrinsic_inclusive_scan:
+   case nir_intrinsic_exclusive_scan:
+      if (options->lower_to_scalar && intrin->num_components > 1)
+         return lower_subgroup_op_to_scalar(b, intrin);
+      break;
+
    default:
       break;
    }
diff --git a/src/compiler/nir/nir_print.c b/src/compiler/nir/nir_print.c
index fcc8025346e..7888dbd3384 100644
--- a/src/compiler/nir/nir_print.c
+++ b/src/compiler/nir/nir_print.c
@@ -619,6 +619,8 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
       [NIR_INTRINSIC_BINDING] = "binding",
       [NIR_INTRINSIC_COMPONENT] = "component",
       [NIR_INTRINSIC_INTERP_MODE] = "interp_mode",
+      [NIR_INTRINSIC_REDUCTION_OP] = "reduction_op",
+      [NIR_INTRINSIC_CLUSTER_SIZE] = "cluster_size",
    };
    for (unsigned idx = 1; idx < NIR_INTRINSIC_NUM_INDEX_FLAGS; idx++) {
       if (!info->index_map[idx])
@@ -631,6 +633,9 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
          for (unsigned i = 0; i < 4; i++)
             if ((wrmask >> i) & 1)
                fprintf(fp, "%c", "xyzw"[i]);
+      } else if (idx == NIR_INTRINSIC_REDUCTION_OP) {
+         nir_op reduction_op = nir_intrinsic_reduction_op(instr);
+         fprintf(fp, " reduction_op=%s", nir_op_infos[reduction_op].name);
       } else {
          unsigned off = info->index_map[idx] - 1;
          assert(index_name[idx]);  /* forgot to update index_name table? */