Improve type constraints for AIEVec_MulElemOp #1487

muradq-amd · 2024-05-15T08:44:47Z

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 · 2024-05-15T12:20:48Z

include/aie/Dialect/AIEVec/IR/AIEVecTypeConstraints.td

+def AIE2MulElemLHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>], 
+            "a vector compatible with a lhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;
+
+def AIE2MulElemRHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>],
+            "a vector compatible with a rhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;


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.

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

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?

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.

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 · 2024-05-15T12:23:36Z

include/aie/Dialect/AIEVec/IR/AIEVecTypeConstraints.td

+class IsValidAIE2MulElemShapeAndType<string lhs, string rhs, string acc> :
+  PredOpTrait<lhs # " x " # rhs # " = " # acc # " is a valid AIE2 " #
+              "element-wise multiply and accumulate op",
+              Or<[VectorTypesMatch<lhs, VectorOfShapeAndType<[32], I8>,
+                                   rhs, VectorOfShapeAndType<[32], I8>,
+                                   acc, VectorOfShapeAndType<[32], I8>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[32], I8>,
+                                   rhs, VectorOfShapeAndType<[32], I8>,
+                                   acc, VectorOfShapeAndType<[32], I32>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[32], I16>,
+                                   rhs, VectorOfShapeAndType<[32], I16>,
+                                   acc, VectorOfShapeAndType<[32], I16>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[32], I16>,
+                                   rhs, VectorOfShapeAndType<[32], I16>,
+                                   acc, VectorOfShapeAndType<[32], I32>>,
+
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[16], I32>,
+                                   rhs, VectorOfShapeAndType<[16], I32>,
+                                   acc, VectorOfShapeAndType<[16], I32>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[16], BF16>,
+                                   rhs, VectorOfShapeAndType<[16], BF16>,
+                                   acc, VectorOfShapeAndType<[16], BF16>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[16], BF16>,
+                                   rhs, VectorOfShapeAndType<[16], BF16>,
+                                   acc, VectorOfShapeAndType<[16], F32>>,
+                  VectorTypesMatch<lhs, VectorOfShapeAndType<[16], F32>,
+                                   rhs, VectorOfShapeAndType<[16], F32>,
+                                   acc, VectorOfShapeAndType<[16], F32>>]>>;


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 · 2024-05-15T22:02:46Z

include/aie/Dialect/AIEVec/IR/AIEVecOps.td

+         `vector<32xi8>`    | `vector<32xi8>`    | `vector<32xi8>`
+         `vector<32xi8>`    | `vector<32xi8>`    | `vector<32xi32>`
+         `vector<32xi16>`   | `vector<32xi16>`   | `vector<32xi16>`
+         `vector<32xi16>`   | `vector<32xi16>`   | `vector<32xi32>`
+         `vector<16xi32>`   | `vector<16xi32>`   | `vector<16xi32>`
+         `vector<16xbf16>`  | `vector<16xbf16>`  | `vector<16xbf16>`
+         `vector<16xbf16>`  | `vector<16xbf16>`  | `vector<16xf32>`
+         `vector<16xf32>`   | `vector<16xf32>`   | `vector<16xf32>`


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 · 2024-05-15T22:22:35Z

include/aie/Dialect/AIEVec/IR/AIEVecTypeConstraints.td

+def AIE2MulElemLHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>], 
+            "a vector compatible with a lhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;
+
+def AIE2MulElemRHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>],
+            "a vector compatible with a rhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;


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 · 2024-05-15T22:26:45Z

include/aie/Dialect/AIEVec/IR/AIEVecTypeConstraints.td

+def AIE2MulElemLHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>], 
+            "a vector compatible with a lhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;
+
+def AIE2MulElemRHS :
+  AnyTypeOf<[VectorOfShapeAndType<[32], I8>,
+             VectorOfShapeAndType<[32], I16>,
+             VectorOfShapeAndType<[16], I32>,
+             VectorOfShapeAndType<[16], BF16>,
+             VectorOfShapeAndType<[16], F32>],
+            "a vector compatible with a rhs operand of element-wise multiply and "
+            # "accumulate",
+            "::mlir::VectorType">;


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 · 2024-05-20T21:16:13Z

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 · 2024-05-21T17:02:31Z

platforms/boards

This shouldn't be in the commit.

jsetoain · 2024-05-21T17:05:09Z

include/aie/Dialect/AIEVec/IR/AIEVecTypeConstraints.td

@@ -101,6 +101,10 @@ class ShapesCompatibleWithContraction<string lhs, string rhs, string acc> :
 class VectorType<string name> : StrFunc<"cast<VectorType>($" # name #
                                        ".getType())">;

+class VectorElementType<string name> : 


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 · 2024-05-21T17:06:55Z

include/aie/Dialect/AIEVec/IR/AIEVecOps.td

+  Arguments<(ins AnyVector:$lhs,
+                 AnyVector:$rhs)>,
+  Results<(outs AnyVector:$acc)> {


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

muradq-amd requested review from jsetoain, david-vc and jamestcl-amd May 15, 2024 08:44

muradq-amd requested review from makslevental, jackl-xilinx and jgmelber as code owners May 15, 2024 08:44

jsetoain requested changes May 15, 2024

View reviewed changes

jamestcl-amd reviewed May 15, 2024

View reviewed changes

muradq-amd requested a review from eddierichter-amd as a code owner May 17, 2024 19:50

muradq-amd requested review from jsetoain and jamestcl-amd May 20, 2024 21:17

jsetoain requested changes May 21, 2024

View reviewed changes

muradq-amd closed this May 22, 2024

muradq-amd force-pushed the 1171-improve-type-constraints-for-aievecmul_elem-and-co branch from a38d63d to 37ef519 Compare May 22, 2024 18:57

muradq-amd mentioned this pull request May 23, 2024

[aievec] Adding type constraints for aievec::mul_elem operation #1514

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve type constraints for AIEVec_MulElemOp #1487

Improve type constraints for AIEVec_MulElemOp #1487

muradq-amd commented May 15, 2024

jsetoain left a comment

jsetoain May 15, 2024

jamestcl-amd May 15, 2024

jamestcl-amd May 15, 2024

jsetoain May 16, 2024

jsetoain May 16, 2024

jsetoain May 15, 2024

jamestcl-amd May 15, 2024

jamestcl-amd May 15, 2024

jamestcl-amd May 15, 2024

muradq-amd commented May 20, 2024

jsetoain May 21, 2024

jsetoain May 21, 2024 •

edited

Loading

jsetoain May 21, 2024

	// 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");

	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();
	}

Improve type constraints for AIEVec_MulElemOp #1487

Improve type constraints for AIEVec_MulElemOp #1487

Conversation

muradq-amd commented May 15, 2024

jsetoain left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

muradq-amd commented May 20, 2024

Choose a reason for hiding this comment

jsetoain May 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jsetoain May 21, 2024 •

edited

Loading