[llvm] [InstCombine] Remove scalable vector extracts to and from the same type (PR #69702)

[Kerry McLaughlin via llvm-commits]

https://github.com/kmclaughlin-arm updated https://github.com/llvm/llvm-project/pull/69702

>From 8ee1ab40ffb696178777cbb30d66be8aa0058025 Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Thu, 19 Oct 2023 16:12:10 +0000
Subject: [PATCH 1/2] [InstCombine] Remove scalable vector extracts to and from
 the same type

visitCallInst already looks for fixed width vector extracts where the
source and destination types are equal. This patch modifies the function
to also identify scalable extracts which can be removed.
---
 .../InstCombine/InstCombineCalls.cpp           | 18 ++++++++++--------
 .../InstCombine/canonicalize-vector-extract.ll | 10 ++++++++++
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index e29fb869686ca0b..00505b5edb58dca 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -2997,24 +2997,26 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         return replaceOperand(CI, 0, InsertTuple);
     }
 
-    auto *DstTy = dyn_cast<FixedVectorType>(ReturnType);
-    auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
+    auto *DstTy = dyn_cast<VectorType>(ReturnType);
+    auto *VecTy = dyn_cast<VectorType>(Vec->getType());
 
-    // Only canonicalize if the destination vector and Vec are fixed
-    // vectors.
     if (DstTy && VecTy) {
-      unsigned DstNumElts = DstTy->getNumElements();
-      unsigned VecNumElts = VecTy->getNumElements();
+      auto DstEltCnt = DstTy->getElementCount();
       unsigned IdxN = cast<ConstantInt>(Idx)->getZExtValue();
 
       // Extracting the entirety of Vec is a nop.
-      if (VecNumElts == DstNumElts) {
+      if (DstEltCnt == VecTy->getElementCount()) {
         replaceInstUsesWith(CI, Vec);
         return eraseInstFromFunction(CI);
       }
 
+      // Only canonicalize to shufflevector if the destination vector and
+      // Vec are fixed vectors.
+      if (dyn_cast<ScalableVectorType>(VecTy) || DstEltCnt.isScalable())
+        break;
+
       SmallVector<int, 8> Mask;
-      for (unsigned i = 0; i != DstNumElts; ++i)
+      for (unsigned i = 0; i != DstEltCnt.getKnownMinValue(); ++i)
         Mask.push_back(IdxN + i);
 
       Value *Shuffle = Builder.CreateShuffleVector(Vec, Mask);
diff --git a/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll b/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
index 26af633698f3552..0b7f1ebc040c2a1 100644
--- a/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
+++ b/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
@@ -10,6 +10,7 @@ declare <3 x i32> @llvm.vector.extract.v3i32.v8i32(<8 x i32> %vec, i64 %idx)
 declare <4 x i32> @llvm.vector.extract.v4i32.nxv4i32(<vscale x 4 x i32> %vec, i64 %idx)
 declare <4 x i32> @llvm.vector.extract.v4i32.v8i32(<8 x i32> %vec, i64 %idx)
 declare <8 x i32> @llvm.vector.extract.v8i32.v8i32(<8 x i32> %vec, i64 %idx)
+declare <vscale x 8 x i32> @llvm.vector.extract.nxv8i32.nxv8i32(<vscale x 8 x i32> %vec, i64 %idx)
 
 ; ============================================================================ ;
 ; Trivial cases
@@ -24,6 +25,15 @@ define <8 x i32> @trivial_nop(<8 x i32> %vec) {
   ret <8 x i32> %1
 }
 
+define <vscale x 8 x i32> @trivial_nop_scalable(<vscale x 8 x i32> %vec) {
+; CHECK-LABEL: define <vscale x 8 x i32> @trivial_nop_scalable(
+; CHECK-SAME: <vscale x 8 x i32> [[VEC:%.*]]) {
+; CHECK-NEXT:    ret <vscale x 8 x i32> [[VEC]]
+;
+  %ext = call <vscale x 8 x i32> @llvm.vector.extract.nxv8i32.nxv8i32(<vscale x 8 x i32> %vec, i64 0)
+  ret <vscale x 8 x i32> %ext
+}
+
 ; ============================================================================ ;
 ; Valid canonicalizations
 ; ============================================================================ ;

>From 1e5b3c8ec3ddd8a65f7dd44bb03bb1bc944572f3 Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Mon, 23 Oct 2023 09:24:39 +0000
Subject: [PATCH 2/2] - Replaced dyn_cast<ScalableVectorType>(VecTy) with
 VecEltCnt.isScalable()

---
 llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 00505b5edb58dca..60233f01e0bde2b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3002,6 +3002,7 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
 
     if (DstTy && VecTy) {
       auto DstEltCnt = DstTy->getElementCount();
+      auto VecEltCnt = VecTy->getElementCount();
       unsigned IdxN = cast<ConstantInt>(Idx)->getZExtValue();
 
       // Extracting the entirety of Vec is a nop.
@@ -3012,7 +3013,7 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
 
       // Only canonicalize to shufflevector if the destination vector and
       // Vec are fixed vectors.
-      if (dyn_cast<ScalableVectorType>(VecTy) || DstEltCnt.isScalable())
+      if (VecEltCnt.isScalable() || DstEltCnt.isScalable())
         break;
 
       SmallVector<int, 8> Mask;