Don't attempt to create vectors with complex element types.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

48

getLegalVectorType / getValidVectorType plz, otherwise it's unclear why you're duplicating standard infra.

900

auto maybeReadType = ...

followed by

auto readType = *maybeReadType

and avoid propagating the FailureOr / ptr changes forward plz.

1220

This is a second different way of doing the same thing, please make it uniform and single entry point.
Either use the new method you introduced everywhere or use this.

Actually using this check everywhere is less footprint, so please just:

if (!VectorType::isValidElementType(a) || !VectorType::isValidElementType(b) ... )
  return failure();

everywhere plz and drop the getVectorType altogether.

The whole PR should be 5-10 lines of change.

[nicolasvasilache]

Please address my comments as post-commit

[dcaballe]

Thanks for the contribution! It would be great to give more time for other reviewers to chime in. Do you plan to create a new diff for this?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
 /// Try to vectorize `convOp` as a convolution.
 static FailureOr<Operation *> vectorizeConvolution(RewriterBase &rewriter,
                                                    LinalgOp convOp);
 
+/// Return the given vector type if `elementType` is valid.
+static FailureOr<VectorType> getVectorType(ArrayRef<int64_t> shape,
+                                           Type elementType) {
+  if (!VectorType::isValidElementType(elementType)) {
+    return failure();
+  }
+  return VectorType::get(shape, elementType);
+}
+
+/// Cast the given type to a vector type if its element type is valid.
+static FailureOr<VectorType> getVectorType(ShapedType type) {
+  return getVectorType(type.getShape(), type.getElementType());
+}
+
 /// Return the unique instance of OpType in `block` if it is indeed unique.
 /// Return null if none or more than 1 instances exist.
 template <typename OpType>
 static OpType getSingleOpOfType(Block &block) {
   OpType res;
   block.walk([&](OpType op) {
     if (res) {
       res = nullptr;
static Value broadcastIfNeeded(OpBuilder &b, Value value,
   if (shape.empty())
     return value;
   VectorType targetVectorType =
       VectorType::get(shape, getElementTypeOrSelf(value));
   if (vector::isBroadcastableTo(value.getType(), targetVectorType) !=
       vector::BroadcastableToResult::Success)
     return value;
   Location loc = b.getInsertionPoint()->getLoc();
-  return b.createOrFold<vector::BroadcastOp>(loc, targetVectorType,
-                                                    value);
+  return b.createOrFold<vector::BroadcastOp>(loc, targetVectorType, value);
 }
 
 /// Create MultiDimReductionOp to compute the reduction for `reductionOp`. This
 /// assumes that `reductionOp` has two operands and one of them is the reduction
 /// initial value.buildMultiDimReduce
 // Note: this is a true builder that notifies the OpBuilder listener.
 // TODO: Consider moving as a static helper on the ReduceOp.
 static Operation *buildMultiDimReduce(OpBuilder &b,
if (linalgOp.isDpsInput(opOperand)) {
       // 3.a.ii. For output reads (iteration-carried dependence, e.g.,
       // reductions), the vector shape is computed by mapping the canonical
       // vector shape to the output domain and back to the canonical domain.
       readMap = inversePermutation(reindexIndexingMap(indexingMap));
       readVecShape =
           readMap.compose(indexingMap.compose(state.getCanonicalVecShape()));
     }
 
     auto readType =
-        VectorType::get(readVecShape, getElementTypeOrSelf(opOperand->get()));
+        getVectorType(readVecShape, getElementTypeOrSelf(opOperand->get()));
+    if (!succeeded(readType))
+      return failure();
     SmallVector<Value> indices(linalgOp.getShape(opOperand).size(), zero);
 
     Operation *read = rewriter.create<vector::TransferReadOp>(
-        loc, readType, opOperand->get(), indices, readMap);
+        loc, *readType, opOperand->get(), indices, readMap);
     read = state.maskOperation(rewriter, read, linalgOp, maskingMap);
     Value readValue = read->getResult(0);
 
     // 3.b. If masked, set in-bounds to true. Masking guarantees that the access
     // will be in-bounds.
     if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(read)) {
-      SmallVector<bool> inBounds(readType.getRank(), true);
+      SmallVector<bool> inBounds(readType->getRank(), true);
       cast<vector::TransferReadOp>(maskOp.getMaskableOp())
           .setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
     }
 
     // 3.c. Not all ops support 0-d vectors, extract the scalar for now.
     // TODO: remove this.
     if (readValue.getType().cast<VectorType>().getRank() == 0)
       readValue = rewriter.create<vector::ExtractElementOp>(loc, readValue);

 LogicalResult mlir::linalg::vectorizeCopy(RewriterBase &rewriter,
                                           memref::CopyOp copyOp) {
 
   auto srcType = copyOp.getSource().getType().cast<MemRefType>();
   auto dstType = copyOp.getTarget().getType().cast<MemRefType>();
   if (!srcType.hasStaticShape() || !dstType.hasStaticShape())
     return failure();
 
-  auto readType =
-      VectorType::get(srcType.getShape(), getElementTypeOrSelf(srcType));
-  auto writeType =
-      VectorType::get(dstType.getShape(), getElementTypeOrSelf(dstType));
+  auto readType = getVectorType(srcType);
+  auto writeType = getVectorType(dstType);
+  if (!(succeeded(readType) && succeeded(writeType)))
+    return failure();
 
   Location loc = copyOp->getLoc();
   Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   SmallVector<Value> indices(srcType.getRank(), zero);
 
   Value readValue = rewriter.create<vector::TransferReadOp>(
-      loc, readType, copyOp.getSource(), indices,
+      loc, *readType, copyOp.getSource(), indices,
       rewriter.getMultiDimIdentityMap(srcType.getRank()));
   if (readValue.getType().cast<VectorType>().getRank() == 0) {
     readValue = rewriter.create<vector::ExtractElementOp>(loc, readValue);
-    readValue = rewriter.create<vector::BroadcastOp>(loc, writeType, readValue);
+    readValue =
+        rewriter.create<vector::BroadcastOp>(loc, *writeType, readValue);
   }
   Operation *writeValue = rewriter.create<vector::TransferWriteOp>(
       loc, readValue, copyOp.getTarget(), indices,
       rewriter.getMultiDimIdentityMap(srcType.getRank()));
   rewriter.replaceOp(copyOp, writeValue->getResults());
   return success();
 }
 
struct GenericPadOpVectorizationPattern : public GeneralizePadOpPattern {
   /// Vectorize the copying of a tensor::PadOp's source. This is possible if
   /// each dimension size is statically know in the source type or the result
   /// type (or both).
   static LogicalResult tryVectorizeCopy(PatternRewriter &rewriter,
                                         tensor::PadOp padOp, Value dest) {
     auto sourceType = padOp.getSourceType();
     auto resultType = padOp.getResultType();
 
+    // Complex is not a valid vector element type.
+    if (!VectorType::isValidElementType(sourceType.getElementType()))
+      return failure();
+
     // Copy cannot be vectorized if pad value is non-constant and source shape
     // is dynamic. In case of a dynamic source shape, padding must be appended
     // by TransferReadOp, but TransferReadOp supports only constant padding.
     auto padValue = padOp.getConstantPaddingValue();
     if (!padValue) {
       if (!sourceType.hasStaticShape())
         return failure();
       // Create dummy padding value.
if (!padValue)
       return failure();
     // Dynamic shapes not supported.
     if (!padOp.getResult().getType().cast<ShapedType>().hasStaticShape())
       return failure();
     // Pad result not used as destination.
     if (insertOp.getDest() == padOp.getResult())
       return failure();
 
-    auto vecType = VectorType::get(padOp.getType().getShape(),
-                                   padOp.getType().getElementType());
-    unsigned vecRank = vecType.getRank();
+    auto vecType = getVectorType(padOp.getType());
+    if (!succeeded(vecType))
+      return failure();
+    unsigned vecRank = vecType->getRank();
     unsigned tensorRank = insertOp.getType().getRank();
 
     // Check if sizes match: Insert the entire tensor into most minor dims.
     // (No permutations allowed.)
     SmallVector<int64_t> expectedSizes(tensorRank - vecRank, 1);
-    expectedSizes.append(vecType.getShape().begin(), vecType.getShape().end());
+    expectedSizes.append(vecType->getShape().begin(),
+                         vecType->getShape().end());
     if (!llvm::all_of(
             llvm::zip(insertOp.getMixedSizes(), expectedSizes), [](auto it) {
               return getConstantIntValue(std::get<0>(it)) == std::get<1>(it);
             }))
       return failure();
 
     // Insert the TransferReadOp and TransferWriteOp at the position of the
     // InsertSliceOp.
     rewriter.setInsertionPoint(insertOp);
 
     // Generate TransferReadOp: Read entire source tensor and add high
     // padding.
     SmallVector<Value> readIndices(
         vecRank, rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
     auto read = rewriter.create<vector::TransferReadOp>(
-        padOp.getLoc(), vecType, padOp.getSource(), readIndices, padValue);
+        padOp.getLoc(), *vecType, padOp.getSource(), readIndices, padValue);
 
     // Generate TransferWriteOp: Write to InsertSliceOp's dest tensor at
     // specified offsets. Write is fully in-bounds because a InsertSliceOp's
     // source must fit into the destination at the specified offsets.
     auto writeIndices =
         ofrToIndexValues(rewriter, padOp.getLoc(), insertOp.getMixedOffsets());
     SmallVector<bool> inBounds(vecRank, true);
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
Conv1DGenerator(RewriterBase &rewriter, LinalgOp linalgOp, int strideW,
     if (!maybeKind || (*maybeKind != vector::CombiningKind::ADD &&
                        (oper != Pool || !isSupportedPoolKind(*maybeKind)))) {
       return;
     }
 
     auto rhsRank = rhsShapedType.getRank();
     switch (oper) {
     case Conv:
-      if (rhsRank != 2 && rhsRank!= 3)
+      if (rhsRank != 2 && rhsRank != 3)
         return;
       break;
     case Pool:
       if (rhsRank != 1)
         return;
       break;
     }
     // The op is now known to be valid.
FailureOr<Operation *> conv(Conv1DOpOrder conv1DOpOrder) {
     // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
     // When strideW == 1, we can batch the contiguous loads and avoid
     // unrolling
     int64_t wSizeStep = strideW == 1 ? wSize : 1;
 
     Type lhsEltType = lhsShapedType.getElementType();
     Type rhsEltType = rhsShapedType.getElementType();
     Type resEltType = resShapedType.getElementType();
-    auto lhsType = VectorType::get(lhsShape, lhsEltType);
-    auto rhsType = VectorType::get(rhsShape, rhsEltType);
-    auto resType = VectorType::get(resShape, resEltType);
+    auto lhsType = getVectorType(lhsShape, lhsEltType);
+    auto rhsType = getVectorType(rhsShape, rhsEltType);
+    auto resType = getVectorType(resShape, resEltType);
+    if (!(succeeded(lhsType) && succeeded(rhsType) && succeeded(resType)))
+      return failure();
     // Read lhs slice of size {w * strideW + kw * dilationW, c, f} @ [0, 0,
     // 0].
     Value lhs = rewriter.create<vector::TransferReadOp>(
-        loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
+        loc, *lhsType, lhsShaped, ValueRange{zero, zero, zero});
     // Read rhs slice of size {kw, c, f} @ [0, 0, 0].
     // This is needed only for Conv.
     Value rhs = nullptr;
     if (oper == Conv)
       rhs = rewriter.create<vector::TransferReadOp>(
-          loc, rhsType, rhsShaped, ValueRange{zero, zero, zero});
+          loc, *rhsType, rhsShaped, ValueRange{zero, zero, zero});
     // Read res slice of size {n, w, f} @ [0, 0, 0].
     Value res = rewriter.create<vector::TransferReadOp>(
-        loc, resType, resShaped, ValueRange{zero, zero, zero});
+        loc, *resType, resShaped, ValueRange{zero, zero, zero});
 
     // The base vectorization case is input: {n,w,c}, weight: {kw,c,f}, output:
     // {n,w,f}. To reuse the base pattern vectorization case, we do pre
     // transpose on input, weight, and output.
     switch (conv1DOpOrder) {
     case Conv1DOpOrder::Nwc:
       // Base case, so no transposes necessary.
       break;
FailureOr<Operation *> depthwiseConv() {
     // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
     // When strideW == 1, we can batch the contiguous loads and avoid
     // unrolling
     int64_t wSizeStep = strideW == 1 ? wSize : 1;
 
     Type lhsEltType = lhsShapedType.getElementType();
     Type rhsEltType = rhsShapedType.getElementType();
     Type resEltType = resShapedType.getElementType();
-    VectorType lhsType = VectorType::get(
+    auto lhsType = getVectorType(
         {nSize,
          // iw = ow * sw + kw *  dw - 1
          //   (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
          ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) - 1,
          cSize},
         lhsEltType);
-    VectorType rhsType = VectorType::get({kwSize, cSize}, rhsEltType);
-    VectorType resType = VectorType::get({nSize, wSize, cSize}, resEltType);
+    auto rhsType = getVectorType({kwSize, cSize}, rhsEltType);
+    auto resType = getVectorType({nSize, wSize, cSize}, resEltType);
+    if (!(succeeded(lhsType) && succeeded(rhsType) && succeeded(resType)))
+      return failure();
 
     // Read lhs slice of size {n, w * strideW + kw * dilationW, c} @ [0, 0,
     // 0].
     Value lhs = rewriter.create<vector::TransferReadOp>(
-        loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
+        loc, *lhsType, lhsShaped, ValueRange{zero, zero, zero});
     // Read rhs slice of size {kw, c} @ [0, 0].
-    Value rhs = rewriter.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
-                                                        ValueRange{zero, zero});
+    Value rhs = rewriter.create<vector::TransferReadOp>(
+        loc, *rhsType, rhsShaped, ValueRange{zero, zero});
     // Read res slice of size {n, w, c} @ [0, 0, 0].
     Value res = rewriter.create<vector::TransferReadOp>(
-        loc, resType, resShaped, ValueRange{zero, zero, zero});
+        loc, *resType, resShaped, ValueRange{zero, zero, zero});
 
     //===------------------------------------------------------------------===//
     // Begin vector-only rewrite part
     //===------------------------------------------------------------------===//
     // Unroll along kw and read slices of lhs and rhs.
     SmallVector<Value> lhsVals, rhsVals, resVals;
     // Extract lhs slice of size {n, wSizeStep, c}
     //   @ [0, sw * w + dw * kw, 0].

mlir/test/Dialect/Linalg/vectorize-convolution.mlir

 // CHECK: %[[V2:.+]] = vector.transpose %[[V0]], [0, 2, 1] : vector<4x2x5xf32> to vector<4x5x2xf32>
 // CHECK: %[[V3:.+]] = vector.transpose %[[V1]], [0, 2, 1] : vector<4x2x3xf32> to vector<4x3x2xf32>
 // CHECK: %[[V4:.+]] = vector.extract_strided_slice %[[V2]] {offsets = [0, 0, 0], sizes = [4, 3, 2], strides = [1, 1, 1]} : vector<4x5x2xf32> to vector<4x3x2xf32>
 // CHECK: %[[V5:.+]] = vector.extract_strided_slice %[[V2]] {offsets = [0, 2, 0], sizes = [4, 3, 2], strides = [1, 1, 1]} : vector<4x5x2xf32> to vector<4x3x2xf32>
 // CHECK: %[[V6:.+]] = arith.addf %[[V4]], %[[V3]] : vector<4x3x2xf32>
 // CHECK: %[[V7:.+]] = arith.addf %[[V5]], %[[V6]] : vector<4x3x2xf32>
 // CHECK: %[[V8:.+]] = vector.transpose %[[V7]], [0, 2, 1] : vector<4x3x2xf32> to vector<4x2x3xf32>
 // CHECK: vector.transfer_write %[[V8:.+]], %[[OUTPUT]][%[[Vc0]], %[[Vc0]], %[[Vc0]]] {in_bounds = [true, true, true]} : vector<4x2x3xf32>, memref<4x2x3xf32>
+
+
+// -----
+
+func.func @pooling_ncw_sum_memref_complex(%input: memref<4x2x5xcomplex<f32>>,
+     %filter: memref<2xcomplex<f32>>, %output: memref<4x2x3xcomplex<f32>>) {
+  linalg.pooling_ncw_sum
+    {dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
+    ins(%input, %filter : memref<4x2x5xcomplex<f32>>, memref<2xcomplex<f32>>)
+    outs(%output : memref<4x2x3xcomplex<f32>>)
+  return
+}
+
+// Regression test: just check that this lowers successfully
+// CHECK-LABEL: @pooling_ncw_sum_memref_complex