Gather and Scatter intrinsics in the Loop Vectorizer

[delena]

Summary

Loop vectorizer now knows to vectorize GEP and create masked gather and scatter intrinsics for random memory access.

The feature is enabled on AVX-512 target.
the following loops are vectorized on AVX-512:

struct In {

float a;
float b;

};

void foo1 (float * restrict in, float * restrict out, int * restrict trigger, int * restrict index) {

for (int i=0; i<SIZE; ++i) {
  if (trigger[i] > 0) {
    out[i] = in[index[i]] + (float) 0.5;
  }
}

}
void foo2 (In * restrict in, float * restrict out, int * restrict trigger) {

for (int i=0; i<SIZE; ++i) {
  if (trigger[i] > 0) {
    out[i] = in[i].b + (float) 0.5;
  }
}

}

struct Out {

float a;
float b;

};
void foo3 (In * restrict in, Out * restrict out, int * restrict trigger) {

for (int i=0; i<SIZE; ++i) {
  if (trigger[i] > 0) {
    out[i].b = in[i].b + (float) 0.5;
  }
}

}

[delena]

Updated diff after commit of gather-scatter cost

[spatel]

I don't know the vectorizer or scatter/gather well enough for a complete review, so I've just pointed out some small fixes.

Also, I think it's fine to have scatter and gather in one patch for review, but they should be split into separate patches for commit to make it easier to investigate bugs/regressions.

../lib/IR/IRBuilder.cpp

261–263

'dyn_cast' can be a plain 'cast' since we're asserting.

265

This can be 'auto' to match the earlier cast.

285–286

Use plain 'cast' since we're asserting.

293

Use 'auto' for consistency.

304

I don't understand the comment. Please explain more or delete.

../lib/Transforms/Vectorize/LoopVectorize.cpp

2397

typo: supported

2468

Use complete sentence for comment.

2469

Extra spaces.

2470–2493

The function is getting too long / indented. I would prefer to see it broken up with helper functions.

2511–2512

Formatting - 'else' should be on the same line as the brace.

5331

typo: transferred

../test/Transforms/LoopVectorize/X86/gather_scatter.ll

20–23

Please check the full text for any masked ops that are created in this test and others. We do not want to miss any bugs/regressions resulting from changes in the data types or number of instructions produced.

../test/Transforms/LoopVectorize/X86/masked_load_store.ll

283

Please add positive checks for the expected / correct mask ops that will now be produced.

[Ayal]

../include/llvm/IR/IRBuilder.h

436

Set default Mask to 0 similar to scatter below?

../lib/IR/IRBuilder.cpp

250

Load >> Gather

253

a[n] vector

265

As Sanjay suggested, you could do

auto *PtrsTy = cast<...;
auto *PtrTy = cast<...;

and rely on these casts asserting for you. Also good to place these two lookalike variables together.

278

Val >> Data

279

ditto

281

a[n] vector

298

Inco[m]patible.
Two separate asserts?

../lib/Transforms/Vectorize/LoopVectorize.cpp

2401–2402

Can fold into

bool CreateGatherScatter = (!ConsecutiveStride && ((LI && ...) || (SI && ...)));

2470

comment that in vectorizing Gep, across UF parts, we want to keep each loop-invariant base or index of Gep scalar.

2483

looks better to first set GEPBasePtr and then set the GEP Ops.

5331

transfer[r]ed to >> translated into

5599

can assert ConsecutiveStride here to be sure

[delena]

../lib/Transforms/Vectorize/LoopVectorize.cpp

2470–2493

May be after review. Otherwise the diff will be inconvenient for reviewers.

../test/Transforms/LoopVectorize/X86/gather_scatter.ll

20–23

I can add types. But the vector factor at this point is less than I want.
I received VF=8 instead of 16.
The problem is in the cost model. I tried to fix the cost model, but the patch was rejected.
So, I'm adding v8f32 and hope that it will be fixed later.

[delena]

Code changes according to Sanjay's comments.

[delena]

Minor fixes, according to Ayal comments.

[nadav]

../lib/Transforms/Vectorize/LoopVectorize.cpp

2506–2514

Please initialize to null.

2557

Please don't clean the code as part of this change. It makes it more difficult to review the code.

5335

Please use Doxygen style comment and document each of the arguments with \p.

../unittests/Support/AlignOfTest.cpp

93 ↗

