[LFRic] [PSyAD] Change in operator index ordering does not match PSyAD's matmul standard

[DrTVockerodtMO]

Previously, an valid operator from PSyAD's perspective would be of the form operator( : , : , idx ), i.e.: the first two indices are full ranges. LFRic is now changing the index accessing to be like operator( idx, : , : ) which does not fit the standard imposed by PSyAD. This causes the following failure:

psyclone.psyir.transformations.transformation_error.TransformationError: Transformation Error: To use matmul2code_trans on matmul, the first two indices of the 1st argument 'operator' must be full ranges.

[DrTVockerodtMO]

This is related to #2766

[arporter]

This is a limitation of the Matmul2Code transformation:

PSyclone/src/psyclone/psyir/transformations/intrinsics/matmul2code_trans.py

Lines 261 to 271 in d865146

	# The first two indices should be ranges. This is a
	# limitation of this transformation, not the PSyIR, as it
	# would be valid to have any two indices being
	# ranges. Further, this transformation is currently
	# limited to Ranges which specify the full extent of the
	# dimension.
	if not (matrix1.is_full_range(0) and matrix1.is_full_range(1)):
	raise TransformationError(
	f"To use matmul2code_trans on matmul, the first two "
	f"indices of the 1st argument '{matrix1.name}' must be "
	f"full ranges.")

As the comment says, it should be possible to generalise it so that the ':'s appear at any location in the array reference.

[DrTVockerodtMO]

I'm thinking you might just be able to count the number of full range indices, if it's not equal to 2 then you throw a modified error message.

If it is simple enough to get the number of matrix indices out then I don't mind having a go at this!

Looks like len(matrix1.children) can give the span of the index set, iterating over that range may also allow us to remove the code beneath it checking that the remaining indices must not be full ranges so this works quite nicely to consolidate the code!

[arporter]

This is just the validation. You will need to update the implementation too:

PSyclone/src/psyclone/psyir/transformations/intrinsics/matmul2code_trans.py

Lines 462 to 469 in d865146

	# Create "matrix2(ii,j)"
	m2_array_reference = _create_matrix_ref(matrix2.symbol,
	[ii_loop_sym, j_loop_sym],
	matrix2.children[2:])
	# Create "matrix1(i,ii)"
	m1_array_reference = _create_matrix_ref(matrix1.symbol,
	[i_loop_sym, ii_loop_sym],
	matrix1.children[2:])

We could take the opportunity to remove the requirement for full ranges too but I suggest that be a separate PR as it's not required to fix this problem (and presumably it needs fixing for 3.0?).

[DrTVockerodtMO]

I think the validation would still throw an error since the forward code will use operators that fail that validation. Of course, only fixing the validation ought to run into errors if it is only operating on matrix1.children[2:] say (since some of those ranges will be full whereas currently they are not), so the resulting adjoint ought to not be correct. Both would need to be fixed to run PSyAD without needing to patch it.

For 3.0 it may be part of the work of #2766, so far I have only encountered this once when trying to use PSyAD but not sure how prevalent it is. It can be patched currently so it's less pressing but as always ideally does need to be fixed.

[arporter]

For an array reference, the indices method is what you want. You'd need to count the number of Range nodes.

[arporter]

I don't think it should be part of #2766 (although may be required for it). It's a self-contained change to the transformation.