nir/search: Respect the bit_size parameter on nir_search_value

Reviewed-by: Samuel Iglesias Gonsálvez <[email protected]>
Reviewed-by: Iago Toral Quiroga <[email protected]>

2 files changed, 18 insertions, 1 deletions

diff --git a/src/compiler/nir/nir_opt_algebraic.py b/src/compiler/nir/nir_opt_algebraic.py
index fa7d7c7e96f..fc439f08363 100644
--- a/src/compiler/nir/nir_opt_algebraic.py
+++ b/src/compiler/nir/nir_opt_algebraic.py
@@ -56,7 +56,10 @@ d = 'd'
 #
 # All expression types can have a bit-size specified.  For opcodes, this
 # looks like "op@32", for variables it is "a@32" or "a@uint32" to specify a
-# type and size, and for literals, you can write "2.0@32".
+# type and size, and for literals, you can write "2.0@32".  In the search half
+# of the expression this indicates that it should only match that particular
+# bit-size.  In the replace half of the expression this indicates that the
+# constructed value should have that bit-size.
 
 optimizations = [
    (('fneg', ('fneg', a)), a),
diff --git a/src/compiler/nir/nir_search.c b/src/compiler/nir/nir_search.c
index df80c0f8e14..2c2fd9208c2 100644
--- a/src/compiler/nir/nir_search.c
+++ b/src/compiler/nir/nir_search.c
@@ -88,6 +88,11 @@ match_value(const nir_search_value *value, nir_alu_instr *instr, unsigned src,
    for (unsigned i = 0; i < num_components; ++i)
       new_swizzle[i] = instr->src[src].swizzle[swizzle[i]];
 
+   /* If the value has a specific bit size and it doesn't match, bail */
+   if (value->bit_size &&
+       nir_src_bit_size(instr->src[src].src) != value->bit_size)
+      return false;
+
    switch (value->type) {
    case nir_search_value_expression:
       if (!instr->src[src].src.is_ssa)
@@ -243,6 +248,10 @@ match_expression(const nir_search_expression *expr, nir_alu_instr *instr,
 
    assert(instr->dest.dest.is_ssa);
 
+   if (expr->value.bit_size &&
+       instr->dest.dest.ssa.bit_size != expr->value.bit_size)
+      return false;
+
    state->inexact_match = expr->inexact || state->inexact_match;
    state->has_exact_alu = instr->exact || state->has_exact_alu;
    if (state->inexact_match && state->has_exact_alu)
@@ -353,6 +362,11 @@ build_bitsize_tree(void *mem_ctx, struct match_state *state,
    }
    }
 
+   if (value->bit_size) {
+      assert(!tree->is_dest_sized || tree->dest_size == value->bit_size);
+      tree->common_size = value->bit_size;
+   }
+
    return tree;
 }