[LoweringStrategy] Use a more general method to fetch input dims and sizes

[yzhang93]

There is a bug in the existing codes to get M/N/K size from a matmul-like op, i.e.,

const uint64_t M = initShape[0];
const uint64_t N = initShape[1];
const uint64_t K = lhsShape[1];

Apparently, K shouldn't be lhsShape[1] if it's a matmul-tranpose-a op.
It's hard to infer the shape if the input matmul-like ops are transposed and in linalg.generic form. Or even the input has a higher number of dimensions such as mmt4d ops.

In addition, the indexing maps of matmul-like linalg.generic ops can be transposed during dispatch generation by default because of TransposeGenericOpsPass.
Example:
parallel parallel reduction parallel parallel reduction will become
parallel parallel parallel parallel reduction reduction after dispatch generation.
So if we still put the pack size/tile size at the dim for the former indexing map, it will generate wrong size for tiling.

This PR uses the upstream method linalg::inferContractionDims to infer the dim indices of M/N/K for all contraction ops.

[Abhishek-Varma]

Nice! A few comments to address.

[Abhishek-Varma]

Do we really need this check ?
This seems like something which the upstream utility should handle, if not already.

Therefore if at all we need to ensure this check, I think we can simply assert on this.

[Abhishek-Varma]

Suggested change

	for (auto i : idx) sizes.push_back(shapes[i]);
	for (unsigned i : idx) sizes.push_back(shapes[i]);

[Abhishek-Varma]

We can use linalgOp.getNumLoops(); within each function ?
In that case sending numLoops as a function argument can be dropped here and elsewhere.