AlexIndustrial/mesa
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

radv/nir/lower_cmat: split up larger nested switches

Browse Source 
This has been annoying me for quite some while, the level of indention
makes reviewing code changes in Gitlab harder.

I think now is a good time to change this before more cmat lowering is added.

Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37002>
This commit is contained in:
Georg Lehmann
2025-08-26 09:20:41 +02:00
committed by Marge Bot
parent 047b95a8c3
commit 5a10142a9f
1 changed files with 380 additions and 319 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/amd/vulkan/nir/radv_nir_lower_cooperative_matrix.c
+380 -319
View File
@@ -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, &params) / radv_nir_cmat_length_mul(desc, &params);
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, &params);
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, &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_rewrite_uses(&intr->def, elem);
nir_instr_remove(instr);
progress = true;
case nir_intrinsic_cmat_extract:
progress |= lower_cmat_extract(&b, intr, &params);
break;
}
case nir_intrinsic_cmat_insert: {
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(instr);
progress = true;
case nir_intrinsic_cmat_insert:
progress |= lower_cmat_insert(&b, intr, &params);
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, &params));
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, &params);
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, &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 (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 (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, &params);
break;
}
case nir_intrinsic_cmat_muladd: {
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(instr);
progress = true;
case nir_intrinsic_cmat_muladd:
progress |= lower_cmat_muladd(&b, intr, &params);
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, &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(instr);
progress = true;
case nir_intrinsic_cmat_convert:
progress |= lower_cmat_convert_transpose(&b, intr, &params);
break;
}
case nir_intrinsic_cmat_unary_op: {
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(instr);
progress = true;
case nir_intrinsic_cmat_unary_op:
progress |= lower_cmat_unary_op(&b, intr, &params);
break;
}
case nir_intrinsic_cmat_scalar_op: {
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(instr);
progress = true;
case nir_intrinsic_cmat_scalar_op:
progress |= lower_cmat_scalar_op(&b, intr, &params);
break;
}
case nir_intrinsic_cmat_binary_op: {
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(instr);
progress = true;
case nir_intrinsic_cmat_binary_op:
progress |= lower_cmat_binary_op(&b, intr, &params);
break;
}
case nir_intrinsic_cmat_bitcast: {
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(instr);
progress = true;
case nir_intrinsic_cmat_bitcast:
progress |= lower_cmat_bitcast(&b, intr, &params);
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, &params);
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, &params);
break;
}
default:
continue;
}