InstSimplify: Add a hook for shufflevector

Summary:
Add a hook for simplification of shufflevector's with the following rules:
- Constant folding - NFC, as it was already being done by the default handler.
-  If only one of the operands is constant, constant fold the shuffle if the
    mask does not select elements from the variable operand -  to show the hook is firing and affecting the test-cases.

Reviewers: RKSimon, craig.topper, spatel, sanjoy, nlopes, majnemer

Reviewed By: spatel

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D31525

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@299393 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/Analysis/ConstantFolding.h

@@ -100,6 +100,12 @@ Constant *ConstantFoldExtractValueInstruction(Constant *Agg,
 /// successful; if not, null is returned.
 Constant *ConstantFoldExtractElementInstruction(Constant *Val, Constant *Idx);
 
+/// \brief Attempt to constant fold a shufflevector instruction with the
+/// specified operands and indices.  The constant result is returned if
+/// successful; if not, null is returned.
+Constant *ConstantFoldShuffleVectorInstruction(Constant *V1, Constant *V2,
+                                               Constant *Mask);
+
 /// ConstantFoldLoadFromConstPtr - Return the value that a load from C would
 /// produce if it is constant and determinable.  If this is not determinable,
 /// return null.

include/llvm/Analysis/InstructionSimplify.h

@@ -247,6 +247,14 @@ namespace llvm {
                           AssumptionCache *AC = nullptr,
                           const Instruction *CxtI = nullptr);
 
+  /// Given operands for a ShuffleVectorInst, fold the result or return null.
+  Value *SimplifyShuffleVectorInst(Value *Op0, Value *Op1, Constant *Mask,
+                                   Type *RetTy, const DataLayout &DL,
+                                   const TargetLibraryInfo *TLI = nullptr,
+                                   const DominatorTree *DT = nullptr,
+                                   AssumptionCache *AC = nullptr,
+                                   const Instruction *CxtI = nullptr);
+
   //=== Helper functions for higher up the class hierarchy.
 
 

lib/Analysis/InstructionSimplify.cpp

@@ -4081,6 +4081,50 @@ Value *llvm::SimplifyCastInst(unsigned CastOpc, Value *Op, Type *Ty,
                             RecursionLimit);
 }
 
+static Value *SimplifyShuffleVectorInst(Value *Op0, Value *Op1, Constant *Mask,
+                                        Type *RetTy, const Query &Q,
+                                        unsigned MaxRecurse) {
+  unsigned MaskNumElts = Mask->getType()->getVectorNumElements();
+  unsigned InVecNumElts = Op0->getType()->getVectorNumElements();
+
+  auto *Op0Const = dyn_cast<Constant>(Op0);
+  auto *Op1Const = dyn_cast<Constant>(Op1);
+
+  // If all operands are constant, constant fold the shuffle.
+  if (Op0Const && Op1Const)
+    return ConstantFoldShuffleVectorInstruction(Op0Const, Op1Const, Mask);
+
+  // If only one of the operands is constant, constant fold the shuffle if the
+  // mask does not select elements from the variable operand.
+  bool MaskSelects0 = false, MaskSelects1 = false;
+  for (unsigned i = 0; i != MaskNumElts; ++i) {
+    int Idx = ShuffleVectorInst::getMaskValue(Mask, i);
+    if (Idx == -1)
+      continue;
+    if ((unsigned)Idx < InVecNumElts)
+      MaskSelects0 = true;
+    else
+      MaskSelects1 = true;
+  }
+  if (!MaskSelects0 && Op1Const)
+    return ConstantFoldShuffleVectorInstruction(UndefValue::get(Op0->getType()),
+                                                Op1Const, Mask);
+  if (!MaskSelects1 && Op0Const)
+    return ConstantFoldShuffleVectorInstruction(
+        Op0Const, UndefValue::get(Op0->getType()), Mask);
+
+  return nullptr;
+}
+
+/// Given operands for a ShuffleVectorInst, fold the result or return null.
+Value *llvm::SimplifyShuffleVectorInst(
+    Value *Op0, Value *Op1, Constant *Mask, Type *RetTy,
+    const DataLayout &DL, const TargetLibraryInfo *TLI, const DominatorTree *DT,
+    AssumptionCache *AC, const Instruction *CxtI) {
+  return ::SimplifyShuffleVectorInst(
+      Op0, Op1, Mask, RetTy, Query(DL, TLI, DT, AC, CxtI), RecursionLimit);
+}
+
 //=== Helper functions for higher up the class hierarchy.
 
 /// Given operands for a BinaryOperator, see if we can fold the result.
