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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (agorenstein-nvidia)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+6) 
- (added) llvm/test/Transforms/InstCombine/vec_extract_through_broadcast.ll (+12) 


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

``````````

</details>


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