nir/lower_io: Add support for lowering deref_mode_is

The guts are still missing so it will blow up if it sees any
deref_mode_is intrinsic that it can't constant-fold from the mode.

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

src/compiler/nir/nir_intrinsics.py

@@ -216,6 +216,8 @@ intrinsic("get_ubo_size", src_comp=[-1], dest_comp=1,
 # mode checks, a pointer can only have exactly one mode at runtime.
 intrinsic("deref_mode_is", src_comp=[-1], dest_comp=1,
           indices=[MEMORY_MODES], flags=[CAN_ELIMINATE, CAN_REORDER])
+intrinsic("addr_mode_is", src_comp=[-1], dest_comp=1,
+          indices=[MEMORY_MODES], flags=[CAN_ELIMINATE, CAN_REORDER])
 
 # a barrier is an intrinsic with no inputs/outputs but which can't be moved
 # around/optimized in general

src/compiler/nir/nir_lower_io.c

@@ -898,6 +898,18 @@ build_addr_for_var(nir_builder *b, nir_variable *var,
    }
 }
 
+static nir_ssa_def *
+build_runtime_addr_mode_check(nir_builder *b, nir_ssa_def *addr,
+                              nir_address_format addr_format,
+                              nir_variable_mode mode)
+{
+   /* The compile-time check failed; do a run-time check */
+   switch (addr_format) {
+   default:
+      unreachable("Unsupported address mode");
+   }
+}
+
 static nir_ssa_def *
 addr_to_index(nir_builder *b, nir_ssa_def *addr,
               nir_address_format addr_format)
@@ -1708,6 +1720,32 @@ lower_explicit_io_array_length(nir_builder *b, nir_intrinsic_instr *intrin,
    nir_instr_remove(&intrin->instr);
 }
 
+static void
+lower_explicit_io_mode_check(nir_builder *b, nir_intrinsic_instr *intrin,
+                             nir_address_format addr_format)
+{
+   if (addr_format_is_global(addr_format, 0)) {
+      /* If the address format is always global, then the driver can use
+       * global addresses regardless of the mode.  In that case, don't create
+       * a check, just whack the intrinsic to addr_mode_is and delegate to the
+       * driver lowering that.
+       */
+      intrin->intrinsic = nir_intrinsic_addr_mode_is;
+      return;
+   }
+
+   assert(intrin->src[0].is_ssa);
+   nir_ssa_def *addr = intrin->src[0].ssa;
+
+   b->cursor = nir_instr_remove(&intrin->instr);
+
+   nir_ssa_def *is_mode =
+      build_runtime_addr_mode_check(b, addr, addr_format,
+                                    nir_intrinsic_memory_modes(intrin));
+
+   nir_ssa_def_rewrite_uses(&intrin->dest.ssa, nir_src_for_ssa(is_mode));
+}
+
 static bool
 nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
                            nir_address_format addr_format)
@@ -1769,6 +1807,15 @@ nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
                break;
             }
 
+            case nir_intrinsic_deref_mode_is: {
+               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+               if (nir_deref_mode_is_in_set(deref, modes)) {
+                  lower_explicit_io_mode_check(&b, intrin, addr_format);
+                  progress = true;
+               }
+               break;
+            }
+
             default:
                break;
             }

src/compiler/nir/nir_validate.c

@@ -671,6 +671,7 @@ validate_intrinsic_instr(nir_intrinsic_instr *instr, validate_state *state)
       break;
 
    case nir_intrinsic_deref_mode_is:
+   case nir_intrinsic_addr_mode_is:
       validate_assert(state,
          util_bitcount(nir_intrinsic_memory_modes(instr)) == 1);
       break;