/*
 * Copyright © 2023 Bas Nieuwenhuizen
 *
 * SPDX-License-Identifier: MIT
 */

#include "nir_builder.h"
#include "radv_nir.h"

/* This pass lowers cooperative matrix.
 *
 * On GFX11, the A&B matrices needs to be replicated, lanes 0..15 are replicated
 * to 16..31 and for wave64 also into lanes 32..47 and 48..63. A&B matrices are
 * always vectors of 16 elements.
 *
 * On GFX12, there is no data replication and the matrices layout is described
 * as below:
 *
 * Wave32:
 *         0..15  | 16..31 (lanes)
 * v0 lo:  row 0  | row 8
 * v0 hi:  row 1  | row 9
 * v1 lo:  row 2  | row 10
 * v1 hi:  row 3  | row 11
 * v2 lo:  row 4  | row 12
 * v2 hi:  row 5  | row 13
 * v3 lo:  row 6  | row 14
 * v3 hi:  row 7  | row 15
 *
 * Wave64:
 *         0..15 | 16..31 | 32..47 | 48..63 (lanes)
 * v0 lo:  row 0 | row 8  | row 4  | row 12
 * v0 hi:  row 1 | row 9  | row 5  | row 13
 * v1 lo:  row 2 | row 10 | row 6  | row 14
 * v1 hi:  row 3 | row 11 | row 7  | row 15
 *
 * Note that the GFX12 ISA doc describes other layouts for A/B, but they are identical
 * to the C layout with the exception of the order of the rows (columns for A).
 * And as long as these are swapped in the same way for both A and B, the muladd
 * result will be the same. So we use the C layout for all uses.
 */

typedef struct {
   enum amd_gfx_level gfx_level;
   unsigned wave_size;
} lower_cmat_params;

static unsigned
radv_nir_cmat_bits(struct glsl_cmat_description desc)
{
   return glsl_base_type_bit_size(desc.element_type);
}

static unsigned
radv_nir_cmat_length(struct glsl_cmat_description desc, const lower_cmat_params *params)
{
   if (params->gfx_level >= GFX12) {
      assert(desc.cols == 16 && desc.rows == 16);
      return 256 / params->wave_size;
   } else if (desc.use != GLSL_CMAT_USE_ACCUMULATOR) {
      return 16;
   } else {
      return desc.cols * desc.rows / params->wave_size * (radv_nir_cmat_bits(desc) == 16 ? 2 : 1);
   }
}

static unsigned
radv_nir_cmat_length_mul(struct glsl_cmat_description desc, const lower_cmat_params *params)
{
   if (params->gfx_level >= GFX12 || desc.use != GLSL_CMAT_USE_ACCUMULATOR) {
      return 1;
   } else {
      /* For  GFX11 C matrices we have 1 VGPR per element even if the element type is
       * 16bits. So with 8 fp16 elements we implement that with a f16vec16.
       * We then use the coefficient generated by this function to figure out
       * how many elements we really have.
       */
      return radv_nir_cmat_bits(desc) == 16 ? 2 : 1;
   }
}

static nir_def *
radv_nir_load_cmat(nir_builder *b, const lower_cmat_params *params, nir_def *src)
{
   nir_deref_instr *deref = nir_def_as_deref(src);
   struct glsl_cmat_description desc = *glsl_get_cmat_description(deref->type);
   return nir_build_load_deref(b, radv_nir_cmat_length(desc, params), radv_nir_cmat_bits(desc), src, 0);
}

static const struct glsl_type *
radv_nir_translate_matrix_type(const struct glsl_type *orig_type, struct hash_table *type_map,
                               const lower_cmat_params *params)
{
   struct hash_entry *entry = _mesa_hash_table_search(type_map, orig_type);
   if (entry) {
      return entry->data;
   } else if (glsl_type_is_cmat(orig_type)) {
      struct glsl_cmat_description desc = *glsl_get_cmat_description(orig_type);
      unsigned length = radv_nir_cmat_length(desc, params);

      return glsl_vector_type(desc.element_type, length);
   } else if (glsl_type_is_array(orig_type)) {
      const struct glsl_type *elem_type = glsl_get_array_element(orig_type);
      const struct glsl_type *new_elem_type = radv_nir_translate_matrix_type(elem_type, type_map, params);

      if (elem_type == new_elem_type)
         return orig_type;

      return glsl_array_type(new_elem_type, glsl_get_length(orig_type), glsl_get_explicit_stride(orig_type));
   } else if (glsl_type_is_struct(orig_type)) {
      unsigned num_fields = glsl_get_length(orig_type);

      bool change = false;
      for (unsigned i = 0; i < num_fields; ++i) {
         const struct glsl_type *field_type = glsl_get_struct_field(orig_type, i);
         const struct glsl_type *new_field_type = radv_nir_translate_matrix_type(field_type, type_map, params);

         if (field_type != new_field_type) {
            change = true;
            break;
         }
      }

      if (!change)
         return orig_type;

      struct glsl_struct_field *fields = malloc(sizeof(struct glsl_struct_field) * num_fields);

      for (unsigned i = 0; i < num_fields; ++i) {
         fields[i] = *glsl_get_struct_field_data(orig_type, i);

         fields[i].type = radv_nir_translate_matrix_type(fields[i].type, type_map, params);
      }

      const struct glsl_type *ret =
         glsl_struct_type(fields, num_fields, glsl_get_type_name(orig_type), glsl_struct_type_is_packed(orig_type));
      free(fields);

      _mesa_hash_table_insert(type_map, orig_type, (void *)ret);
      return ret;
   } else
      return orig_type;
}

static nir_def *
radv_get_base_row(nir_builder *b, struct glsl_cmat_description desc, const lower_cmat_params *params,
                  nir_def *local_idx)
{
   nir_def *base_row;

   if (params->gfx_level >= GFX12) {
      base_row = nir_udiv_imm(b, local_idx, 16);

      if (params->wave_size == 64) {
         /* Switch rows from lanes 16..31 to 32..47 */
         base_row = nir_ushr_imm(b, nir_bitfield_reverse(b, base_row), 30);
         base_row = nir_imul_imm(b, base_row, 4);
      } else {
         base_row = nir_imul_imm(b, base_row, 8);
      }
   } else {
      base_row = desc.use == GLSL_CMAT_USE_ACCUMULATOR ? nir_udiv_imm(b, local_idx, 16) : nir_imm_int(b, 0);
   }

   return base_row;
}

static nir_def *
shuffle_xor_imm(nir_builder *b, nir_def *data, unsigned imm)
{
   assert(imm < 64);
   if (imm == 32) {
      /* v_permlane64_b32 */
      return nir_rotate(b, data, nir_imm_int(b, 32), .cluster_size = 64);
   } else if (imm < 32) {
      /* All of these map to a single DPP/v_permlanex16 instruction */
      return nir_masked_swizzle_amd(b, data, .swizzle_mask = 0x1f | (imm << 10), .fetch_inactive = 1);
   } else {
      /* There isn't a single instruction that can do this, but for cooperative matrix
       * opcodes all invocations must be active.
       * So we can split the operation into the two previous cases instead of
       * having to use full width shuffle.
       */
      data = shuffle_xor_imm(b, data, 32);
      return shuffle_xor_imm(b, data, imm & ~32);
   }
}

