Change getTgtMemIntrinsic interface with returning a pair

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -32,6 +32,7 @@
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/InstructionCost.h"
+#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include <functional>
 #include <optional>
 #include <utility>
@@ -1696,8 +1697,9 @@ class TargetTransformInfo {
   /// will contain additional information - whether the intrinsic may write
   /// or read to memory, volatility and the pointer.  Info is undefined
   /// if false is returned.
-  LLVM_ABI bool getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                   MemIntrinsicInfo &Info) const;
+  LLVM_ABI std::pair<std::optional<MemIntrinsicInfo>,
+                     std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const;
 
   /// \returns The maximum element size, in bytes, for an element
   /// unordered-atomic memory intrinsic.

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -975,9 +975,10 @@ class TargetTransformInfoImplBase {
     return 0;
   }
 
-  virtual bool getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                  MemIntrinsicInfo &Info) const {
-    return false;
+  virtual std::pair<std::optional<MemIntrinsicInfo>,
+                    std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const {
+    return std::make_pair(std::nullopt, std::nullopt);
   }
 
   virtual unsigned getAtomicMemIntrinsicMaxElementSize() const {

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
@@ -1285,9 +1286,10 @@ TargetTransformInfo::getCostOfKeepingLiveOverCall(ArrayRef<Type *> Tys) const {
   return TTIImpl->getCostOfKeepingLiveOverCall(Tys);
 }
 
-bool TargetTransformInfo::getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                             MemIntrinsicInfo &Info) const {
-  return TTIImpl->getTgtMemIntrinsic(Inst, Info);
+std::pair<std::optional<MemIntrinsicInfo>,
+          std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+TargetTransformInfo::getTgtMemIntrinsic(IntrinsicInst *Inst) const {
+  return TTIImpl->getTgtMemIntrinsic(Inst);
 }
 
 unsigned TargetTransformInfo::getAtomicMemIntrinsicMaxElementSize() const {

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -5125,8 +5125,10 @@ Value *AArch64TTIImpl::getOrCreateResultFromMemIntrinsic(IntrinsicInst *Inst,
   }
 }
 
-bool AArch64TTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                        MemIntrinsicInfo &Info) const {
+std::pair<std::optional<MemIntrinsicInfo>,
+          std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+AArch64TTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst) const {
+  MemIntrinsicInfo Info;
   switch (Inst->getIntrinsicID()) {
   default:
     break;
@@ -5148,7 +5150,7 @@ bool AArch64TTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
 
   switch (Inst->getIntrinsicID()) {
   default:
-    return false;
+    return std::make_pair(std::nullopt, std::nullopt);
   case Intrinsic::aarch64_neon_ld2:
   case Intrinsic::aarch64_neon_st2:
     Info.MatchingId = VECTOR_LDST_TWO_ELEMENTS;
@@ -5162,7 +5164,7 @@ bool AArch64TTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
     Info.MatchingId = VECTOR_LDST_FOUR_ELEMENTS;
     break;
   }
-  return true;
+  return std::make_pair(Info, std::nullopt);
 }
 
 /// See if \p I should be considered for address type promotion. We check if \p

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

@@ -273,8 +273,9 @@ class AArch64TTIImpl final : public BasicTTIImplBase<AArch64TTIImpl> {
   getOrCreateResultFromMemIntrinsic(IntrinsicInst *Inst, Type *ExpectedType,
                                     bool CanCreate = true) const override;
 
-  bool getTgtMemIntrinsic(IntrinsicInst *Inst,
-                          MemIntrinsicInfo &Info) const override;
+  std::pair<std::optional<MemIntrinsicInfo>,
+            std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const override;
 
   bool isElementTypeLegalForScalableVector(Type *Ty) const override {
     if (Ty->isPointerTy())

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

@@ -500,29 +500,32 @@ unsigned GCNTTIImpl::getMaxInterleaveFactor(ElementCount VF) const {
   return 8;
 }
 
-bool GCNTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                       MemIntrinsicInfo &Info) const {
+std::pair<std::optional<MemIntrinsicInfo>,
+          std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+GCNTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst) const {
   switch (Inst->getIntrinsicID()) {
   case Intrinsic::amdgcn_ds_ordered_add:
   case Intrinsic::amdgcn_ds_ordered_swap: {
     auto *Ordering = dyn_cast<ConstantInt>(Inst->getArgOperand(2));
     auto *Volatile = dyn_cast<ConstantInt>(Inst->getArgOperand(4));
     if (!Ordering || !Volatile)
-      return false; // Invalid.
+      return std::make_pair(std::nullopt, std::nullopt); // Invalid
 
     unsigned OrderingVal = Ordering->getZExtValue();
-    if (OrderingVal > static_cast<unsigned>(AtomicOrdering::SequentiallyConsistent))
-      return false;
+    if (OrderingVal >
+        static_cast<unsigned>(AtomicOrdering::SequentiallyConsistent))
+      return std::make_pair(std::nullopt, std::nullopt);
 
+    MemIntrinsicInfo Info;
     Info.PtrVal = Inst->getArgOperand(0);
     Info.Ordering = static_cast<AtomicOrdering>(OrderingVal);
     Info.ReadMem = true;
     Info.WriteMem = true;
     Info.IsVolatile = !Volatile->isZero();
-    return true;
+    return std::make_pair(Info, std::nullopt);
   }
   default:
-    return false;
+    return std::make_pair(std::nullopt, std::nullopt);
   }
 }
 

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

@@ -151,8 +151,9 @@ class GCNTTIImpl final : public BasicTTIImplBase<GCNTTIImpl> {
       std::optional<uint32_t> AtomicCpySize) const override;
   unsigned getMaxInterleaveFactor(ElementCount VF) const override;
 
-  bool getTgtMemIntrinsic(IntrinsicInst *Inst,
-                          MemIntrinsicInfo &Info) const override;
+  std::pair<std::optional<MemIntrinsicInfo>,
+            std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const override;
 
   InstructionCost getArithmeticInstrCost(
       unsigned Opcode, Type *Ty, TTI::TargetCostKind CostKind,

llvm/lib/Target/PowerPC/PPCTargetTransformInfo.cpp

@@ -975,8 +975,10 @@ bool PPCTTIImpl::shouldBuildRelLookupTables() const {
   return BaseT::shouldBuildRelLookupTables();
 }
 
-bool PPCTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
-                                    MemIntrinsicInfo &Info) const {
+std::pair<std::optional<MemIntrinsicInfo>,
+          std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+PPCTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst) const {
+  MemIntrinsicInfo Info;
   switch (Inst->getIntrinsicID()) {
   case Intrinsic::ppc_altivec_lvx:
   case Intrinsic::ppc_altivec_lvxl:
@@ -993,7 +995,7 @@ bool PPCTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
     Info.PtrVal = Inst->getArgOperand(0);
     Info.ReadMem = true;
     Info.WriteMem = false;
-    return true;
+    return std::make_pair(Info, std::nullopt);
   }
   case Intrinsic::ppc_altivec_stvx:
   case Intrinsic::ppc_altivec_stvxl:
@@ -1010,7 +1012,7 @@ bool PPCTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
     Info.PtrVal = Inst->getArgOperand(1);
     Info.ReadMem = false;
     Info.WriteMem = true;
-    return true;
+    return std::make_pair(Info, std::nullopt);
   }
   case Intrinsic::ppc_stbcx:
   case Intrinsic::ppc_sthcx:
@@ -1019,13 +1021,12 @@ bool PPCTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst,
     Info.PtrVal = Inst->getArgOperand(0);
     Info.ReadMem = false;
     Info.WriteMem = true;
-    return true;
+    return std::make_pair(Info, std::nullopt);
   }
   default:
     break;
   }
-
-  return false;
+  return std::make_pair(std::nullopt, std::nullopt);
 }
 
 bool PPCTTIImpl::supportsTailCallFor(const CallBase *CB) const {

llvm/lib/Target/PowerPC/PPCTargetTransformInfo.h

@@ -69,8 +69,9 @@ class PPCTTIImpl final : public BasicTTIImplBase<PPCTTIImpl> {
   bool canSaveCmp(Loop *L, BranchInst **BI, ScalarEvolution *SE, LoopInfo *LI,
                   DominatorTree *DT, AssumptionCache *AC,
                   TargetLibraryInfo *LibInfo) const override;
-  bool getTgtMemIntrinsic(IntrinsicInst *Inst,
-                          MemIntrinsicInfo &Info) const override;
+  std::pair<std::optional<MemIntrinsicInfo>,
+            std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const override;
   void getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                                TTI::UnrollingPreferences &UP,
                                OptimizationRemarkEmitter *ORE) const override;

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

@@ -15,6 +15,7 @@
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicsRISCV.h"
 #include "llvm/IR/PatternMatch.h"
 #include <cmath>
 #include <optional>
@@ -2673,6 +2674,85 @@ void RISCVTTIImpl::getPeelingPreferences(Loop *L, ScalarEvolution &SE,
   BaseT::getPeelingPreferences(L, SE, PP);
 }
 
+std::pair<std::optional<MemIntrinsicInfo>,
+          std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+RISCVTTIImpl::getTgtMemIntrinsic(IntrinsicInst *Inst) const {
+  const DataLayout &DL = getDataLayout();
+  Intrinsic::ID IID = Inst->getIntrinsicID();
+  LLVMContext &C = Inst->getContext();
+  bool HasMask = false;
+  switch (IID) {
+  case Intrinsic::riscv_vle_mask:
+  case Intrinsic::riscv_vse_mask:
+    HasMask = true;
+    [[fallthrough]];
+  case Intrinsic::riscv_vle:
+  case Intrinsic::riscv_vse: {
+    // Intrinsic interface:
+    // riscv_vle(merge, ptr, vl)
+    // riscv_vle_mask(merge, ptr, mask, vl, policy)
+    // riscv_vse(val, ptr, vl)
+    // riscv_vse_mask(val, ptr, mask, vl, policy)
+    SmallVector<InterestingMemoryOperand, 1> Interesting;
+    bool IsWrite = Inst->getType()->isVoidTy();
+    Type *Ty = IsWrite ? Inst->getArgOperand(0)->getType() : Inst->getType();
+    const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
+    unsigned VLIndex = RVVIInfo->VLOperand;
+    unsigned PtrOperandNo = VLIndex - 1 - HasMask;
+    MaybeAlign Alignment =
+        Inst->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
+    Type *MaskType = Ty->getWithNewType(Type::getInt1Ty(C));
+    Value *Mask = ConstantInt::getTrue(MaskType);
+    if (HasMask)
+      Mask = Inst->getArgOperand(VLIndex - 1);
+    Value *EVL = Inst->getArgOperand(VLIndex);
+    Interesting.emplace_back(Inst, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
+                             EVL);
+    return std::make_pair(std::nullopt, Interesting);
+  }
+  case Intrinsic::riscv_vlse_mask:
+  case Intrinsic::riscv_vsse_mask:
+    HasMask = true;
+    [[fallthrough]];
+  case Intrinsic::riscv_vlse:
+  case Intrinsic::riscv_vsse: {
+    // Intrinsic interface:
+    // riscv_vlse(merge, ptr, stride, vl)
+    // riscv_vlse_mask(merge, ptr, stride, mask, vl, policy)
+    // riscv_vsse(val, ptr, stride, vl)
+    // riscv_vsse_mask(val, ptr, stride, mask, vl, policy)
+    SmallVector<InterestingMemoryOperand, 1> Interesting;
+    bool IsWrite = Inst->getType()->isVoidTy();
+    Type *Ty = IsWrite ? Inst->getArgOperand(0)->getType() : Inst->getType();
+    const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
+    unsigned VLIndex = RVVIInfo->VLOperand;
+    unsigned PtrOperandNo = VLIndex - 2 - HasMask;
+    MaybeAlign Alignment =
+        Inst->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
+
+    Value *Stride = Inst->getArgOperand(PtrOperandNo + 1);
+    // Use the pointer alignment as the element alignment if the stride is a
+    // multiple of the pointer alignment. Otherwise, the element alignment
+    // should be the greatest common divisor of pointer alignment and stride.
+    // For simplicity, just consider unalignment for elements.
+    unsigned PointerAlign = Alignment.valueOrOne().value();
+    if (!isa<ConstantInt>(Stride) ||
+        cast<ConstantInt>(Stride)->getZExtValue() % PointerAlign != 0)
+      Alignment = Align(1);
+
+    Type *MaskType = Ty->getWithNewType(Type::getInt1Ty(C));
+    Value *Mask = ConstantInt::getTrue(MaskType);
+    if (HasMask)
+      Mask = Inst->getArgOperand(VLIndex - 1);
+    Value *EVL = Inst->getArgOperand(VLIndex);
+    Interesting.emplace_back(Inst, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
+                             EVL, Stride);
+    return std::make_pair(std::nullopt, Interesting);
+  }
+  }
+  return std::make_pair(std::nullopt, std::nullopt);
+}
+
 unsigned RISCVTTIImpl::getRegUsageForType(Type *Ty) const {
   if (Ty->isVectorTy()) {
     // f16 with only zvfhmin and bf16 will be promoted to f32

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h

@@ -158,6 +158,10 @@ class RISCVTTIImpl final : public BasicTTIImplBase<RISCVTTIImpl> {
   void getPeelingPreferences(Loop *L, ScalarEvolution &SE,
                              TTI::PeelingPreferences &PP) const override;
 
+  std::pair<std::optional<MemIntrinsicInfo>,
+            std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+  getTgtMemIntrinsic(IntrinsicInst *Inst) const override;
+
   unsigned getMinVectorRegisterBitWidth() const override {
     return ST->useRVVForFixedLengthVectors() ? 16 : 0;
   }

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -29,6 +29,7 @@
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/BinaryFormat/MachO.h"
 #include "llvm/Demangle/Demangle.h"
@@ -803,7 +804,8 @@ struct AddressSanitizer {
 
   bool ignoreAccess(Instruction *Inst, Value *Ptr);
   void getInterestingMemoryOperands(
-      Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
+      Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting,
+      const TargetTransformInfo *TTI);
 
   void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
                      InterestingMemoryOperand &O, bool UseCalls,
@@ -843,7 +845,7 @@ struct AddressSanitizer {
   void instrumentMemIntrinsic(MemIntrinsic *MI, RuntimeCallInserter &RTCI);
   Value *memToShadow(Value *Shadow, IRBuilder<> &IRB);
   bool suppressInstrumentationSiteForDebug(int &Instrumented);
-  bool instrumentFunction(Function &F, const TargetLibraryInfo *TLI);
+  bool instrumentFunction(Function &F, const TargetLibraryInfo *TLI, const TargetTransformInfo *TTI);
   bool maybeInsertAsanInitAtFunctionEntry(Function &F);
   bool maybeInsertDynamicShadowAtFunctionEntry(Function &F);
   void markEscapedLocalAllocas(Function &F);
@@ -1312,7 +1314,8 @@ PreservedAnalyses AddressSanitizerPass::run(Module &M,
         Options.MaxInlinePoisoningSize, Options.CompileKernel, Options.Recover,
         Options.UseAfterScope, Options.UseAfterReturn);
     const TargetLibraryInfo &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
-    Modified |= FunctionSanitizer.instrumentFunction(F, &TLI);
+    const TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
+    Modified |= FunctionSanitizer.instrumentFunction(F, &TLI, &TTI);
   }
   Modified |= ModuleSanitizer.instrumentModule();
   if (!Modified)
@@ -1450,7 +1453,8 @@ bool AddressSanitizer::ignoreAccess(Instruction *Inst, Value *Ptr) {
 }
 
 void AddressSanitizer::getInterestingMemoryOperands(
-    Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting) {
+    Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting,
+    const TargetTransformInfo *TTI) {
   // Do not instrument the load fetching the dynamic shadow address.
   if (LocalDynamicShadow == I)
     return;
@@ -1568,6 +1572,15 @@ void AddressSanitizer::getInterestingMemoryOperands(
       break;
     }
     default:
+      if (auto *II = dyn_cast<IntrinsicInst>(I)) {
+        std::pair<std::optional<MemIntrinsicInfo>,
+                  std::optional<SmallVector<InterestingMemoryOperand, 1>>>
+            MemInfo;
+        MemInfo = TTI->getTgtMemIntrinsic(II);
+        if (MemInfo.second != std::nullopt)
+          Interesting = *MemInfo.second;
+        return;
+      }
       for (unsigned ArgNo = 0; ArgNo < CI->arg_size(); ArgNo++) {
         if (!ClInstrumentByval || !CI->isByValArgument(ArgNo) ||
             ignoreAccess(I, CI->getArgOperand(ArgNo)))
@@ -2983,7 +2996,8 @@ bool AddressSanitizer::suppressInstrumentationSiteForDebug(int &Instrumented) {
 }
 
 bool AddressSanitizer::instrumentFunction(Function &F,
-                                          const TargetLibraryInfo *TLI) {
+                                          const TargetLibraryInfo *TLI,
+					  const TargetTransformInfo *TTI) {
   bool FunctionModified = false;
 
   // Do not apply any instrumentation for naked functions.
@@ -3036,7 +3050,7 @@ bool AddressSanitizer::instrumentFunction(Function &F,
       if (Inst.hasMetadata(LLVMContext::MD_nosanitize))
         continue;
       SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
-      getInterestingMemoryOperands(&Inst, InterestingOperands);
+      getInterestingMemoryOperands(&Inst, InterestingOperands, TTI);
 
       if (!InterestingOperands.empty()) {
         for (auto &Operand : InterestingOperands) {