[CodeGenPrepare][X86] Teach optimizeGatherScatterInst to turn a splat pointer into GEP with scalar base and 0 index

[craig.topper]

Summary

This helps SelectionDAGBuilder recognize the splat can be used as a uniform base.

[craig.topper]

Ping

[spatel]

llvm/lib/Analysis/VectorUtils.cpp

345

Inconsistent use of "llvm::" ? Not sure if there's some guideline/pref, but seems odd to see one Value use the namespace but the other does not.

llvm/lib/CodeGen/CodeGenPrepare.cpp

5403

I know we're doing this cast in the other clause too, but is there anything guarding this function from crashing with scalable vectors?
cc @ctetreau

llvm/test/CodeGen/X86/masked_gather_scatter.ll

3323

pre-commit baseline tests?

[craig.topper]

-Pre-commit tests
-Remove hopefully unnecessary llvm::
-Add early bailout and FIXME for scalable vetors.

[efriedma]

If the address is a splat, generating a gather/scatter seems a little weird. I guess if the mask is non-trivial, there isn't any obvious way to rewrite it in target-independent IR. though.

No testcases for gather?

[craig.topper]

Add gather test cases

[RKSimon]

In D86371#2247925, @efriedma wrote:

If the address is a splat, generating a gather/scatter seems a little weird. I guess if the mask is non-trivial, there isn't any obvious way to rewrite it in target-independent IR. though.

Do we have any constant mask test cases?

[craig.topper]

In D86371#2249125, @RKSimon wrote:

In D86371#2247925, @efriedma wrote:

If the address is a splat, generating a gather/scatter seems a little weird. I guess if the mask is non-trivial, there isn't any obvious way to rewrite it in target-independent IR. though.

Do we have any constant mask test cases?

No, but I don't think it really matters for this transform. If the mask is constant and the pointer is a splat, we should probably replace the gather/scatter with scalar load/store and shuffle/extractelement in instcombine.

[RKSimon]

OK that makes sense, LGTM

llvm/include/llvm/Analysis/VectorUtils.h

 /// If all non-negative \p Mask elements are the same value, return that value.
 /// If all elements are negative (undefined) or \p Mask contains different
 /// non-negative values, return -1.
 int getSplatIndex(ArrayRef<int> Mask);
 
 /// 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.
-const Value *getSplatValue(const Value *V);
+Value *getSplatValue(const Value *V);
 
 /// Return true if each element of the vector value \p V is poisoned or equal to
 /// every other non-poisoned element. If an index element is specified, either
 /// every element of the vector is poisoned or the element at that index is not
 /// poisoned and equal to every other non-poisoned element.
 /// This may be more powerful than the related getSplatValue() because it is
 /// not limited by finding a scalar source value to a splatted vector.
 bool isSplatValue(const Value *V, int Index = -1, unsigned Depth = 0);

llvm/lib/Analysis/VectorUtils.cpp

int llvm::getSplatIndex(ArrayRef<int> Mask) {
   assert((SplatIndex == -1 || SplatIndex >= 0) && "Negative index?");
   return SplatIndex;
 }
 
 /// Get splat value if the input is a splat vector or return nullptr.
 /// This function is not fully general. It checks only 2 cases:
 /// the input value is (1) a splat constant vector or (2) a sequence
 /// of instructions that broadcasts a scalar at element 0.