static bool
lower_cmat_length(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   struct glsl_cmat_description desc = nir_intrinsic_cmat_desc(intr);
   unsigned len = radv_nir_cmat_length(desc, params) / radv_nir_cmat_length_mul(desc, params);
   nir_def_replace(&intr->def, nir_imm_int(b, len));
   return true;
}

static bool
lower_cmat_extract(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   nir_deref_instr *src_deref = nir_src_as_deref(intr->src[0]);
   struct glsl_cmat_description desc = *glsl_get_cmat_description(src_deref->type);
   nir_def *src0 = radv_nir_load_cmat(b, params, intr->src[0].ssa);

   nir_def *index = intr->src[1].ssa;
   index = nir_imul_imm(b, index, radv_nir_cmat_length_mul(desc, params));

   nir_def *elem = nir_vector_extract(b, src0, index);

   nir_def_replace(&intr->def, elem);
   return true;
}

static bool
lower_cmat_insert(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   nir_def *src1 = radv_nir_load_cmat(b, params, intr->src[2].ssa);
   nir_deref_instr *dst_deref = nir_src_as_deref(intr->src[0]);
   struct glsl_cmat_description desc = *glsl_get_cmat_description(dst_deref->type);
   nir_def *index = intr->src[3].ssa;
   index = nir_imul_imm(b, index, radv_nir_cmat_length_mul(desc, params));

   nir_def *elem = intr->src[1].ssa;
   nir_def *r = nir_vector_insert(b, src1, elem, index);
   nir_store_deref(b, dst_deref, r, nir_component_mask(r->num_components));
   nir_instr_remove(&intr->instr);
   return true;
}

static bool
lower_cmat_construct(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   nir_deref_instr *dst_deref = nir_src_as_deref(intr->src[0]);
   struct glsl_cmat_description desc = *glsl_get_cmat_description(dst_deref->type);
   nir_def *elem = intr->src[1].ssa;

   nir_def *r = nir_replicate(b, elem, radv_nir_cmat_length(desc, params));

   nir_store_deref(b, dst_deref, r, nir_component_mask(r->num_components));
   nir_instr_remove(&intr->instr);
   return true;
}

static bool
lower_cmat_load_store(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   const bool is_load = intr->intrinsic == nir_intrinsic_cmat_load;

   nir_deref_instr *cmat_deref = nir_src_as_deref(intr->src[!is_load]);
   struct glsl_cmat_description desc = *glsl_get_cmat_description(cmat_deref->type);
   enum glsl_matrix_layout layout = nir_intrinsic_matrix_layout(intr);

   nir_deref_instr *deref = nir_src_as_deref(intr->src[is_load]);
   nir_def *stride = intr->src[2].ssa;

   const uint32_t ptr_stride = glsl_get_bit_size(deref->type) / 8 * glsl_get_vector_elements(deref->type);
   deref = nir_build_deref_cast(b, &deref->def, deref->modes, deref->type, ptr_stride);

   nir_def *local_idx = nir_load_subgroup_invocation(b);
   nir_def *inner_idx = nir_iand_imm(b, local_idx, 15);

   bool load_acc_as_b = is_load && params->gfx_level < GFX12 && desc.use == GLSL_CMAT_USE_ACCUMULATOR &&
                        radv_nir_cmat_bits(desc) == 8 && params->wave_size == 32 &&
                        layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR;
   if (load_acc_as_b)
      desc.use = GLSL_CMAT_USE_B;

   /* A input is transposed */
   if (desc.use == GLSL_CMAT_USE_A)
      layout =
         layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR ? GLSL_MATRIX_LAYOUT_ROW_MAJOR : GLSL_MATRIX_LAYOUT_COLUMN_MAJOR;

   unsigned length = radv_nir_cmat_length(desc, params);
   unsigned mul = radv_nir_cmat_length_mul(desc, params);
   unsigned lanes_per_iter = desc.use == GLSL_CMAT_USE_ACCUMULATOR ? params->wave_size : 16;
   nir_def *vars[16];
   if (is_load) {
      if (mul > 1) {
         for (unsigned i = 0; i < length; ++i)
            if (i % mul != 0)
               vars[i] = nir_undef(b, 1, radv_nir_cmat_bits(desc));
      }
   } else {
      if (params->gfx_level < GFX12 && desc.use != GLSL_CMAT_USE_ACCUMULATOR)
         nir_push_if(b, nir_ilt_imm(b, local_idx, 16));

      nir_def *src = radv_nir_load_cmat(b, params, &cmat_deref->def);
      for (unsigned i = 0; i < length; ++i)
         vars[i] = nir_channel(b, src, i);
   }

   unsigned idx_bits = deref->def.bit_size;
   nir_def *base_row = radv_get_base_row(b, desc, params, local_idx);

   /* VUID-RuntimeSpirv-OpCooperativeMatrixLoadKHR-08986:
    * For OpCooperativeMatrixLoadKHR and OpCooperativeMatrixStoreKHR instructions,
    * the Pointer and Stride operands must be aligned to at least the lesser of 16 bytes
    * or the natural alignment of a row or column (depending on ColumnMajor) of the matrix
    * (where the natural alignment is the number of columns/rows multiplied by the component size).
    */
   unsigned align_mul = 0;
   if (layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR)
      align_mul = MIN2(16, radv_nir_cmat_bits(desc) * desc.rows / 8);

   if (params->gfx_level >= GFX12)
      align_mul /= params->wave_size / 16;
   else if (desc.use == GLSL_CMAT_USE_ACCUMULATOR)
      align_mul = 0;

   for (unsigned i = 0; i < length / mul; ++i) {
      nir_def *col_offset = inner_idx;
      nir_def *row_offset;
      uint32_t row_iter;

      if (params->gfx_level >= GFX12) {
         row_iter = i;
      } else {
         row_iter = i * lanes_per_iter / 16;
      }

      row_offset = nir_iadd_imm(b, base_row, row_iter);

      if (layout == GLSL_MATRIX_LAYOUT_ROW_MAJOR) {
         SWAP(col_offset, row_offset);
      }

      col_offset = nir_imul(b, col_offset, stride);

      col_offset = nir_u2uN(b, col_offset, idx_bits);
      row_offset = nir_u2uN(b, row_offset, idx_bits);

      nir_deref_instr *iter_deref = nir_build_deref_ptr_as_array(b, deref, col_offset);
      iter_deref = nir_build_deref_cast(b, &iter_deref->def, deref->modes, glsl_scalar_type(desc.element_type),
                                        radv_nir_cmat_bits(desc) / 8);
      iter_deref = nir_build_deref_ptr_as_array(b, iter_deref, row_offset);

      if (align_mul) {
         unsigned align_offset = row_iter * radv_nir_cmat_bits(desc) / 8 % align_mul;
         iter_deref = nir_build_deref_cast_with_alignment(b, &iter_deref->def, deref->modes, iter_deref->type,
                                                          radv_nir_cmat_bits(desc) / 8, align_mul, align_offset);
      }

      if (is_load) {
         vars[i * mul] = nir_load_deref(b, iter_deref);
      } else {
         nir_store_deref(b, iter_deref, vars[i * mul], 1);
      }
   }

   if (is_load) {
      nir_def *mat = nir_vec(b, vars, length);

      if (load_acc_as_b) {
         /* Shift data in the high bits for the first 16 lanes (it's already there for the other 16.) */
         nir_def *lo = nir_ult_imm(b, local_idx, 16);
         nir_def *shift = nir_bcsel(b, lo, nir_imm_int(b, 8), nir_imm_int(b, 0));
         for (unsigned i = 0; i < mat->num_components / 4; i++) {
            nir_def *packed = nir_pack_32_4x8(b, nir_channels(b, mat, 0xfu << (4 * i)));
            packed = nir_ishl(b, packed, shift);
            nir_def *unpacked = nir_unpack_32_4x8(b, packed);
            vars[i * 2] = nir_channel(b, unpacked, 1);
            vars[i * 2 + 1] = nir_channel(b, unpacked, 3);
         }

         mat = nir_vec(b, vars, length / 2);
      }

      nir_store_deref(b, cmat_deref, mat, nir_component_mask(mat->num_components));
   } else if (params->gfx_level < GFX12 && desc.use != GLSL_CMAT_USE_ACCUMULATOR) {
      nir_pop_if(b, NULL);
   }
   nir_instr_remove(&intr->instr);
   return true;
}

