nir: Add a helper to get the live set at a cursor

Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7356>

src/compiler/nir/nir.h

@@ -5064,6 +5064,8 @@ bool nir_normalize_cubemap_coords(nir_shader *shader);
 
 void nir_live_ssa_defs_impl(nir_function_impl *impl);
 
+const BITSET_WORD *nir_get_live_ssa_defs(nir_cursor cursor, void *mem_ctx);
+
 void nir_loop_analyze_impl(nir_function_impl *impl,
                            nir_variable_mode indirect_mask);
 

src/compiler/nir/nir_liveness.c

@@ -214,6 +214,64 @@ nir_live_ssa_defs_impl(nir_function_impl *impl)
    nir_block_worklist_fini(&state.worklist);
 }
 
+/** Return the live set at a cursor
+ *
+ * Note: The bitset returned may be the live_in or live_out from the block in
+ *       which the instruction lives.  Do not ralloc_free() it directly;
+ *       instead, provide a mem_ctx and free that.
+ */
+const BITSET_WORD *
+nir_get_live_ssa_defs(nir_cursor cursor, void *mem_ctx)
+{
+   nir_block *block = nir_cursor_current_block(cursor);
+   nir_function_impl *impl = nir_cf_node_get_function(&block->cf_node);
+   assert(impl->valid_metadata & nir_metadata_live_ssa_defs);
+
+   switch (cursor.option) {
+   case nir_cursor_before_block:
+      return cursor.block->live_in;
+
+   case nir_cursor_after_block:
+      return cursor.block->live_out;
+
+   case nir_cursor_before_instr:
+      if (cursor.instr == nir_block_first_instr(cursor.instr->block))
+         return cursor.instr->block->live_in;
+      break;
+
+   case nir_cursor_after_instr:
+      if (cursor.instr == nir_block_last_instr(cursor.instr->block))
+         return cursor.instr->block->live_out;
+      break;
+   }
+
+   /* If we got here, we're an instruction cursor mid-block */
+   const unsigned bitset_words = BITSET_WORDS(impl->ssa_alloc);
+   BITSET_WORD *live = ralloc_array(mem_ctx, BITSET_WORD, bitset_words);
+   memcpy(live, block->live_out, bitset_words * sizeof(BITSET_WORD));
+
+   nir_foreach_instr_reverse(instr, block) {
+      if (cursor.option == nir_cursor_after_instr && instr == cursor.instr)
+         break;
+
+      /* If someone asked for liveness in the middle of a bunch of phis,
+       * that's an error.  Since we are going backwards and they are at the
+       * beginning, we can just blow up as soon as we see one.
+       */
+      assert(instr->type != nir_instr_type_phi);
+      if (instr->type == nir_instr_type_phi)
+         break;
+
+      nir_foreach_ssa_def(instr, set_ssa_def_dead, live);
+      nir_foreach_src(instr, set_src_live, live);
+
+      if (cursor.option == nir_cursor_before_instr && instr == cursor.instr)
+         break;
+   }
+
+   return live;
+}
+
 static bool
 src_does_not_use_def(nir_src *src, void *def)
 {