[llvm] [Instcombine] Combine extractelement from a vector_extract at index 0 (PR #151491)

[Kerry McLaughlin via llvm-commits]

https://github.com/kmclaughlin-arm created https://github.com/llvm/llvm-project/pull/151491

Extracting any element from a subvector starting at index 0 is
equivalent to extracting from the original vector, i.e.
  extract_elt(vector_extract(x, 0), y) -> extract_elt(x, y)

>From 03413139be5f006859c933203922c7f63bfd2a60 Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Thu, 17 Jul 2025 08:49:24 +0000
Subject: [PATCH 1/2] - Add a test for extractelement from scalable
 vector_extract at index 0

---
 .../InstCombine/scalable-extract-subvec-elt.ll     | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll

diff --git a/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll b/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll
new file mode 100644
index 0000000000000..d0cfc48420b43
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll
@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=instcombine < %s | FileCheck %s
+
+define i1 @scalable_test(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: define i1 @scalable_test(
+; CHECK-SAME: <vscale x 4 x i1> [[A:%.*]]) {
+; CHECK-NEXT:    [[SUBVEC:%.*]] = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1(<vscale x 4 x i1> [[A]], i64 0)
+; CHECK-NEXT:    [[ELT:%.*]] = extractelement <vscale x 2 x i1> [[SUBVEC]], i64 1
+; CHECK-NEXT:    ret i1 [[ELT]]
+;
+  %subvec = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %a, i64 0)
+  %elt = extractelement <vscale x 2 x i1> %subvec, i32 1
+  ret i1 %elt
+}

>From d93bdead6c6d36907c3ed539407736070792f92b Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Thu, 17 Jul 2025 08:51:47 +0000
Subject: [PATCH 2/2] [Instcombine] Combine extractelement from a
 vector_extract at index 0

Extracting any element from a subvector starting at index 0 is
equivalent to extracting from the original vector, i.e.

  extract_elt(vector_extract(x, 0), y) -> extract_elt(x, y)
---
 llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp    | 6 +++++-
 .../Transforms/InstCombine/scalable-extract-subvec-elt.ll   | 3 +--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 00b877b8a07ef..6f2adba9e3f6b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -444,7 +444,11 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
         else
           Idx = PoisonValue::get(Ty);
         return replaceInstUsesWith(EI, Idx);
-      }
+      } else if (IID == Intrinsic::vector_extract)
+        // If II is a subvector starting at index 0, extract from the wider
+        // source vector
+        if (cast<ConstantInt>(II->getArgOperand(1))->getZExtValue() == 0)
+          return ExtractElementInst::Create(II->getArgOperand(0), Index);
     }
 
     // InstSimplify should handle cases where the index is invalid.
diff --git a/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll b/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll
index d0cfc48420b43..38cbeb0df98bd 100644
--- a/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll
+++ b/llvm/test/Transforms/InstCombine/scalable-extract-subvec-elt.ll
@@ -4,8 +4,7 @@
 define i1 @scalable_test(<vscale x 4 x i1> %a) {
 ; CHECK-LABEL: define i1 @scalable_test(
 ; CHECK-SAME: <vscale x 4 x i1> [[A:%.*]]) {
-; CHECK-NEXT:    [[SUBVEC:%.*]] = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1(<vscale x 4 x i1> [[A]], i64 0)
-; CHECK-NEXT:    [[ELT:%.*]] = extractelement <vscale x 2 x i1> [[SUBVEC]], i64 1
+; CHECK-NEXT:    [[ELT:%.*]] = extractelement <vscale x 4 x i1> [[A]], i64 1
 ; CHECK-NEXT:    ret i1 [[ELT]]
 ;
   %subvec = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %a, i64 0)