[llvm] [InstCombine] Pull extract through broadcast (PR #143380)

[via llvm-commits]

https://github.com/agorenstein-nvidia created https://github.com/llvm/llvm-project/pull/143380

The change adds a new instcombine pattern, and associated test, for patterns like this:

```
  %3 = shufflevector <2 x float> %1, <2 x float> poison, <4 x i32> zeroinitializer
  %4 = extractelement <4 x float> %3, i64 %idx
```

The shufflevector has an all-zeros mask, so the extractelement simply must be the first element of %1, so we transform this to

```
  %2 = extractelement <2 x float> %1, i64 0
```

>From 34c1222e9119655c5d5286bd01be13752f300bbd Mon Sep 17 00:00:00 2001
From: Aaron Gorenstein <agorenstein at nvidia.com>
Date: Fri, 6 Jun 2025 09:55:52 -0400
Subject: [PATCH] Initial changes

---
 .../Transforms/InstCombine/InstCombineVectorOps.cpp  |  6 ++++++
 .../InstCombine/vec_extract_through_broadcast.ll     | 12 ++++++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/vec_extract_through_broadcast.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index f946c3856948b..f2fb26c3ffae8 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -542,6 +542,12 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
         }
       }
     } else if (auto *SVI = dyn_cast<ShuffleVectorInst>(I)) {
+      // extractelt (shufflevector %v1, %v2, zeroinitializer) ->
+      // extractelt %v1, 0
+      if (isa<FixedVectorType>(SVI->getType()))
+        if (all_of(SVI->getShuffleMask(), [](int Elt) { return Elt == 0; }))
+          return ExtractElementInst::Create(SVI->getOperand(0), Builder.getInt64(0));
+
       // If this is extracting an element from a shufflevector, figure out where
       // it came from and extract from the appropriate input element instead.
       // Restrict the following transformation to fixed-length vector.
diff --git a/llvm/test/Transforms/InstCombine/vec_extract_through_broadcast.ll b/llvm/test/Transforms/InstCombine/vec_extract_through_broadcast.ll
new file mode 100644
index 0000000000000..ec80fbf01a3c9
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/vec_extract_through_broadcast.ll
@@ -0,0 +1,12 @@
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+define float @extract_from_zero_init_shuffle(<2 x float> %1, i64 %idx) {
+; CHECK-LABEL: @extract_from_zero_init_shuffle(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x float> [[W:%.*]], i64 0
+; CHECK-NEXT:    ret float [[TMP1]]
+;
+  %3 = shufflevector <2 x float> %1, <2 x float> poison, <4 x i32> zeroinitializer
+  %4 = extractelement <4 x float> %3, i64 %idx
+  ret float %4
+}
+