@@ -4569,6 +4613,13 @@ Value *llvm::SimplifyInstruction(Instruction *I, const DataLayout &DL,
         EEI->getVectorOperand(), EEI->getIndexOperand(), DL, TLI, DT, AC, I);
     break;
   }
+  case Instruction::ShuffleVector: {
+    auto *SVI = cast<ShuffleVectorInst>(I);
+    Result = SimplifyShuffleVectorInst(SVI->getOperand(0), SVI->getOperand(1),
+                                       SVI->getMask(), SVI->getType(), DL, TLI,
+                                       DT, AC, I);
+    break;
+  }
   case Instruction::PHI:
     Result = SimplifyPHINode(cast<PHINode>(I), Query(DL, TLI, DT, AC, I));
     break;

test/Transforms/InstSimplify/shufflevector.ll

@@ -3,13 +3,29 @@
 
 define <4 x i32> @const_folding(<4 x i32> %x) {
 ; CHECK-LABEL: @const_folding(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> zeroinitializer, <4 x i32> <i32 5, i32 4, i32 5, i32 4>
+; CHECK-NEXT:    ret <4 x i32> zeroinitializer
-; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> zeroinitializer, <4 x i32> <i32 5, i32 4, i32 5, i32 4>
   ret <4 x i32> %shuf
 }
 
+define <4 x i32> @const_folding1(<4 x i32> %x) {
+; CHECK-LABEL: @const_folding1(
+; CHECK-NEXT:    ret <4 x i32> <i32 5, i32 5, i32 5, i32 5>
+;
+  %shuf = shufflevector <4 x i32> <i32 5, i32 4, i32 5, i32 4>, <4 x i32> %x, <4 x i32> zeroinitializer
+  ret <4 x i32> %shuf
+}
+
+define <4 x i32> @const_folding_negative(<3 x i32> %x) {
+; CHECK-LABEL: @const_folding_negative(
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <3 x i32> [[X:%.*]], <3 x i32> zeroinitializer, <4 x i32> <i32 2, i32 4, i32 5, i32 4>
+; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
+;
+  %shuf = shufflevector <3 x i32> %x, <3 x i32> zeroinitializer, <4 x i32> <i32 2, i32 4, i32 5, i32 4>
+  ret <4 x i32> %shuf
+}
+
 define <4 x i32> @splat_operand(<4 x i32> %x) {
 ; CHECK-LABEL: @splat_operand(
 ; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> undef, <4 x i32> zeroinitializer
@@ -23,8 +39,7 @@ define <4 x i32> @splat_operand(<4 x i32> %x) {
 
 define <4 x i32> @undef_mask(<4 x i32> %x) {
 ; CHECK-LABEL: @undef_mask(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> undef, <4 x i32> undef
+; CHECK-NEXT:    ret <4 x i32> undef
-; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> undef
   ret <4 x i32> %shuf
@@ -59,8 +74,7 @@ define <4 x i32> @pseudo_identity_mask(<4 x i32> %x) {
 
 define <4 x i32> @const_operand(<4 x i32> %x) {
 ; CHECK-LABEL: @const_operand(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> <i32 42, i32 43, i32 44, i32 45>, <4 x i32> [[X:%.*]], <4 x i32> <i32 0, i32 3, i32 2, i32 1>
+; CHECK-NEXT:    ret <4 x i32> <i32 42, i32 45, i32 44, i32 43>
-; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
 ;
   %shuf = shufflevector <4 x i32> <i32 42, i32 43, i32 44, i32 45>, <4 x i32> %x, <4 x i32> <i32 0, i32 3, i32 2, i32 1>
   ret <4 x i32> %shuf