nir/builder: Add a nir_trim_vector helper

This pattern pops up a bunch and the semantics of nir_channels() aren't
very convenient much of the time.  Let's add a nir_trim_vector() which
matches nir_pad_vector().

Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16309>

src/amd/vulkan/radv_meta_bufimage.c

@@ -60,7 +60,7 @@ build_nir_itob_compute_shader(struct radv_device *dev, bool is_3d)
    tex->sampler_dim = dim;
    tex->op = nir_texop_txf;
    tex->src[0].src_type = nir_tex_src_coord;
-   tex->src[0].src = nir_src_for_ssa(nir_channels(&b, img_coord, is_3d ? 0x7 : 0x3));
+   tex->src[0].src = nir_src_for_ssa(nir_trim_vector(&b, img_coord, 2 + is_3d));
    tex->src[1].src_type = nir_tex_src_lod;
    tex->src[1].src = nir_src_for_ssa(nir_imm_int(&b, 0));
    tex->src[2].src_type = nir_tex_src_texture_deref;
@@ -570,7 +570,7 @@ build_nir_itoi_compute_shader(struct radv_device *dev, bool is_3d, int samples)
       tex->sampler_dim = dim;
       tex->op = is_multisampled ? nir_texop_txf_ms : nir_texop_txf;
       tex->src[0].src_type = nir_tex_src_coord;
-      tex->src[0].src = nir_src_for_ssa(nir_channels(&b, src_coord, is_3d ? 0x7 : 0x3));
+      tex->src[0].src = nir_src_for_ssa(nir_trim_vector(&b, src_coord, 2 + is_3d));
       tex->src[1].src_type = nir_tex_src_lod;
       tex->src[1].src = nir_src_for_ssa(nir_imm_int(&b, 0));
       tex->src[2].src_type = nir_tex_src_texture_deref;

src/compiler/nir/nir_builder.h

@@ -1021,6 +1021,16 @@ nir_bitcast_vector(nir_builder *b, nir_ssa_def *src, unsigned dest_bit_size)
    return nir_extract_bits(b, &src, 1, 0, dest_num_components, dest_bit_size);
 }
 
+static inline nir_ssa_def *
+nir_trim_vector(nir_builder *b, nir_ssa_def *src, unsigned num_components)
+{
+   assert(src->num_components >= num_components);
+   if (src->num_components == num_components)
+      return src;
+
+   return nir_channels(b, src, nir_component_mask(num_components));
+}
+
 /**
  * Pad a value to N components with undefs of matching bit size.
  * If the value already contains >= num_components, it is returned without change.

src/compiler/nir/nir_lower_readonly_images_to_tex.c

@@ -147,8 +147,7 @@ lower_readonly_image_op(nir_builder *b, nir_instr *instr, void *context)
    case nir_intrinsic_image_deref_load: {
       assert(intrin->src[1].is_ssa);
       nir_ssa_def *coord =
-         nir_channels(b, intrin->src[1].ssa,
+         nir_trim_vector(b, intrin->src[1].ssa, coord_components);
-                      (1 << coord_components) - 1);
       tex->src[1].src_type = nir_tex_src_coord;
       tex->src[1].src = nir_src_for_ssa(coord);
       tex->coord_components = coord_components;
@@ -188,7 +187,7 @@ lower_readonly_image_op(nir_builder *b, nir_instr *instr, void *context)
    nir_ssa_def *res = &tex->dest.ssa;
    if (res->num_components != intrin->dest.ssa.num_components) {
       unsigned num_components = intrin->dest.ssa.num_components;
-      res = nir_channels(b, res, (1 << num_components) - 1);
+      res = nir_trim_vector(b, res, num_components);
    }
 
    return res;

src/compiler/nir/nir_lower_subgroups.c

@@ -98,7 +98,7 @@ uint_to_ballot_type(nir_builder *b, nir_ssa_def *value,
     * have enough ballot bits.
     */
    if (value->num_components > num_components)
-      value = nir_channels(b, value, nir_component_mask(num_components));
+      value = nir_trim_vector(b, value, num_components);
 
    return value;
 }

src/compiler/nir/nir_opt_shrink_stores.c

@@ -48,7 +48,7 @@ opt_shrink_vectors_image_store(nir_builder *b, nir_intrinsic_instr *instr)
    if (components >= instr->num_components)
       return false;
 
-   nir_ssa_def *data = nir_channels(b, instr->src[3].ssa, BITSET_MASK(components));
+   nir_ssa_def *data = nir_trim_vector(b, instr->src[3].ssa, components);
    nir_instr_rewrite_src(&instr->instr, &instr->src[3], nir_src_for_ssa(data));
    instr->num_components = components;
 
