diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp --- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp @@ -279,6 +279,18 @@ return false; } + // No vectorization for ops with result types that are not vectorizable. + auto types = matcher::Op([](Operation &op) -> bool { + return llvm::any_of(op.getResultTypes(), [](Type type) { + return !VectorType::isValidElementType(type); + }); + }); + SmallVector opsMatched; + types.match(forOp, &opsMatched); + if (!opsMatched.empty()) { + return false; + } + // No vectorization across unknown regions. auto regions = matcher::Op([](Operation &op) -> bool { return op.getNumRegions() != 0 && !isa(op); 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 +}