[Mlir-commits] [mlir] [mlir][vector] Skip uniform vectorization for non scalar type (PR #128294)

[Kai Sasaki]

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

>From f67d45133220f0b821c337d7067535fd66ed33e7 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Sat, 22 Feb 2025 16:19:33 +0900
Subject: [PATCH] [mlir][vector] Skip uniform vectorization for non scalar type

The code to vectorize the uniform type effectively assume the scalar
type such as integer, float and index type. We should not regard other
types as vectorizable uniform type because the following vectorization
cannot treat the type properlty for now.

See: https://github.com/llvm/llvm-project/issues/128277
---
 .../Dialect/Affine/Transforms/SuperVectorize.cpp |  4 ++++
 .../Affine/SuperVectorize/vectorize_1d.mlir      | 16 ++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 71e9648a5e00f..eaaafaf68767e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1106,6 +1106,10 @@ static bool isUniformDefinition(Value value,
     if (!loop.isDefinedOutsideOfLoop(value))
       return false;
   }
+
+  if (!value.getType().isIntOrIndexOrFloat())
+    return false;
+
   return true;
 }
 
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
index 9244604128cb7..81b04ccceaf27 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
@@ -684,3 +684,19 @@ func.func @vec_vecdim_reduction_rejected(%in: memref<256x512xf32>, %out: memref<
 
 // CHECK-LABEL: @vec_vecdim_reduction_rejected
 // CHECK-NOT: vector
+
+
+// -----
+
+// Non scalar type is not regarded as the uniform in the vectorization
+func.func @vec_non_scalar_type() {
+  %idx0 = index.constant 0
+  %alloc_82 = memref.alloc() : memref<1xi64>
+  affine.for %arg0 = 0 to 78 {
+    %dim_191 = memref.dim %alloc_82, %idx0 : memref<1xi64>
+  }
+  return
+}
+
+// CHECK-LABEL: @vec_non_scalar_type
+// CHECK-NOT: vector