@@ -83,8 +83,7 @@ opt_shrink_store_instr(nir_builder *b, nir_intrinsic_instr *instr, bool shrink_i
    unsigned write_mask = nir_intrinsic_write_mask(instr);
    unsigned last_bit = util_last_bit(write_mask);
    if (last_bit < instr->num_components && instr->src[0].is_ssa) {
-      nir_ssa_def *def = nir_channels(b, instr->src[0].ssa,
+      nir_ssa_def *def = nir_trim_vector(b, instr->src[0].ssa, last_bit);
-                                      BITSET_MASK(last_bit));
       nir_instr_rewrite_src(&instr->instr,
                             &instr->src[0],
                             nir_src_for_ssa(def));

src/compiler/spirv/spirv_to_nir.c

@@ -2874,8 +2874,7 @@ vtn_handle_texture(struct vtn_builder *b, SpvOp opcode,
        */
       vtn_fail_if(coord->num_components < coord_components,
                   "Coordinate value passed has fewer components than sampler dimensionality.");
-      p->src = nir_src_for_ssa(nir_channels(&b->nb, coord,
+      p->src = nir_src_for_ssa(nir_trim_vector(&b->nb, coord, coord_components));
-                                            (1 << coord_components) - 1));
 
       /* OpenCL allows integer sampling coordinates */
       if (glsl_type_is_integer(coord_val->type) &&
@@ -3156,8 +3155,8 @@ vtn_handle_texture(struct vtn_builder *b, SpvOp opcode,
       struct vtn_ssa_value *dest = vtn_create_ssa_value(b, struct_type->type);
       unsigned result_size = glsl_get_vector_elements(ret_type->type);
       dest->elems[0]->def = nir_channel(&b->nb, &instr->dest.ssa, result_size);
-      dest->elems[1]->def = nir_channels(&b->nb, &instr->dest.ssa,
+      dest->elems[1]->def = nir_trim_vector(&b->nb, &instr->dest.ssa,
-                                         nir_component_mask(result_size));
+                                              result_size);
       vtn_push_ssa_value(b, w[2], dest);
    } else {
       vtn_push_nir_ssa(b, w[2], &instr->dest.ssa);
@@ -3574,9 +3573,8 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
 
       nir_builder_instr_insert(&b->nb, &intrin->instr);
 
-      nir_ssa_def *result = &intrin->dest.ssa;
+      nir_ssa_def *result = nir_trim_vector(&b->nb, &intrin->dest.ssa,
-      if (nir_intrinsic_dest_components(intrin) != dest_components)
+                                              dest_components);
-         result = nir_channels(&b->nb, result, (1 << dest_components) - 1);
 
       if (opcode == SpvOpImageQuerySize ||
           opcode == SpvOpImageQuerySizeLod)
@@ -3588,8 +3586,7 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
          dest->elems[0]->def = nir_channel(&b->nb, result, res_type_size);
          if (intrin->dest.ssa.bit_size != 32)
             dest->elems[0]->def = nir_u2u32(&b->nb, dest->elems[0]->def);
-         dest->elems[1]->def = nir_channels(&b->nb, result,
+         dest->elems[1]->def = nir_trim_vector(&b->nb, result, res_type_size);
-                                            nir_component_mask(res_type_size));
          vtn_push_ssa_value(b, w[2], dest);
       } else {
          vtn_push_nir_ssa(b, w[2], result);

src/freedreno/ir3/ir3_nir_analyze_ubo_ranges.c

@@ -621,7 +621,7 @@ ir3_nir_lower_load_const_instr(nir_builder *b, nir_instr *in_instr, void *data)
 
    if (nir_dest_bit_size(instr->dest) == 16) {
       result = nir_bitcast_vector(b, result, 16);
-      result = nir_channels(b, result, BITSET_MASK(instr->num_components));
+      result = nir_trim_vector(b, result, instr->num_components);
    }
 
    return result;

src/intel/blorp/blorp_blit.c

@@ -940,7 +940,7 @@ bit_cast_color(struct nir_builder *b, nir_ssa_def *color,
       /* Restrict to only the channels we actually have */
       const unsigned src_channels =
          isl_format_get_num_channels(key->src_format);
-      color = nir_channels(b, color, (1 << src_channels) - 1);
+      color = nir_trim_vector(b, color, src_channels);
 
       color = nir_format_bitcast_uvec_unmasked(b, color, src_bpc, dst_bpc);
    }

src/intel/compiler/brw_nir_lower_storage_image.c

@@ -100,7 +100,7 @@ image_address(nir_builder *b, const struct intel_device_info *devinfo,
                           nir_channel(b, coord, 1));
    } else {
       unsigned dims = glsl_get_sampler_coordinate_components(deref->type);
-      coord = nir_channels(b, coord, (1 << dims) - 1);
+      coord = nir_trim_vector(b, coord, dims);
    }
 
    nir_ssa_def *offset = load_image_param(b, deref, OFFSET);
@@ -457,7 +457,7 @@ convert_color_for_store(nir_builder *b, const struct intel_device_info *devinfo,
    struct format_info image = get_format_info(image_fmt);
    struct format_info lower = get_format_info(lower_fmt);
 
-   color = nir_channels(b, color, (1 << image.chans) - 1);
+   color = nir_trim_vector(b, color, image.chans);
 
    if (image_fmt == lower_fmt)
       return color;