Extracting a splat value from vector

[delena]

Summary

We have the following interface in IRBuilder.h:
Value *CreateVectorSplat( unsigned NumElts , Value * V , const Twine & Name = "" ) { ..}

This function creates 2 instructions - “insertelement” and “shuffle” with all-zero mask.

Now I'm adding
Value *getSplatValue(Value *Val)

This function recognizes the pattern - insertelement+shuffle and returns the splat value (or nullptr).
It also returns a splat value form ConstantDataVector, for completeness.

( I don't know why functions in this file are not wrapped with "namespace llvm")

[jroelofs]

Probably best to commit this with a use of it... otherwise this is just dead code.

[delena]

I use the new getSplatValue() interface in DAG building.

[delena]

Ping*

[hfinkel]

lib/Analysis/VectorUtils.cpp

419

Does this handle the case where some of the mask entries are undef?

[delena]

Addressed Hal's comments: added more verification for shuffle mask - zero or undef.
Added index verification for InsertElement instruction.

[hfinkel]

This LGTM. Please either commit this with a test case covering the masks with undef (and the insertelement index verification), or add such a test case to D11121 before you commit.

lib/Analysis/VectorUtils.cpp

428

fisrt -> first

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3092

I'd highly prefer that Gep were GEP, GepPtrs -> GEPPtrs, etc. to match our general naming convention. But that should be done in a separate commit regardless.

include/llvm/Analysis/VectorUtils.h

 /// strides "a[i*stride]". Returns the symbolic stride, or null otherwise.
 Value *getStrideFromPointer(Value *Ptr, ScalarEvolution *SE, Loop *Lp);
 
 /// \brief Given a vector and an element number, see if the scalar value is
 /// already around as a register, for example if it were inserted then extracted
 /// from the vector.
 Value *findScalarElement(Value *V, unsigned EltNo);
 
+/// \brief Get splat value if the input is a splat vector or return nullptr.
+/// The value may be extracted from a splat constants vector or from
+/// a sequence of instructions that broadcast a single value into a vector.
+Value *getSplatValue(Value *V);
+
 } // llvm namespace
 
 #endif

lib/Analysis/VectorUtils.cpp

 
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Value.h"
+#include "llvm/IR/Constants.h"
 
 /// \brief Identify if the intrinsic is trivially vectorizable.
 /// This method returns true if the intrinsic's argument types are all
 /// scalars for the scalar form of the intrinsic and all vectors for
 /// the vector form of the intrinsic.
 bool llvm::isTriviallyVectorizable(Intrinsic::ID ID) {
   switch (ID) {
   case Intrinsic::sqrt:
if (match(V,
                                       llvm::PatternMatch::m_Constant(Con)))) {
     if (Con->getAggregateElement(EltNo)->isNullValue())
       return findScalarElement(Val, EltNo);
   }
 
   // Otherwise, we don't know.
   return nullptr;
 }
+
+/// \brief Get splat value if the input is a splat vector or return nullptr.
+/// The value may be extracted from a splat constants vector or from
+/// a sequence of instructions that broadcast a single value into a vector.
+llvm::Value *llvm::getSplatValue(Value *V) {
+  llvm::ConstantDataVector *CV = dyn_cast<llvm::ConstantDataVector>(V);
+  if (CV)
+    return CV->getSplatValue();
+  llvm::ShuffleVectorInst *ShuffleInst = dyn_cast<llvm::ShuffleVectorInst>(V);
+  if (!ShuffleInst)
+    return nullptr;
+  // All-zero (our undef) shuffle mask elements.
+  for (int i : ShuffleInst->getShuffleMask())
+    if (i != 0 && i != -1)
+      return nullptr;
+  // The fisrt shuffle source is 'insertelement' with index 0.
+  llvm::InsertElementInst *InsertEltInst =
+    dyn_cast<llvm::InsertElementInst>(ShuffleInst->getOperand(0));
+  if (!InsertEltInst || !isa<ConstantInt>(InsertEltInst->getOperand(2)) ||
+      !cast<ConstantInt>(InsertEltInst->getOperand(2))->isNullValue())
+    return nullptr;
+
+  return InsertEltInst->getOperand(1);
+}

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/VectorUtils.h"
 #include "llvm/CodeGen/FastISel.h"
 #include "llvm/CodeGen/FunctionLoweringInfo.h"
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/GCStrategy.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"

 // Gather/scatter receive a vector of pointers.
 // This vector of pointers may be represented as a base pointer + vector of 
 // indices, it depends on GEP and instruction preceeding GEP
 // that calculates indices
 static bool getUniformBase(Value *& Ptr, SDValue& Base, SDValue& Index,
                            SelectionDAGBuilder* SDB) {
 
   assert (Ptr->getType()->isVectorTy() && "Uexpected pointer type");
   GetElementPtrInst *Gep = dyn_cast<GetElementPtrInst>(Ptr);
   if (!Gep || Gep->getNumOperands() > 2)
     return false;
-  ShuffleVectorInst *ShuffleInst = 
-    dyn_cast<ShuffleVectorInst>(Gep->getPointerOperand());
-  if (!ShuffleInst || !ShuffleInst->getMask()->isNullValue() ||
-      cast<Instruction>(ShuffleInst->getOperand(0))->getOpcode() !=
-      Instruction::InsertElement)
+  Value *GepPtrs = Gep->getPointerOperand();
+  if (!(Ptr = getSplatValue(GepPtrs)))
     return false;
 
-  Ptr = cast<InsertElementInst>(ShuffleInst->getOperand(0))->getOperand(1);
-
   SelectionDAG& DAG = SDB->DAG;
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   // Check is the Ptr is inside current basic block
   // If not, look for the shuffle instruction
   if (SDB->findValue(Ptr))
     Base = SDB->getValue(Ptr);
-  else if (SDB->findValue(ShuffleInst)) {
-    SDValue ShuffleNode = SDB->getValue(ShuffleInst);
-    SDLoc sdl = ShuffleNode;
-    Base = DAG.getNode(
-        ISD::EXTRACT_VECTOR_ELT, sdl,
-        ShuffleNode.getValueType().getScalarType(), ShuffleNode,
-        DAG.getConstant(0, sdl, TLI.getVectorIdxTy(DAG.getDataLayout())));
+  else if (SDB->findValue(GepPtrs)) {
+    SDValue GepPtrsVal = SDB->getValue(GepPtrs);
+    SDLoc sdl = GepPtrsVal;
+    EVT IdxVT = TLI.getVectorIdxTy(DAG.getDataLayout());
+    Base = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, sdl,
+                       GepPtrsVal.getValueType().getScalarType(), GepPtrsVal,
+                       DAG.getConstant(0, sdl, IdxVT));
     SDB->setValue(Ptr, Base);
   }
   else
     return false;
 
   Value *IndexVal = Gep->getOperand(1);
   if (SDB->findValue(IndexVal)) {
     Index = SDB->getValue(IndexVal);