[llvm] InstCombine: Fold shufflevector(select) and shufflevector(phi) (PR #113746)

[Matthias Braun via llvm-commits]

https://github.com/MatzeB updated https://github.com/llvm/llvm-project/pull/113746

>From f5ceabaa8ecc7053e2a312a1cced49fd7c85fc30 Mon Sep 17 00:00:00 2001
From: Matthias Braun <matze at braunis.de>
Date: Fri, 4 Oct 2024 15:45:09 -0700
Subject: [PATCH 1/2] InstCombine: Fold shufflevector(select) and
 shufflevector(phi)

- Transform `shufflevector(select(c, x, y), C)` to
  `select(c, shufflevector(x, C), shufflevector(y, C))` by re-using
  the `FoldOpIntoSelect` helper.
- Transform `shufflevector(phi(x, y), C)` to
  `phi(shufflevector(x, C), shufflevector(y, C))` by re-using the
  `foldOpInotPhi` helper.
---
 .../InstCombine/InstCombineVectorOps.cpp      | 11 +++++
 .../Transforms/InstCombine/vec_shuffle.ll     | 44 +++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index d68ae64f08aa90..7989af5e0e2146 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2900,6 +2900,17 @@ Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   if (Instruction *I = foldIdentityPaddedShuffles(SVI))
     return I;
 
+  if (Constant *C = dyn_cast<Constant>(RHS)) {
+    if (SelectInst *SI = dyn_cast<SelectInst>(LHS)) {
+      if (Instruction *I = FoldOpIntoSelect(SVI, SI, /*FoldWIthMultiUse=*/false))
+        return I;
+    }
+    if (PHINode *PN = dyn_cast<PHINode>(LHS)) {
+      if (Instruction *I = foldOpIntoPhi(SVI, PN))
+        return I;
+    }
+  }
+
   if (match(RHS, m_Poison()) && canEvaluateShuffled(LHS, Mask)) {
     Value *V = evaluateInDifferentElementOrder(LHS, Mask, Builder);
     return replaceInstUsesWith(SVI, V);
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index 75a84e51279b80..d050cf10849e3c 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -2377,3 +2377,47 @@ define <2 x i32> @not_splat_shuffle2(i32 %x) {
   %shuf = shufflevector <2 x i32> %vec, <2 x i32> undef, <2 x i32> <i32 1, i32 3>
   ret <2 x i32> %shuf
 }
+define <2 x i32> @foldselect0(i1 %c) {
+; CHECK-LABEL: @foldselect0(
+; CHECK-NEXT:    [[SHUF:%.*]] = select i1 [[C:%.*]], <2 x i32> <i32 7, i32 42>, <2 x i32> <i32 1, i32 0>
+; CHECK-NEXT:    ret <2 x i32> [[SHUF]]
+;
+  %sel = select i1 %c, <2 x i32> <i32 42, i32 7>, <2 x i32> <i32 0, i32 1>
+  %shuf = shufflevector <2 x i32> %sel, <2 x i32> poison, <2 x i32> <i32 1, i32 0>
+  ret <2 x i32> %shuf
+}
+
+declare i1 @cond()
+declare <4 x i32> @value()
+
+define <4 x i32> @foldphi1() {
+; CHECK-LABEL: @foldphi1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[V:%.*]] = phi <4 x i32> [ zeroinitializer, [[ENTRY:%.*]] ], [ [[XOR:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[VAL:%.*]] = call <4 x i32> @value()
+; CHECK-NEXT:    [[XOR]] = xor <4 x i32> [[V]], [[VAL]]
+; CHECK-NEXT:    [[C:%.*]] = call i1 @cond()
+; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[SHUF1:%.*]] = shufflevector <4 x i32> [[XOR]], <4 x i32> poison, <4 x i32> <i32 3, i32 0, i32 1, i32 2>
+; CHECK-NEXT:    ret <4 x i32> [[SHUF1]]
+;
+entry:
+  br label %loop
+
+loop:
+  %v = phi <4 x i32> [zeroinitializer, %entry], [%shuf1, %loop]
+
+  %shuf0 = shufflevector <4 x i32> %v, <4 x i32> poison, <4 x i32> <i32 1, i32 2, i32 3, i32 0>
+  %val = call <4 x i32> @value()
+  %xor = xor <4 x i32> %shuf0, %val
+  %shuf1 = shufflevector <4 x i32> %xor, <4 x i32> poison, <4 x i32> <i32 3, i32 0, i32 1, i32 2>
+
+  %c = call i1 @cond()
+  br i1 %c, label %loop, label %exit
+
+exit:
+  ret <4 x i32> %shuf1
+}

>From c788c65090e3f949fdb04ef2bbbfe8ca0721b31e Mon Sep 17 00:00:00 2001
From: Matthias Braun <matze at braunis.de>
Date: Fri, 25 Oct 2024 17:33:59 -0700
Subject: [PATCH 2/2] clang-format

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

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 7989af5e0e2146..286d556f0eb362 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2902,7 +2902,8 @@ Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
 
   if (Constant *C = dyn_cast<Constant>(RHS)) {
     if (SelectInst *SI = dyn_cast<SelectInst>(LHS)) {
-      if (Instruction *I = FoldOpIntoSelect(SVI, SI, /*FoldWIthMultiUse=*/false))
+      if (Instruction *I =
+              FoldOpIntoSelect(SVI, SI, /*FoldWIthMultiUse=*/false))
         return I;
     }
     if (PHINode *PN = dyn_cast<PHINode>(LHS)) {