static bool
lower_cmat_muladd(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   nir_def *A = radv_nir_load_cmat(b, params, intr->src[1].ssa);
   nir_def *B = radv_nir_load_cmat(b, params, intr->src[2].ssa);
   nir_def *C = radv_nir_load_cmat(b, params, intr->src[3].ssa);

   nir_deref_instr *a_deref = nir_src_as_deref(intr->src[1]);
   nir_deref_instr *b_deref = nir_src_as_deref(intr->src[2]);
   struct glsl_cmat_description a_desc = *glsl_get_cmat_description(a_deref->type);
   struct glsl_cmat_description b_desc = *glsl_get_cmat_description(b_deref->type);

   const nir_cmat_signed cmat_signed_mask = nir_intrinsic_cmat_signed_mask(intr);

   enum glsl_base_type a_element_type =
      glsl_apply_signedness_to_base_type(a_desc.element_type, cmat_signed_mask & NIR_CMAT_A_SIGNED);
   enum glsl_base_type b_element_type =
      glsl_apply_signedness_to_base_type(b_desc.element_type, cmat_signed_mask & NIR_CMAT_B_SIGNED);

   nir_def *ret = nir_cmat_muladd_amd(b, A, B, C, .saturate = nir_intrinsic_saturate(intr),
                                      .src_base_type = a_element_type, .src_base_type2 = b_element_type);

   nir_deref_instr *dst_deref = nir_src_as_deref(intr->src[0]);
   nir_store_deref(b, dst_deref, ret, nir_component_mask(ret->num_components));
   nir_instr_remove(&intr->instr);
   return true;
}

static nir_def *
convert_base_type(nir_builder *b, nir_def *src, enum glsl_base_type src_type, enum glsl_base_type dst_type, bool sat)
{
   if (dst_type == src_type)
      return src;

   if (src_type == GLSL_TYPE_BFLOAT16) {
      src = nir_bf2f(b, src);
      return convert_base_type(b, src, GLSL_TYPE_FLOAT, dst_type, sat);
   } else if (dst_type == GLSL_TYPE_BFLOAT16) {
      src = convert_base_type(b, src, src_type, GLSL_TYPE_FLOAT, sat);
      return nir_f2bf(b, src);
   } else if (src_type == GLSL_TYPE_FLOAT_E4M3FN) {
      src = nir_e4m3fn2f(b, src);
      return convert_base_type(b, src, GLSL_TYPE_FLOAT, dst_type, sat);
   } else if (dst_type == GLSL_TYPE_FLOAT_E4M3FN) {
      src = convert_base_type(b, src, src_type, GLSL_TYPE_FLOAT, sat);
      if (sat)
         return nir_f2e4m3fn_sat(b, src);
      else
         return nir_f2e4m3fn(b, src);
   } else if (src_type == GLSL_TYPE_FLOAT_E5M2) {
      src = nir_e5m22f(b, src);
      return convert_base_type(b, src, GLSL_TYPE_FLOAT, dst_type, sat);
   } else if (dst_type == GLSL_TYPE_FLOAT_E5M2) {
      src = convert_base_type(b, src, src_type, GLSL_TYPE_FLOAT, sat);
      if (sat)
         return nir_f2e5m2_sat(b, src);
      else
         return nir_f2e5m2(b, src);
   }

   nir_op op = nir_type_conversion_op(nir_get_nir_type_for_glsl_base_type(src_type),
                                      nir_get_nir_type_for_glsl_base_type(dst_type), nir_rounding_mode_undef);

   return nir_build_alu1(b, op, src);
}

