InstSimplify: Add a hook for shufflevector

[zvi]

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.

[spatel]

Why would a shuffle with undef mask automatically return undef? For example:
shufflevector <2 x i8> <i8 42, i8 42>, <2 x i8> <i8 42, i8 42>, <2 x i32> undef

Should this be constant folded to <i8 42, i8 42> , <2 x i8> undef or something else?

lib/Analysis/InstructionSimplify.cpp

4091–4092

Prefer 'auto *' to make the pointer-ness of the cast explicit.

[zvi]

In D31525#715234, @spatel wrote:

Why would a shuffle with undef mask automatically return undef? For example:
shufflevector <2 x i8> <i8 42, i8 42>, <2 x i8> <i8 42, i8 42>, <2 x i32> undef

Should this be constant folded to <i8 42, i8 42> , <2 x i8> undef or something else?

I was asking myself the same question when i saw this in both InstCombine and in ConstanFolding. SimplifyDemandedVectorElts behaves the same.
I couldn't deduce this is allowed from reading the language reference.

[zvi]

auto->auto*

[zvi]

Backing out of the undef mask simpflication which as per pr32486 is probably wrong. Replaced with a different simplification to show the hook is firing.

[spatel]

lib/Analysis/InstructionSimplify.cpp

4097–4099

Why do we need to copy this list of ints? Just read them with ShuffleVectorInst::getMaskValue() and set the MaskSelects0/1 flags right here.

[zvi]

Fixed Sanjay's comment by removing unnecessary container of mask indices.

[spatel]

LGTM.

include/llvm/Analysis/ConstantFolding.h

 Constant *ConstantFoldExtractValueInstruction(Constant *Agg,
                                               ArrayRef<unsigned> Idxs);
 
 /// \brief Attempt to constant fold an extractelement instruction with the
 /// specified operands and indices.  The constant result is returned if
 /// 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.
 Constant *ConstantFoldLoadFromConstPtr(Constant *C, Type *Ty, const DataLayout &DL);
 
 /// ConstantFoldLoadThroughGEPConstantExpr - Given a constant and a
 /// getelementptr constantexpr, return the constant value being addressed by the
 /// constant expression, or null if something is funny and we can't decide.

include/llvm/Analysis/InstructionSimplify.h

namespace llvm {
   /// Given operands for a CastInst, fold the result or return null.
   Value *SimplifyCastInst(unsigned CastOpc, Value *Op, Type *Ty,
                           const DataLayout &DL,
                           const TargetLibraryInfo *TLI = nullptr,
                           const DominatorTree *DT = nullptr,
                           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.
 
 
   /// Given operands for a CmpInst, fold the result or return null.
   Value *SimplifyCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                          const DataLayout &DL,
                          const TargetLibraryInfo *TLI = nullptr,
                          const DominatorTree *DT = nullptr,

lib/Analysis/InstructionSimplify.cpp

Value *llvm::SimplifyCastInst(unsigned CastOpc, Value *Op, Type *Ty,
                               const DataLayout &DL,
                               const TargetLibraryInfo *TLI,
                               const DominatorTree *DT, AssumptionCache *AC,
                               const Instruction *CxtI) {
   return ::SimplifyCastInst(CastOpc, Op, Ty, Query(DL, TLI, DT, AC, CxtI),
                             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);
+
+  SmallVector<int, 32> Idxs;
+  for (unsigned i = 0; i != MaskNumElts; ++i)
+    Idxs.push_back(ShuffleVectorInst::getMaskValue(Mask, i));
+
+  // 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 (auto Idx : Idxs) {
+    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.
 /// If not, this returns null.
 static Value *SimplifyBinOp(unsigned Opcode, Value *LHS, Value *RHS,
                             const Query &Q, unsigned MaxRecurse) {
   switch (Opcode) {
   case Instruction::Add:
case Instruction::ExtractValue: {
     break;
   }
   case Instruction::ExtractElement: {
     auto *EEI = cast<ExtractElementInst>(I);
     Result = SimplifyExtractElementInst(
         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;
   case Instruction::Call: {
     CallSite CS(cast<CallInst>(I));
     Result = SimplifyCall(CS.getCalledValue(), CS.arg_begin(), CS.arg_end(), DL,
                           TLI, DT, AC, I);
     break;

test/Transforms/InstSimplify/shufflevector.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instsimplify -S | FileCheck %s
 
 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> [[SHUF]]
+; CHECK-NEXT:    ret <4 x i32> zeroinitializer
 ;
   %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
 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[SPLAT]], <4 x i32> undef, <4 x i32> <i32 0, i32 3, i32 2, i32 1>
 ; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
 ;
   %splat = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> zeroinitializer
   %shuf = shufflevector <4 x i32> %splat, <4 x i32> undef, <4 x i32> <i32 0, i32 3, i32 2, i32 1>
   ret <4 x i32> %shuf
 }
 
 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> [[SHUF]]
+; CHECK-NEXT:    ret <4 x i32> undef
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> undef
   ret <4 x i32> %shuf
 }
 
 define <4 x i32> @identity_mask_0(<4 x i32> %x) {
 ; CHECK-LABEL: @identity_mask_0(
 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>

 ; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> %x, <4 x i32> <i32 0, i32 1, i32 2, i32 7>
   ret <4 x i32> %shuf
 }
 
 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> [[SHUF]]
+; CHECK-NEXT:    ret <4 x i32> <i32 42, i32 45, i32 44, i32 43>
 ;
   %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
 }
 
 define <4 x i32> @merge(<4 x i32> %x) {
 ; CHECK-LABEL: @merge(
 ; CHECK-NEXT:    [[LOWER:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> undef, <2 x i32> <i32 1, i32 0>