Skip to content

Conversation

@bjacob
Copy link
Collaborator

@bjacob bjacob commented Oct 29, 2025

This PR does multiple things that were easier done all at once:

  1. Harmonize *ScaledMMAAttr operand order:
    • The operand order of ScaledMMAAttr was lhs, lhs_scale, rhs, rhs_scale, while the operand order of DataTiledScaledMMAAttr was lhs, rhs, lhs_scale, rhs_scale.
    • This PR changes ScaledMMAAttr to match the DataTiledScaledMMAAttr convention. This propagates to a change of operand order in the enclosing inner_tiled ops.
  2. Drop MMAFragment:
    • There used to be a TableGen enum MMAFragment, that had unclear semantics: the enum values were sometimes used as opaque symbolic enums to refer to operand by "role", e.g. "Lhs", and sometimes used as the underlying integer values as operand indices, e.g. "Rhs == 1". This was originally OK as all MMA-like ops had the same 3 operands Lhs, Rhs, Acc. But when ScaledMMAAttr was introduced, that didn't... scale: now the preexisting enum value "Rhs == 1" didn't equal anymore the corresponding operand index under the lhs, lhs_scale, rhs, rhs_scale convention (1 != 2) and even regardless of convention, the enum value "Acc ==2" never corresponded to operand index anymore (2 != 4).
    • This PR drops MMAFragment and generalizes the use of plain integer operandIndex.
    • This is made reasonable to implement by the harmonization of operand orders (above 1.). If we tried doing this without 1., then we would feel more of a need for opaque "role" enums to replace MMAFragment without having to fix up operand indices between two conventions.
  3. Drop some logic that only existed to make up for the discrepancy between convensions and the fuzzy semantics of MMAFragment.

@kuhar
Copy link
Member

kuhar commented Oct 29, 2025

Very harmonious, this PR restores tranquility.

@kuhar kuhar removed their request for review October 29, 2025 20:24
@bjacob bjacob force-pushed the drop-mmafragment-take2 branch from 079ee69 to 305f0c2 Compare October 29, 2025 20:29
@bjacob bjacob requested a review from krzysz00 October 29, 2025 20:33
Signed-off-by: Benoit Jacob <[email protected]>
@bjacob bjacob force-pushed the drop-mmafragment-take2 branch from 305f0c2 to 1a009cf Compare October 30, 2025 18:19
Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did not participate in the discussion about what the order should be, as there are reordering changes in lit tests. The change itself looks good to me.

Comment on lines +210 to +216
using MMAIntrinsicTy = decltype(mma.getIntrinsic());
const bool isScaled = std::is_same<MMAIntrinsicTy, ScaledMMAIntrinsic>::value;
const bool isLhs = isIntrinsicLhs<MMAIntrinsicTy>(operandIdx);
const bool isRhs = isIntrinsicRhs<MMAIntrinsicTy>(operandIdx);
const bool isAcc = isIntrinsicAcc<MMAIntrinsicTy>(operandIdx);
const bool isLhsScale = isIntrinsicLhsScale<MMAIntrinsicTy>(operandIdx);
const bool isRhsScale = isIntrinsicRhsScale<MMAIntrinsicTy>(operandIdx);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks much better than magic numbers to me.

Comment on lines +1101 to +1110
ValueRange inputs = linalgOp->getOperands();

SmallVector<Type> eltTypes;
smmaKind.getElementTypes(eltTypes);
if (cast<RankedTensorType>(inputs[0].getType()).getElementType() !=
eltTypes[0] ||
cast<RankedTensorType>(inputs[2].getType()).getElementType() !=
eltTypes[2] ||
cast<RankedTensorType>(inputs[4].getType()).getElementType() !=
eltTypes[4]) {
return failure();
for (int i :
{kScaledMMAOperandLhs, kScaledMMAOperandRhs, kScaledMMAOperandAcc}) {
if (cast<RankedTensorType>(inputs[i].getType()).getElementType() !=
eltTypes[i]) {
return failure();
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks much better to me!

Copy link
Contributor

@krzysz00 krzysz00 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@bjacob bjacob merged commit b0103f6 into iree-org:main Oct 30, 2025
35 of 46 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants