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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Kai Sasaki (Lewuathe)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/128294.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+4) 
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir (+16) 


``````````diff
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

``````````

</details>


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