static nir_def *
convert_use(nir_builder *b, nir_def *src, enum glsl_cmat_use src_use, enum glsl_cmat_use dst_use,
            const lower_cmat_params *params)
{
   if (src_use == dst_use)
      return src;
   if (params->gfx_level >= GFX12) {
      if (src_use == GLSL_CMAT_USE_B && dst_use == GLSL_CMAT_USE_ACCUMULATOR)
         return src;
      if (src_use == GLSL_CMAT_USE_ACCUMULATOR && dst_use == GLSL_CMAT_USE_B)
         return src;
   }

   if (src_use == GLSL_CMAT_USE_A && dst_use == GLSL_CMAT_USE_ACCUMULATOR) {
      src = convert_use(b, src, GLSL_CMAT_USE_A, GLSL_CMAT_USE_B, params);
      return convert_use(b, src, GLSL_CMAT_USE_B, GLSL_CMAT_USE_ACCUMULATOR, params);
   } else if (src_use == GLSL_CMAT_USE_ACCUMULATOR && dst_use == GLSL_CMAT_USE_A) {
      src = convert_use(b, src, GLSL_CMAT_USE_ACCUMULATOR, GLSL_CMAT_USE_B, params);
      return convert_use(b, src, GLSL_CMAT_USE_B, GLSL_CMAT_USE_A, params);
   }

   nir_def *components[NIR_MAX_VEC_COMPONENTS] = {NULL};

   unsigned num_comps = src->num_components;
   for (unsigned i = 0; i < num_comps; i++)
      components[i] = nir_channel(b, src, i);

   if (src_use == GLSL_CMAT_USE_ACCUMULATOR && dst_use == GLSL_CMAT_USE_B) {
      assert(params->gfx_level < GFX12);
      nir_def *tmp[NIR_MAX_VEC_COMPONENTS];

      if (src->bit_size == 32) {
         for (int cross32 = params->wave_size == 64; cross32 >= 0; cross32--) {
            uint64_t low_mask = cross32 ? UINT32_MAX : 0xffff0000ffffull;
            nir_def *low_lanes = nir_inverse_ballot_imm(b, low_mask, params->wave_size);
            for (int i = 0; i < num_comps; i++) {
               nir_def *comp = components[i];
               nir_def *half_swap = shuffle_xor_imm(b, comp, cross32 ? 32 : 16);

               tmp[i * 2] = nir_bcsel(b, low_lanes, comp, half_swap);
               tmp[i * 2 + 1] = nir_bcsel(b, low_lanes, half_swap, comp);
            }
            num_comps *= 2;
            memcpy(components, tmp, sizeof(components));
         }
      } else {
         /* Same as above, but operate on 32 bits at once, using byte_perm as vectorized bcsel. */
         assert(src->bit_size < 32);
         nir_def *packed = nir_extract_bits(b, components, num_comps, 0, num_comps / (32 / src->bit_size), 32);
         num_comps /= 32 / src->bit_size;
         for (unsigned i = 0; i < num_comps; i++)
            components[i] = nir_channel(b, packed, i);

         nir_def *low_sel = nir_imm_int(b, src->bit_size == 8 ? 0x05010400 : 0x05040100);
         nir_def *high_sel = nir_imm_int(b, src->bit_size == 8 ? 0x01050004 : 0x01000504);

         for (int cross32 = params->wave_size == 64; cross32 >= 0; cross32--) {
            uint64_t low_mask = cross32 ? UINT32_MAX : 0xffff0000ffffull;
            nir_def *low_lanes = nir_inverse_ballot_imm(b, low_mask, params->wave_size);

            nir_def *first_perm = nir_bcsel(b, low_lanes, low_sel, high_sel);
            nir_def *second_perm = nir_ior_imm(b, first_perm, 0x02020202);
            for (int i = 0; i < num_comps; i++) {
               nir_def *comp = components[i];
               nir_def *half_swap = shuffle_xor_imm(b, comp, cross32 ? 32 : 16);

               tmp[i * 2] = nir_byte_perm_amd(b, half_swap, comp, first_perm);
               tmp[i * 2 + 1] = nir_byte_perm_amd(b, half_swap, comp, second_perm);
            }
            num_comps *= 2;
            memcpy(components, tmp, sizeof(components));
         }

         nir_def *unpacked =
            nir_extract_bits(b, components, num_comps, 0, num_comps * (32 / src->bit_size), src->bit_size);
         num_comps *= 32 / src->bit_size;
         for (unsigned i = 0; i < num_comps; i++)
            components[i] = nir_channel(b, unpacked, i);
      }

      assert(num_comps == 16);
   } else if (src_use == GLSL_CMAT_USE_B && dst_use == GLSL_CMAT_USE_ACCUMULATOR) {
      assert(params->gfx_level < GFX12);
      assert(num_comps == 16);
      if (src->bit_size == 32) {
         for (unsigned keep32 = 0; keep32 < ((params->wave_size == 64) ? 2 : 1); keep32++) {
            nir_def *keep = nir_inverse_ballot_imm(b, keep32 ? UINT32_MAX : 0xffff0000ffffull, params->wave_size);
            num_comps /= 2;
            for (unsigned i = 0; i < num_comps; i++) {
               components[i] = nir_bcsel(b, keep, components[i * 2], components[i * 2 + 1]);
            }
         }
      } else {
         /* Same as above, but operate on 32 bits at once, using byte_perm as vectorized bcsel. */
         assert(src->bit_size < 32);
         nir_def *packed = nir_extract_bits(b, components, num_comps, 0, num_comps / (32 / src->bit_size), 32);
         num_comps /= 32 / src->bit_size;
         for (unsigned i = 0; i < num_comps; i++)
            components[i] = nir_channel(b, packed, i);

         nir_def *low_sel = nir_imm_int(b, src->bit_size == 8 ? 0x06040200 : 0x05040100);
         nir_def *high_sel = nir_imm_int(b, src->bit_size == 8 ? 0x07050301 : 0x07060302);

         for (unsigned keep32 = 0; keep32 < ((params->wave_size == 64) ? 2 : 1); keep32++) {
            nir_def *keep = nir_inverse_ballot_imm(b, keep32 ? UINT32_MAX : 0xffff0000ffffull, params->wave_size);
            nir_def *perm = nir_bcsel(b, keep, low_sel, high_sel);
            num_comps /= 2;
            for (unsigned i = 0; i < num_comps; i++) {
               components[i] = nir_byte_perm_amd(b, components[i * 2 + 1], components[i * 2], perm);
            }
         }

         nir_def *unpacked =
            nir_extract_bits(b, components, num_comps, 0, num_comps * (32 / src->bit_size), src->bit_size);
         num_comps *= 32 / src->bit_size;
         for (unsigned i = 0; i < num_comps; i++)
            components[i] = nir_channel(b, unpacked, i);
      }
   } else if ((src_use == GLSL_CMAT_USE_A && dst_use == GLSL_CMAT_USE_B) ||
              (src_use == GLSL_CMAT_USE_B && dst_use == GLSL_CMAT_USE_A)) {
      /* Transpose is a mess... */
      for (unsigned x_mask = 1; x_mask < num_comps; x_mask *= 2) {
         /* Use separate masks to always keep the masked_swizzle on the first source of v_cndmask. */
         uint64_t mask = 0;
         for (unsigned i = 0; i < 64; i += 2 * x_mask) {
            mask |= BITFIELD64_MASK(x_mask) << i;
         }

         nir_def *even = nir_inverse_ballot_imm(b, mask, params->wave_size);
         nir_def *odd = nir_inverse_ballot_imm(b, mask << x_mask, params->wave_size);

         for (unsigned i = 0; i < num_comps; i += 2 * x_mask) {
            for (unsigned j = 0; j < x_mask; j++) {
               unsigned pos0 = i + j;
               unsigned pos1 = pos0 + x_mask;
               nir_def *comp0 = components[pos0];
               nir_def *comp1 = components[pos1];

               nir_def *comp0x = shuffle_xor_imm(b, comp0, x_mask);
               nir_def *comp1x = shuffle_xor_imm(b, comp1, x_mask);
               components[pos0] = nir_bcsel(b, even, comp0, comp1x);
               components[pos1] = nir_bcsel(b, odd, comp1, comp0x);
            }
         }
      }

      assert(num_comps == 16 || params->gfx_level >= GFX12);

      if (params->gfx_level >= GFX12) {
         /* One component contains 2/4 rows in wave32/64, so we must transpose inside it. */
         for (int cross32 = params->wave_size == 64; cross32 >= 0; cross32--) {
            uint64_t even = cross32 ? 0xf0f0f0f00f0f0f0f : 0xff0000ffff0000ff;
            nir_def *cond = nir_inverse_ballot_imm(b, even, params->wave_size);
            unsigned x_mask = cross32 ? 0x24 : 0x18;
            for (unsigned i = 0; i < num_comps; i++) {
               nir_def *comp = components[i];
               nir_def *compx = shuffle_xor_imm(b, comp, x_mask);
               components[i] = nir_bcsel(b, cond, comp, compx);
            }
         }
      }
   }

   return nir_vec(b, components, num_comps);
}

