From 3974ecb7ec05ce63e8075b7bc8de921dd135823a Mon Sep 17 00:00:00 2001
From: Joshua Cao <cao.joshua@yahoo.com>
Date: Wed, 5 Apr 2023 21:47:59 -0700
Subject: [PATCH] [mlir][affine] SuperVectorizer only widen ops with valid
 types

fixes https://github.com/llvm/llvm-project/issues/61309

Differential Revision: https://reviews.llvm.org/D147679
---
 mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp  | 19 +++++++++++++
 .../Affine/SuperVectorize/invalid_type.mlir        | 31 ++++++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir

diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 81d0655..6b028b7 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -279,6 +279,25 @@ isVectorizableLoopBodyWithOpCond(AffineForOp loop,
     return false;
   }
 
+  // No vectorization for ops with operand or result types that are not
+  // vectorizable.
+  auto types = matcher::Op([](Operation &op) -> bool {
+    if (llvm::any_of(op.getOperandTypes(), [](Type type) {
+          if (MemRefType t = dyn_cast<MemRefType>(type))
+            return !VectorType::isValidElementType(t.getElementType());
+          return !VectorType::isValidElementType(type);
+        }))
+      return true;
+    return llvm::any_of(op.getResultTypes(), [](Type type) {
+      return !VectorType::isValidElementType(type);
+    });
+  });
+  SmallVector<NestedMatch, 8> 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<AffineIfOp, AffineForOp>(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
index 0000000..65df957
--- /dev/null
+++ b/mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @invalid_operand
+func.func @invalid_operand(%a : vector<4xf32>, %b : vector<4xf32>) {
+// CHECK: affine.for %{{.*}} = 0 to 10
+// CHECK:   %{{.*}} = vector.reduction <add>, %{{.*}} : vector<4xf32> into f32
+// CHECK: }
+// CHECK: return
+  affine.for %j = 0 to 10 {
+    %1 = vector.reduction <add>, %a : vector<4xf32> into f32
+  }
+  return
+}
+
+// CHECK-LABEL: func @invalid_result
+func.func @invalid_result(%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
+}
-- 
2.7.4