intel/fs: teach ubo range analysis pass about resource_intel

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Kenneth Graunke <kenneth@whitecape.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21645>

src/intel/compiler/brw_fs_nir.cpp

@@ -4629,34 +4629,32 @@ fs_visitor::nir_emit_intrinsic(const fs_builder &bld, nir_intrinsic_instr *instr
           */
          const unsigned type_size = type_sz(dest.type);
          const unsigned load_offset = nir_src_as_uint(instr->src[1]);
+         const unsigned ubo_block =
+            brw_nir_ubo_surface_index_get_push_block(instr->src[0]);
+         const unsigned offset_256b = load_offset / 32;
 
          /* See if we've selected this as a push constant candidate */
-         if (nir_src_is_const(instr->src[0])) {
+         fs_reg push_reg;
-            const unsigned ubo_block = nir_src_as_uint(instr->src[0]);
+         for (int i = 0; i < 4; i++) {
-            const unsigned offset_256b = load_offset / 32;
+            const struct brw_ubo_range *range = &prog_data->ubo_ranges[i];
+            if (range->block == ubo_block &&
+                offset_256b >= range->start &&
+                offset_256b < range->start + range->length) {
 
-            fs_reg push_reg;
+               push_reg = fs_reg(UNIFORM, UBO_START + i, dest.type);
-            for (int i = 0; i < 4; i++) {
+               push_reg.offset = load_offset - 32 * range->start;
-               const struct brw_ubo_range *range = &prog_data->ubo_ranges[i];
-               if (range->block == ubo_block &&
-                   offset_256b >= range->start &&
-                   offset_256b < range->start + range->length) {
-
-                  push_reg = fs_reg(UNIFORM, UBO_START + i, dest.type);
-                  push_reg.offset = load_offset - 32 * range->start;
-                  break;
-               }
-            }
-
-            if (push_reg.file != BAD_FILE) {
-               for (unsigned i = 0; i < instr->num_components; i++) {
-                  bld.MOV(offset(dest, bld, i),
-                          byte_offset(push_reg, i * type_size));
-               }
                break;
             }
          }
 
+         if (push_reg.file != BAD_FILE) {
+            for (unsigned i = 0; i < instr->num_components; i++) {
+               bld.MOV(offset(dest, bld, i),
+                       byte_offset(push_reg, i * type_size));
+            }
+            break;
+         }
+
          prog_data->has_ubo_pull = true;
 
          const unsigned block_sz = 64; /* Fetch one cacheline at a time. */

src/intel/compiler/brw_nir.h

@@ -103,6 +103,64 @@ struct brw_nir_compiler_opts {
    unsigned input_vertices;
 };
 
+/* UBO surface index can come in 2 flavors :
+ *    - nir_intrinsic_resource_intel
+ *    - anything else
+ *
+ * In the first case, checking that the surface index is const requires
+ * checking resource_intel::src[1]. In any other case it's a simple
+ * nir_src_is_const().
+ *
+ * This function should only be called on src[0] of load_ubo intrinsics.
+ */
+static inline bool
+brw_nir_ubo_surface_index_is_pushable(nir_src src)
+{
+   nir_intrinsic_instr *intrin = src.is_ssa &&
+      src.ssa->parent_instr->type == nir_instr_type_intrinsic ?
+      nir_instr_as_intrinsic(src.ssa->parent_instr) : NULL;
+
+   if (intrin && intrin->intrinsic == nir_intrinsic_resource_intel) {
+      return (nir_intrinsic_resource_access_intel(intrin) &
+              nir_resource_intel_pushable) &&
+             nir_src_is_const(intrin->src[1]);
+   }
+
+   return nir_src_is_const(src);
+}
+
+static inline unsigned
+brw_nir_ubo_surface_index_get_push_block(nir_src src)
+{
+   if (nir_src_is_const(src))
+      return nir_src_as_uint(src);
+
+   if (!brw_nir_ubo_surface_index_is_pushable(src))
+      return UINT32_MAX;
+
+   assert(src.ssa->parent_instr->type == nir_instr_type_intrinsic);
+
+   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(src.ssa->parent_instr);
+   assert(intrin->intrinsic == nir_intrinsic_resource_intel);
+
+   return nir_intrinsic_resource_block_intel(intrin);
+}
+
+static inline unsigned
+brw_nir_ubo_surface_index_get_bti(nir_src src)
+{
+   if (nir_src_is_const(src))
+      return nir_src_as_uint(src);
+
+   assert(src.ssa->parent_instr->type == nir_instr_type_intrinsic);
+
+   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(src.ssa->parent_instr);
+   assert(intrin->intrinsic == nir_intrinsic_resource_intel);
+   assert(nir_src_is_const(intrin->src[1]));
+
+   return nir_src_as_uint(intrin->src[1]);
+}
+
 void brw_preprocess_nir(const struct brw_compiler *compiler,
                         nir_shader *nir,
                         const struct brw_nir_compiler_opts *opts);

src/intel/compiler/brw_nir_analyze_ubo_ranges.c

@@ -141,9 +141,9 @@ analyze_ubos_block(struct ubo_analysis_state *state, nir_block *block)
          continue; /* Not a uniform or UBO intrinsic */
       }
 
-      if (nir_src_is_const(intrin->src[0]) &&
+      if (brw_nir_ubo_surface_index_is_pushable(intrin->src[0]) &&
           nir_src_is_const(intrin->src[1])) {
-         const int block = nir_src_as_uint(intrin->src[0]);
+         const int block = brw_nir_ubo_surface_index_get_push_block(intrin->src[0]);
          const unsigned byte_offset = nir_src_as_uint(intrin->src[1]);
          const int offset = byte_offset / 32;