anv: add a shrinking push constant loading pass

Shaders will often contains things like this :

con 32    %469 = @load_push_constant (%468 (0x30)) (base=0, range=128, align_mul=256, align_offset=48)

We don't need 128 bytes of push constants to do that load.

This will become important when we rely more on base/range in the next
commit to promote things to inline parameters (only 32B of space
available).

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/39405>

src/intel/vulkan/anv_nir.h

@@ -118,6 +118,8 @@ struct anv_nir_push_layout_info {
    bool mesh_dynamic;
 };
 
+bool anv_nir_shrink_push_constant_ranges(nir_shader *nir);
+
 bool anv_nir_compute_push_layout(nir_shader *nir,
                                  const struct anv_physical_device *pdevice,
                                  enum brw_robustness_flags robust_flags,

src/intel/vulkan/anv_nir_compute_push_layout.c

@@ -620,6 +620,34 @@ anv_nir_compute_push_layout(nir_shader *nir,
    return progress;
 }
 
+static bool
+shrink_push_constant_range_instr(nir_builder *b,
+                                 nir_intrinsic_instr *intrin,
+                                 void *data)
+{
+   if (!((intrin->intrinsic == nir_intrinsic_load_push_constant ||
+          intrin->intrinsic == nir_intrinsic_load_push_data_intel) &&
+         nir_src_is_const(intrin->src[0])))
+      return false;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+
+   nir_intrinsic_set_base(intrin, nir_intrinsic_base(intrin) +
+                                  nir_src_as_uint(intrin->src[0]));
+   nir_intrinsic_set_range(intrin,
+                           intrin->def.num_components * intrin->def.bit_size / 8);
+
+   nir_src_rewrite(&intrin->src[0], nir_imm_zero(b, 1, 32));
+   return true;
+}
+
+bool
+anv_nir_shrink_push_constant_ranges(nir_shader *shader)
+{
+   return nir_shader_intrinsics_pass(shader, shrink_push_constant_range_instr,
+                                     nir_metadata_control_flow, NULL);
+}
+
 void
 anv_nir_validate_push_layout(const struct anv_physical_device *pdevice,
                              struct brw_stage_prog_data *prog_data,

src/intel/vulkan/anv_shader_compile.c

@@ -1575,6 +1575,8 @@ anv_shader_lower_nir(struct anv_device *device,
 
    NIR_PASS(_, nir, anv_nir_update_resource_intel_block);
 
+   NIR_PASS(_, nir, anv_nir_shrink_push_constant_ranges);
+
    NIR_PASS(_, nir, anv_nir_compute_push_layout,
                pdevice, shader_data->key.base.robust_flags,
                &(struct anv_nir_push_layout_info) {