[Mlir-commits] [mlir] [MLIR] Avoid `vector.extract_strided_slice` when not needed (PR #115941)

[llvmlistbot at llvm.org]

https://github.com/lialan created https://github.com/llvm/llvm-project/pull/115941

In `staticallyExtractSubvector`, When the extracting slice is the same as source vector, do not need to emit `vector.extract_strided_slice`.

This fixes the lit test case `@vector_store_i4` in
`mlir\test\Dialect\Vector\vector-emulate-narrow-type.mlir`, where converting from `vector<8xi4>` to `vector<4xi8>` does not need slice extraction.

>From 025fc2f920bfe8d2178d84f6a4db24a855318151 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 12 Nov 2024 16:10:40 -0500
Subject: [PATCH] [MLIR] Avoid `vector.extract_strided_slice` when not needed

In `staticallyExtractSubvector`, When the extracting slice is the same as source vector,
do not need to emit `vector.extract_strided_slice`.

This fixes the lit test `@vector_store_i4` in
`mlir\test\Dialect\Vector\vector-emulate-narrow-type.mlir`, where converting
from `vector<8xi4>` to `vector<4xi8>` does not need slice extraction.

Signed-off-by: Alan Li <me at alanli.org>
---
 .../Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp  | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 7578aadee23a6e..6fd6ccc04e8045 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -150,6 +150,13 @@ static Value staticallyExtractSubvector(OpBuilder &rewriter, Location loc,
   assert((vectorType.getRank() == 1 && extractType.getRank() == 1) &&
          "expected 1-D source and destination types");
   (void)vectorType;
+  assert(frontOffset + subvecSize <= vectorType.getNumElements() &&
+         "subvector out of bounds");
+
+  // do not need extraction if the subvector size is the same as the source
+  if (vectorType.getNumElements() == subvecSize)
+    return source;
+
   auto offsets = rewriter.getI64ArrayAttr({frontOffset});
   auto sizes = rewriter.getI64ArrayAttr({subvecSize});
   auto strides = rewriter.getI64ArrayAttr({1});