diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp --- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp @@ -1395,9 +1395,15 @@ /// operands. static Operation *widenOp(Operation *op, VectorizationState &state) { SmallVector vectorTypes; - for (Value result : op->getResults()) + for (Value result : op->getResults()) { + if (!VectorType::isValidElementType(result.getType())) { + LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ an operation failed vectorize " + "because it has an invalid result type\n"); + return nullptr; + } vectorTypes.push_back( VectorType::get(state.strategy->vectorSizes, result.getType())); + } SmallVector vectorOperands; for (Value operand : op->getOperands()) { diff --git a/mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir b/mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir @@ -0,0 +1,19 @@ +// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s + +// CHECK-LABEL: func @vector_loop +func.func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>) { +// CHECK: affine.for %{{.*}} = 0 to 10 +// CHECK: affine.for %{{.*}} = 0 to 5 +// CHECK: %{{.*}} = affine.vector_load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x20xf32>, vector<4xf32> +// CHECK: affine.vector_store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x20xf32>, vector<4xf32> +// CHECK: } +// CHECK: } +// CHECK: return + affine.for %j = 0 to 10 { + affine.for %i = 0 to 5 { + %ld0 = affine.vector_load %a[%j, %i] : memref<10x20xf32>, vector<4xf32> + affine.vector_store %ld0, %b[%j, %i] : memref<10x20xf32>, vector<4xf32> + } + } + return +}