static bool
lower_cmat_convert_transpose(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr);
   nir_deref_instr *dst_deref = nir_src_as_deref(intr->src[0]);
   nir_deref_instr *src_deref = nir_src_as_deref(intr->src[1]);
   struct glsl_cmat_description dst_desc = *glsl_get_cmat_description(dst_deref->type);
   struct glsl_cmat_description src_desc = *glsl_get_cmat_description(src_deref->type);
   nir_def *src = radv_nir_load_cmat(b, params, intr->src[1].ssa);

   bool sat = false;
   const bool transpose = intr->intrinsic == nir_intrinsic_cmat_transpose;

   enum glsl_cmat_use dst_use = dst_desc.use;
   enum glsl_cmat_use src_use = src_desc.use;

   enum glsl_base_type dst_element_type = dst_desc.element_type;
   enum glsl_base_type src_element_type = src_desc.element_type;

   if (transpose) {
      /* NV_cmat2 only support acc -> b transpose, but we can handle any transpose except acc -> acc. */
      if (dst_use == GLSL_CMAT_USE_A) {
         dst_use = GLSL_CMAT_USE_B;
      } else if (dst_use == GLSL_CMAT_USE_B) {
         dst_use = GLSL_CMAT_USE_A;
      } else if (dst_use == GLSL_CMAT_USE_ACCUMULATOR) {
         if (src_use == GLSL_CMAT_USE_A)
            src_use = GLSL_CMAT_USE_B;
         else if (src_use == GLSL_CMAT_USE_B)
            src_use = GLSL_CMAT_USE_A;
         else
            UNREACHABLE("unsupported transpose");
      }
   } else {
      sat = nir_intrinsic_saturate(intr);
      nir_cmat_signed cmat_signed_mask = nir_intrinsic_cmat_signed_mask(intr);

      dst_element_type =
         glsl_apply_signedness_to_base_type(dst_element_type, cmat_signed_mask & NIR_CMAT_RESULT_SIGNED);
      src_element_type = glsl_apply_signedness_to_base_type(src_element_type, cmat_signed_mask & NIR_CMAT_A_SIGNED);
   }

   unsigned dst_mul = radv_nir_cmat_length_mul(dst_desc, params);
   unsigned src_mul = radv_nir_cmat_length_mul(src_desc, params);

   if (src_mul > dst_mul) {
      nir_def *components[NIR_MAX_VEC_COMPONENTS];
      unsigned scale = src_mul / dst_mul;
      for (unsigned i = 0; i * scale < src->num_components; ++i) {
         components[i] = nir_channel(b, src, i * scale);
      }
      src = nir_vec(b, components, src->num_components / scale);
   }

   if (radv_nir_cmat_bits(src_desc) <= radv_nir_cmat_bits(dst_desc))
      src = convert_use(b, src, src_use, dst_use, params);

   nir_def *ret = convert_base_type(b, src, src_element_type, dst_element_type, sat);

   if (radv_nir_cmat_bits(src_desc) > radv_nir_cmat_bits(dst_desc))
      ret = convert_use(b, ret, src_use, dst_use, params);

   if (dst_mul > src_mul) {
      nir_def *components[NIR_MAX_VEC_COMPONENTS];
      unsigned scale = dst_mul / src_mul;
      for (unsigned i = 0; i < ret->num_components; ++i) {
         components[i * scale] = nir_channel(b, ret, i);
         for (unsigned j = 1; j < scale; j++)
            components[i * scale + j] = nir_undef(b, 1, ret->bit_size);
      }
      ret = nir_vec(b, components, ret->num_components * scale);
   }

   nir_store_deref(b, dst_deref, ret, nir_component_mask(ret->num_components));
   nir_instr_remove(&intr->instr);
   b->fp_math_ctrl = nir_fp_fast_math;
   return true;
}

static bool
lower_cmat_unary_op(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr);
   nir_def *src = radv_nir_load_cmat(b, params, intr->src[1].ssa);
   nir_op op = nir_intrinsic_alu_op(intr);
   nir_def *ret = nir_build_alu1(b, op, src);
   nir_store_deref(b, nir_src_as_deref(intr->src[0]), ret, nir_component_mask(ret->num_components));
   nir_instr_remove(&intr->instr);
   b->fp_math_ctrl = nir_fp_fast_math;
   return true;
}

static bool
lower_cmat_scalar_op(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr);
   nir_def *src1 = radv_nir_load_cmat(b, params, intr->src[1].ssa);
   nir_op op = nir_intrinsic_alu_op(intr);
   nir_def *ret = nir_build_alu2(b, op, src1, intr->src[2].ssa);
   nir_store_deref(b, nir_src_as_deref(intr->src[0]), ret, nir_component_mask(ret->num_components));
   nir_instr_remove(&intr->instr);
   b->fp_math_ctrl = nir_fp_fast_math;
   return true;
}

static bool
lower_cmat_binary_op(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr);
   nir_def *src1 = radv_nir_load_cmat(b, params, intr->src[1].ssa);
   nir_def *src2 = radv_nir_load_cmat(b, params, intr->src[2].ssa);
   nir_op op = nir_intrinsic_alu_op(intr);
   nir_def *ret = nir_build_alu2(b, op, src1, src2);
   nir_store_deref(b, nir_src_as_deref(intr->src[0]), ret, nir_component_mask(ret->num_components));
   nir_instr_remove(&intr->instr);
   b->fp_math_ctrl = nir_fp_fast_math;
   return true;
}

static bool
lower_cmat_bitcast(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params)
{
   nir_def *src1 = radv_nir_load_cmat(b, params, intr->src[1].ssa);
   nir_store_deref(b, nir_src_as_deref(intr->src[0]), src1, nir_component_mask(src1->num_components));
   nir_instr_remove(&intr->instr);
   return true;
}

static bool
lower_cmat_copy(nir_builder *b, nir_intrinsic_instr *intr)
{
   nir_build_copy_deref(b, intr->src[0].ssa, intr->src[1].ssa);
   nir_instr_remove(&intr->instr);
   return true;
}

static bool
lower_cmat_reduce_finish_call(nir_builder *b, nir_cmat_call_instr *call, const lower_cmat_params *params)
{
   nir_deref_instr *dst_deref = nir_src_as_deref(call->params[0]);
   nir_deref_instr *src0_deref = nir_src_as_deref(call->params[1]);
   nir_function *fnptr = call->callee;
   nir_cmat_reduce reduce = nir_cmat_call_reduce_flags(call);
   struct glsl_cmat_description src_desc = *glsl_get_cmat_description(src0_deref->type);

   nir_def *src0 = radv_nir_load_cmat(b, params, call->params[1].ssa);
   nir_def *src1 = radv_nir_load_cmat(b, params, call->params[2].ssa);

   unsigned length = radv_nir_cmat_length(src_desc, params);
   unsigned mul = radv_nir_cmat_length_mul(src_desc, params);
   nir_def *vars[NIR_MAX_VEC_COMPONENTS];

   assert(src_desc.use == GLSL_CMAT_USE_ACCUMULATOR);

   if (mul > 1) {
      for (unsigned i = 0; i < length; ++i)
         if (i % mul != 0)
            vars[i] = nir_undef(b, 1, radv_nir_cmat_bits(src_desc));
   }

   if (reduce & NIR_CMAT_REDUCE_COLUMN) {
      nir_variable *col_tmp = nir_local_variable_create(b->impl, glsl_get_bare_type(fnptr->params[0].type), "col_tmp");
      /* All of the rows contains the same data, so just reduce both first rows. */
      nir_def *row_accum0 = nir_channel(b, src0, 0);
      nir_def *row_accum1 = nir_channel(b, src1, 0);

      nir_deref_instr *col_tmp_deref = nir_build_deref_var(b, col_tmp);

      nir_call(b, fnptr, &col_tmp_deref->def, row_accum0, row_accum1);

      nir_def *first_col = nir_load_deref(b, col_tmp_deref);

      for (unsigned i = 0; i < length / mul; i++)
         vars[i * mul] = first_col;
   } else if (reduce & NIR_CMAT_REDUCE_ROW) {
      /* All of the columns contain the same data, but not the rows, so across all components. */
      for (unsigned i = 0; i < length / mul; ++i) {
         nir_def *row0_accum = nir_channel(b, src0, i * mul);
         nir_def *row1_accum = nir_channel(b, src1, i * mul);

         nir_variable *row_tmp = nir_local_variable_create(b->impl, glsl_get_bare_type(fnptr->params[0].type), "row_tmp");
         nir_deref_instr *row_tmp_deref = nir_build_deref_var(b, row_tmp);

         nir_call(b, fnptr, &row_tmp_deref->def, row0_accum, row1_accum);

         nir_def *row = nir_load_deref(b, row_tmp_deref);
         vars[i * mul] = row;
      }
   }

   nir_def *mat = nir_vec(b, vars, length);
   nir_store_deref(b, dst_deref, mat, nir_component_mask(mat->num_components));
   nir_instr_remove(&call->instr);
   return true;
}

