Improve type constraints for AIEVec_MulElemOp

[muradq-amd]

This PR includes adding constrains to aievec::mul_elem on the number of lanes and types of operands and results.

Changes:

• Add constraints to the LHS and RHS operands and result (AIE2MulElemLHS, AIE2MulElemRHS, and AIE2MulElemACC)
• Add PredOpTrait constrains to allow only supported type/number of lanes combinations.

[jsetoain]

I understand you've followed the constraint scheme for aievec.matmul, but that's is a very bad template for this use case because that one was a very irregular type combination, while this one is pretty regular.

You need to check that lhs/rhs match, and that either the element type of lhs/rhs matches acc, or acc is the "wide type" (i32 for integrals, f32 for floats).

[jsetoain]

This is overkill. Lhs & Rhs types have to match, you don't need to spell out every single valid type for each of the operands.

[jamestcl-amd]

The op verifier checks that the lanes and datatypes are the same for operands. See

mlir-aie/lib/Dialect/AIEVec/IR/AIEVecOps.cpp

Lines 872 to 894 in aba1887

	// Additional checks for FMAElem op
	// Get the width of the underlying scalars of all the vectors
	Type ltype = lhsType.getElementType();
	Type rtype = rhsType.getElementType();
	Type atype = resultType.getElementType();
	unsigned ltypeWidth = ltype.getIntOrFloatBitWidth();
	unsigned rtypeWidth = rtype.getIntOrFloatBitWidth();
	unsigned atypeWidth = atype.getIntOrFloatBitWidth();
	// Checks on the number of lanes
	unsigned rhsLanes = getVectorLaneSize(rhsType);
	unsigned lhsLanes = getVectorLaneSize(lhsType);
	// lane size must match
	if (lhsLanes != rhsLanes) {
	return op.emitError("The number of lanes in lhs operand "
	"must be the same as rhs operand");
	}
	// lhs and rhs vector's element type must match
	if (ltype != rtype)
	return op.emitError("The element type of lhs and rhs "
	"operand vectors must match");

This can probably be replaced by SameTypeOperands and SameOperandsAndResultShape traits....

[jamestcl-amd]

However, I'm not sure how to better constraint the fact that the acc is i32 for i8/i16 operands, i64 for i32 operands, and f32 for bf16/f32 operands. Maybe just constraints it in the op verifier in cpp code?

[jsetoain]

The op verifier checks that the lanes and datatypes are the same for operands. See

mlir-aie/lib/Dialect/AIEVec/IR/AIEVecOps.cpp

Lines 872 to 894 in aba1887

	// Additional checks for FMAElem op
	// Get the width of the underlying scalars of all the vectors
	Type ltype = lhsType.getElementType();
	Type rtype = rhsType.getElementType();
	Type atype = resultType.getElementType();
	unsigned ltypeWidth = ltype.getIntOrFloatBitWidth();
	unsigned rtypeWidth = rtype.getIntOrFloatBitWidth();
	unsigned atypeWidth = atype.getIntOrFloatBitWidth();
	// Checks on the number of lanes
	unsigned rhsLanes = getVectorLaneSize(rhsType);
	unsigned lhsLanes = getVectorLaneSize(lhsType);
	// lane size must match
	if (lhsLanes != rhsLanes) {
	return op.emitError("The number of lanes in lhs operand "
	"must be the same as rhs operand");
	}
	// lhs and rhs vector's element type must match
	if (ltype != rtype)
	return op.emitError("The element type of lhs and rhs "
	"operand vectors must match");

Notice that the C++ verifier will be replaced by the automatically generated one.

This can probably be replaced by SameTypeOperands and SameOperandsAndResultShape traits....

Indeed.

[jsetoain]

However, I'm not sure how to better constraint the fact that the acc is i32 for i8/i16 operands, i64 for i32 operands, and f32 for bf16/f32 operands. Maybe just constraints it in the op verifier in cpp code?

