InstCombine: Fold shufflevector(select) and shufflevector(phi)

[MatzeB]

• 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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Matthias Braun (MatzeB)

Changes

• 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.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+11)
• (modified) llvm/test/Transforms/InstCombine/vec_shuffle.ll (+44)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.