diff --git a/src/amd/vulkan/nir/radv_nir_lower_cooperative_matrix.c b/src/amd/vulkan/nir/radv_nir_lower_cooperative_matrix.c index ac06d918915..8625ea35333 100644 --- a/src/amd/vulkan/nir/radv_nir_lower_cooperative_matrix.c +++ b/src/amd/vulkan/nir/radv_nir_lower_cooperative_matrix.c @@ -165,6 +165,221 @@ radv_get_base_row(nir_builder *b, struct glsl_cmat_description desc, const lower return base_row; } +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, + iter_deref->cast.ptr_stride, 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) { @@ -400,6 +615,145 @@ convert_use(nir_builder *b, nir_def *src, enum glsl_cmat_use src_use, enum glsl_ 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) +{ + 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); + return true; +} + +static bool +lower_cmat_unary_op(nir_builder *b, nir_intrinsic_instr *intr, const lower_cmat_params *params) +{ + 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); + return true; +} + +static bool +lower_cmat_scalar_op(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_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); + return true; +} + +static bool +lower_cmat_binary_op(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_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); + 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_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; +} + bool radv_nir_lower_cooperative_matrix(nir_shader *shader, enum amd_gfx_level gfx_level, unsigned wave_size) { @@ -443,345 +797,52 @@ radv_nir_lower_cooperative_matrix(nir_shader *shader, enum amd_gfx_level gfx_lev case nir_instr_type_intrinsic: { nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); switch (intr->intrinsic) { - case nir_intrinsic_cmat_length: { - struct glsl_cmat_description desc = nir_intrinsic_cmat_desc(intr); - unsigned len = radv_nir_cmat_length(desc, ¶ms) / radv_nir_cmat_length_mul(desc, ¶ms); - nir_def_rewrite_uses(&intr->def, nir_imm_int(&b, len)); - nir_instr_remove(instr); - progress = true; + case nir_intrinsic_cmat_length: + progress |= lower_cmat_length(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_extract: { - 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, ¶ms, intr->src[0].ssa); - - nir_def *index = intr->src[1].ssa; - index = nir_imul_imm(&b, index, radv_nir_cmat_length_mul(desc, ¶ms)); - - nir_def *elem = nir_vector_extract(&b, src0, index); - - nir_def_rewrite_uses(&intr->def, elem); - nir_instr_remove(instr); - progress = true; + case nir_intrinsic_cmat_extract: + progress |= lower_cmat_extract(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_insert: { - nir_def *src1 = radv_nir_load_cmat(&b, ¶ms, 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, ¶ms)); - - 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(instr); - progress = true; + case nir_intrinsic_cmat_insert: + progress |= lower_cmat_insert(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_construct: { - 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, ¶ms)); - - nir_store_deref(&b, dst_deref, r, nir_component_mask(r->num_components)); - nir_instr_remove(instr); - progress = true; + case nir_intrinsic_cmat_construct: + progress |= lower_cmat_construct(&b, intr, ¶ms); break; - } case nir_intrinsic_cmat_load: - case nir_intrinsic_cmat_store: { - 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 && gfx_level < GFX12 && desc.use == GLSL_CMAT_USE_ACCUMULATOR && - radv_nir_cmat_bits(desc) == 8 && 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, ¶ms); - unsigned mul = radv_nir_cmat_length_mul(desc, ¶ms); - 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 (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, ¶ms, &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, ¶ms, 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 (gfx_level >= GFX12) - align_mul /= 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 (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, - iter_deref->cast.ptr_stride, 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 (gfx_level < GFX12 && desc.use != GLSL_CMAT_USE_ACCUMULATOR) { - nir_pop_if(&b, NULL); - } - nir_instr_remove(instr); - progress = true; + case nir_intrinsic_cmat_store: + progress |= lower_cmat_load_store(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_muladd: { - nir_def *A = radv_nir_load_cmat(&b, ¶ms, intr->src[1].ssa); - nir_def *B = radv_nir_load_cmat(&b, ¶ms, intr->src[2].ssa); - nir_def *C = radv_nir_load_cmat(&b, ¶ms, 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(instr); - progress = true; + case nir_intrinsic_cmat_muladd: + progress |= lower_cmat_muladd(&b, intr, ¶ms); break; - } case nir_intrinsic_cmat_transpose: - case nir_intrinsic_cmat_convert: { - 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, ¶ms, 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, ¶ms); - unsigned src_mul = radv_nir_cmat_length_mul(src_desc, ¶ms); - - 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, ¶ms); - - 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, ¶ms); - - 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(instr); - progress = true; + case nir_intrinsic_cmat_convert: + progress |= lower_cmat_convert_transpose(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_unary_op: { - nir_def *src = radv_nir_load_cmat(&b, ¶ms, 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(instr); - progress = true; + case nir_intrinsic_cmat_unary_op: + progress |= lower_cmat_unary_op(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_scalar_op: { - nir_def *src1 = radv_nir_load_cmat(&b, ¶ms, 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(instr); - progress = true; + case nir_intrinsic_cmat_scalar_op: + progress |= lower_cmat_scalar_op(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_binary_op: { - nir_def *src1 = radv_nir_load_cmat(&b, ¶ms, intr->src[1].ssa); - nir_def *src2 = radv_nir_load_cmat(&b, ¶ms, 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(instr); - progress = true; + case nir_intrinsic_cmat_binary_op: + progress |= lower_cmat_binary_op(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_bitcast: { - nir_def *src1 = radv_nir_load_cmat(&b, ¶ms, 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(instr); - progress = true; + case nir_intrinsic_cmat_bitcast: + progress |= lower_cmat_bitcast(&b, intr, ¶ms); break; - } - case nir_intrinsic_cmat_copy: { - nir_build_copy_deref(&b, intr->src[0].ssa, intr->src[1].ssa); - nir_instr_remove(instr); - progress = true; + case nir_intrinsic_cmat_copy: + progress |= lower_cmat_copy(&b, intr); break; - } default: continue; } break; } - case nir_instr_type_deref: { - nir_deref_instr *deref = nir_instr_as_deref(instr); - const struct glsl_type *new_type = radv_nir_translate_matrix_type(deref->type, type_map, ¶ms); - if (new_type != deref->type) { - deref->type = new_type; - progress = true; - } + case nir_instr_type_deref: + progress |= lower_cmat_deref(nir_instr_as_deref(instr), type_map, ¶ms); break; - } default: continue; }