static bool
lower_cmat_reduce_call(nir_builder *b, nir_cmat_call_instr *call, const lower_cmat_params *params)
{
   nir_deref_instr *dst_deref = nir_src_as_deref(call->params[0]);
   nir_deref_instr *src_deref = nir_src_as_deref(call->params[1]);
   struct glsl_cmat_description src_desc = *glsl_get_cmat_description(src_deref->type);
   nir_def *src = radv_nir_load_cmat(b, params, call->params[1].ssa);
   nir_function *fnptr = call->callee;

   nir_cmat_reduce reduce = nir_cmat_call_reduce_flags(call);

   unsigned length = radv_nir_cmat_length(src_desc, params);
   unsigned mul = radv_nir_cmat_length_mul(src_desc, params);
   nir_def *vars[16];

   assert(src_desc.use == GLSL_CMAT_USE_ACCUMULATOR);

   for (unsigned i = 0; i < length; ++i) {
      if (mul > 1 && (i % mul != 0))
         vars[i] = nir_undef(b, 1, radv_nir_cmat_bits(src_desc));
      else
         vars[i] = nir_channel(b, src, i);
   }

   if (reduce & NIR_CMAT_REDUCE_COLUMN) {
      nir_variable *col_tmp = nir_local_variable_create(b->impl, glsl_get_bare_type(fnptr->params[0].type), "col_tmp");

      nir_deref_instr *col_tmp_deref = nir_build_deref_var(b, col_tmp);
      nir_store_deref(b, col_tmp_deref, vars[0], 1);

      for (unsigned i = 1; i < (length / mul); i++) {
         nir_def *col_accum_val = nir_load_deref(b, col_tmp_deref);
         nir_call(b, fnptr, &col_tmp_deref->def, col_accum_val, vars[i * mul]);
      }

      for (unsigned x_mask = 16; x_mask < params->wave_size; x_mask *= 2) {
         nir_def *col_accum_val = nir_load_deref(b, col_tmp_deref);
         nir_def *other_row = shuffle_xor_imm(b, col_accum_val, x_mask);
         nir_call(b, fnptr, &col_tmp_deref->def, col_accum_val, other_row);
      }
      for (unsigned i = 0; i < length / mul; i++)
         vars[i * mul] = nir_load_deref(b, col_tmp_deref);
   }

   if (reduce & NIR_CMAT_REDUCE_ROW) {
      for (unsigned i = 0; i < length / mul; ++i) {
         nir_def *row_accum = vars[i * mul];

         nir_variable *row_tmp = nir_local_variable_create(b->impl, glsl_get_bare_type(fnptr->params[0].type), "row_tmp");
         nir_deref_instr *row_tmp_deref = nir_build_deref_var(b, row_tmp);
         nir_store_deref(b, row_tmp_deref, row_accum, 1);

         /* 4 steps for 16 wide reduction */
         for (unsigned j = 1; j < 16; j *= 2) {
            nir_def *prev_row_accum_val = nir_load_deref(b, row_tmp_deref);

            nir_def *this_row = shuffle_xor_imm(b, prev_row_accum_val, j);

            nir_call(b, fnptr, &row_tmp_deref->def, prev_row_accum_val, this_row);
         }
         row_tmp_deref = nir_build_deref_var(b, row_tmp);
         vars[i * mul] = nir_load_deref(b, row_tmp_deref);
      }
   }

   /* this should be lowered earlier */
   assert(!(reduce & NIR_CMAT_REDUCE_2X2));
   nir_def *mat = nir_vec(b, vars, length);
   nir_store_deref(b, dst_deref, mat, nir_component_mask(mat->num_components));
   nir_instr_remove(&call->instr);
   return true;
}

