Optimization for Gather/Scatter with uniform base

[delena]

Summary

Vector 'getelementptr' with scalar base is an opportunity for gather/scatter intrinsic to generate a better sequence.
While looking for uniform base, we want to use the scalar base pointer of GEP, if exists.

[delena]

Vector 'getelementptr' with scalar base is an opportunity for gather/scatter intrinsic to generate a better sequence.
While looking for uniform base, we want to use the scalar base pointer of GEP, if exists.

[ab]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3083

unform -> uniform

3083–3084

receive -> take, perhaps?

3085–3086

Perhaps IR examples for both representations might be helpful?

3141–3143

else on the same line?

3147–3148

Should we check that the insertelement index is 0? (or even generalize and check that the broadcast and insertelement have the same index, though that sounds unnecessary)

3150

inside -> inside the

3157

I'm a bit uneasy with changing IndexVal outside the if(): when the latter fails, IndexVal will be incorrect for subsequent users.

Not sure it's better, but what about something like:

Value *OrigIndexVal = ...;
if (SDB->findValue(OrigIndexVal)) {
  IndexVal = OrigIndexVal;
  ...
}

3158

unnecessary end-of-line whitespace?

3163

Why not EVT?

[delena]

Ahmed,

Thanks a lot for the review. I addressed all your comments and uploaded a new patch.

[delena]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3083–3084

Vector of pointers is the first argument of Gather/Scatter

3147–3148

I submitted another patch with getSplatValue() as vector utility and planned to use it here, but it is stuck on review.
http://reviews.llvm.org/D11124

3157

Actually, I don't need the IndexVal any more, but I changed the code.

3163

The same interface in EVT requires Context. But I changed, the result is the same.

[hfinkel]

This LGTM, but let's get getSplatValue figured out, commit that, and then commit this patch using getSplatValue.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3122

I suppose you'll replace this logic with getSplatValue when that lands.

[delena]

I use the new interface getSplatValue().
I added more tests, Hal and James asked for.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

getMachineMemOperand(MachinePointerInfo(PtrOperand),
                           MachineMemOperand::MOStore,  VT.getStoreSize(),
                           Alignment, AAInfo);
   SDValue StoreNode = DAG.getMaskedStore(getRoot(), sdl, Src0, Ptr, Mask, VT,
                                          MMO, false);
   DAG.setRoot(StoreNode);
   setValue(&I, StoreNode);
 }
 
-// Gather/scatter receive a vector of pointers.
+// Get a uniform base for the Gather/Scatter input.
+// The first argument of the Gather/Scatter is 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
+// indices. Usually, the vector of pointers comes from a 'getelementptr'
+// instruction. Extracting a uniform base depends on the GEP and the instruction
+// preceeding GEP that calculates indices.
+// Example:
+//   %gep.ptr = getelementptr i32, <8 x i32*> %vptr, <8 x i32> %ind
+//  or
+//   %gep.ptr = getelementptr i32, i32* %ptr,        <8 x i32> %ind
+// %res = call <8 x i32> @llvm.masked.gather.v8i32(<8 x i32*> %gep.ptr, ..
+
 static bool getUniformBase(Value *& Ptr, SDValue& Base, SDValue& Index,
                            SelectionDAGBuilder* SDB) {
 
+  SelectionDAG& DAG = SDB->DAG;
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  LLVMContext &Context = *DAG.getContext();
+
   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());
+
+  Value *GepBasePtr = Gep->getPointerOperand();
+  Value *IndexVal = Gep->getOperand(1);
+  // The operands of the GEP may be defined in another basic block.
+  // In this case we'll not find nodes for the operands.
+  if (!SDB->findValue(GepBasePtr) || !SDB->findValue(IndexVal))
+    return false;
+
+  // If GEP base is scalar - this is the uniform base we are looking for.
+  if (!GepBasePtr->getType()->isVectorTy()) {
+    Ptr = GepBasePtr;
+    Base = SDB->getValue(Ptr);
+  } else {
+    // The base is a vector. But may be this vector is splat.
+    // Try to find a preceding broadcast.
+    ShuffleVectorInst *ShuffleInst = dyn_cast<ShuffleVectorInst>(GepBasePtr);
     if (!ShuffleInst || !ShuffleInst->getMask()->isNullValue() ||
-      cast<Instruction>(ShuffleInst->getOperand(0))->getOpcode() !=
-      Instruction::InsertElement)
+        !isa<InsertElementInst>(ShuffleInst->getOperand(0)))
       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
+    // Check if the Ptr is inside the current basic block.
     // If not, look for the shuffle instruction
     if (SDB->findValue(Ptr))
       Base = SDB->getValue(Ptr);
-  else if (SDB->findValue(ShuffleInst)) {
+    else {
       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())));
+      SDLoc sdl(ShuffleNode);
+      EVT IdxVT = TLI.getVectorIdxTy(DAG.getDataLayout());
+      Base = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, sdl,
+                         ShuffleNode.getValueType().getScalarType(),
+                         ShuffleNode, 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);
 
-    if (SExtInst* Sext = dyn_cast<SExtInst>(IndexVal)) {
+  // Suppress sign extension.
+    if (SDB->findValue(Sext->getOperand(0))) {
       IndexVal = Sext->getOperand(0);
-      if (SDB->findValue(IndexVal))
-        Index = SDB->getValue(IndexVal);
     }
-    return true;
   }
