[aievec] to-llvm flow for aievec.broadcast_scalar op (#1446)

* This PR add the support for aievec.broadcast_scalar op going through the to-llvm flow.
* Add aievec-to-llvm conversion pattern/tests for the aievec.broadcast_scalar op.
* Add op verifier for the aievec.broadcast_scalar op.
* Add target external llvm translation tests.

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

@@ -352,7 +352,7 @@ def AIEVec_BroadcastScalarOp:
   AIEVec_Op<"broadcast_scalar", [
     Pure
   ]>,
-  Arguments<(ins AnyType:$source)>,
+  Arguments<(ins AnyTypeOf<[BF16, F32, I32, I16, I8]>:$source)>,
   Results<(outs AnyVector:$result)> {
   let summary = "AIE-ML broadcast scalar";
   let description = [{

include/aie/Dialect/XLLVM/IR/XLLVMAIE2IntrOps.td

@@ -134,17 +134,27 @@ def Vector16AccFloatToV16BF16IntrOp :
 def VectorBroadcast8I512IntrOp :
     AIEVec2_IntrOp<"vbroadcast8.I512",
         [TypeIs<"res", VectorOfLengthAndType<[64], [I8]>>]>,
-    Arguments<(ins I32:$value)>;
+    Arguments<(ins I32:$src)>;
+
+def VectorBroadcast16I512IntrOp :
+    AIEVec2_IntrOp<"vbroadcast16.I512",
+        [TypeIs<"res", VectorOfLengthAndType<[32], [I16]>>]>,
+    Arguments<(ins I32:$src)>;
 
 def VectorBroadcast32I512IntrOp :
     AIEVec2_IntrOp<"vbroadcast32.I512",
         [TypeIs<"res", VectorOfLengthAndType<[16], [I32]>>]>,
-    Arguments<(ins I32:$value)>;
+    Arguments<(ins I32:$src)>;
 
 def VectorBroadcast16BF512IntrOp :
     AIEVec2_IntrOp<"vbroadcast16.bf512",
         [TypeIs<"res", VectorOfLengthAndType<[32], [BF16]>>]>,
-    Arguments<(ins BF16:$value)>;
+    Arguments<(ins BF16:$src)>;
+
+def VectorBroadcastfloatI512IntrOp :
+    AIEVec2_IntrOp<"vbroadcastfloat.I512",
+        [TypeIs<"res", VectorOfLengthAndType<[16], [F32]>>]>,
+    Arguments<(ins F32:$src)>;
 
 // ----- EXT ----- 
 

lib/Conversion/AIEVecToLLVM/AIEVecToLLVM.cpp

@@ -1249,6 +1249,15 @@ class BroadcastScalarOpConversion
     VectorType resultType = cast<VectorType>(result.getType());
     Type resultScaTy = resultType.getElementType();
     unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
+    int resultLanes = getVectorLaneSize(resultType);
+    int resultVectorSize = resultBitWidth * resultLanes;
+
+    if (resultVectorSize != 512) {
+      op.emitWarning()
+          << "aievec.broadcast_scalar conversion with result vector size "
+          << resultVectorSize << " is not implemented.\n";
+      return failure();
+    }
 
     // Integer types
     if (llvm::isa<IntegerType>(resultScaTy)) {
@@ -1257,16 +1266,15 @@ class BroadcastScalarOpConversion
       unsigned srcBitWidth = srcType.getIntOrFloatBitWidth();
 
       if (srcBitWidth < 32) {
-        src = rewriter.create<LLVM::SExtOp>(loc, rewriter.getI32Type(),
-                                            adaptor.getSource());
-      } else if (srcBitWidth > 32) {
-        src = rewriter.create<LLVM::TruncOp>(loc, rewriter.getI32Type(),
-                                             adaptor.getSource());
+        src = rewriter.create<LLVM::SExtOp>(loc, rewriter.getI32Type(), src);
       }
 
       if (resultBitWidth == 8) {
         rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast8I512IntrOp>(
             op, VectorType::get({64}, rewriter.getI8Type()), src);
+      } else if (resultBitWidth == 16) {
+        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast16I512IntrOp>(
+            op, VectorType::get({32}, rewriter.getI16Type()), src);
       } else if (resultBitWidth == 32) {
         rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast32I512IntrOp>(
             op, VectorType::get({16}, rewriter.getI32Type()), src);
@@ -1282,6 +1290,10 @@ class BroadcastScalarOpConversion
         rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast16BF512IntrOp>(
             op, VectorType::get({32}, rewriter.getBF16Type()),
             adaptor.getSource());
+      } else if (resultBitWidth == 32) {
+        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcastfloatI512IntrOp>(
+            op, VectorType::get({16}, rewriter.getF32Type()),
+            adaptor.getSource());
       } else {
         op.emitWarning()
             << "aievec.broadcast_scalar conversion with result bitwidth "

lib/Dialect/AIEVec/IR/AIEVecOps.cpp

@@ -541,8 +541,14 @@ LogicalResult BroadcastScalarOp::verify() {
   if (!resultType)
     return emitError("requires vector type");
 
-  if (!sourceType)
-    return emitError("requires source type");
+  if (!sourceType.isa<IntegerType, FloatType>())
+    return emitError("requires source type to be integer or float");
+
+  Type resultElemType = resultType.getElementType();
+  if (sourceType != resultElemType) {
+    return emitError("the element type of result vector must be the same as "
+                     "the source type");
+  }
 
   return success();
 }

test/Conversion/AIEVecToLLVM/broadcast_scalar.mlir

@@ -15,6 +15,21 @@ func.func @i8_broadcast_scalar(%arg0 : i8) -> vector<64xi8> {
 
 // -----
 
+func.func @i16_broadcast_scalar(%arg0 : i16) -> vector<32xi16> {
+  %0 = aievec.broadcast_scalar %arg0 : i16, vector<32xi16>
+  return %0 : vector<32xi16>
+}
+
+// CHECK-LABEL: @i16_broadcast_scalar
+// CHECK-SAME: %[[ARG0:.*]]: i16
+// CHECK: %[[VAL:.*]]  = llvm.sext %[[ARG0]] : i16 to i32
+// CHECK-NEXT: %[[VBROADCAST:.*]] = "xllvm.intr.aie2.vbroadcast16.I512"(
+// CHECK-SAME: %[[VAL]]) : 
+// CHECK-SAME: (i32) -> vector<32xi16>
+// CHECK-NEXT: return %[[VBROADCAST]] : vector<32xi16>
+
+// -----
+
 func.func @i32_broadcast_scalar(%arg0 : i32) -> vector<16xi32> {
   %0 = aievec.broadcast_scalar %arg0 : i32, vector<16xi32>
   return %0 : vector<16xi32>
@@ -39,4 +54,18 @@ func.func @bf16_broadcast_scalar(%arg0 : bf16) -> vector<32xbf16> {
 // CHECK: %[[VBROADCAST:.*]] = "xllvm.intr.aie2.vbroadcast16.bf512"(
 // CHECK-SAME: %[[ARG0]]) : 
 // CHECK-SAME: (bf16) -> vector<32xbf16>
-// CHECK-NEXT: return %[[VBROADCAST]] : vector<32xbf16>
+// CHECK-NEXT: return %[[VBROADCAST]] : vector<32xbf16>
+
+// -----
+
+func.func @f32_broadcast_scalar(%arg0 : f32) -> vector<16xf32> {
+  %0 = aievec.broadcast_scalar %arg0 : f32, vector<16xf32>
+  return %0 : vector<16xf32>
+}
+
+// CHECK-LABEL: @f32_broadcast_scalar
+// CHECK-SAME: %[[ARG0:.*]]: f32
+// CHECK: %[[VBROADCAST:.*]] = "xllvm.intr.aie2.vbroadcastfloat.I512"(
+// CHECK-SAME: %[[ARG0]]) : 
+// CHECK-SAME: (f32) -> vector<16xf32>
+// CHECK-NEXT: return %[[VBROADCAST]] : vector<16xf32>

test/Target/LLVMIR/aievec.mlir

@@ -152,6 +152,14 @@ llvm.func @vbroadcast8_i512(%val : i32) -> vector<64xi8> {
     llvm.return %0 : vector<64xi8>
 }
 
+// CHECK-LABEL: define <32 x i16> @vbroadcast16_i512
+llvm.func @vbroadcast16_i512(%val : i32) -> vector<32xi16> {
+    // CHECK: call <32 x i16> @llvm.aie2.vbroadcast16.I512(
+    // CHECK-SAME: i32 %{{[0-9]+}})
+    %0 = "xllvm.intr.aie2.vbroadcast16.I512"(%val) : (i32) -> vector<32xi16>
+    llvm.return %0 : vector<32xi16>
+}
+
 // CHECK-LABEL: define <16 x i32> @vbroadcast32_i512
 llvm.func @vbroadcast32_i512(%val : i32) -> vector<16xi32> {
     // CHECK: call <16 x i32> @llvm.aie2.vbroadcast32.I512(
@@ -168,6 +176,14 @@ llvm.func @vbroadcast16_bf512(%val : bf16) -> vector<32xbf16> {
     llvm.return %0 : vector<32xbf16>
 }
 
+// CHECK-LABEL: define <16 x float> @vbroadcastfloat_i512
+llvm.func @vbroadcastfloat_i512(%val : f32) -> vector<16xf32> {
+    // CHECK: call <16 x float> @llvm.aie2.vbroadcastfloat.I512(
+    // CHECK-SAME: float %{{[0-9]+}})
+    %0 = "xllvm.intr.aie2.vbroadcastfloat.I512"(%val) : (f32) -> vector<16xf32>
+    llvm.return %0 : vector<16xf32>
+}
+
 // ----- EXT ----- 
 
 // CHECK-LABEL: define <8 x i32> @ext_i256_i512