From b5792c1a3421572c86667b6fbb9878531caed4de Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcin=20=C5=9Alusarz?= <marcin.slusarz@intel.com>
Date: Fri, 31 Mar 2023 11:51:16 +0200
Subject: [PATCH] nir: extract try_lower_id_to_index_1d
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/22334>
---
 src/compiler/nir/nir_lower_system_values.c | 36 +++++++++++++---------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/compiler/nir/nir_lower_system_values.c b/src/compiler/nir/nir_lower_system_values.c
index d91366c82d0..5253efd0098 100644
--- a/src/compiler/nir/nir_lower_system_values.c
+++ b/src/compiler/nir/nir_lower_system_values.c
@@ -423,6 +423,24 @@ lower_compute_system_value_filter(const nir_instr *instr, const void *_state)
    return instr->type == nir_instr_type_intrinsic;
 }
 
+static nir_ssa_def *
+try_lower_id_to_index_1d(nir_builder *b, nir_ssa_def *index, const uint16_t *size)
+{
+   /* size_x = 1, size_y = 1, therefore Z = local index */
+   if (size[0] == 1 && size[1] == 1)
+      return nir_vec3(b, nir_imm_int(b, 0), nir_imm_int(b, 0), index);
+
+   /* size_x = 1, size_z = 1, therefore Y = local index */
+   if (size[0] == 1 && size[2] == 1)
+      return nir_vec3(b, nir_imm_int(b, 0), index, nir_imm_int(b, 0));
+
+   /* size_y = 1, size_z = 1, therefore X = local index */
+   if (size[1] == 1 && size[2] == 1)
+      return nir_vec3(b, index, nir_imm_int(b, 0), nir_imm_int(b, 0));
+
+   return NULL;
+}
+
 static nir_ssa_def *
 lower_compute_system_value_instr(nir_builder *b,
                                  nir_instr *instr, void *_state)
@@ -454,20 +472,10 @@ lower_compute_system_value_instr(nir_builder *b,
              * this way we don't leave behind extra ALU instrs.
              */
 
-            /* size_x = 1, size_y = 1, therefore Z = local index */
-            if (b->shader->info.workgroup_size[0] == 1 &&
-                b->shader->info.workgroup_size[1] == 1)
-               return nir_vec3(b, nir_imm_int(b, 0), nir_imm_int(b, 0), local_index);
-
-            /* size_x = 1, size_z = 1, therefore Y = local index */
-            if (b->shader->info.workgroup_size[0] == 1 &&
-                b->shader->info.workgroup_size[2] == 1)
-               return nir_vec3(b, nir_imm_int(b, 0), local_index, nir_imm_int(b, 0));
-
-            /* size_y = 1, size_z = 1, therefore X = local index */
-            if (b->shader->info.workgroup_size[1] == 1 &&
-                b->shader->info.workgroup_size[2] == 1)
-               return nir_vec3(b, local_index, nir_imm_int(b, 0), nir_imm_int(b, 0));
+            nir_ssa_def *val = try_lower_id_to_index_1d(b, local_index,
+                  b->shader->info.workgroup_size);
+            if (val)
+               return val;
          }
 
          nir_ssa_def *local_size = nir_load_workgroup_size(b);