nir: add intrinsics for AMD_shader_ballot

Reviewed-by: Connor Abbott <[email protected]>

3 files changed, 31 insertions, 0 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index e179481b8ea..3aed1c29945 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1306,6 +1306,12 @@ typedef enum {
     */
    NIR_INTRINSIC_TYPE = 20,
 
+   /**
+    * The swizzle mask for the instructions
+    * SwizzleInvocationsAMD and SwizzleInvocationsMaskedAMD
+    */
+   NIR_INTRINSIC_SWIZZLE_MASK = 20,
+
    NIR_INTRINSIC_NUM_INDEX_FLAGS,
 
 } nir_intrinsic_index_flag;
@@ -1411,6 +1417,7 @@ INTRINSIC_IDX_ACCESSORS(align_mul, ALIGN_MUL, unsigned)
 INTRINSIC_IDX_ACCESSORS(align_offset, ALIGN_OFFSET, unsigned)
 INTRINSIC_IDX_ACCESSORS(desc_type, DESC_TYPE, unsigned)
 INTRINSIC_IDX_ACCESSORS(type, TYPE, nir_alu_type)
+INTRINSIC_IDX_ACCESSORS(swizzle_mask, SWIZZLE_MASK, unsigned)
 
 static inline void
 nir_intrinsic_set_align(nir_intrinsic_instr *intrin,
diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py
index 0ae33a17fe0..155a548d1b8 100644
--- a/src/compiler/nir/nir_intrinsics.py
+++ b/src/compiler/nir/nir_intrinsics.py
@@ -120,6 +120,8 @@ ALIGN_OFFSET = "NIR_INTRINSIC_ALIGN_OFFSET"
 DESC_TYPE = "NIR_INTRINSIC_DESC_TYPE"
 # The nir_alu_type of a uniform/input/output
 TYPE = "NIR_INTRINSIC_TYPE"
+# The swizzle mask for quad_swizzle_amd & masked_swizzle_amd
+SWIZZLE_MASK = "NIR_INTRINSIC_SWIZZLE_MASK"
 
 #
 # Possible flags:
@@ -267,6 +269,14 @@ intrinsic("inclusive_scan", src_comp=[0], dest_comp=0, indices=[REDUCTION_OP],
 intrinsic("exclusive_scan", src_comp=[0], dest_comp=0, indices=[REDUCTION_OP],
           flags=[CAN_ELIMINATE])
 
+# AMD shader ballot operations
+intrinsic("quad_swizzle_amd", src_comp=[0], dest_comp=0, indices=[SWIZZLE_MASK],
+          flags=[CAN_ELIMINATE])
+intrinsic("masked_swizzle_amd", src_comp=[0], dest_comp=0, indices=[SWIZZLE_MASK],
+          flags=[CAN_ELIMINATE])
+intrinsic("write_invocation_amd", src_comp=[0, 0, 1], dest_comp=0, flags=[CAN_ELIMINATE])
+intrinsic("mbcnt_amd", src_comp=[1], dest_comp=1, flags=[CAN_ELIMINATE])
+
 # Basic Geometry Shader intrinsics.
 #
 # emit_vertex implements GLSL's EmitStreamVertex() built-in.  It takes a single
diff --git a/src/compiler/nir/nir_print.c b/src/compiler/nir/nir_print.c
index 03509e1cc65..d67c6c79c17 100644
--- a/src/compiler/nir/nir_print.c
+++ b/src/compiler/nir/nir_print.c
@@ -776,6 +776,7 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
       [NIR_INTRINSIC_ALIGN_OFFSET] = "align_offset",
       [NIR_INTRINSIC_DESC_TYPE] = "desc_type",
       [NIR_INTRINSIC_TYPE] = "type",
+      [NIR_INTRINSIC_SWIZZLE_MASK] = "swizzle_mask",
    };
    for (unsigned idx = 1; idx < NIR_INTRINSIC_NUM_INDEX_FLAGS; idx++) {
       if (!info->index_map[idx])
@@ -827,6 +828,19 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
             fprintf(fp, " type=%s%u", name, size);
          else
             fprintf(fp, " type=%s", name);
+      } else if (idx == NIR_INTRINSIC_SWIZZLE_MASK) {
+         fprintf(fp, " swizzle_mask=");
+         unsigned mask = nir_intrinsic_swizzle_mask(instr);
+         if (instr->intrinsic == nir_intrinsic_quad_swizzle_amd) {
+            for (unsigned i = 0; i < 4; i++)
+               fprintf(fp, "%d", (mask >> (i * 2) & 3));
+         } else if (instr->intrinsic == nir_intrinsic_masked_swizzle_amd) {
+            fprintf(fp, "((id & %d) | %d) ^ %d", mask & 0x1F,
+                                                (mask >> 5) & 0x1F,
+                                                (mask >> 10) & 0x1F);
+         } else {
+            fprintf(fp, "%d", mask);
+         }
       } else {
          unsigned off = info->index_map[idx] - 1;
          assert(index_name[idx]);  /* forgot to update index_name table? */