static bool
lower_cmat_reduce_2x2_call(nir_builder *b, nir_cmat_call_instr *call, const lower_cmat_params *params)
{
   nir_deref_instr *dst_deref = nir_src_as_deref(call->params[0]);
   nir_deref_instr *src_deref = nir_src_as_deref(call->params[1]);
   struct glsl_cmat_description src_desc = *glsl_get_cmat_description(src_deref->type);
   nir_function *fnptr = call->callee;

   unsigned length = radv_nir_cmat_length(src_desc, params);
   unsigned mul = radv_nir_cmat_length_mul(src_desc, params);

   assert(src_desc.use == GLSL_CMAT_USE_ACCUMULATOR);

   nir_def *src_components[4][NIR_MAX_VEC_COMPONENTS];
   for (unsigned m = 0; m < 4; m++) {
      nir_def *src = radv_nir_load_cmat(b, params, call->params[m + 1].ssa);
      for (unsigned i = 0; i < length / mul; i++) {
         src_components[m][i] = nir_channel(b, src, i * mul);
      }
   }

   nir_variable *qd_tmp = nir_local_variable_create(b->impl, glsl_get_bare_type(fnptr->params[0].type), "qd_tmp");
   nir_deref_instr *qd_tmp_deref = nir_build_deref_var(b, qd_tmp);

   nir_def *low16 = nir_inverse_ballot_imm(b, 0xffff0000ffff, params->wave_size);
   for (unsigned m = 0; m < 4; m++) {
      for (unsigned i = 0; i < length / mul / 2; i++) {
         if (params->gfx_level >= GFX12) {
            /* The neighboring row is in the VGPR next to us */
            nir_call(b, fnptr, &qd_tmp_deref->def, src_components[m][i * 2], src_components[m][i * 2 + 1]);
            src_components[m][i] = nir_load_deref(b, qd_tmp_deref);
         } else {
            /* The neighboring row is in the upper 16 lanes */
            for (unsigned j = 0; j < 2; j++) {
               nir_def *other_row = shuffle_xor_imm(b, src_components[m][i * 2 + j], 16);
               nir_call(b, fnptr, &qd_tmp_deref->def, src_components[m][i * 2 + j], other_row);
               src_components[m][i * 2 + j] = nir_load_deref(b, qd_tmp_deref);
            }

            /* Fill the duplicated space with the rows from the next VGPR */
            src_components[m][i] = nir_bcsel(b, low16, src_components[m][i * 2], src_components[m][i * 2 + 1]);
         }

         /* Get data from the neighboring column */
         nir_def *other_col = shuffle_xor_imm(b, src_components[m][i], 1);
         nir_call(b, fnptr, &qd_tmp_deref->def, src_components[m][i], other_col);
         src_components[m][i] = nir_load_deref(b, qd_tmp_deref);
      }
   }

   nir_def *vars[NIR_MAX_VEC_COMPONENTS];

   /* Combine the matrices into one matrix. */
   nir_def *even = nir_inverse_ballot_imm(b, 0x5555555555555555, params->wave_size);
   for (unsigned m = 0; m < 2; m++) {
      for (unsigned i = 0; i < length / mul / 2; i++) {
         nir_def *m0_comp = src_components[m * 2][i];
         nir_def *m1_comp = src_components[m * 2 + 1][i];

         nir_def *combined = nir_bcsel(b, even, m0_comp, m1_comp);
         vars[m * (length / mul / 2) + i] = combined;
      }
   }

   /* In each row, every second element needs to be moved to the end, and the others to the front.
    * This can be done with v_permlane16_b32.
    */
   nir_def *perm_low = nir_imm_int(b, 0xeca86420);
   nir_def *perm_high = nir_imm_int(b, 0xfdb97531);
   for (unsigned i = 0; i < length / mul; i++)
      vars[i] = nir_lane_permute_16_amd(b, vars[i], perm_low, perm_high);

   if (params->gfx_level >= GFX12) {
      /* For GFX12, we still have to swap the row(s) in upper half coming from the bottom two
       * matrices with low row(s) in the from other two matrices.
       */
      nir_def *low = nir_inverse_ballot_imm(b, BITFIELD64_MASK(params->wave_size / 2), params->wave_size);
      for (unsigned i = 0; i < length / mul / 2; i++) {
         nir_def **comp0 = &vars[i];
         nir_def **comp1 = &vars[length / mul / 2 + i];
         nir_def *comp0x = shuffle_xor_imm(b, *comp0, params->wave_size / 2);
         nir_def *comp1x = shuffle_xor_imm(b, *comp1, params->wave_size / 2);
         *comp0 = nir_bcsel(b, low, *comp0, comp1x);
         *comp1 = nir_bcsel(b, low, comp0x, *comp1);
      }
   }

   if (params->wave_size == 64) {
      /* For wave64, we also need to swap the inner two rows. */
      nir_def *keep = nir_inverse_ballot_imm(b, 0xffff00000000ffff, params->wave_size);
      for (unsigned i = 0; i < length / mul; i++)
         vars[i] = nir_bcsel(b, keep, vars[i], shuffle_xor_imm(b, vars[i], 0x30));
   }

   for (int i = length / mul - 1; i >= 0; i--) {
      vars[i * mul] = vars[i];
      for (unsigned j = 1; j < mul; j++)
         vars[i * mul + j] = nir_undef(b, 1, radv_nir_cmat_bits(src_desc));
   }

   nir_def *mat = nir_vec(b, vars, length);
   nir_store_deref(b, dst_deref, mat, nir_component_mask(mat->num_components));
   nir_instr_remove(&call->instr);
   return true;
}

static bool
lower_cmat_deref(nir_deref_instr *deref, struct hash_table *type_map, const lower_cmat_params *params)
{
   const struct glsl_type *new_type = radv_nir_translate_matrix_type(deref->type, type_map, params);
   if (new_type != deref->type) {
      deref->type = new_type;
      return true;
   }
   return false;
}

static bool
lower_cmat_per_element_op(nir_builder *b, nir_cmat_call_instr *call, const lower_cmat_params *params)
{
   nir_def *src = radv_nir_load_cmat(b, params, call->params[3].ssa);
   nir_deref_instr *dst_deref = nir_src_as_deref(call->params[0]);
   nir_function *fnptr = call->callee;

   struct glsl_cmat_description desc = *glsl_get_cmat_description(dst_deref->type);

   nir_variable *elem_tmp = nir_local_variable_create(b->impl, glsl_get_cmat_element(dst_deref->type), "elemtmp");
   nir_deref_instr *elem_deref = nir_build_deref_var(b, elem_tmp);
   nir_def *local_idx = nir_load_subgroup_invocation(b);
   nir_def *inner_idx = nir_iand_imm(b, local_idx, 15);
   unsigned length = radv_nir_cmat_length(desc, params);
   unsigned mul = radv_nir_cmat_length_mul(desc, params);
   unsigned lanes_per_iter = desc.use == GLSL_CMAT_USE_ACCUMULATOR ? params->wave_size : 16;
   nir_def *base_row = radv_get_base_row(b, desc, params, local_idx);
   nir_def *vars[16];

   if (mul > 1) {
      for (unsigned i = 0; i < length; ++i)
         if (i % mul != 0)
            vars[i] = nir_undef(b, 1, radv_nir_cmat_bits(desc));
   }

   for (unsigned i = 0; i < length / mul; i++) {
      nir_def *src_elem = nir_channel(b, src, i * mul);
      nir_call_instr *new_call = nir_call_instr_create(b->shader, fnptr);
      uint32_t row_iter;

      if (params->gfx_level >= GFX12) {
         row_iter = i;
      } else {
         row_iter = i * lanes_per_iter / 16;
      }

      nir_def *row_val = nir_iadd_imm(b, base_row, row_iter);
      nir_def *col_val = inner_idx;

      if (desc.use == GLSL_CMAT_USE_A)
         SWAP(col_val, row_val);

      row_val = nir_iadd(b, call->params[1].ssa, row_val);
      col_val = nir_iadd(b, call->params[2].ssa, col_val);

      new_call->params[0] = nir_src_for_ssa(&elem_deref->def);
      new_call->params[1] = nir_src_for_ssa(row_val);
      new_call->params[2] = nir_src_for_ssa(col_val);
      new_call->params[3] = nir_src_for_ssa(src_elem);

      for (unsigned p = 4; p < call->num_params; p++) {
         nir_deref_instr *deref = nir_src_as_deref(call->params[p]);
         nir_def *def = call->params[p].ssa;
         if (deref) {
            if (glsl_type_is_cmat(deref->type)) {
               def = nir_build_load_deref(b, radv_nir_cmat_length(desc, params), radv_nir_cmat_bits(desc), def);
               def = nir_channel(b, def, i * mul);
            }
         }
         new_call->params[p] = nir_src_for_ssa(def);
      }
      nir_builder_instr_insert(b, &new_call->instr);
      vars[i * mul] = nir_build_load_deref(b, 1, radv_nir_cmat_bits(desc), &elem_deref->def, 0);
   }

   nir_def *mat = nir_vec(b, vars, length);
   nir_store_deref(b, dst_deref, mat, nir_component_mask(src->num_components));

   nir_instr_remove(&call->instr);
   return true;
}

