[NFC][TTI] Add Alignment for isLegalMasked[Load/Store]

[samparker]

Summary

Add an extra parameter so the backend can take the alignment into consideration. I'll need this to continue with D68337.

[dmgreen]

Seems very sensible to me. There is an Align object recently added to llvm. Should we make use of it here?

[samparker]

I saw that, but wasn't sure how to use it! The alignment is just an unsigned value at all the call sites, so that's why I went for that.

[dmgreen]

lib/Target/X86/X86TargetTransformInfo.h

189

", unsigned Alignment" :)

Also the isLegalNTLoad below use the Align. I think they should be fairly simple to use (but haven't used them myself yet, just fixed merge conflicts).

[samparker]

Yes, it turns out Align is simple to use, the constructor just takes the unsigned value... but it will not accept zero and so causes assertion failures.

[dmgreen]

I think that's what "MaybeAlign" is for, if I read it correctly. It is essentially an Optional<Align> with 0 alignments being the "false" optional state.

[samparker]

Looks like it! I'll try it.

[samparker]

• Now using MaybeAlign
• Now using an alignment helper function in the vectorizer.

[dmgreen]

Thanks! LGTM

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4543 ↗

(On Diff #223809)

Does this need to be a MaybeAlign(Alignment)? Because the constructor is explicit, as far as I understand (hence, needs to look like the others).

[samparker]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4543 ↗

(On Diff #223809)

These are the internal helper functions, so they're constructed explicitly there. lines 1172 and 1179.

include/llvm/Analysis/TargetTransformInfo.h

   /// \return True is LSR should make efforts to create/preserve post-inc
   /// addressing mode expressions.
   bool shouldFavorPostInc() const;
 
   /// Return true if LSR should make efforts to generate indexed addressing
   /// modes that operate across loop iterations.
   bool shouldFavorBackedgeIndex(const Loop *L) const;
 
-  /// Return true if the target supports masked load.
-  bool isLegalMaskedStore(Type *DataType) const;
   /// Return true if the target supports masked store.
-  bool isLegalMaskedLoad(Type *DataType) const;
+  bool isLegalMaskedStore(Type *DataType, unsigned Alignment) const;
+  /// Return true if the target supports masked load.
+  bool isLegalMaskedLoad(Type *DataType, unsigned Alignment) const;
 
   /// Return true if the target supports nontemporal store.
   bool isLegalNTStore(Type *DataType, Align Alignment) const;
   /// Return true if the target supports nontemporal load.
   bool isLegalNTLoad(Type *DataType, Align Alignment) const;
 
   /// Return true if the target supports masked scatter.
   bool isLegalMaskedScatter(Type *DataType) const;

   virtual bool isLSRCostLess(TargetTransformInfo::LSRCost &C1,
                              TargetTransformInfo::LSRCost &C2) = 0;
   virtual bool canMacroFuseCmp() = 0;
   virtual bool canSaveCmp(Loop *L, BranchInst **BI, ScalarEvolution *SE,
                           LoopInfo *LI, DominatorTree *DT, AssumptionCache *AC,
                           TargetLibraryInfo *LibInfo) = 0;
   virtual bool shouldFavorPostInc() const = 0;
   virtual bool shouldFavorBackedgeIndex(const Loop *L) const = 0;
-  virtual bool isLegalMaskedStore(Type *DataType) = 0;
-  virtual bool isLegalMaskedLoad(Type *DataType) = 0;
+  virtual bool isLegalMaskedStore(Type *DataType, unsigned Alignment) = 0;
+  virtual bool isLegalMaskedLoad(Type *DataType, unsigned Alignment) = 0;
   virtual bool isLegalNTStore(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalNTLoad(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedScatter(Type *DataType) = 0;
   virtual bool isLegalMaskedGather(Type *DataType) = 0;
   virtual bool isLegalMaskedCompressStore(Type *DataType) = 0;
   virtual bool isLegalMaskedExpandLoad(Type *DataType) = 0;
   virtual bool hasDivRemOp(Type *DataType, bool IsSigned) = 0;
   virtual bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) = 0;

     return Impl.canSaveCmp(L, BI, SE, LI, DT, AC, LibInfo);
   }
   bool shouldFavorPostInc() const override {
     return Impl.shouldFavorPostInc();
   }
   bool shouldFavorBackedgeIndex(const Loop *L) const override {
     return Impl.shouldFavorBackedgeIndex(L);
   }
-  bool isLegalMaskedStore(Type *DataType) override {
-    return Impl.isLegalMaskedStore(DataType);
+  bool isLegalMaskedStore(Type *DataType, unsigned Alignment) override {
+    return Impl.isLegalMaskedStore(DataType, Alignment);
   }
-  bool isLegalMaskedLoad(Type *DataType) override {
-    return Impl.isLegalMaskedLoad(DataType);
+  bool isLegalMaskedLoad(Type *DataType, unsigned Alignment) override {
+    return Impl.isLegalMaskedLoad(DataType, Alignment);
   }
   bool isLegalNTStore(Type *DataType, Align Alignment) override {
     return Impl.isLegalNTStore(DataType, Alignment);
   }
   bool isLegalNTLoad(Type *DataType, Align Alignment) override {
     return Impl.isLegalNTLoad(DataType, Alignment);
   }
   bool isLegalMaskedScatter(Type *DataType) override {

include/llvm/Analysis/TargetTransformInfoImpl.h

bool canSaveCmp(Loop *L, BranchInst **BI, ScalarEvolution *SE, LoopInfo *LI,
                   TargetLibraryInfo *LibInfo) {
     return false;
   }
 
   bool shouldFavorPostInc() const { return false; }
 
   bool shouldFavorBackedgeIndex(const Loop *L) const { return false; }
 
-  bool isLegalMaskedStore(Type *DataType) { return false; }
+  bool isLegalMaskedStore(Type *DataType, unsigned Alignment) { return false; }
 
-  bool isLegalMaskedLoad(Type *DataType) { return false; }
+  bool isLegalMaskedLoad(Type *DataType, unsigned Alignment) { return false; }
 
   bool isLegalNTStore(Type *DataType, Align Alignment) {
     // By default, assume nontemporal memory stores are available for stores
     // that are aligned and have a size that is a power of 2.
     unsigned DataSize = DL.getTypeStoreSize(DataType);
     return Alignment >= DataSize && isPowerOf2_32(DataSize);
   }
 

lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::shouldFavorPostInc() const {
   return TTIImpl->shouldFavorPostInc();
 }
 
 bool TargetTransformInfo::shouldFavorBackedgeIndex(const Loop *L) const {
   return TTIImpl->shouldFavorBackedgeIndex(L);
 }
 
-bool TargetTransformInfo::isLegalMaskedStore(Type *DataType) const {
-  return TTIImpl->isLegalMaskedStore(DataType);
+bool TargetTransformInfo::isLegalMaskedStore(Type *DataType,
+                                             unsigned Alignment) const {
+  return TTIImpl->isLegalMaskedStore(DataType, Alignment);
 }
 
-bool TargetTransformInfo::isLegalMaskedLoad(Type *DataType) const {
-  return TTIImpl->isLegalMaskedLoad(DataType);
+bool TargetTransformInfo::isLegalMaskedLoad(Type *DataType,
+                                            unsigned Alignment) const {
+  return TTIImpl->isLegalMaskedLoad(DataType, Alignment);
 }
 
 bool TargetTransformInfo::isLegalNTStore(Type *DataType,
                                          Align Alignment) const {
   return TTIImpl->isLegalNTStore(DataType, Alignment);
 }
 
 bool TargetTransformInfo::isLegalNTLoad(Type *DataType, Align Alignment) const {

lib/CodeGen/ScalarizeMaskedMemIntrin.cpp

 
 bool ScalarizeMaskedMemIntrin::optimizeCallInst(CallInst *CI,
                                                 bool &ModifiedDT) {
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI);
   if (II) {
     switch (II->getIntrinsicID()) {
     default:
       break;
-    case Intrinsic::masked_load:
+    case Intrinsic::masked_load: {
+      auto Alignment = cast<ConstantInt>(CI->getArgOperand(1));
       // Scalarize unsupported vector masked load
-      if (TTI->isLegalMaskedLoad(CI->getType()))
+      if (TTI->isLegalMaskedLoad(CI->getType(), Alignment->getZExtValue()))
         return false;
       scalarizeMaskedLoad(CI, ModifiedDT);
       return true;
-    case Intrinsic::masked_store:
-      if (TTI->isLegalMaskedStore(CI->getArgOperand(0)->getType()))
+    }
+    case Intrinsic::masked_store: {
+      auto Alignment = cast<ConstantInt>(CI->getArgOperand(2));
+      if (TTI->isLegalMaskedStore(CI->getArgOperand(0)->getType(),
+                                  Alignment->getZExtValue()))
         return false;
       scalarizeMaskedStore(CI, ModifiedDT);
       return true;
+    }
     case Intrinsic::masked_gather:
       if (TTI->isLegalMaskedGather(CI->getType()))
         return false;
       scalarizeMaskedGather(CI, ModifiedDT);
       return true;
     case Intrinsic::masked_scatter:
       if (TTI->isLegalMaskedScatter(CI->getArgOperand(0)->getType()))
         return false;

lib/Target/ARM/ARMTargetTransformInfo.h

unsigned getRegisterBitWidth(bool Vector) const {
 
     return 32;
   }
 
   unsigned getMaxInterleaveFactor(unsigned VF) {
     return ST->getMaxInterleaveFactor();
   }
 
-  bool isLegalMaskedLoad(Type *DataTy);
-  bool isLegalMaskedStore(Type *DataTy) { return isLegalMaskedLoad(DataTy); }
+  bool isLegalMaskedLoad(Type *DataTy, unsigned Alignment);
+  bool isLegalMaskedStore(Type *DataTy, unsigned Alignment) {
+    return isLegalMaskedLoad(DataTy, Alignment);
+  }
 
   int getMemcpyCost(const Instruction *I);
 
   int getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index, Type *SubTp);
 
   bool useReductionIntrinsic(unsigned Opcode, Type *Ty,
                              TTI::ReductionFlags Flags) const;
 

lib/Target/ARM/ARMTargetTransformInfo.cpp

if (ST->hasNEON()) {
 
     // In many cases the address computation is not merged into the instruction
     // addressing mode.
     return 1;
   }
   return BaseT::getAddressComputationCost(Ty, SE, Ptr);
 }
 
-bool ARMTTIImpl::isLegalMaskedLoad(Type *DataTy) {
+bool ARMTTIImpl::isLegalMaskedLoad(Type *DataTy, unsigned Alignment) {
   if (!EnableMaskedLoadStores || !ST->hasMVEIntegerOps())
     return false;
 
   if (DataTy->isVectorTy()) {
     // We don't yet support narrowing or widening masked loads/stores. Expand
     // them for the moment.
     unsigned VecWidth = DataTy->getPrimitiveSizeInBits();
     if (VecWidth != 128)

lib/Target/X86/X86TargetTransformInfo.h

public:
   unsigned getUserCost(const User *U, ArrayRef<const Value *> Operands);
 
   int getIntImmCost(unsigned Opcode, unsigned Idx, const APInt &Imm, Type *Ty);
   int getIntImmCost(Intrinsic::ID IID, unsigned Idx, const APInt &Imm,
                     Type *Ty);
   bool isLSRCostLess(TargetTransformInfo::LSRCost &C1,
                      TargetTransformInfo::LSRCost &C2);
   bool canMacroFuseCmp();
-  bool isLegalMaskedLoad(Type *DataType);
+  bool isLegalMaskedLoad(Type *DataType, unsigned Alignment);
   bool isLegalMaskedStore(Type *DataType);
   bool isLegalNTLoad(Type *DataType, Align Alignment);
   bool isLegalNTStore(Type *DataType, Align Alignment);
   bool isLegalMaskedGather(Type *DataType);
   bool isLegalMaskedScatter(Type *DataType);
   bool isLegalMaskedExpandLoad(Type *DataType);
   bool isLegalMaskedCompressStore(Type *DataType);
   bool hasDivRemOp(Type *DataType, bool IsSigned);
   bool isFCmpOrdCheaperThanFCmpZero(Type *Ty);

lib/Target/X86/X86TargetTransformInfo.cpp

   VectorType *SrcVTy = dyn_cast<VectorType>(SrcTy);
   if (!SrcVTy)
     // To calculate scalar take the regular cost, without mask
     return getMemoryOpCost(Opcode, SrcTy, Alignment, AddressSpace);
 
   unsigned NumElem = SrcVTy->getVectorNumElements();
   VectorType *MaskTy =
       VectorType::get(Type::getInt8Ty(SrcVTy->getContext()), NumElem);
-  if ((IsLoad && !isLegalMaskedLoad(SrcVTy)) ||
-      (IsStore && !isLegalMaskedStore(SrcVTy)) || !isPowerOf2_32(NumElem)) {
+  if ((IsLoad && !isLegalMaskedLoad(SrcVTy, Alignment)) ||
+      (IsStore && !isLegalMaskedStore(SrcVTy, Alignment)) ||
+      !isPowerOf2_32(NumElem)) {
     // Scalarization
     int MaskSplitCost = getScalarizationOverhead(MaskTy, false, true);
     int ScalarCompareCost = getCmpSelInstrCost(
         Instruction::ICmp, Type::getInt8Ty(SrcVTy->getContext()), nullptr);
     int BranchCost = getCFInstrCost(Instruction::Br);
     int MaskCmpCost = NumElem * (BranchCost + ScalarCompareCost);
 
     int ValueSplitCost = getScalarizationOverhead(SrcVTy, IsLoad, IsStore);

                     C2.NumIVMuls, C2.NumBaseAdds,
                     C2.ScaleCost, C2.ImmCost, C2.SetupCost);
 }
 
 bool X86TTIImpl::canMacroFuseCmp() {
   return ST->hasMacroFusion() || ST->hasBranchFusion();
 }
 
-bool X86TTIImpl::isLegalMaskedLoad(Type *DataTy) {
+bool X86TTIImpl::isLegalMaskedLoad(Type *DataTy, unsigned Alignment) {
   if (!ST->hasAVX())
     return false;
 
   // The backend can't handle a single element vector.
   if (isa<VectorType>(DataTy) && DataTy->getVectorNumElements() == 1)
     return false;
   Type *ScalarTy = DataTy->getScalarType();
 
   if (ScalarTy->isPointerTy())
     return true;
 
   if (ScalarTy->isFloatTy() || ScalarTy->isDoubleTy())
     return true;
 
   if (!ScalarTy->isIntegerTy())
     return false;
 
   unsigned IntWidth = ScalarTy->getIntegerBitWidth();
   return IntWidth == 32 || IntWidth == 64 ||
          ((IntWidth == 8 || IntWidth == 16) && ST->hasBWI());
 }
 
-bool X86TTIImpl::isLegalMaskedStore(Type *DataType) {
-  return isLegalMaskedLoad(DataType);
+bool X86TTIImpl::isLegalMaskedStore(Type *DataType, unsigned Alignment) {
+  return isLegalMaskedLoad(DataType, Alignment);
 }
 
 bool X86TTIImpl::isLegalNTLoad(Type *DataType, Align Alignment) {
   unsigned DataSize = DL.getTypeStoreSize(DataType);
   // The only supported nontemporal loads are for aligned vectors of 16 or 32
   // bytes.  Note that 32-byte nontemporal vector loads are supported by AVX2
   // (the equivalent stores only require AVX).
   if (Alignment >= DataSize && (DataSize == 16 || DataSize == 32))

lib/Transforms/Vectorize/LoopVectorize.cpp

       return;
     setCostBasedWideningDecision(VF);
     collectLoopUniforms(VF);
     collectLoopScalars(VF);
   }
 
   /// Returns true if the target machine supports masked store operation
   /// for the given \p DataType and kind of access to \p Ptr.
-  bool isLegalMaskedStore(Type *DataType, Value *Ptr) {
-    return Legal->isConsecutivePtr(Ptr) && TTI.isLegalMaskedStore(DataType);
+  bool isLegalMaskedStore(Type *DataType, unsigned Alignment, Value *Ptr) {
+    return Legal->isConsecutivePtr(Ptr) &&
+           TTI.isLegalMaskedStore(DataType, Alignment);
   }
 
   /// 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 Legal->isConsecutivePtr(Ptr) && TTI.isLegalMaskedLoad(DataType);
+  bool isLegalMaskedLoad(Type *DataType, unsigned Alignment, Value *Ptr) {
+    return Legal->isConsecutivePtr(Ptr) &&
+           TTI.isLegalMaskedLoad(DataType, Alignment);
   }
 
   /// Returns true if the target machine supports masked scatter operation
   /// for the given \p DataType.
   bool isLegalMaskedScatter(Type *DataType) {
     return TTI.isLegalMaskedScatter(DataType);
   }
 

     // result.
     if (VF > 1) {
       InstWidening WideningDecision = getWideningDecision(I, VF);
       assert(WideningDecision != CM_Unknown &&
              "Widening decision should be ready at this moment");
       return WideningDecision == CM_Scalarize;
     }
     return isa<LoadInst>(I) ?
-        !(isLegalMaskedLoad(Ty, Ptr)  || isLegalMaskedGather(Ty))
-      : !(isLegalMaskedStore(Ty, Ptr) || isLegalMaskedScatter(Ty));
+        !(isLegalMaskedLoad(Ty, cast<LoadInst>(I)->getAlignment(), Ptr) ||
+          isLegalMaskedGather(Ty))
+      : !(isLegalMaskedStore(Ty, cast<StoreInst>(I)->getAlignment(), Ptr) ||
+          isLegalMaskedScatter(Ty));
   }
   case Instruction::UDiv:
   case Instruction::SDiv:
   case Instruction::SRem:
   case Instruction::URem:
     return mayDivideByZero(*I);
   }
   return false;

 
   // If masked interleaving is required, we expect that the user/target had
   // enabled it, because otherwise it either wouldn't have been created or
   // it should have been invalidated by the CostModel.
   assert(useMaskedInterleavedAccesses(TTI) &&
          "Masked interleave-groups for predicated accesses are not enabled.");
 
   auto *Ty = getMemInstValueType(I);
-  return isa<LoadInst>(I) ? TTI.isLegalMaskedLoad(Ty)
-                          : TTI.isLegalMaskedStore(Ty);
+  return isa<LoadInst>(I) ?
+    TTI.isLegalMaskedLoad(Ty, cast<LoadInst>(I)->getAlignment()) :
+    TTI.isLegalMaskedStore(Ty, cast<StoreInst>(I)->getAlignment());
 }
 
 bool LoopVectorizationCostModel::memoryInstructionCanBeWidened(Instruction *I,
                                                                unsigned VF) {
   // Get and ensure we have a valid memory instruction.
   LoadInst *LI = dyn_cast<LoadInst>(I);
   StoreInst *SI = dyn_cast<StoreInst>(I);
   assert((LI || SI) && "Invalid memory instruction");