(On Diff #46244)

Why is this change needed?

[delena]

Updated diff following Nadav's comments.

[Ayal]

Check your documentation to follow Nadav's comment: "Please use Doxygen style comment and document each of the arguments with \p."

../lib/IR/IRBuilder.cpp

262–263

This assert is redundant as it is taken care of by the two 'cast's. Unless you want to issue this particular error message?

282

"where the Val elements" >> "where the \p Data element"

291–292

These two asserts are mostly dead code as they are taken care of by the three 'cast's.

../lib/Transforms/Vectorize/LoopVectorize.cpp

2554

Formatting - 'else' should be on the same line as the brace.

[delena]

Changes according to Ayal's comments.

[delena]

Fixed one more style issue.

[apilipenko]

I expect this code to fail trying to vectorize non-default address space accesses. See D17270 for details. We should either support arbitrary address space for gather/scatter intrinsics or at least check that we are vectorizing a default address space accesses.

[delena]

Style changes according to Ayal comments.

[Ayal]

This looks good to me, thanks for accommodating the changes.

../include/llvm/IR/IRBuilder.h

public:
   /// \brief Create a call to Masked Load intrinsic
   CallInst *CreateMaskedLoad(Value *Ptr, unsigned Align, Value *Mask,
                              Value *PassThru = nullptr, const Twine &Name = "");
 
   /// \brief Create a call to Masked Store intrinsic
   CallInst *CreateMaskedStore(Value *Val, Value *Ptr, unsigned Align,
                               Value *Mask);
 
+  /// \brief Create a call to Masked Gather intrinsic
+  CallInst *CreateMaskedGather(Value *Ptrs, unsigned Align, Value *Mask,
+                               Value *PassThru = 0, const Twine& Name = "");
+
+  /// \brief Create a call to Masked Scatter intrinsic
+  CallInst *CreateMaskedScatter(Value *Val, Value *Ptrs, unsigned Align,
+                                Value *Mask = 0);
+
   /// \brief Create an assume intrinsic call that allows the optimizer to
   /// assume that the provided condition will be true.
   CallInst *CreateAssumption(Value *Cond);
 
   /// \brief Create a call to the experimental.gc.statepoint intrinsic to
   /// start a new statepoint sequence.
   CallInst *CreateGCStatepointCall(uint64_t ID, uint32_t NumPatchBytes,
                                    Value *ActualCallee,

../lib/IR/IRBuilder.cpp

CallInst *IRBuilderBase::CreateMaskedIntrinsic(Intrinsic::ID Id,
                                                Type *DataTy,
                                                const Twine &Name) {
   Module *M = BB->getParent()->getParent();
   Type *OverloadedTypes[] = { DataTy };
   Value *TheFn = Intrinsic::getDeclaration(M, Id, OverloadedTypes);
   return createCallHelper(TheFn, Ops, this, Name);
 }
 
+/// Create a call to a Masked Load intrinsic.
+/// Ptrs     - the vector of pointers for loading
+/// Align    - alignment for one element
+/// Mask     - an vector of booleans which indicates what vector lanes should
+///            be accessed in memory
+/// PassThru - a pass-through value that is used to fill the masked-off lanes
+///            of the result
+/// Name     - name of the result variable
+CallInst *IRBuilderBase::CreateMaskedGather(Value *Ptrs, unsigned Align,
+                                            Value *Mask,  Value *PassThru,
+                                            const Twine& Name) {
+  auto *PtrsTy = dyn_cast<VectorType>(Ptrs->getType());
+  assert(PtrsTy && PtrsTy->getElementType()->isPointerTy() &&
+    "Expected vector of pointers for gather");
+
+  PointerType *PtrTy = cast<PointerType>(PtrsTy->getElementType());
+  unsigned NumElts = PtrsTy->getVectorNumElements();
+  Type *DataTy = VectorType::get(PtrTy->getElementType(), NumElts);
+
+  if (!Mask)
+    Mask = Constant::getAllOnesValue(VectorType::get(Type::getInt1Ty(Context),
+                                     NumElts));
+
+  Value * Ops[] = {Ptrs, getInt32(Align), Mask, UndefValue::get(DataTy)};
+  return CreateMaskedIntrinsic(Intrinsic::masked_gather, Ops, DataTy, Name);
+}
+
+/// Create a call to a Masked Scatter intrinsic.
+/// Val   - the data to be stored,
+/// Ptrs  - the vector of pointers, where the Val elements should be stored
+/// Align - alignment for one element
+/// Mask  - an vector of booleans which indicates what vector lanes should
+///         be accessed in memory
+CallInst *IRBuilderBase::CreateMaskedScatter(Value *Data, Value *Ptrs,
+                                             unsigned Align, Value *Mask) {
+  auto *PtrsTy = dyn_cast<VectorType>(Ptrs->getType());
+  auto *DataTy = dyn_cast<VectorType>(Data->getType());
+
+  assert(PtrsTy && PtrsTy->getElementType()->isPointerTy() &&
+    "Expected vector of pointers for scatter");
+
+  assert(DataTy && "Unexpected data type for scatter");
+
+  PointerType *PtrTy = cast<PointerType>(PtrsTy->getElementType());
+  unsigned NumElts = PtrsTy->getVectorNumElements();
+
+  assert(NumElts == DataTy->getVectorNumElements() &&
+         PtrTy->getElementType() ==  DataTy->getElementType() &&
+        "Incopatible pointer and data types");
+
+  if (!Mask)
+    Mask = Constant::getAllOnesValue(VectorType::get(Type::getInt1Ty(Context),
+                                     NumElts));
+  Value * Ops[] = {Data, Ptrs, getInt32(Align), Mask};
+  // Type of the data to be stored - the only one overloaded type
+  return CreateMaskedIntrinsic(Intrinsic::masked_scatter, Ops, DataTy);
+}
+
 template <typename T0, typename T1, typename T2, typename T3>
 static std::vector<Value *>
 getStatepointArgs(IRBuilderBase &B, uint64_t ID, uint32_t NumPatchBytes,
                   Value *ActualCallee, uint32_t Flags, ArrayRef<T0> CallArgs,
                   ArrayRef<T1> TransitionArgs, ArrayRef<T2> DeoptArgs,
                   ArrayRef<T3> GCArgs) {
   std::vector<Value *> Args;
   Args.push_back(B.getInt64(ID));

../lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   bool isLegalMaskedStore(Type *DataType, Value *Ptr) {
     return isConsecutivePtr(Ptr) && TTI->isLegalMaskedStore(DataType);
   }
   /// Returns true if the target machine supports masked load operation
   /// for the given \p DataType and kind of access to \p Ptr.
   bool isLegalMaskedLoad(Type *DataType, Value *Ptr) {
     return isConsecutivePtr(Ptr) && TTI->isLegalMaskedLoad(DataType);
   }
+  /// Returns true if the target machine supports masked scatter operation
+  /// for the given \p DataType.
+  bool isLegalMaskedScatter(Type *DataType) {
+    return TTI->isLegalMaskedScatter(DataType);
+  }
+  /// Returns true if the target machine supports masked gather operation
+  /// for the given \p DataType.
+  bool isLegalMaskedGather(Type *DataType) {
+    return TTI->isLegalMaskedGather(DataType);
+  }
+
   /// Returns true if vector representation of the instruction \p I
   /// requires mask.
   bool isMaskRequired(const Instruction* I) {
     return (MaskedOp.count(I) != 0);
   }
   unsigned getNumStores() const {
     return LAI->getNumStores();
   }
void InnerLoopVectorizer::vectorizeMemoryInstruction(Instruction *Instr) {
 
   if (SI && Legal->blockNeedsPredication(SI->getParent()) &&
       !Legal->isMaskRequired(SI))
     return scalarizeInstruction(Instr, true);
 
   if (ScalarAllocatedSize != VectorElementSize)
     return scalarizeInstruction(Instr);
 
-  // If the pointer is loop invariant or if it is non-consecutive,
-  // scalarize the load.
+  // If the pointer is loop invariant scalarize the load.
+  if (LI && Legal->isUniform(Ptr))
+    return scalarizeInstruction(Instr);
+
+  // If the pointer is non-consecutive and gather/scatter is not suppored
+  // scalarize the instruction.
   int ConsecutiveStride = Legal->isConsecutivePtr(Ptr);
   bool Reverse = ConsecutiveStride < 0;
-  bool UniformLoad = LI && Legal->isUniform(Ptr);
-  if (!ConsecutiveStride || UniformLoad)
+  bool CreateGatherScatter = false;
+  if (!ConsecutiveStride)
+    CreateGatherScatter =
+      (LI && Legal->isLegalMaskedGather(ScalarDataTy)) ||
+      (SI && Legal->isLegalMaskedScatter(ScalarDataTy));
+ 
+  if (!ConsecutiveStride && !CreateGatherScatter)
     return scalarizeInstruction(Instr);
 
   Constant *Zero = Builder.getInt32(0);
   VectorParts &Entry = WidenMap.get(Instr);
+  VectorParts VectorGep;
 
   // Handle consecutive loads/stores.
   GetElementPtrInst *Gep = getGEPInstruction(Ptr);
+  if (ConsecutiveStride) {
     if (Gep && Legal->isInductionVariable(Gep->getPointerOperand())) {
       setDebugLocFromInst(Builder, Gep);
         Value *PtrOperand = Gep->getPointerOperand();
         Value *FirstBasePtr = getVectorValue(PtrOperand)[0];
         FirstBasePtr = Builder.CreateExtractElement(FirstBasePtr, Zero);
 
         // Create the new GEP with the new induction variable.
         GetElementPtrInst *Gep2 = cast<GetElementPtrInst>(Gep->clone());
         Gep2->setOperand(0, FirstBasePtr);
         Gep2->setName("gep.indvar.base");
         Ptr = Builder.Insert(Gep2);
       } else if (Gep) {
        setDebugLocFromInst(Builder, Gep);
        assert(PSE.getSE()->isLoopInvariant(PSE.getSCEV(Gep->getPointerOperand()),
                                              OrigLoop) &&
                              "Base ptr must be invariant");
-
         // The last index does not have to be the induction. It can be
         // consecutive and be a function of the index. For example A[I+1];
         unsigned NumOperands = Gep->getNumOperands();
         unsigned InductionOperand = getGEPInductionOperand(Gep);
         // Create the new GEP with the new induction variable.
         GetElementPtrInst *Gep2 = cast<GetElementPtrInst>(Gep->clone());
 
         for (unsigned i = 0; i < NumOperands; ++i) {
           Value *GepOperand = Gep->getOperand(i);
           Instruction *GepOperandInst = dyn_cast<Instruction>(GepOperand);
 
           // Update last index or loop invariant instruction anchored in loop.
           if (i == InductionOperand ||
               (GepOperandInst && OrigLoop->contains(GepOperandInst))) {
             assert((i == InductionOperand ||
                     PSE.getSE()->isLoopInvariant(PSE.getSCEV(GepOperandInst),
                                                  OrigLoop)) &&
                     "Must be last index or loop invariant");
 
             VectorParts &GEPParts = getVectorValue(GepOperand);
             Value *Index = GEPParts[0];
             Index = Builder.CreateExtractElement(Index, Zero);
             Gep2->setOperand(i, Index);
             Gep2->setName("gep.indvar.idx");
           }
         }
         Ptr = Builder.Insert(Gep2);
-  } else {
+      } else { // No GEP
         // Use the induction element ptr.
         assert(isa<PHINode>(Ptr) && "Invalid induction ptr");
         setDebugLocFromInst(Builder, Ptr);
         VectorParts &PtrVal = getVectorValue(Ptr);
         Ptr = Builder.CreateExtractElement(PtrVal[0], Zero);
       }
+    } else {
+      // Supported random access (gather / scatter)
+      assert(CreateGatherScatter  &&  "The instruction should be scalarized");
+      if (Gep) {
+        SmallVector<VectorParts, 4> OpsV;
+        for (Value *Op : Gep->operands()) {
+          if (PSE.getSE()->isLoopInvariant(PSE.getSCEV(Op), OrigLoop))
+            OpsV.push_back(VectorParts(UF, Op));
+          else
+            OpsV.push_back(getVectorValue(Op));
+        }
+
+        for (unsigned Part = 0; Part < UF; ++Part) {
+          SmallVector<Value*, 4> Ops;
+          for (unsigned i = 1; i < Gep->getNumOperands(); i++)
+            Ops.push_back(OpsV[i][Part]);
+          Value *GEPBasePtr = OpsV[0][Part];
+          Value *NewGep = Builder.CreateGEP(nullptr, GEPBasePtr, Ops,
+                                            "VectorGep");
+          assert(NewGep->getType()->isVectorTy() && "Expected vector GEP");
+          NewGep = Builder.CreateBitCast(NewGep,
+                                         VectorType::get(Ptr->getType(), VF));
+          VectorGep.push_back(NewGep);
+        }
+      } else
+        VectorGep = getVectorValue(Ptr);
+    }
 
   VectorParts Mask = createBlockInMask(Instr->getParent());
   // Handle Stores:
   if (SI) {
     assert(!Legal->isUniform(SI->getPointerOperand()) &&
            "We do not allow storing to uniform addresses");
     setDebugLocFromInst(Builder, SI);
     // We don't want to update the value in the map as it might be used in
     // another expression. So don't use a reference type for "StoredVal".
     VectorParts StoredVal = getVectorValue(SI->getValueOperand());
 
     for (unsigned Part = 0; Part < UF; ++Part) {
+      Instruction *NewSI;
+      if (CreateGatherScatter) {
+        Value *MaskPart = Legal->isMaskRequired(SI) ? Mask[Part] : nullptr;
+        NewSI = Builder.CreateMaskedScatter(StoredVal[Part], VectorGep[Part],
+                                            Alignment, MaskPart);
+      }
+      else {
         // Calculate the pointer for the specific unroll-part.
         Value *PartPtr =
           Builder.CreateGEP(nullptr, Ptr, Builder.getInt32(Part * VF));
 
         if (Reverse) {
           // If we store to reverse consecutive memory locations, then we need
           // to reverse the order of elements in the stored value.
           StoredVal[Part] = reverseVector(StoredVal[Part]);
           // If the address is consecutive but reversed, then the
           // wide store needs to start at the last vector element.
           PartPtr = Builder.CreateGEP(nullptr, Ptr, Builder.getInt32(-Part * VF));
           PartPtr = Builder.CreateGEP(nullptr, PartPtr, Builder.getInt32(1 - VF));
           Mask[Part] = reverseVector(Mask[Part]);
         }
 
         Value *VecPtr = Builder.CreateBitCast(PartPtr,
                                               DataTy->getPointerTo(AddressSpace));
 
-      Instruction *NewSI;
         if (Legal->isMaskRequired(SI))
           NewSI = Builder.CreateMaskedStore(StoredVal[Part], VecPtr, Alignment,
                                             Mask[Part]);
         else 
-        NewSI = Builder.CreateAlignedStore(StoredVal[Part], VecPtr, Alignment);
+          NewSI = Builder.CreateAlignedStore(StoredVal[Part], VecPtr,
+                                             Alignment);
+      }
       propagateMetadata(NewSI, SI);
     }
     return;
   }
 
   // Handle loads.
   assert(LI && "Must have a load instruction");
   setDebugLocFromInst(Builder, LI);
   for (unsigned Part = 0; Part < UF; ++Part) {
+    Instruction* NewLI;
+    if (CreateGatherScatter) {
+      Value *MaskPart = Legal->isMaskRequired(LI) ? Mask[Part] : nullptr;
+      NewLI = Builder.CreateMaskedGather(VectorGep[Part], Alignment,
+                                         MaskPart, 0,
+                                         "wide.masked.gather");
+      Entry[Part] = NewLI;
+    }
+    else {
       // Calculate the pointer for the specific unroll-part.
-    Value *PartPtr =
-        Builder.CreateGEP(nullptr, Ptr, Builder.getInt32(Part * VF));
+      Value *PartPtr = Builder.CreateGEP(nullptr, Ptr, Builder.getInt32(Part * VF));
 
       if (Reverse) {
         // If the address is consecutive but reversed, then the
         // wide load needs to start at the last vector element.
         PartPtr = Builder.CreateGEP(nullptr, Ptr, Builder.getInt32(-Part * VF));
         PartPtr = Builder.CreateGEP(nullptr, PartPtr, Builder.getInt32(1 - VF));
         Mask[Part] = reverseVector(Mask[Part]);
       }
 
-    Instruction* NewLI;
       Value *VecPtr = Builder.CreateBitCast(PartPtr,
                                             DataTy->getPointerTo(AddressSpace));
       if (Legal->isMaskRequired(LI))
         NewLI = Builder.CreateMaskedLoad(VecPtr, Alignment, Mask[Part],
                                          UndefValue::get(DataTy),
                                          "wide.masked.load");
       else
         NewLI = Builder.CreateAlignedLoad(VecPtr, Alignment, "wide.load");
-    propagateMetadata(NewLI, LI);
       Entry[Part] = Reverse ? reverseVector(NewLI) :  NewLI;
     }
+    propagateMetadata(NewLI, LI);
+  }
 }
 
 void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr,
                                                bool IfPredicateStore) {
   assert(!Instr->getType()->isAggregateType() && "Can't handle vectors");
   // Holds vector parameters or scalars, in case of uniform vals.
   SmallVector<VectorParts, 4> Params;
 
for (Instruction::op_iterator OI = it->op_begin(), OE = it->op_end();
           return false;
     }
     // We might be able to hoist the load.
     if (it->mayReadFromMemory()) {
       LoadInst *LI = dyn_cast<LoadInst>(it);
       if (!LI)
         return false;
       if (!SafePtrs.count(LI->getPointerOperand())) {
-        if (isLegalMaskedLoad(LI->getType(), LI->getPointerOperand())) {
+        if (isLegalMaskedLoad(LI->getType(), LI->getPointerOperand()) ||
+            isLegalMaskedGather(LI->getType())) {
           MaskedOp.insert(LI);
           continue;
         }
         return false;
       }
     }
 
     // We don't predicate stores at the moment.
     if (it->mayWriteToMemory()) {
       StoreInst *SI = dyn_cast<StoreInst>(it);
       // We only support predication of stores in basic blocks with one
       // predecessor.
       if (!SI)
         return false;
 
       bool isSafePtr = (SafePtrs.count(SI->getPointerOperand()) != 0);
       bool isSinglePredecessor = SI->getParent()->getSinglePredecessor();
       
       if (++NumPredStores > NumberOfStoresToPredicate || !isSafePtr ||
           !isSinglePredecessor) {
         // Build a masked store if it is legal for the target, otherwise
         // scalarize the block.
         bool isLegalMaskedOp =
           isLegalMaskedStore(SI->getValueOperand()->getType(),
-                             SI->getPointerOperand());
+                             SI->getPointerOperand()) ||
+          isLegalMaskedScatter(SI->getValueOperand()->getType());
         if (isLegalMaskedOp) {
           --NumPredStores;
           MaskedOp.insert(SI);
           continue;
         }
         return false;
       }
     }
if (VF == 1 && Legal->blockNeedsPredication(*bb))
       BlockCost /= 2;
 
     Cost += BlockCost;
   }
 
   return Cost;
 }
 
+/// \brief Check if the load/store instruction may be transfered to
+/// gather/scatter during vectorization.
+///
+/// It is possible when gather/scatter is supported by target.
+static bool isGatherOrScatterLegal(Instruction *I, Value *Ptr,
+                                   LoopVectorizationLegality *Legal) {
+  Type *DataTy = cast<PointerType>(Ptr->getType())->getElementType();
+  return (isa<LoadInst>(I) && Legal->isLegalMaskedGather(DataTy)) ||
+    (isa<StoreInst>(I) && Legal->isLegalMaskedScatter(DataTy));
+}
+
 /// \brief Check whether the address computation for a non-consecutive memory
 /// access looks like an unlikely candidate for being merged into the indexing
 /// mode.
 ///
 /// We look for a GEP which has one index that is an induction variable and all
 /// other indices are loop invariant. If the stride of this access is also
 /// within a small bound we decide that this address computation can likely be
 /// merged into the addressing mode.
if (Legal->isAccessInterleaved(I)) {
       // FIXME: The interleaved load group with a huge gap could be even more
       // expensive than scalar operations. Then we could ignore such group and
       // use scalar operations instead.
       return Cost;
     }
 
     // Scalarized loads/stores.
     int ConsecutiveStride = Legal->isConsecutivePtr(Ptr);
+    bool UseGatherOrScatter = (ConsecutiveStride == 0) &&
+      isGatherOrScatterLegal(I, Ptr, Legal);
+
     bool Reverse = ConsecutiveStride < 0;
     const DataLayout &DL = I->getModule()->getDataLayout();
     unsigned ScalarAllocatedSize = DL.getTypeAllocSize(ValTy);
     unsigned VectorElementSize = DL.getTypeStoreSize(VectorTy) / VF;
-    if (!ConsecutiveStride || ScalarAllocatedSize != VectorElementSize) {
+    if ((!ConsecutiveStride && !UseGatherOrScatter) ||
+        ScalarAllocatedSize != VectorElementSize) {
       bool IsComplexComputation =
         isLikelyComplexAddressComputation(Ptr, Legal, SE, TheLoop);
       unsigned Cost = 0;
       // The cost of extracting from the value vector and pointer vector.
       Type *PtrTy = ToVectorTy(Ptr->getType(), VF);
       for (unsigned i = 0; i < VF; ++i) {
         //  The cost of extracting the pointer operand.
         Cost += TTI.getVectorInstrCost(Instruction::ExtractElement, PtrTy, i);
         // In case of STORE, the cost of ExtractElement from the vector.
         // In case of LOAD, the cost of InsertElement into the returned
         // vector.
         Cost += TTI.getVectorInstrCost(SI ? Instruction::ExtractElement :
                                             Instruction::InsertElement,
                                             VectorTy, i);
       }
 
       // The cost of the scalar loads/stores.
       Cost += VF * TTI.getAddressComputationCost(PtrTy, IsComplexComputation);
       Cost += VF * TTI.getMemoryOpCost(I->getOpcode(), ValTy->getScalarType(),
                                        Alignment, AS);
       return Cost;
     }
 
-    // Wide load/stores.
     unsigned Cost = TTI.getAddressComputationCost(VectorTy);
+    if (UseGatherOrScatter)
+      return Cost +
+        TTI.getGatherScatterOpCost(I->getOpcode(), VectorTy, Ptr,
+                                   Legal->isMaskRequired(I), Alignment);
+
+    // Wide load/stores.
     if (Legal->isMaskRequired(I))
       Cost += TTI.getMaskedMemoryOpCost(I->getOpcode(), VectorTy, Alignment,
                                         AS);
     else
       Cost += TTI.getMemoryOpCost(I->getOpcode(), VectorTy, Alignment, AS);
 
     if (Reverse)
       Cost += TTI.getShuffleCost(TargetTransformInfo::SK_Reverse,

../test/Transforms/LoopVectorize/X86/gather_scatter.ll

+; RUN: opt < %s  -O3 -mcpu=knl -S | FileCheck %s -check-prefix=AVX512
+
+;AVX1-NOT: llvm.masked
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc_linux"
+
+; The source code:
+;
+;void foo1(float * __restrict__ in, float * __restrict__ out, int * __restrict__ trigger, int * __restrict__ index) {
+;
+;  for (int i=0; i < SIZE; ++i) {
+;    if (trigger[i] > 0) {
+;      out[i] = in[index[i]] + (float) 0.5;
+;    }
+;  }
+;}
+
+;AVX512-LABEL: @foo1
+;AVX512:  llvm.masked.load
+;AVX512: llvm.masked.gather
+;AVX512: llvm.masked.store
+;AVX512: ret void
+
+; Function Attrs: nounwind uwtable
+define void @foo1(float* noalias %in, float* noalias %out, i32* noalias %trigger, i32* noalias %index) {
+entry:
+  %in.addr = alloca float*, align 8
+  %out.addr = alloca float*, align 8
+  %trigger.addr = alloca i32*, align 8
+  %index.addr = alloca i32*, align 8
+  %i = alloca i32, align 4
+  store float* %in, float** %in.addr, align 8
+  store float* %out, float** %out.addr, align 8
+  store i32* %trigger, i32** %trigger.addr, align 8
+  store i32* %index, i32** %index.addr, align 8
+  store i32 0, i32* %i, align 4
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+  %0 = load i32, i32* %i, align 4
+  %cmp = icmp slt i32 %0, 4096
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond
+  %1 = load i32, i32* %i, align 4
+  %idxprom = sext i32 %1 to i64
+  %2 = load i32*, i32** %trigger.addr, align 8
+  %arrayidx = getelementptr inbounds i32, i32* %2, i64 %idxprom
+  %3 = load i32, i32* %arrayidx, align 4
+  %cmp1 = icmp sgt i32 %3, 0
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+  %4 = load i32, i32* %i, align 4
+  %idxprom2 = sext i32 %4 to i64
+  %5 = load i32*, i32** %index.addr, align 8
+  %arrayidx3 = getelementptr inbounds i32, i32* %5, i64 %idxprom2
+  %6 = load i32, i32* %arrayidx3, align 4
+  %idxprom4 = sext i32 %6 to i64
+  %7 = load float*, float** %in.addr, align 8
+  %arrayidx5 = getelementptr inbounds float, float* %7, i64 %idxprom4
+  %8 = load float, float* %arrayidx5, align 4
+  %add = fadd float %8, 5.000000e-01
+  %9 = load i32, i32* %i, align 4
+  %idxprom6 = sext i32 %9 to i64
+  %10 = load float*, float** %out.addr, align 8
+  %arrayidx7 = getelementptr inbounds float, float* %10, i64 %idxprom6
+  store float %add, float* %arrayidx7, align 4
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %for.body
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.end
+  %11 = load i32, i32* %i, align 4
+  %inc = add nsw i32 %11, 1
+  store i32 %inc, i32* %i, align 4
+  br label %for.cond
+
+for.end:                                          ; preds = %for.cond
+  ret void
+}
+
+; The source code
+;void foo2 (In * __restrict__ in, float * __restrict__ out, int * __restrict__ trigger) {
+;
+;  for (int i=0; i<SIZE; ++i) {
+;    if (trigger[i] > 0) {
+;      out[i] = in[i].b + (float) 0.5;
+;    }
+;  }
+;}
+
+%struct.In = type { float, float }
+
+;AVX512-LABEL: @foo2
+;AVX512: getelementptr %struct.In, %struct.In* %in, <16 x i64> %induction, i32 1
+;AVX512: llvm.masked.gather.v16f32
+;AVX512: llvm.masked.store
+;AVX512: ret void
+define void @foo2(%struct.In* noalias %in, float* noalias %out, i32* noalias %trigger, i32* noalias %index) #0 {
+entry:
+  %in.addr = alloca %struct.In*, align 8
+  %out.addr = alloca float*, align 8
+  %trigger.addr = alloca i32*, align 8
+  %index.addr = alloca i32*, align 8
+  %i = alloca i32, align 4
+  store %struct.In* %in, %struct.In** %in.addr, align 8
+  store float* %out, float** %out.addr, align 8
+  store i32* %trigger, i32** %trigger.addr, align 8
+  store i32* %index, i32** %index.addr, align 8
+  store i32 0, i32* %i, align 4
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+  %0 = load i32, i32* %i, align 4
+  %cmp = icmp slt i32 %0, 4096
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond
+  %1 = load i32, i32* %i, align 4
+  %idxprom = sext i32 %1 to i64
+  %2 = load i32*, i32** %trigger.addr, align 8
+  %arrayidx = getelementptr inbounds i32, i32* %2, i64 %idxprom
+  %3 = load i32, i32* %arrayidx, align 4
+  %cmp1 = icmp sgt i32 %3, 0
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+  %4 = load i32, i32* %i, align 4
+  %idxprom2 = sext i32 %4 to i64
+  %5 = load %struct.In*, %struct.In** %in.addr, align 8
+  %arrayidx3 = getelementptr inbounds %struct.In, %struct.In* %5, i64 %idxprom2
+  %b = getelementptr inbounds %struct.In, %struct.In* %arrayidx3, i32 0, i32 1
+  %6 = load float, float* %b, align 4
+  %add = fadd float %6, 5.000000e-01
+  %7 = load i32, i32* %i, align 4
+  %idxprom4 = sext i32 %7 to i64
+  %8 = load float*, float** %out.addr, align 8
+  %arrayidx5 = getelementptr inbounds float, float* %8, i64 %idxprom4
+  store float %add, float* %arrayidx5, align 4
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %for.body
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.end
+  %9 = load i32, i32* %i, align 4
+  %inc = add nsw i32 %9, 1
+  store i32 %inc, i32* %i, align 4
+  br label %for.cond
+
+for.end:                                          ; preds = %for.cond
+  ret void
+}
+
+; The source code
+;struct Out {
+;  float a;
+;  float b;
+;};
+;void foo3 (In * __restrict__ in, Out * __restrict__ out, int * __restrict__ trigger) {
+;
+;  for (int i=0; i<SIZE; ++i) {
+;    if (trigger[i] > 0) {
+;      out[i].b = in[i].b + (float) 0.5;
+;    }
+;  }
+;}
+
+;AVX512-LABEL: @foo3
+;AVX512: getelementptr %struct.In, %struct.In* %in, <16 x i64> %induction, i32 1
+;AVX512: llvm.masked.gather.v16f32
+;AVX512: fadd <16 x float>
+;AVX512: getelementptr %struct.Out, %struct.Out* %out, <16 x i64> %induction, i32 1
+;AVX512: llvm.masked.scatter.v16f32
+;AVX512: ret void
+
+%struct.Out = type { float, float }
+
+define void @foo3(%struct.In* noalias %in, %struct.Out* noalias %out, i32* noalias %trigger) {
+entry:
+  %in.addr = alloca %struct.In*, align 8
+  %out.addr = alloca %struct.Out*, align 8
+  %trigger.addr = alloca i32*, align 8
+  %i = alloca i32, align 4
+  store %struct.In* %in, %struct.In** %in.addr, align 8
+  store %struct.Out* %out, %struct.Out** %out.addr, align 8
+  store i32* %trigger, i32** %trigger.addr, align 8
+  store i32 0, i32* %i, align 4
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+  %0 = load i32, i32* %i, align 4
+  %cmp = icmp slt i32 %0, 4096
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond
+  %1 = load i32, i32* %i, align 4
+  %idxprom = sext i32 %1 to i64
+  %2 = load i32*, i32** %trigger.addr, align 8
+  %arrayidx = getelementptr inbounds i32, i32* %2, i64 %idxprom
+  %3 = load i32, i32* %arrayidx, align 4
+  %cmp1 = icmp sgt i32 %3, 0
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+  %4 = load i32, i32* %i, align 4
+  %idxprom2 = sext i32 %4 to i64
+  %5 = load %struct.In*, %struct.In** %in.addr, align 8
+  %arrayidx3 = getelementptr inbounds %struct.In, %struct.In* %5, i64 %idxprom2
+  %b = getelementptr inbounds %struct.In, %struct.In* %arrayidx3, i32 0, i32 1
+  %6 = load float, float* %b, align 4
+  %add = fadd float %6, 5.000000e-01
+  %7 = load i32, i32* %i, align 4
+  %idxprom4 = sext i32 %7 to i64
+  %8 = load %struct.Out*, %struct.Out** %out.addr, align 8
+  %arrayidx5 = getelementptr inbounds %struct.Out, %struct.Out* %8, i64 %idxprom4
+  %b6 = getelementptr inbounds %struct.Out, %struct.Out* %arrayidx5, i32 0, i32 1
+  store float %add, float* %b6, align 4
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %for.body
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.end
+  %9 = load i32, i32* %i, align 4
+  %inc = add nsw i32 %9, 1
+  store i32 %inc, i32* %i, align 4
+  br label %for.cond
+
+for.end:                                          ; preds = %for.cond
+  ret void
+}
+declare void @llvm.masked.scatter.v16f32(<16 x float>, <16 x float*>, i32, <16 x i1>)

../test/Transforms/LoopVectorize/X86/masked_load_store.ll

 ;  }
 ;}
 
 ;AVX2-LABEL: @foo4
 ;AVX2-NOT: llvm.masked
 ;AVX2: ret void
 
 ;AVX512-LABEL: @foo4
-;AVX512-NOT: llvm.masked
+;AVX512-NOT: llvm.masked.load
 ;AVX512: ret void
 
 ; Function Attrs: nounwind uwtable
 define void @foo4(double* %A, double* %B, i32* %trigger)  {
 entry:
   %A.addr = alloca double*, align 8
   %B.addr = alloca double*, align 8
   %trigger.addr = alloca i32*, align 8