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

[Kai Sasaki]

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

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.

Fixes https://github.com/llvm/llvm-project/issues/128277

>From 05ceaf00337f536999c19f9abe45408c8446a15f 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..7e69bce34aea4 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 (!isa<IndexType, IntegerType, FloatType>(value.getType()))
+    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..388d9ff423626 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
@@ -682,5 +682,21 @@ func.func @vec_vecdim_reduction_rejected(%in: memref<256x512xf32>, %out: memref<
  return
 }
 
+// -----
+
+// Non scalar type is not regarded as the uniform in the vectorization
+
 // CHECK-LABEL: @vec_vecdim_reduction_rejected
 // CHECK-NOT: vector
+
+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