-const llvm::Value *llvm::getSplatValue(const Value *V) {
+llvm::Value *llvm::getSplatValue(const Value *V) {
   if (isa<VectorType>(V->getType()))
     if (auto *C = dyn_cast<Constant>(V))
       return C->getSplatValue();
 
   // shuf (inselt ?, Splat, 0), ?, <0, undef, 0, ...>
   Value *Splat;
   if (match(V,
             m_Shuffle(m_InsertElt(m_Value(), m_Value(Splat), m_ZeroInt()),

llvm/lib/CodeGen/CodeGenPrepare.cpp

 /// zeroes, we can replace it with 2 GEPs where the second has 2 operands.
 ///
 /// If the final index isn't a vector or is a splat, we can emit a scalar GEP
 /// followed by a GEP with an all zeroes vector index. This will enable
 /// SelectionDAGBuilder to use a the scalar GEP as the uniform base and have a
 /// zero index.
 bool CodeGenPrepare::optimizeGatherScatterInst(Instruction *MemoryInst,
                                                Value *Ptr) {
-  const GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Ptr);
-  if (!GEP || !GEP->hasIndices())
+  Value *NewAddr;
+
+  if (const auto *GEP = dyn_cast<GetElementPtrInst>(Ptr)) {
+    // Don't optimize GEPs that don't have indices.
+    if (!GEP->hasIndices())
       return false;
 
     // If the GEP and the gather/scatter aren't in the same BB, don't optimize.
     // FIXME: We should support this by sinking the GEP.
     if (MemoryInst->getParent() != GEP->getParent())
       return false;
 
     SmallVector<Value *, 2> Ops(GEP->op_begin(), GEP->op_end());
 
     bool RewriteGEP = false;
 
     if (Ops[0]->getType()->isVectorTy()) {
-    Ops[0] = const_cast<Value *>(getSplatValue(Ops[0]));
+      Ops[0] = getSplatValue(Ops[0]);
       if (!Ops[0])
         return false;
       RewriteGEP = true;
     }
 
     unsigned FinalIndex = Ops.size() - 1;
 
     // Ensure all but the last index is 0.
     // FIXME: This isn't strictly required. All that's required is that they are
     // all scalars or splats.
     for (unsigned i = 1; i < FinalIndex; ++i) {
       auto *C = dyn_cast<Constant>(Ops[i]);
       if (!C)
         return false;
       if (isa<VectorType>(C->getType()))
         C = C->getSplatValue();
       auto *CI = dyn_cast_or_null<ConstantInt>(C);
       if (!CI || !CI->isZero())
         return false;
       // Scalarize the index if needed.
       Ops[i] = CI;
     }
 
     // Try to scalarize the final index.
     if (Ops[FinalIndex]->getType()->isVectorTy()) {
-    if (Value *V = const_cast<Value *>(getSplatValue(Ops[FinalIndex]))) {
+      if (Value *V = getSplatValue(Ops[FinalIndex])) {
         auto *C = dyn_cast<ConstantInt>(V);
         // Don't scalarize all zeros vector.
         if (!C || !C->isZero()) {
           Ops[FinalIndex] = V;
           RewriteGEP = true;
         }
       }
     }
 
     // If we made any changes or the we have extra operands, we need to generate
     // new instructions.
     if (!RewriteGEP && Ops.size() == 2)
       return false;
 
     unsigned NumElts = cast<FixedVectorType>(Ptr->getType())->getNumElements();
 
     IRBuilder<> Builder(MemoryInst);
 
     Type *ScalarIndexTy = DL->getIndexType(Ops[0]->getType()->getScalarType());
 
-  Value *NewAddr;
-
     // If the final index isn't a vector, emit a scalar GEP containing all ops
     // and a vector GEP with all zeroes final index.
     if (!Ops[FinalIndex]->getType()->isVectorTy()) {
       NewAddr = Builder.CreateGEP(Ops[0], makeArrayRef(Ops).drop_front());
       auto *IndexTy = FixedVectorType::get(ScalarIndexTy, NumElts);
       NewAddr = Builder.CreateGEP(NewAddr, Constant::getNullValue(IndexTy));
     } else {
       Value *Base = Ops[0];
       Value *Index = Ops[FinalIndex];
 
       // Create a scalar GEP if there are more than 2 operands.
       if (Ops.size() != 2) {
         // Replace the last index with 0.
         Ops[FinalIndex] = Constant::getNullValue(ScalarIndexTy);
         Base = Builder.CreateGEP(Base, makeArrayRef(Ops).drop_front());
       }
 
       // Now create the GEP with scalar pointer and vector index.
       NewAddr = Builder.CreateGEP(Base, Index);
     }
+  } else if (!isa<Constant>(Ptr)) {
+    // Not a GEP, maybe its a splat and we can create a GEP to enable
+    // SelectionDAGBuilder to use it as a uniform base.
+    Value *V = getSplatValue(Ptr);
+    if (!V)
+      return false;
+
+    unsigned NumElts = cast<FixedVectorType>(Ptr->getType())->getNumElements();
+
+    IRBuilder<> Builder(MemoryInst);
+
+    // Emit a vector GEP with a scalar pointer and all 0s vector index.
+    Type *ScalarIndexTy = DL->getIndexType(V->getType()->getScalarType());
+    auto *IndexTy = FixedVectorType::get(ScalarIndexTy, NumElts);
+    NewAddr = Builder.CreateGEP(V, Constant::getNullValue(IndexTy));
+  } else {
+    // Constant, SelectionDAGBuilder knows to check if its a splat.
+    return false;
+  }
 
   MemoryInst->replaceUsesOfWith(Ptr, NewAddr);
 
   // If we have no uses, recursively delete the value and all dead instructions
   // using it.
   if (Ptr->use_empty())
     RecursivelyDeleteTriviallyDeadInstructions(
         Ptr, TLInfo, nullptr,

llvm/test/CodeGen/X86/masked_gather_scatter.ll

 ; SKX_32-NEXT:    vmovdqa64 {{.*#+}} zmm0 = [0,4294967294,1,4294967288,10,20,50,65536,16777215,2147483647,100,4294965296,2147483648,76897723,7,4227069609]
 ; SKX_32-NEXT:    vpscatterdd %zmm1, (%eax,%zmm0,4) {%k1}
 ; SKX_32-NEXT:    vzeroupper
 ; SKX_32-NEXT:    retl
   %gep = getelementptr i32, i32* %ptr, <16 x i64> <i64 0, i64 -2, i64 1, i64 -8, i64 10, i64 20, i64 50, i64 65536, i64 16777215, i64 2147483647, i64 100, i64 -2000, i64 -2147483648, i64 76897723, i64 7, i64 -67897687>
   call void @llvm.masked.scatter.v16i32.v16p0i32(<16 x i32> %src0, <16 x i32*> %gep, i32 4, <16 x i1> %mask)
   ret void
 }
+
+define void @splat_ptr(i32* %ptr, <4 x i1> %mask, <4 x i32> %val) {
+; CHECK-LABEL: @splat_ptr(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[PTR:%.*]], <4 x i64> zeroinitializer
+; CHECK-NEXT:    call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> [[VAL:%.*]], <4 x i32*> [[TMP1]], i32 4, <4 x i1> [[MASK:%.*]])
+; CHECK-NEXT:    ret void
+;
+; KNL_64-LABEL: splat_ptr:
+; KNL_64:       # %bb.0:
+; KNL_64-NEXT:    # kill: def $xmm1 killed $xmm1 def $zmm1
+; KNL_64-NEXT:    vpslld $31, %xmm0, %xmm0
+; KNL_64-NEXT:    vptestmd %zmm0, %zmm0, %k0
+; KNL_64-NEXT:    kshiftlw $12, %k0, %k0
+; KNL_64-NEXT:    kshiftrw $12, %k0, %k1
+; KNL_64-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; KNL_64-NEXT:    vpscatterdd %zmm1, (%rdi,%zmm0,4) {%k1}
+; KNL_64-NEXT:    vzeroupper
+; KNL_64-NEXT:    retq
+;
+; KNL_32-LABEL: splat_ptr:
+; KNL_32:       # %bb.0:
+; KNL_32-NEXT:    # kill: def $xmm1 killed $xmm1 def $zmm1
+; KNL_32-NEXT:    vpslld $31, %xmm0, %xmm0
+; KNL_32-NEXT:    vptestmd %zmm0, %zmm0, %k0
+; KNL_32-NEXT:    kshiftlw $12, %k0, %k0
+; KNL_32-NEXT:    kshiftrw $12, %k0, %k1
+; KNL_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; KNL_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; KNL_32-NEXT:    vpscatterdd %zmm1, (%eax,%zmm0,4) {%k1}
+; KNL_32-NEXT:    vzeroupper
+; KNL_32-NEXT:    retl
+;
+; SKX-LABEL: splat_ptr:
+; SKX:       # %bb.0:
+; SKX-NEXT:    vpslld $31, %xmm0, %xmm0
+; SKX-NEXT:    vpmovd2m %xmm0, %k1
+; SKX-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; SKX-NEXT:    vpscatterdd %xmm1, (%rdi,%xmm0,4) {%k1}
+; SKX-NEXT:    retq
+;
+; SKX_32-LABEL: splat_ptr:
+; SKX_32:       # %bb.0:
+; SKX_32-NEXT:    vpslld $31, %xmm0, %xmm0
+; SKX_32-NEXT:    vpmovd2m %xmm0, %k1
+; SKX_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; SKX_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; SKX_32-NEXT:    vpscatterdd %xmm1, (%eax,%xmm0,4) {%k1}
+; SKX_32-NEXT:    retl
+  %1 = insertelement <4 x i32*> undef, i32* %ptr, i32 0
+  %2 = shufflevector <4 x i32*> %1, <4 x i32*> undef, <4 x i32> zeroinitializer
+  call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> %val, <4 x i32*> %2, i32 4, <4 x i1> %mask)
+  ret void
+}
+

llvm/test/Transforms/CodeGenPrepare/X86/gather-scatter-opt.ll

 ;
   %1 = insertelement <4 x %struct.a*> undef, %struct.a* @c, i32 0
   %2 = shufflevector <4 x %struct.a*> %1, <4 x %struct.a*> undef, <4 x i32> zeroinitializer
   %3 = getelementptr %struct.a, <4 x %struct.a*> %2, <4 x i64> zeroinitializer, i32 1
   %4 = call <4 x i32> @llvm.masked.gather.v4i32.v4p0i32(<4 x i32*> %3, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, <4 x i32> undef)
   ret <4 x i32> %4
 }
 
+define void @splat_ptr(i32* %ptr, <4 x i1> %mask, <4 x i32> %val) {
+; CHECK-LABEL: @splat_ptr(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[PTR:%.*]], <4 x i64> zeroinitializer
+; CHECK-NEXT:    call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> [[VAL:%.*]], <4 x i32*> [[TMP1]], i32 4, <4 x i1> [[MASK:%.*]])
+; CHECK-NEXT:    ret void
+;
+  %1 = insertelement <4 x i32*> undef, i32* %ptr, i32 0
+  %2 = shufflevector <4 x i32*> %1, <4 x i32*> undef, <4 x i32> zeroinitializer
+  call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> %val, <4 x i32*> %2, i32 4, <4 x i1> %mask)
+  ret void
+}
+
 declare <4 x i32> @llvm.masked.gather.v4i32.v4p0i32(<4 x i32*>, i32, <4 x i1>, <4 x i32>)
+declare void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32>, <4 x i32*>, i32, <4 x i1>)