[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.

[newling]

Great improvement!

Mostly nits, but a question about using ContractionDimensions's inferred batch dimension(s).

[newling]

Are these asserts needed, where is this assumption of <= 7 used?

[yzhang93]

Kind of. I was thinking to restrict the contraction ops to only include 2D inputs like matmul ops or 4D inputs like mmt4d (or their batch versions) because the current pack strategy is specially designed for those ops. However, I think there's barely input ops with more than 7 loops in real application. I can delete this assert for now.

[newling]

When might this be empty after you've already not received failure from linalg::inferContractionDims ?

[yzhang93]

I'm not fully sure about this, but I've seen this check everytime this function is called in IREE, e.g. https://github.com/iree-org/iree/blob/624a9fae4f8dd03becfaae69110411c84c4ed6e8/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp#L51.

I think even if linalg::inferContractionDims returns success, it doesn't guarantee all the m, n, k are non-empty. We want to make sure all these dimensions are non-empty in order to proceed.

[newling]

At the link above, it's doing the check just before calling back() on each of the vectors, so at a point where it knows that the intrinsic it's targeting has all 3. Your implicit assumption here is that you're doing some kind of matmul, but maybe we'll want to reuse this for matvec or something where there is no n dimension.

Perhaps add a comment, change the class name, or move the checks later.

[newling]

Personal preference: Use SmallVector<T> instead of SmallVector<T, 2>.

[yzhang93]

I'd prefer to use the same type as in upstream. In addition it seems not work to use SmallVector when calling

SmallVector<unsigned> mDims = contractionDims.m;

It has error: no viable conversion from 'SmallVector<[...], 2>' to 'SmallVector<[...], (default) 12>'

[newling]

Ok, I didn't realize this is the way upstream does it, that's fine then.

FWIW you could do

SmallVector mDims {contractionDims.m.begin(), contractionDims.m.end()};

[jtuyls]

LGTM

[MaheshRavishankar]

Nice!

[newling]

Could assert that batchDims is not empty, like the others.

[yzhang93]

Done.

[newling]

At the link above, it's doing the check just before calling back() on each of the vectors, so at a point where it knows that the intrinsic it's targeting has all 3. Your implicit assumption here is that you're doing some kind of matmul, but maybe we'll want to reuse this for matvec or something where there is no n dimension.

Perhaps add a comment, change the class name, or move the checks later.

[Abhishek-Varma]

LGTM!