nir/algebraic: allow swizzle in nir_algebraic replace expression

This is to allow optimizations in nir_opt_algebraic not otherwise possible

Signed-off-by: Jonathan Marek <[email protected]>
Reviewed-by: Ian Romanick <[email protected]>
Acked-by: Matt Turner <[email protected]>

4 files changed, 22 insertions, 6 deletions

diff --git a/src/compiler/nir/nir_algebraic.py b/src/compiler/nir/nir_algebraic.py
index 9b8af1a7470..3b136f800d4 100644
--- a/src/compiler/nir/nir_algebraic.py
+++ b/src/compiler/nir/nir_algebraic.py
@@ -198,6 +198,7 @@ class Value(object):
    ${'true' if val.is_constant else 'false'},
    ${val.type() or 'nir_type_invalid' },
    ${val.cond if val.cond else 'NULL'},
+   ${val.swizzle()},
 % elif isinstance(val, Expression):
    ${'true' if val.inexact else 'false'},
    ${val.comm_expr_idx}, ${val.comm_exprs},
@@ -284,7 +285,8 @@ class Constant(Value):
 
 _var_name_re = re.compile(r"(?P<const>#)?(?P<name>\w+)"
                           r"(?:@(?P<type>int|uint|bool|float)?(?P<bits>\d+)?)?"
-                          r"(?P<cond>\([^\)]+\))?")
+                          r"(?P<cond>\([^\)]+\))?"
+                          r"(?P<swiz>\.[xyzw]+)?")
 
 class Variable(Value):
    def __init__(self, val, name, varset):
@@ -306,6 +308,7 @@ class Variable(Value):
       self.cond = m.group('cond')
       self.required_type = m.group('type')
       self._bit_size = int(m.group('bits')) if m.group('bits') else None
+      self.swiz = m.group('swiz')
 
       if self.required_type == 'bool':
          if self._bit_size is not None:
@@ -339,6 +342,12 @@ class Variable(Value):
 
       return self.index == other.index
 
+   def swizzle(self):
+      if self.swiz is not None:
+         swizzles = {'x' : 0, 'y' : 1, 'z' : 2, 'w': 3}
+         return '{' + ', '.join([str(swizzles[c]) for c in self.swiz[1:]]) + '}'
+      return '{0, 1, 2, 3}'
+
 _opcode_re = re.compile(r"(?P<inexact>~)?(?P<opcode>\w+)(?:@(?P<bits>\d+))?"
                         r"(?P<cond>\([^\)]+\))?")
 
diff --git a/src/compiler/nir/nir_opt_algebraic.py b/src/compiler/nir/nir_opt_algebraic.py
index 3ffa061c01e..00145fc80f5 100644
--- a/src/compiler/nir/nir_opt_algebraic.py
+++ b/src/compiler/nir/nir_opt_algebraic.py
@@ -50,11 +50,12 @@ e = 'e'
 # however, be used for backend-requested lowering operations as those need to
 # happen regardless of precision.
 #
-# Variable names are specified as "[#]name[@type][(cond)]" where "#" inicates
-# that the given variable will only match constants and the type indicates that
-# the given variable will only match values from ALU instructions with the
-# given output type, and (cond) specifies an additional condition function
-# (see nir_search_helpers.h).
+# Variable names are specified as "[#]name[@type][(cond)][.swiz]" where:
+# "#" indicates that the given variable will only match constants,
+# type indicates that the given variable will only match values from ALU
+#    instructions with the given output type,
+# (cond) specifies an additional condition function (see nir_search_helpers.h),
+# swiz is a swizzle applied to the variable (only in the <replace> expression)
 #
 # For constants, you have to be careful to make sure that it is the right
 # type because python is unaware of the source and destination types of the
diff --git a/src/compiler/nir/nir_search.c b/src/compiler/nir/nir_search.c
index 8e3c3529a74..c62de4ad7fe 100644
--- a/src/compiler/nir/nir_search.c
+++ b/src/compiler/nir/nir_search.c
@@ -506,6 +506,9 @@ construct_value(nir_builder *build,
                        (void *)build->shader);
       assert(!var->is_constant);
 
+      for (unsigned i = 0; i < NIR_MAX_VEC_COMPONENTS; i++)
+         val.swizzle[i] = state->variables[var->variable].swizzle[var->swizzle[i]];
+
       return val;
    }
 
diff --git a/src/compiler/nir/nir_search.h b/src/compiler/nir/nir_search.h
index 526a498cd47..ac37f1ac653 100644
--- a/src/compiler/nir/nir_search.h
+++ b/src/compiler/nir/nir_search.h
@@ -95,6 +95,9 @@ typedef struct {
     */
    bool (*cond)(nir_alu_instr *instr, unsigned src,
                 unsigned num_components, const uint8_t *swizzle);
+
+	/** Swizzle (for replace only) */
+	uint8_t swizzle[NIR_MAX_VEC_COMPONENTS];
 } nir_search_variable;
 
 typedef struct {