Mesa (master): nir/opt_uniform_atomics: recognize more complicated invocation comparisons

[GitLab Mirror]

Module: Mesa
Branch: master
Commit: eb70c52abef44f183f2a768fb3a7e5b55fe828e8
URL:    http://cgit.freedesktop.org/mesa/mesa/commit/?id=eb70c52abef44f183f2a768fb3a7e5b55fe828e8

Author: Rhys Perry <pendingchaos02 at gmail.com>
Date:   Thu Jan 21 17:27:13 2021 +0000

nir/opt_uniform_atomics: recognize more complicated invocation comparisons

For example, gl_LocalInvocationID.x + gl_LocalInvocationID.y * 8.

fossil-db (GFX10.3):
Totals from 8 (0.01% of 139391) affected shaders:
CodeSize: 15224 -> 14800 (-2.79%)
Instrs: 2880 -> 2798 (-2.85%)
Cycles: 44556 -> 44204 (-0.79%)
VMEM: 407 -> 473 (+16.22%); split: +17.69%, -1.47%
Copies: 491 -> 483 (-1.63%)
Branches: 200 -> 192 (-4.00%)

Signed-off-by: Rhys Perry <pendingchaos02 at gmail.com>
Reviewed-by: Daniel Schürmann <daniel at schuermann.dev>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/8654>

---

 src/compiler/nir/nir_opt_uniform_atomics.c | 73 ++++++++++++++++++++----------
 1 file changed, 48 insertions(+), 25 deletions(-)

diff --git a/src/compiler/nir/nir_opt_uniform_atomics.c b/src/compiler/nir/nir_opt_uniform_atomics.c
index ce5090b0d83..1b7d6708a5e 100644
--- a/src/compiler/nir/nir_opt_uniform_atomics.c
+++ b/src/compiler/nir/nir_opt_uniform_atomics.c
@@ -80,6 +80,48 @@ parse_atomic_op(nir_intrinsic_op op, unsigned *offset_src, unsigned *data_src)
    }
 }
 
+static unsigned
+get_dim(nir_ssa_scalar scalar)
+{
+   if (!scalar.def->divergent)
+      return 0;
+
+   if (scalar.def->parent_instr->type == nir_instr_type_intrinsic) {
+      nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(scalar.def->parent_instr);
+      if (intrin->intrinsic == nir_intrinsic_load_subgroup_invocation)
+         return 0x8;
+      else if (intrin->intrinsic == nir_intrinsic_load_local_invocation_index)
+         return 0x7;
+      else if (intrin->intrinsic == nir_intrinsic_load_local_invocation_id)
+         return 1 << scalar.comp;
+      else if (intrin->intrinsic == nir_intrinsic_load_global_invocation_index)
+         return 0x7;
+      else if (intrin->intrinsic == nir_intrinsic_load_global_invocation_id)
+         return 1 << scalar.comp;
+   } else if (nir_ssa_scalar_is_alu(scalar)) {
+      if (nir_ssa_scalar_alu_op(scalar) == nir_op_iadd ||
+          nir_ssa_scalar_alu_op(scalar) == nir_op_imul) {
+         nir_ssa_scalar src0 = nir_ssa_scalar_chase_alu_src(scalar, 0);
+         nir_ssa_scalar src1 = nir_ssa_scalar_chase_alu_src(scalar, 1);
+
+         unsigned src0_dim = get_dim(src0);
+         if (!src0_dim && src0.def->divergent)
+            return 0;
+         unsigned src1_dim = get_dim(src1);
+         if (!src1_dim && src1.def->divergent)
+            return 0;
+
+         return src0_dim | src1_dim;
+      } else if (nir_ssa_scalar_alu_op(scalar) == nir_op_ishl) {
+         nir_ssa_scalar src0 = nir_ssa_scalar_chase_alu_src(scalar, 0);
+         nir_ssa_scalar src1 = nir_ssa_scalar_chase_alu_src(scalar, 1);
+         return src1.def->divergent ? 0 : get_dim(src0);
+      }
+   }
+
+   return 0;
+}
+
 /* Returns a bitmask of invocation indices that are compared against a subgroup
  * uniform value.
  */
@@ -93,36 +135,17 @@ match_invocation_comparison(nir_ssa_scalar scalar)
       return match_invocation_comparison(nir_ssa_scalar_chase_alu_src(scalar, 0)) |
              match_invocation_comparison(nir_ssa_scalar_chase_alu_src(scalar, 1));
    } else if (nir_ssa_scalar_alu_op(scalar) == nir_op_ieq) {
-      unsigned dims = 0;
-      for (unsigned i = 0; i < 2; i++) {
-         nir_ssa_scalar src = nir_ssa_scalar_chase_alu_src(scalar, i);
-         if (src.def->parent_instr->type != nir_instr_type_intrinsic)
-            continue;
-         if (nir_ssa_scalar_chase_alu_src(scalar, !i).def->divergent)
-            continue;
-
-         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(src.def->parent_instr);
-         if (intrin->intrinsic == nir_intrinsic_load_subgroup_invocation)
-            dims = 0x8;
-         else if (intrin->intrinsic == nir_intrinsic_load_local_invocation_index)
-            dims = 0x7;
-         else if (intrin->intrinsic == nir_intrinsic_load_local_invocation_id)
-            dims = 1 << src.comp;
-         else if (intrin->intrinsic == nir_intrinsic_load_global_invocation_index)
-            dims = 0x7;
-         else if (intrin->intrinsic == nir_intrinsic_load_global_invocation_id)
-            dims = 1 << src.comp;
-      }
-
-      return dims;
+      if (!nir_ssa_scalar_chase_alu_src(scalar, 0).def->divergent)
+         return get_dim(nir_ssa_scalar_chase_alu_src(scalar, 1));
+      if (!nir_ssa_scalar_chase_alu_src(scalar, 1).def->divergent)
+         return get_dim(nir_ssa_scalar_chase_alu_src(scalar, 0));
    } else if (scalar.def->parent_instr->type == nir_instr_type_intrinsic) {
       nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(scalar.def->parent_instr);
       if (intrin->intrinsic == nir_intrinsic_elect)
          return 0x8;
-      return 0;
-   } else {
-      return 0;
    }
+
+   return 0;
 }
 
 /* Returns true if the intrinsic is already conditional so that at most one