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

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

@@ -313,22 +313,33 @@ def AIEVec_MulOp:
 }
 
 def AIEVec_MulElemOp:
-  AIEVec_Op<"mul_elem", [
-    Pure
-  ]>,
-  Arguments<(ins AnyVector:$lhs, AnyVector:$rhs)>,
-  Results<(outs AnyVector:$result)> {
+  AIEVec_Op<"mul_elem", [	 
+    Pure,
+    SameTypeOperands,
+    SameOperandsShape,
+    SameOperandsAndResultShape,
+    isOperandResultTypePairValidForAIE2MulElem<"lhs", "rhs", "result">
+  ]>,	  
+  Arguments<(ins
+    VectorOfLengthAndType<[16, 32], [I8, I16, I32, BF16, F32]>:$lhs, 
+    VectorOfLengthAndType<[16, 32], [I8, I16, I32, BF16, F32]>:$rhs)>,
+  Results<(outs
+    VectorOfLengthAndType<[16, 32], [I32, I64, F32]>:$result)> {
   let summary = "AIE2 vector element-wise multiply";
   let description = [{
     AMD-specific multiply operation that multiplies two 1-D vectors in the same channel.
     The vector sizes are at least 512 bits.
     `$result = `$lhs * $rhs`.
+    Currently, the following are the supported type combinations:
+        lhs                | rhs                | Accumulator
+      :------------------:|:------------------:|:-----------------:
+        `vector<32xi8>`    | `vector<32xi8>`    | `vector<32xi32>`
+        `vector<32xi16>`   | `vector<32xi16>`   | `vector<32xi32>`
+        `vector<16xi32>`   | `vector<16xi32>`   | `vector<16xi64>`
+        `vector<16xbf16>`  | `vector<16xbf16>`  | `vector<16xf32>`
+        `vector<16xf32>`   | `vector<16xf32>`   | `vector<16xf32>`'
   }];
-  let builders = [
-    OpBuilder<(ins "mlir::Value":$lhs, "mlir::Value":$rhs, "mlir::Type":$accType),
-    [{build($_builder, $_state, accType, lhs, rhs);}]>
-  ];
-}
+} 
 
 def AIEVec_BroadcastOp:
   AIEVec_Op<"broadcast", [

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

@@ -146,4 +146,18 @@
                                    rhs, VectorOfShapeAndType<[8, 4], BF16>,
                                    acc, VectorOfShapeAndType<[4, 4], F32>>]>>;
 
+
+
+class isOperandResultTypePairValidForAIE2MulElem<string lhs, string rhs, string acc> :
+  PredOpTrait<acc # " type is a valid accumulator type given the type of the" #
+              " operands.",
+    Or<[  
+      And <[ElementTypeIsPred<acc,I32>, ElementTypeIsPred<lhs,I8>]>, 
+      And <[ElementTypeIsPred<acc,I32>, ElementTypeIsPred<lhs,I16>]>, 
+      And <[ElementTypeIsPred<acc,I64>, ElementTypeIsPred<lhs,I32>]>, 
+      And <[ElementTypeIsPred<acc,F32>, ElementTypeIsPred<lhs,BF16>]>, 
+      And <[ElementTypeIsPred<acc,F32>, ElementTypeIsPred<lhs,F32>]>
+    ]>>;
+
+
 #endif

lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp

@@ -30,6 +30,7 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 
+#include "llvm/Support/raw_ostream.h"
 #include <bitset>
 #include <optional>
 #include <tuple>
@@ -678,21 +679,37 @@ struct ConvertMulFToAIEVecMulElemOpPattern
     }
 
     // Prepare lhr/rhs for the aievec.mul_elem op
-    VectorType targetInputType =
-        createVectorType(512 / lBitWidth, lSrcType.getElementType());
-    if (rBitWidth > lBitWidth)
-      targetInputType =
-          createVectorType(512 / rBitWidth, rSrcType.getElementType());
-    auto lValConverted = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(),
-                                                      lval, targetInputType);
-    auto rValConverted = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(),
-                                                      rval, targetInputType);
-    if (!lValConverted || !rValConverted)
+    unsigned bitWidth = (rBitWidth > lBitWidth) ? rBitWidth : lBitWidth;
+    Type srcElemType = (rBitWidth > lBitWidth) ? rSrcType.getElementType()
+                                               : lSrcType.getElementType();
+    unsigned numLanes = 0;
+    if (isa<FloatType>(srcElemType) && (bitWidth == 16 || bitWidth == 32)) {
+      numLanes = 16;
+    } else if (isa<IntegerType>(srcElemType) &&
+               (bitWidth == 8 || bitWidth == 16)) {
+      numLanes = 32;
+    } else if (isa<IntegerType>(srcElemType) && (bitWidth == 32)) {
+      numLanes = 16;
+    } else {
+      return failure();
+    }
+    VectorType targetInputType = createVectorType(numLanes, srcElemType);
+    if (targetInputType != lSrcType) {
+      lval = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(), lval,
+                                          targetInputType)
+                 .value();
+    }
+    if (targetInputType != rSrcType) {
+      rval = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(), rval,
+                                          targetInputType)
+                 .value();
+    }
+    if (!lval || !rval)
       return failure();
 
     // Create an aievec.mul_elem op
-    auto mulElemOp = rewriter.create<aievec::MulElemOp>(
-        mulOp.getLoc(), accType, *lValConverted, *rValConverted);
+    auto mulElemOp =
+        rewriter.create<aievec::MulElemOp>(mulOp.getLoc(), accType, lval, rval);
 
     // Create an aievec.cast or an aievec.srs op
     auto mulElemResultType = mulElemOp.getType();
@@ -752,8 +769,10 @@ struct ConvertMulIToAIEVecMulElemOpPattern
     // Decide the accType for aievec.mul_elem based on mulOp's lhs & rhs
     auto lval = adaptor.getLhs();
     auto rval = adaptor.getRhs();
+
     lval = getSourceOfWideningOp(lval).value_or(lval);
     rval = getSourceOfWideningOp(rval).value_or(rval);
+
     auto lSrcType = cast<VectorType>(lval.getType());
     auto rSrcType = cast<VectorType>(rval.getType());
     unsigned lBitWidth = lSrcType.getElementType().getIntOrFloatBitWidth();
@@ -764,22 +783,37 @@ struct ConvertMulIToAIEVecMulElemOpPattern
     }
 
     // Prepare lhr/rhs for the aievec.mul_elem op
-    VectorType targetInputType =
-        createVectorType(512 / lBitWidth, lSrcType.getElementType());
-    if (rBitWidth > lBitWidth) {
-      targetInputType =
-          createVectorType(512 / rBitWidth, rSrcType.getElementType());
+    unsigned bitWidth = (rBitWidth > lBitWidth) ? rBitWidth : lBitWidth;
+    Type srcElemType = (rBitWidth > lBitWidth) ? rSrcType.getElementType()
+                                               : lSrcType.getElementType();
+    unsigned numLanes = 0;
+    if (isa<FloatType>(srcElemType) && (bitWidth == 16 || bitWidth == 32)) {
+      numLanes = 16;
+    } else if (isa<IntegerType>(srcElemType) &&
+               (bitWidth == 8 || bitWidth == 16)) {
+      numLanes = 32;
+    } else if (isa<IntegerType>(srcElemType) && (bitWidth == 32)) {
+      numLanes = 16;
+    } else {
+      return failure();
+    }
+    VectorType targetInputType = createVectorType(numLanes, srcElemType);
+    if (targetInputType != lSrcType) {
+      lval = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(), lval,
+                                          targetInputType)
+                 .value();
+    }
+    if (targetInputType != rSrcType) {
+      rval = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(), rval,
+                                          targetInputType)
+                 .value();
     }
-    auto lValConverted = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(),
-                                                      lval, targetInputType);
-    auto rValConverted = convertValueToTargetTypeAIE2(rewriter, mulOp.getLoc(),
-                                                      rval, targetInputType);
-    if (!lValConverted || !rValConverted)
+    if (!lval || !rval)
       return failure();
 
     // Create an aievec.mul_elem op
-    auto mulElemOp = rewriter.create<aievec::MulElemOp>(
-        mulOp.getLoc(), accType, *lValConverted, *rValConverted);
+    auto mulElemOp =
+        rewriter.create<aievec::MulElemOp>(mulOp.getLoc(), accType, lval, rval);
 
     // Create an aievec.cast or an aievec.srs op
     auto mulElemResultType = mulElemOp.getType();

lib/Targets/AIEVecToCpp/TranslateAIEVecToCpp.cpp

@@ -1377,6 +1377,39 @@ static LogicalResult printOperation(CppEmitter &emitter, aievec::MulOp mulOp) {
 
   return success();
 }
+// convert operand to 512 bits
+static std::string printConversionTo512bit(CppEmitter &emitter, Value v) {
+  std::string vName = emitter.getOrCreateName(v).str();
+  auto vTy = cast<VectorType>(v.getType());
+  auto vShape = vTy.getShape();
+  int64_t elemBitWidth = vTy.getElementTypeBitWidth();
+  int64_t numElems = std::accumulate(vShape.begin(), vShape.end(), 1,
+                                     std::multiplies<int64_t>());
+  int64_t vBitWidth = numElems * elemBitWidth;
+  if (vBitWidth >= 512)
+    return vName;
+
+  int64_t newNumElems = 512 / elemBitWidth;
+
+  std::string vNewName = emitter.getNewName();
+  raw_indented_ostream &os = emitter.ostream();
+  auto newVecTy = VectorType::get({512 / elemBitWidth}, vTy.getElementType());
+  auto newTyName = *(
+      emitter.genCppTypeName(newVecTy, /*stdintType=*/false, /*isAcc=*/false));
+  auto oldTyName =
+      *(emitter.genCppTypeName(vTy, /*stdintType=*/false, /*isAcc=*/false));
+
+  os << newTyName << " " << vNewName << " = concat(";
+  if (newNumElems / numElems == 4) {
+    os << "concat(" << vName << ", undef_" << oldTyName << "())";
+    oldTyName = *(emitter.genCppTypeName(
+        VectorType::get({256 / elemBitWidth}, vTy.getElementType())));
+  } else {
+    os << vName;
+  }
+  os << ", undef_" << oldTyName << "());\n";
+  return vNewName;
+}
 
 // Generate the MulElem op
 static LogicalResult printOperation(CppEmitter &emitter,
@@ -1388,6 +1421,9 @@ static LogicalResult printOperation(CppEmitter &emitter,
   if (!emitter.hasValueInScope(lhs) || !emitter.hasValueInScope(rhs))
     return failure();
 
+  auto lhsName = printConversionTo512bit(emitter, lhs);
+  auto rhsName = printConversionTo512bit(emitter, rhs);
+
   std::string opname = "mul_elem";
 
   // Create opname based on the source type
@@ -1417,16 +1453,15 @@ static LogicalResult printOperation(CppEmitter &emitter,
     return failure();
 
   os << opname;
-  os << "(";
-  if (failed(printFMAOrMulElemOperand<aievec::MulElemOp>(emitter, mulElemOp,
-                                                         iType, lsize, 1)))
-    return failure();
-  os << ", ";
-  if (failed(printFMAOrMulElemOperand<aievec::MulElemOp>(emitter, mulElemOp,
-                                                         iType, lsize, 0)))
-    return failure();
+  os << "(" << lhsName;
+  if ((lsize == 32) && iType)
+    os << " ,"
+       << "undef_v16int32()";
+  os << " ," << rhsName;
+  if ((lsize == 32) && iType)
+    os << " , "
+       << "broadcast_zero_s32()";
   os << ")";
-
   return success();
 }
 
@@ -2710,39 +2745,6 @@ static LogicalResult printOperation(CppEmitter &emitter,
   return success();
 }
 
-static std::string printConversionTo512bit(CppEmitter &emitter, Value v) {
-  std::string vName = emitter.getOrCreateName(v).str();
-  auto vTy = cast<VectorType>(v.getType());
-  auto vShape = vTy.getShape();
-  int64_t elemBitWidth = vTy.getElementTypeBitWidth();
-  int64_t numElems = std::accumulate(vShape.begin(), vShape.end(), 1,
-                                     std::multiplies<int64_t>());
-  int64_t vBitWidth = numElems * elemBitWidth;
-  if (vBitWidth >= 512)
-    return vName;
-
-  int64_t newNumElems = 512 / elemBitWidth;
-
-  std::string vNewName = emitter.getNewName();
-  raw_indented_ostream &os = emitter.ostream();
-  auto newVecTy = VectorType::get({512 / elemBitWidth}, vTy.getElementType());
-  auto newTyName = *(
-      emitter.genCppTypeName(newVecTy, /*stdintType=*/false, /*isAcc=*/false));
-  auto oldTyName =
-      *(emitter.genCppTypeName(vTy, /*stdintType=*/false, /*isAcc=*/false));
-
-  os << newTyName << " " << vNewName << " = concat(";
-  if (newNumElems / numElems == 4) {
-    os << "concat(" << vName << ", undef_" << oldTyName << "())";
-    oldTyName = *(emitter.genCppTypeName(
-        VectorType::get({256 / elemBitWidth}, vTy.getElementType())));
-  } else {
-    os << vName;
-  }
-  os << ", undef_" << oldTyName << "());\n";
-  return vNewName;
-}
-
 static LogicalResult printOperation(CppEmitter &emitter,
                                     aievec::MatMulOp matmulOp) {
   auto lhs = matmulOp.getLhs();

test/Conversion/AIEVecToLLVM/mul_elem.mlir

@@ -2,7 +2,6 @@
 // RUN: aie-opt %s -split-input-file -convert-aievec-to-llvm="aie2-fp32-emulation-strategy=accuracy-fast" | FileCheck --check-prefix=FP32FAST %s
 // RUN: aie-opt %s -split-input-file -convert-aievec-to-llvm="aie2-fp32-emulation-strategy=accuracy-low" | FileCheck --check-prefix=FP32LOW %s
 
-
 func.func @i16_i16_i32_mul_elem(%arg0 : vector<32xi16>, %arg1 : vector<32xi16>) -> vector<32xi32> {
   %0 = aievec.mul_elem %arg0, %arg1 : vector<32xi16>, vector<32xi16>, vector<32xi32>
   return %0 : vector<32xi32>
@@ -22,21 +21,30 @@ func.func @i16_i16_i32_mul_elem(%arg0 : vector<32xi16>, %arg1 : vector<32xi16>)
 
 // -----
 
-func.func @i8_i8_i32_mul_elem(%arg0 : vector<64xi8>, %arg1 : vector<64xi8>) -> vector<32xi32> {
-  %0 = aievec.mul_elem %arg0, %arg1 : vector<64xi8>, vector<64xi8>, vector<32xi32>
+func.func @i8_i8_i32_mul_elem(%arg0 : vector<32xi8>, %arg1 : vector<32xi8>) -> vector<32xi32> {
+  %0 = aievec.mul_elem %arg0, %arg1 : vector<32xi8>, vector<32xi8>, vector<32xi32>
   return %0 : vector<32xi32>
 }
 
 // CHECK-LABEL: @i8_i8_i32_mul_elem
-// CHECK-SAME: %[[ARG0:.*]]: vector<64xi8>,
-// CHECK-SAME: %[[ARG1:.*]]: vector<64xi8>
+// CHECK-SAME: %[[ARG0:.*]]: vector<32xi8>,
+// CHECK-SAME: %[[ARG1:.*]]: vector<32xi8>
 // CHECK: %[[CST:.*]] = llvm.mlir.constant(808 : i32) : i32
-// CHECK-NEXT: %[[BITCAST1:.*]] = llvm.bitcast %[[ARG1]] : vector<64xi8> to vector<16xi32>
-// CHECK-NEXT: %[[MULCONF:.*]] = "xllvm.intr.aie2.I512.I512.acc32.mul.conf"(
-// CHECK-SAME: %[[ARG0]], %[[BITCAST1]], %[[CST]]) : 
+// CHECK: %[[C0I32:.*]] = llvm.mlir.constant(0 : i32) : i32 
+// CHECK: %[[BITCAST0:.*]] = llvm.bitcast %[[ARG0]] : vector<32xi8> to vector<8xi32>
+// CHECK: %[[IFA0:.*]] = "xllvm.intr.aie2.set.I512.I256"(%[[BITCAST0]],
+// CHECK-SAME:            %[[C0I32]]) : (vector<8xi32>, i32) -> vector<16xi32>
+// CHECK: %[[BITCAST1:.*]] = llvm.bitcast %[[IFA0]] : vector<16xi32> to vector<64xi8>
+// CHECK: %[[C1I32:.*]] = llvm.mlir.constant(0 : i32) : i32 
+// CHECK: %[[BITCAST2:.*]] = llvm.bitcast %[[ARG1]] : vector<32xi8> to vector<8xi32> 
+// CHECK: %[[IFA1:.*]] = "xllvm.intr.aie2.set.I512.I256"(%[[BITCAST2]],
+// CHECK-SAME:            %[[C1I32]]) : (vector<8xi32>, i32) -> vector<16xi32>
+// CHECK: %[[BITCAST3:.*]] = llvm.bitcast %[[IFA1]] : vector<16xi32> to vector<16xi32>
+// CHECK: %[[MULCONF:.*]] = "xllvm.intr.aie2.I512.I512.acc32.mul.conf"(
+// CHECK-SAME: %[[BITCAST1]], %[[BITCAST3]], %[[CST]]) : 
 // CHECK-SAME: (vector<64xi8>, vector<16xi32>, i32) -> vector<16xi64>
-// CHECK-NEXT: %[[RES:.*]] = llvm.bitcast %[[MULCONF]] : vector<16xi64> to vector<32xi32>
-// CHECK-NEXT: return %[[RES]] : vector<32xi32>
+// CHECK: %[[RES:.*]] = llvm.bitcast %[[MULCONF]] : vector<16xi64> to vector<32xi32> 
+// CHECK-NEXT: return %[[RES]] : vector<32xi32> 
 
 // -----
 
@@ -76,21 +84,30 @@ func.func @i32_i32_i32_mul_elem(%arg0 : vector<16xi32>, %arg1 : vector<16xi32>)
 
 // -----
 
-func.func @bf16_bf16_f32_mul_elem(%arg0 : vector<32xbf16>, %arg1 : vector<32xbf16>) -> vector<16xf32> {
-  %0 = aievec.mul_elem %arg0, %arg1 : vector<32xbf16>, vector<32xbf16>, vector<16xf32>
+func.func @bf16_bf16_f32_mul_elem(%arg0 : vector<16xbf16>, %arg1 : vector<16xbf16>) -> vector<16xf32> {
+  %0 = aievec.mul_elem %arg0, %arg1 : vector<16xbf16>, vector<16xbf16>, vector<16xf32>
   return %0 : vector<16xf32>
 }
 
 // CHECK-LABEL: @bf16_bf16_f32_mul_elem
-// CHECK-SAME: %[[ARG0:.*]]: vector<32xbf16>,
-// CHECK-SAME: %[[ARG1:.*]]: vector<32xbf16>
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xbf16>,
+// CHECK-SAME: %[[ARG1:.*]]: vector<16xbf16>
 // CHECK: %[[CST:.*]] = llvm.mlir.constant(60 : i32) : i32
+// CHECK: %[[C0I32:.*]] = llvm.mlir.constant(0 : i32) : i32 
+// CHECK: %[[BITCAST0:.*]] = llvm.bitcast %[[ARG0]] : vector<16xbf16> to vector<8xi32>
+// CHECK: %[[IFA0:.*]] = "xllvm.intr.aie2.set.I512.I256"(%[[BITCAST0]],
+// CHECK-SAME:            %[[C0I32]]) : (vector<8xi32>, i32) -> vector<16xi32>
+// CHECK: %[[BITCAST1:.*]] = llvm.bitcast %[[IFA0]] : vector<16xi32> to vector<32xbf16>
+// CHECK: %[[C1I32:.*]] = llvm.mlir.constant(0 : i32) : i32 
+// CHECK: %[[BITCAST2:.*]] = llvm.bitcast %[[ARG1]] : vector<16xbf16> to vector<8xi32> 
+// CHECK: %[[IFA1:.*]] = "xllvm.intr.aie2.set.I512.I256"(%[[BITCAST2]],
+// CHECK-SAME:            %[[C1I32]]) : (vector<8xi32>, i32) -> vector<16xi32>
+// CHECK: %[[BITCAST3:.*]] = llvm.bitcast %[[IFA1]] : vector<16xi32> to vector<32xbf16>
 // CHECK-NEXT: %[[MULCONF:.*]] = "xllvm.intr.aie2.bf.mul16.conf"(
-// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[CST]]) : 
+// CHECK-SAME: %[[BITCAST1]], %[[BITCAST3]], %[[CST]]) : 
 // CHECK-SAME: (vector<32xbf16>, vector<32xbf16>, i32) -> vector<8xi64>
 // CHECK-NEXT: %[[RES:.*]] = llvm.bitcast %[[MULCONF]] : vector<8xi64> to vector<16xf32>
 // CHECK-NEXT: return %[[RES]] : vector<16xf32>
-
 // -----
 
 func.func @f32_f32_f32_mul_elem(%arg0 : vector<16xf32>, %arg1 : vector<16xf32>) -> vector<16xf32> {