-  return false;
+  if (!Index.getValueType().isVector()) {
+    unsigned GEPWidth = Gep->getType()->getVectorNumElements();
+    EVT VT = EVT::getVectorVT(Context, Index.getValueType(), GEPWidth);
+    SmallVector<SDValue, 16> Ops(GEPWidth, Index);
+    Index = DAG.getNode(ISD::BUILD_VECTOR, SDLoc(Index), VT, Ops);
+  }
+  return true;
 }
 
 void SelectionDAGBuilder::visitMaskedScatter(const CallInst &I) {
   SDLoc sdl = getCurSDLoc();
 
   // llvm.masked.scatter.*(Src0, Ptrs, alignemt, Mask)
   Value  *Ptr = I.getArgOperand(1);
   SDValue Src0 = getValue(I.getArgOperand(0));
   SDValue Mask = getValue(I.getArgOperand(3));
   EVT VT = Src0.getValueType();
   unsigned Alignment = (cast<ConstantInt>(I.getArgOperand(2)))->getZExtValue();
   if (!Alignment)
     Alignment = DAG.getEVTAlignment(VT);

test/CodeGen/X86/masked_gather_scatter.ll

 ; KNL: vinserti64x4
 define <16 x i32> @test8(<16 x i32*> %ptr.random, <16 x i32> %ind, i16 %mask) {
   %imask = bitcast i16 %mask to <16 x i1>
   %gt1 = call <16 x i32> @llvm.masked.gather.v16i32(<16 x i32*> %ptr.random, i32 4, <16 x i1> %imask, <16 x i32>undef)
   %gt2 = call <16 x i32> @llvm.masked.gather.v16i32(<16 x i32*> %ptr.random, i32 4, <16 x i1> %imask, <16 x i32>%gt1)
   %res = add <16 x i32> %gt1, %gt2
   ret <16 x i32> %res
 }
+
+%struct.RT = type { i8, [10 x [20 x i32]], i8 }
+%struct.ST = type { i32, double, %struct.RT }
+
+; Masked gather for agregate types
+; Test9 and Test10 should give the same result (scalar and vector indices in GEP)
+
+; KNL-LABEL: test9
+; KNL: vpbroadcastq    %rdi, %zmm
+; KNL: vpmovsxdq
+; KNL: vpbroadcastq
+; KNL: vpmuludq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpgatherqd      (,%zmm
+
+define <8 x i32> @test9(%struct.ST* %base, <8 x i64> %ind1, <8 x i32>%ind5) {
+entry:
+  %broadcast.splatinsert = insertelement <8 x %struct.ST*> undef, %struct.ST* %base, i32 0
+  %broadcast.splat = shufflevector <8 x %struct.ST*> %broadcast.splatinsert, <8 x %struct.ST*> undef, <8 x i32> zeroinitializer
+
+  %arrayidx = getelementptr  %struct.ST, <8 x %struct.ST*> %broadcast.splat, <8 x i64> %ind1, <8 x i32> <i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2>, <8 x i32><i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>, <8 x i32> %ind5, <8 x i64> <i64 13, i64 13, i64 13, i64 13, i64 13, i64 13, i64 13, i64 13>
+  %res = call <8 x i32 >  @llvm.masked.gather.v8i32(<8 x i32*>%arrayidx, i32 4, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <8 x i32> undef)
+  ret <8 x i32> %res
+}
+
+; KNL-LABEL: test10
+; KNL: vpbroadcastq    %rdi, %zmm
+; KNL: vpmovsxdq
+; KNL: vpbroadcastq
+; KNL: vpmuludq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpaddq
+; KNL: vpgatherqd      (,%zmm
+define <8 x i32> @test10(%struct.ST* %base, <8 x i64> %i1, <8 x i32>%ind5) {
+entry:
+  %broadcast.splatinsert = insertelement <8 x %struct.ST*> undef, %struct.ST* %base, i32 0
+  %broadcast.splat = shufflevector <8 x %struct.ST*> %broadcast.splatinsert, <8 x %struct.ST*> undef, <8 x i32> zeroinitializer
+
+  %arrayidx = getelementptr  %struct.ST, <8 x %struct.ST*> %broadcast.splat, <8 x i64> %i1, i32 2, i32 1, <8 x i32> %ind5, i64 13
+  %res = call <8 x i32 >  @llvm.masked.gather.v8i32(<8 x i32*>%arrayidx, i32 4, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <8 x i32> undef)
+  ret <8 x i32> %res
+}
+
+; Splat index in GEP, requires broadcast
+; KNL-LABEL: test11
+; KNL: vpbroadcastd    %esi, %zmm
+; KNL: vgatherdps      (%rdi,%zmm
+define <16 x float> @test11(float* %base, i32 %ind) {
+
+  %broadcast.splatinsert = insertelement <16 x float*> undef, float* %base, i32 0
+  %broadcast.splat = shufflevector <16 x float*> %broadcast.splatinsert, <16 x float*> undef, <16 x i32> zeroinitializer
+
+  %gep.random = getelementptr float, <16 x float*> %broadcast.splat, i32 %ind
+
+  %res = call <16 x float> @llvm.masked.gather.v16f32(<16 x float*> %gep.random, i32 4, <16 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <16 x float> undef)
+  ret <16 x float>%res
+}
+
+; We are checking the uniform base here. It is taken directly from input to vgatherdps
+; KNL-LABEL: test12
+; KNL: vgatherdps      (%rdi,%zmm
+define <16 x float> @test12(float* %base, <16 x i32> %ind) {
+
+  %sext_ind = sext <16 x i32> %ind to <16 x i64>
+  %gep.random = getelementptr float, float *%base, <16 x i64> %sext_ind
+
+  %res = call <16 x float> @llvm.masked.gather.v16f32(<16 x float*> %gep.random, i32 4, <16 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <16 x float> undef)
+  ret <16 x float>%res
+}