For that, you can either define a new predicate ('IsValidAccumulatorTypeFor`) that verifies the accumulator type validity using simpler pre-existing predicates, or define a C++ function that makes that check and invoke it from the predicate.

[jsetoain]

This constraint alone already makes all the checks necessary for the supported types, but it's unnecessarily complex for such a regular operation.

[jamestcl-amd]

For aievec.mul_elem the accumulator is acc32 (i32) for i8/i16 and acc64 (i64) for i32. Your Accumulator column here is more like a result type for a arith.mulf/i op from what you see in the e2e unit tests. To handle different result types, we introduce either aievec.cast or aievec.srs to the result of aievec.mul_elem. See

mlir-aie/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp

Lines 702 to 712 in aba1887

	if (mulElemResultElWidth == resultElWidth) {
	rewriter.replaceOpWithNewOp<aievec::CastOp>(
	mulOp, resultType, mulElemOp.getResult(), /*isResAcc*/ false);
	} else if (mulElemResultElWidth > resultElWidth) {
	auto shiftParamOp = rewriter.create<arith::ConstantOp>(
	mulOp.getLoc(), rewriter.getI32IntegerAttr(shiftParam));
	rewriter.replaceOpWithNewOp<aievec::SRSOp>(
	mulOp, resultType, mulElemOp.getResult(), shiftParamOp.getResult());
	} else {
	return failure();
	}

[jamestcl-amd]

The op verifier checks that the lanes and datatypes are the same for operands. See

mlir-aie/lib/Dialect/AIEVec/IR/AIEVecOps.cpp

Lines 872 to 894 in aba1887

	// Additional checks for FMAElem op
	// Get the width of the underlying scalars of all the vectors
	Type ltype = lhsType.getElementType();
	Type rtype = rhsType.getElementType();
	Type atype = resultType.getElementType();
	unsigned ltypeWidth = ltype.getIntOrFloatBitWidth();
	unsigned rtypeWidth = rtype.getIntOrFloatBitWidth();
	unsigned atypeWidth = atype.getIntOrFloatBitWidth();
	// Checks on the number of lanes
	unsigned rhsLanes = getVectorLaneSize(rhsType);
	unsigned lhsLanes = getVectorLaneSize(lhsType);
	// lane size must match
	if (lhsLanes != rhsLanes) {
	return op.emitError("The number of lanes in lhs operand "
	"must be the same as rhs operand");
	}
	// lhs and rhs vector's element type must match
	if (ltype != rtype)
	return op.emitError("The element type of lhs and rhs "
	"operand vectors must match");

This can probably be replaced by SameTypeOperands and SameOperandsAndResultShape traits....

[jamestcl-amd]

However, I'm not sure how to better constraint the fact that the acc is i32 for i8/i16 operands, i64 for i32 operands, and f32 for bf16/f32 operands. Maybe just constraints it in the op verifier in cpp code?

[muradq-amd]

Thanks @jamestcl-amd, @jsetoain for the comments.
I just revised the code changes and added simplified constraints on the types and shapes of aievec::mul_elem's operands/results. Here is the list of supported type/shape combinations:

lhs	rhs	acc
<32xi8>	<32xi8>	<32xi8>
<32xi8>	<32xi8>	<32xi32>
<64xi8>	<64xi8>	<32xi32>
<32xi16>	<32xi16>	<32xi16>
<32xi16>	<32xi16>	<32xi32>
<16xi32>	<16xi32>	<16xi32>
<16xi32>	<16xi32>	<16xi64>
<16xbf16>	<16xbf16>	<16xbf16>
<16xbf16>	<16xbf16>	<16xf32>
<16xf32>	<16xf32>	<16xf32>

Please let me know if you see any thing need to be fixed. thanks.

[jsetoain]

This shouldn't be in the commit.

[jsetoain]

These already exist in upstream mlir. Check llvm-project/mlir/include/mlir/IR/OpBase.td. There's a lot of stuff there you'll find helpful.

[jsetoain]

We don't support AnyVector, you should use an operand type constraint for these, and then your op constraint will be simpler.