[Mesa-dev] [PATCH 1/2] nir/cse: invalidate SSBO loads in presence of ssbo writes or memory barriers

[Iago Toral Quiroga]

In the next commit we are going to allow CSE of SSBO loads, so we need
to make sure that we consider how SSBO stores/atomics invalidate previous
loads.
---
 src/glsl/nir/nir_opt_cse.c | 142 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 142 insertions(+)

diff --git a/src/glsl/nir/nir_opt_cse.c b/src/glsl/nir/nir_opt_cse.c
index 93a6635..4a57eca 100644
--- a/src/glsl/nir/nir_opt_cse.c
+++ b/src/glsl/nir/nir_opt_cse.c
@@ -32,6 +32,143 @@
  * Implements common subexpression elimination
  */
 
+static bool
+is_ssbo_atomic_intrinsic_opcode(unsigned opcode)
+{
+   switch (opcode) {
+   case nir_intrinsic_ssbo_atomic_add:
+   case nir_intrinsic_ssbo_atomic_imin:
+   case nir_intrinsic_ssbo_atomic_umin:
+   case nir_intrinsic_ssbo_atomic_imax:
+   case nir_intrinsic_ssbo_atomic_umax:
+   case nir_intrinsic_ssbo_atomic_and:
+   case nir_intrinsic_ssbo_atomic_or:
+   case nir_intrinsic_ssbo_atomic_xor:
+   case nir_intrinsic_ssbo_atomic_exchange:
+   case nir_intrinsic_ssbo_atomic_comp_swap:
+      return true;
+   default:
+      return false;
+   }
+}
+
+static bool
+is_indirect_ssbo_write(nir_intrinsic_instr *instr)
+{
+   return instr->intrinsic == nir_intrinsic_store_ssbo_indirect ||
+          is_ssbo_atomic_intrinsic_opcode(instr->intrinsic);
+}
+
+static bool
+is_ssbo_write(nir_instr *instr)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intrinsic = nir_instr_as_intrinsic(instr);
+   if (is_ssbo_atomic_intrinsic_opcode(intrinsic->intrinsic))
+      return true;
+
+   return (intrinsic->intrinsic == nir_intrinsic_store_ssbo ||
+           intrinsic->intrinsic == nir_intrinsic_store_ssbo_indirect);
+}
+
+static nir_src *
+get_indirect_ssbo_store_offset(nir_intrinsic_instr *instr)
+{
+   assert(is_indirect_ssbo_write(instr));
+
+   if (instr->intrinsic == nir_intrinsic_store_ssbo_indirect)
+      return &instr->src[2];
+
+   /* atomic */
+   return &instr->src[1];
+}
+
+static bool
+is_memory_barrier(nir_instr *instr)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intrinsic = nir_instr_as_intrinsic(instr);
+   return intrinsic->intrinsic == nir_intrinsic_memory_barrier;
+}
+
+/*
+ * Removes any SSBO load from the set that get invalidated by this
+ * SSBO store/atomic instruction.
+ */
+static void
+invalidate_ssbo_loads(nir_intrinsic_instr *store, struct set *instr_set)
+{
+   nir_src *store_block = &store->src[1];
+   for (struct set_entry *entry = _mesa_set_next_entry(instr_set, NULL);
+        entry != NULL; entry = _mesa_set_next_entry(instr_set, entry)) {
+
+      nir_instr *cached = (nir_instr *) entry->key;
+      if (cached->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *load = nir_instr_as_intrinsic(cached);
+      if (load->intrinsic != nir_intrinsic_load_ssbo &&
+          load->intrinsic != nir_intrinsic_load_ssbo_indirect)
+         continue;
+
+      /* if blocks don't match then the load is still valid */
+      nir_src *load_block = &load->src[0];
+      if ((load_block->ssa->parent_instr->type ==
+           store_block->ssa->parent_instr->type) &&
+          !nir_srcs_equal(*store_block, *load_block))
+         continue;
+
+      /* indirect store / direct load (or viceversa): assume they can conflict */
+      if ((is_indirect_ssbo_write(store) &&
+           load->intrinsic == nir_intrinsic_load_ssbo) ||
+          (store->intrinsic == nir_intrinsic_store_ssbo &&
+           load->intrinsic == nir_intrinsic_load_ssbo_indirect)) {
+         nir_instr_set_remove(instr_set, (nir_instr *)entry->key);
+      }
+      /* indirect store and load: check if src offsets match */
+      else if (is_indirect_ssbo_write(store) &&
+               load->intrinsic == nir_intrinsic_load_ssbo_indirect) {
+         nir_src *store_offset = get_indirect_ssbo_store_offset(store);
+         nir_src *load_offset = &load->src[1];
+         if (nir_srcs_equal(*store_offset, *load_offset)) {
+            nir_instr_set_remove(instr_set, (nir_instr *)entry->key);
+         }
+      }
+      /* direct store and load: check if const_index offsets match */
+      else if (store->intrinsic == nir_intrinsic_store_ssbo &&
+               load->intrinsic == nir_intrinsic_load_ssbo) {
+         unsigned store_offset = store->const_index[0];
+         unsigned load_offset = load->const_index[0];
+         if (store_offset == load_offset) {
+            nir_instr_set_remove(instr_set, (nir_instr *)entry->key);
+         }
+      }
+   }
+}
+
+static void
+invalidate_all_ssbo_loads(struct set *instr_set)
+{
+   for (struct set_entry *entry = _mesa_set_next_entry(instr_set, NULL);
+        entry != NULL;
+        entry = _mesa_set_next_entry(instr_set, entry)) {
+      nir_instr *cached = (nir_instr *) entry->key;
+      if (cached->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *intrinsic = nir_instr_as_intrinsic(cached);
+      if (intrinsic->intrinsic != nir_intrinsic_load_ssbo &&
+          intrinsic->intrinsic != nir_intrinsic_load_ssbo_indirect)
+         continue;
+
+      nir_instr_set_remove(instr_set, (nir_instr *)entry->key);
+   }
+}
+
 /*
  * Visits and CSE's the given block and all its descendants in the dominance
  * tree recursively. Note that the instr_set is guaranteed to only ever
@@ -48,6 +185,11 @@ cse_block(nir_block *block, struct set *instr_set)
          progress = true;
          nir_instr_remove(instr);
       }
+
+      if (is_ssbo_write(instr))
+         invalidate_ssbo_loads(nir_instr_as_intrinsic(instr), instr_set);
+      else if (is_memory_barrier(instr))
+         invalidate_all_ssbo_loads(instr_set);
    }
 
    for (unsigned i = 0; i < block->num_dom_children; i++) {
-- 
1.9.1