bool
radv_nir_lower_cooperative_matrix(nir_shader *shader, enum amd_gfx_level gfx_level, unsigned wave_size)
{
   bool progress = false;

   if (!shader->info.cs.has_cooperative_matrix)
      return false;

   const lower_cmat_params params = {
      .gfx_level = gfx_level,
      .wave_size = wave_size,
   };

   struct nir_function *func = (struct nir_function *)exec_list_get_head_const(&shader->functions);
   struct hash_table *type_map = _mesa_pointer_hash_table_create(NULL);

   nir_foreach_variable_with_modes (var, shader, nir_var_shader_temp) {
      const struct glsl_type *new_type = radv_nir_translate_matrix_type(var->type, type_map, &params);
      if (new_type != var->type) {
         var->type = new_type;
         progress = true;
      }
   }

   nir_foreach_function_temp_variable (var, func->impl) {
      const struct glsl_type *new_type = radv_nir_translate_matrix_type(var->type, type_map, &params);
      if (new_type != var->type) {
         var->type = new_type;
         progress = true;
      }
   }

   nir_builder b = nir_builder_create(func->impl);

   /* Iterate in reverse order so that lowering can still use the matrix types from the derefs before we change it. */
   nir_foreach_block_reverse (block, func->impl) {
      nir_foreach_instr_reverse_safe (instr, block) {
         b.cursor = nir_before_instr(instr);

         switch (instr->type) {
         case nir_instr_type_intrinsic: {
            nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
            switch (intr->intrinsic) {
            case nir_intrinsic_cmat_length:
               progress |= lower_cmat_length(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_extract:
               progress |= lower_cmat_extract(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_insert:
               progress |= lower_cmat_insert(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_construct:
               progress |= lower_cmat_construct(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_load:
            case nir_intrinsic_cmat_store:
               progress |= lower_cmat_load_store(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_muladd:
               progress |= lower_cmat_muladd(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_transpose:
            case nir_intrinsic_cmat_convert:
               progress |= lower_cmat_convert_transpose(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_unary_op:
               progress |= lower_cmat_unary_op(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_scalar_op:
               progress |= lower_cmat_scalar_op(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_binary_op:
               progress |= lower_cmat_binary_op(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_bitcast:
               progress |= lower_cmat_bitcast(&b, intr, &params);
               break;
            case nir_intrinsic_cmat_copy:
               progress |= lower_cmat_copy(&b, intr);
               break;
            default:
               continue;
            }
            break;
         }
         case nir_instr_type_deref:
            progress |= lower_cmat_deref(nir_instr_as_deref(instr), type_map, &params);
            break;
         case nir_instr_type_cmat_call: {
            nir_cmat_call_instr *call = nir_instr_as_cmat_call(instr);
            switch (call->op) {
            case nir_cmat_call_op_reduce_finish:
               progress |= lower_cmat_reduce_finish_call(&b, call, &params);
               break;
            case nir_cmat_call_op_reduce:
               progress |= lower_cmat_reduce_call(&b, call, &params);
               break;
            case nir_cmat_call_op_reduce_2x2:
               progress |= lower_cmat_reduce_2x2_call(&b, call, &params);
               break;
            case nir_cmat_call_op_per_element_op:
               progress |= lower_cmat_per_element_op(&b, call, &params);
               break;
            default:
               break;
            }
            break;
         }
         default:
            continue;
         }
      }
   }

   _mesa_hash_table_destroy(type_map, NULL);

   nir_foreach_function_impl(fnim, shader)
      nir_progress(progress, fnim, 0);

   return nir_progress(progress, func->impl, 0);
}

static bool
apply_component_mods(nir_scalar *comp, unsigned num_comps, unsigned stride, nir_op alu_op)
{
   for (unsigned i = 0; i < num_comps; i++) {
      nir_scalar s = comp[i * stride];
      if (!nir_scalar_is_alu(s) || nir_scalar_alu_op(s) != alu_op)
         return false;
   }

   for (unsigned i = 0; i < num_comps; i++)
      comp[i * stride] = nir_scalar_chase_alu_src(comp[i * stride], 0);

   return true;
}

/* Apply neg_lo/neg_hi modifiers to A/B and neg/abs to C. */
static bool
opt_cmat_modifiers(nir_builder *b, nir_intrinsic_instr *intrin, enum amd_gfx_level gfx_level, unsigned src_idx)
{
   unsigned length_mul = src_idx == 2 && intrin->src[2].ssa->bit_size == 16 && gfx_level < GFX12 ? 2 : 1;
   nir_scalar comp[NIR_MAX_VEC_COMPONENTS] = {0};
   nir_def *src = intrin->src[src_idx].ssa;

   for (unsigned i = 0; i < src->num_components; i += length_mul)
      comp[i] = nir_scalar_resolved(src, i);

   unsigned neg_lo = nir_intrinsic_neg_lo_amd(intrin);
   unsigned neg_hi = nir_intrinsic_neg_hi_amd(intrin);

   bool progress = false;
   if (src_idx == 2) {
      unsigned num_comp = src->num_components / length_mul;
      if (apply_component_mods(comp, num_comp, length_mul, nir_op_fneg)) {
         neg_lo ^= (~neg_hi) & BITFIELD_BIT(src_idx);
         progress = true;
      }
      if (apply_component_mods(comp, num_comp, length_mul, nir_op_fabs)) {
         neg_hi |= BITFIELD_BIT(src_idx);
         progress = true;
      }
   } else {
      unsigned num_comp = src->num_components / 2;
      if (apply_component_mods(comp, num_comp, 2, nir_op_fneg)) {
         neg_lo ^= BITFIELD_BIT(src_idx);
         progress = true;
      }
      if (apply_component_mods(comp + 1, num_comp, 2, nir_op_fneg)) {
         neg_hi ^= BITFIELD_BIT(src_idx);
         progress = true;
      }
   }

   if (!progress)
      return false;

   nir_intrinsic_set_neg_lo_amd(intrin, neg_lo);
   nir_intrinsic_set_neg_hi_amd(intrin, neg_hi);

   /* Avoid creating a new vec if we don't have to. */
   nir_def *new_src = comp[0].def;
   for (unsigned i = 0; i < src->num_components; i += length_mul) {
      if (comp[i].def != new_src || comp[i].comp != i) {
         new_src = NULL;
         break;
      }
   }

   if (!new_src) {
      b->cursor = nir_before_instr(&intrin->instr);
      if (length_mul > 1) {
         nir_scalar undef = nir_get_scalar(nir_undef(b, 1, src->bit_size), 0);
         for (unsigned i = 0; i < src->num_components; i += length_mul) {
            for (unsigned j = 1; j < length_mul; j++)
               comp[i + j] = undef;
         }
      }

      new_src = nir_vec_scalars(b, comp, src->num_components);
   }

   nir_src_rewrite(&intrin->src[src_idx], new_src);
   return true;
}

static bool
opt_cmat(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
{
   enum amd_gfx_level gfx_level = *(enum amd_gfx_level *)data;

   if (intrin->intrinsic != nir_intrinsic_cmat_muladd_amd)
      return false;

   enum glsl_base_type a_type = nir_intrinsic_src_base_type(intrin);

   if (glsl_base_type_is_integer(a_type))
      return false;

   bool progress = false;

   if (a_type == GLSL_TYPE_FLOAT16) {
      for (unsigned i = 0; i < 2; i++)
         progress |= opt_cmat_modifiers(b, intrin, gfx_level, i);
   }

   if (a_type == GLSL_TYPE_FLOAT16 || intrin->def.bit_size == 32)
      progress |= opt_cmat_modifiers(b, intrin, gfx_level, 2);

   return progress;
}

bool
radv_nir_opt_cooperative_matrix(nir_shader *shader, enum amd_gfx_level gfx_level)
{
   return nir_shader_intrinsics_pass(shader, opt_cmat, nir_metadata_control_flow, &gfx_level);
}
