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,12 +4629,11 @@ 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]);
-
-         /* See if we've selected this as a push constant candidate */
-         if (nir_src_is_const(instr->src[0])) {
-            const unsigned ubo_block = nir_src_as_uint(instr->src[0]);
+         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 */
          fs_reg push_reg;
          for (int i = 0; i < 4; i++) {
             const struct brw_ubo_range *range = &prog_data->ubo_ranges[i];
@@ -4655,7 +4654,6 @@ fs_visitor::nir_emit_intrinsic(const fs_builder &bld, nir_intrinsic_instr *instr
             }
             break;
          }
-         }
 
          prog_data->has_ubo_pull = true;
 

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;