nir: use a single set during CSE

Use a single set and ensure dominance by checking after a equivalent
instruction is found.

Besides removing the need to copy a set, this also lets us resize the set
at the start of the pass in the next commit.

ministat (CSE only):
Difference at 95.0% confidence
	-984.956 +/- 28.8559
	-6.90075% +/- 0.190231%
	(Student's t, pooled s = 26.9052)

ministat (entire run):
Difference at 95.0% confidence
	-1246.1 +/- 257.253
	-0.998972% +/- 0.205094%
	(Student's t, pooled s = 239.863)

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

4 files changed, 32 insertions, 36 deletions

diff --git a/src/compiler/nir/nir_instr_set.c b/src/compiler/nir/nir_instr_set.c
index 5a3a0f0cada..86f73d64b79 100644
--- a/src/compiler/nir/nir_instr_set.c
+++ b/src/compiler/nir/nir_instr_set.c
@@ -802,14 +802,20 @@ nir_instr_set_destroy(struct set *instr_set)
 }
 
 bool
-nir_instr_set_add_or_rewrite(struct set *instr_set, nir_instr *instr)
+nir_instr_set_add_or_rewrite(struct set *instr_set, nir_instr *instr,
+                             bool (*cond_function) (const nir_instr *a,
+                                                    const nir_instr *b))
 {
    if (!instr_can_rewrite(instr))
       return false;
 
    struct set_entry *e = _mesa_set_search_or_add(instr_set, instr, NULL);
    nir_instr *match = (nir_instr *) e->key;
-   if (match != instr) {
+   if (match == instr)
+      return false;
+
+   if (!cond_function || cond_function(match, instr)) {
+      /* rewrite instruction if condition is matched */
       nir_ssa_def *def = nir_instr_get_dest_ssa_def(instr);
       nir_ssa_def *new_def = nir_instr_get_dest_ssa_def(match);
 
@@ -822,10 +828,15 @@ nir_instr_set_add_or_rewrite(struct set *instr_set, nir_instr *instr)
          nir_instr_as_alu(match)->exact = true;
 
       nir_ssa_def_rewrite_uses(def, new_def);
+
+      nir_instr_remove(instr);
+
       return true;
+   } else {
+      /* otherwise, replace hashed instruction */
+      e->key = instr;
+      return false;
    }
-
-   return false;
 }
 
 void
diff --git a/src/compiler/nir/nir_instr_set.h b/src/compiler/nir/nir_instr_set.h
index 26817f62e96..9a0442cb40e 100644
--- a/src/compiler/nir/nir_instr_set.h
+++ b/src/compiler/nir/nir_instr_set.h
@@ -46,12 +46,18 @@ struct set *nir_instr_set_create(void *mem_ctx);
 void nir_instr_set_destroy(struct set *instr_set);
 
 /**
- * Adds an instruction to an instruction set if it doesn't exist, or if it
+ * Adds an instruction to an instruction set if it doesn't exist. If it
  * does already exist, rewrites all uses of it to point to the other
  * already-inserted instruction. Returns 'true' if the uses of the instruction
- * were rewritten.
+ * were rewritten. Otherwise, replaces the already-inserted instruction
+ * with the new one.
+ *
+ * If cond_function() is given, only rewrites uses if
+ * cond_function(old_instr, new_instr) returns true.
  */
-bool nir_instr_set_add_or_rewrite(struct set *instr_set, nir_instr *instr);
+bool nir_instr_set_add_or_rewrite(struct set *instr_set, nir_instr *instr,
+                                  bool (*cond_function)(const nir_instr *a,
+                                                        const nir_instr *b));
 
 /**
  * Removes an instruction from an instruction set, so that other instructions
diff --git a/src/compiler/nir/nir_opt_cse.c b/src/compiler/nir/nir_opt_cse.c
index e9ca6889501..80f48410ad0 100644
--- a/src/compiler/nir/nir_opt_cse.c
+++ b/src/compiler/nir/nir_opt_cse.c
@@ -32,33 +32,10 @@
  * Implements common subexpression elimination
  */
 
-/*
- * Visits and CSEs the given block and all its descendants in the dominance
- * tree recursively. Note that the instr_set is guaranteed to only ever
- * contain instructions that dominate the current block.
- */
-
 static bool
-cse_block(nir_block *block, struct set *dominance_set)
+dominates(const nir_instr *old_instr, const nir_instr *new_instr)
 {
-   bool progress = false;
-   struct set *instr_set = _mesa_set_clone(dominance_set, NULL);
-
-   nir_foreach_instr_safe(instr, block) {
-      if (nir_instr_set_add_or_rewrite(instr_set, instr)) {
-         progress = true;
-         nir_instr_remove(instr);
-      }
-   }
-
-   for (unsigned i = 0; i < block->num_dom_children; i++) {
-      nir_block *child = block->dom_children[i];
-      progress |= cse_block(child, instr_set);
-   }
-
-   _mesa_set_destroy(instr_set, NULL);
-
-   return progress;
+   return nir_block_dominates(old_instr->block, new_instr->block);
 }
 
 static bool
@@ -68,7 +45,11 @@ nir_opt_cse_impl(nir_function_impl *impl)
 
    nir_metadata_require(impl, nir_metadata_dominance);
 
-   bool progress = cse_block(nir_start_block(impl), instr_set);
+   bool progress = false;
+   nir_foreach_block(block, impl) {
+      nir_foreach_instr_safe(instr, block)
+         progress |= nir_instr_set_add_or_rewrite(instr_set, instr, dominates);
+   }
 
    if (progress) {
       nir_metadata_preserve(impl, nir_metadata_block_index |
diff --git a/src/compiler/nir/nir_opt_gcm.c b/src/compiler/nir/nir_opt_gcm.c
index 7317ded09ca..891da57fe35 100644
--- a/src/compiler/nir/nir_opt_gcm.c
+++ b/src/compiler/nir/nir_opt_gcm.c
@@ -622,10 +622,8 @@ opt_gcm_impl(nir_function_impl *impl, bool value_number)
    if (value_number) {
       struct set *gvn_set = nir_instr_set_create(NULL);
       foreach_list_typed_safe(nir_instr, instr, node, &state.instrs) {
-         if (nir_instr_set_add_or_rewrite(gvn_set, instr)) {
-            nir_instr_remove(instr);
+         if (nir_instr_set_add_or_rewrite(gvn_set, instr, NULL))
             state.progress = true;
-         }
       }
       nir_instr_set_destroy(gvn_set);
    }