[clang] [mlir][affine] Check the input vector sizes to be greater than 0 (PR #65293)

[Kai Sasaki via cfe-commits]

https://github.com/Lewuathe updated https://github.com/llvm/llvm-project/pull/65293

>From cd1a19a5b12cbdba2654b08b04e54c8313dda696 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Tue, 5 Sep 2023 14:40:44 +0900
Subject: [PATCH 1/2] [mlir][affine] Check the input vector sizes to be greater
 than 0

---
 mlir/include/mlir/Dialect/Affine/Passes.td            | 3 ++-
 mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp | 5 +++++
 mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir  | 9 +++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir

diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td
index 891a6661ca87be7..1036e93a0392409 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Passes.td
@@ -349,7 +349,8 @@ def AffineVectorize : Pass<"affine-super-vectorize", "func::FuncOp"> {
   let dependentDialects = ["vector::VectorDialect"];
   let options = [
     ListOption<"vectorSizes", "virtual-vector-size", "int64_t",
-               "Specify an n-D virtual vector size for vectorization">,
+               "Specify an n-D virtual vector size for vectorization. "
+               "This must be greater than zero.">,
     // Optionally, the fixed mapping from loop to fastest varying MemRef
     // dimension for all the MemRefs within a loop pattern:
     //   the index represents the loop depth, the value represents the k^th
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 072e858220feae3..327616d45c90189 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1732,6 +1732,11 @@ void Vectorize::runOnOperation() {
     return signalPassFailure();
   }
 
+  if (llvm::any_of(vectorSizes, [](int64_t size) { return size == 0; })) {
+    f.emitError("Vectorization factor must be greater than zero.");
+    return signalPassFailure();
+  }
+
   DenseSet<Operation *> parallelLoops;
   ReductionLoopMap reductionLoops;
 
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir b/mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir
new file mode 100644
index 000000000000000..396cc933e25bcf4
--- /dev/null
+++ b/mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir
@@ -0,0 +1,9 @@
+// RUN: mlir-opt %s -verify-diagnostics --affine-super-vectorize=virtual-vector-size=0
+
+// expected-error at +1 {{Vectorization factor must be greater than zero}} 
+func.func @with_zero_vector_size(%arg0: memref<21x12x12xi1>) {
+  affine.for %arg1 = 0 to 84 step 4294967295 {
+  }
+  return
+}
+

>From c0cc18d4d5631ec586d719eb21c37d0610f5ee31 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Tue, 26 Sep 2023 15:25:05 +0900
Subject: [PATCH 2/2] Post review follow-up

---
 mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp           | 2 +-
 .../SuperVectorize/{invalid.mlir => invalid-zero-size.mlir}     | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename mlir/test/Dialect/Affine/SuperVectorize/{invalid.mlir => invalid-zero-size.mlir} (100%)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 327616d45c90189..82623126e49006b 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1732,7 +1732,7 @@ void Vectorize::runOnOperation() {
     return signalPassFailure();
   }
 
-  if (llvm::any_of(vectorSizes, [](int64_t size) { return size == 0; })) {
+  if (llvm::any_of(vectorSizes, [](int64_t size) { return size <= 0; })) {
     f.emitError("Vectorization factor must be greater than zero.");
     return signalPassFailure();
   }
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir b/mlir/test/Dialect/Affine/SuperVectorize/invalid-zero-size.mlir
similarity index 100%
rename from mlir/test/Dialect/Affine/SuperVectorize/invalid.mlir
rename to mlir/test/Dialect/Affine/SuperVectorize/invalid-zero-size.mlir