llvm · HankChang736 · Jun 28, 2024
diff --git a/llvm/include/llvm/Analysis/MemoryRefInfo.h b/llvm/include/llvm/Analysis/MemoryRefInfo.h
@@ -0,0 +1,53 @@
+//===--------- Definition of the MemoryRefInfo class -*- C++ -*------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines MemoryRefInfo class that is used when getting
+// the information of a memory reference instruction.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_ANALYSIS_MEMORYREFINFO_H
+#define LLVM_ANALYSIS_MEMORYREFINFO_H
+
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/Support/TypeSize.h"
+
+namespace llvm {
+class MemoryRefInfo {
+public:
+  Use *PtrUse = nullptr;
+  bool IsWrite;
+  Type *OpType;
+  TypeSize TypeStoreSize = TypeSize::getFixed(0);
+  MaybeAlign Alignment;
+  // The mask Value, if we're looking at a masked load/store.
+  Value *MaybeMask;
+  // The EVL Value, if we're looking at a vp intrinsic.
+  Value *MaybeEVL;
+  // The Stride Value, if we're looking at a strided load/store.
+  Value *MaybeStride;
+
+  MemoryRefInfo() = default;
+  MemoryRefInfo(Instruction *I, unsigned OperandNo, bool IsWrite,
+                class Type *OpType, MaybeAlign Alignment,
+                Value *MaybeMask = nullptr, Value *MaybeEVL = nullptr,
+                Value *MaybeStride = nullptr)
+      : IsWrite(IsWrite), OpType(OpType), Alignment(Alignment),
+        MaybeMask(MaybeMask), MaybeEVL(MaybeEVL), MaybeStride(MaybeStride) {
+    const DataLayout &DL = I->getDataLayout();
+    TypeStoreSize = DL.getTypeStoreSizeInBits(OpType);
+    PtrUse = &I->getOperandUse(OperandNo);
+  }
+
+  Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
+  Value *getPtr() { return PtrUse->get(); }
+  operator bool() { return PtrUse != nullptr; }
+};
+
+} // namespace llvm
+#endif
diff --git a/llvm/include/llvm/Analysis/TargetTransformInfo.h b/llvm/include/llvm/Analysis/TargetTransformInfo.h
@@ -23,6 +23,8 @@
 
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Analysis/MemoryRefInfo.h"
 #include "llvm/IR/FMF.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/PassManager.h"
@@ -998,6 +1000,10 @@ class TargetTransformInfo {
   MemCmpExpansionOptions enableMemCmpExpansion(bool OptSize,
                                                bool IsZeroCmp) const;
 
+  // Add MemoryRefInfo of Intrinsic \p II into array \p Interesting.
+  bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                        IntrinsicInst *II) const;
+
   /// Should the Select Optimization pass be enabled and ran.
   bool enableSelectOptimize() const;
 
@@ -2083,6 +2089,8 @@ class TargetTransformInfo::Concept {
   virtual bool enableAggressiveInterleaving(bool LoopHasReductions) = 0;
   virtual MemCmpExpansionOptions
   enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const = 0;
+  virtual bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                                IntrinsicInst *II) const = 0;
   virtual bool enableSelectOptimize() = 0;
   virtual bool shouldTreatInstructionLikeSelect(const Instruction *I) = 0;
   virtual bool enableInterleavedAccessVectorization() = 0;
@@ -2711,6 +2719,12 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
                                                bool IsZeroCmp) const override {
     return Impl.enableMemCmpExpansion(OptSize, IsZeroCmp);
   }
+
+  bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                        IntrinsicInst *II) const override {
+    return Impl.getMemoryRefInfo(Interesting, II);
+  }
+
   bool enableSelectOptimize() override {
     return Impl.enableSelectOptimize();
   }

diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
@@ -434,6 +434,11 @@ class TargetTransformInfoImplBase {
     return {};
   }
 
+  bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                        IntrinsicInst *II) const {
+    return false;
+  }
+
   bool enableSelectOptimize() const { return true; }
 
   bool shouldTreatInstructionLikeSelect(const Instruction *I) {

diff --git a/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h b/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h
@@ -14,6 +14,7 @@
 #define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
 
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/MemoryRefInfo.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
@@ -22,37 +23,6 @@
 
 namespace llvm {
 
-class InterestingMemoryOperand {
-public:
-  Use *PtrUse;
-  bool IsWrite;
-  Type *OpType;
-  TypeSize TypeStoreSize = TypeSize::getFixed(0);
-  MaybeAlign Alignment;
-  // The mask Value, if we're looking at a masked load/store.
-  Value *MaybeMask;
-  // The EVL Value, if we're looking at a vp intrinsic.
-  Value *MaybeEVL;
-  // The Stride Value, if we're looking at a strided load/store.
-  Value *MaybeStride;
-
-  InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
-                           class Type *OpType, MaybeAlign Alignment,
-                           Value *MaybeMask = nullptr,
-                           Value *MaybeEVL = nullptr,
-                           Value *MaybeStride = nullptr)
-      : IsWrite(IsWrite), OpType(OpType), Alignment(Alignment),
-        MaybeMask(MaybeMask), MaybeEVL(MaybeEVL), MaybeStride(MaybeStride) {
-    const DataLayout &DL = I->getDataLayout();
-    TypeStoreSize = DL.getTypeStoreSizeInBits(OpType);
-    PtrUse = &I->getOperandUse(OperandNo);
-  }
-
-  Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
-
-  Value *getPtr() { return PtrUse->get(); }
-};
-
 // Get AddressSanitizer parameters.
 void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
                                bool IsKasan, uint64_t *ShadowBase,

diff --git a/llvm/lib/Analysis/TargetTransformInfo.cpp b/llvm/lib/Analysis/TargetTransformInfo.cpp
@@ -660,6 +660,11 @@ TargetTransformInfo::enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const {
   return TTIImpl->enableMemCmpExpansion(OptSize, IsZeroCmp);
 }
 
+bool TargetTransformInfo::getMemoryRefInfo(
+    SmallVectorImpl<MemoryRefInfo> &Interesting, IntrinsicInst *II) const {
+  return TTIImpl->getMemoryRefInfo(Interesting, II);
+}
+
 bool TargetTransformInfo::enableSelectOptimize() const {
   return TTIImpl->enableSelectOptimize();
 }

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp
@@ -219,7 +219,7 @@ void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
 
 void getInterestingMemoryOperands(
     Module &M, Instruction *I,
-    SmallVectorImpl<InterestingMemoryOperand> &Interesting) {
+    SmallVectorImpl<MemoryRefInfo> &Interesting) {
   const DataLayout &DL = M.getDataLayout();
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
     Interesting.emplace_back(I, LI->getPointerOperandIndex(), false,

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h
@@ -52,8 +52,7 @@ void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
 /// that needs to be instrumented
 void getInterestingMemoryOperands(
     Module &M, Instruction *I,
-    SmallVectorImpl<InterestingMemoryOperand> &Interesting);
-
+    SmallVectorImpl<MemoryRefInfo> &Interesting);
 } // end namespace AMDGPU
 } // end namespace llvm
 

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -1261,9 +1261,9 @@ bool AMDGPUSwLowerLDS::run() {
   }
 
   if (AsanInstrumentLDS) {
-    SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
+    SmallVector<MemoryRefInfo, 16> OperandsToInstrument;
     for (Instruction *Inst : AsanInfo.Instructions) {
-      SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
+      SmallVector<MemoryRefInfo, 1> InterestingOperands;
       getInterestingMemoryOperands(M, Inst, InterestingOperands);
       for (auto &Operand : InterestingOperands) {
         OperandsToInstrument.push_back(Operand);

diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp b/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>
@@ -37,6 +38,83 @@ static cl::opt<unsigned> SLPMaxVF(
         "exclusively by SLP vectorizer."),
     cl::Hidden);
 
+bool RISCVTTIImpl::getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                                    IntrinsicInst *II) const {
+  const DataLayout &DL = getDataLayout();
+  Intrinsic::ID IID = II->getIntrinsicID();
+  LLVMContext &C = II->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)
+    bool IsWrite = II->getType()->isVoidTy();
+    Type *Ty = IsWrite ? II->getArgOperand(0)->getType() : II->getType();
+    const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
+    unsigned VLIndex = RVVIInfo->VLOperand;
+    unsigned PtrOperandNo = VLIndex - 1 - HasMask;
+    MaybeAlign Alignment =
+        II->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
+    Type *MaskType = Ty->getWithNewType(Type::getInt1Ty(C));
+    Value *Mask = ConstantInt::getTrue(MaskType);
+    if (HasMask)
+      Mask = II->getArgOperand(VLIndex - 1);
+    Value *EVL = II->getArgOperand(VLIndex);
+    Interesting.emplace_back(II, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
+                             EVL);
+    return true;
+  }
+  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)
+    bool IsWrite = II->getType()->isVoidTy();
+    Type *Ty = IsWrite ? II->getArgOperand(0)->getType() : II->getType();
+    const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
+    unsigned VLIndex = RVVIInfo->VLOperand;
+    unsigned PtrOperandNo = VLIndex - 2 - HasMask;
+    MaybeAlign Alignment =
+        II->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
+
+    Value *Stride = II->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 = II->getArgOperand(VLIndex - 1);
+    Value *EVL = II->getArgOperand(VLIndex);
+    Interesting.emplace_back(II, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
+                             EVL, Stride);
+    return true;
+  }
+  }
+  return false;
+}
+
 InstructionCost
 RISCVTTIImpl::getRISCVInstructionCost(ArrayRef<unsigned> OpCodes, MVT VT,
                                       TTI::TargetCostKind CostKind) {
@@ -2796,3 +2874,18 @@ RISCVTTIImpl::enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const {
   }
   return Options;
 }
+
+bool RISCVTTIImpl::areInlineCompatible(const Function *Caller,
+                                       const Function *Callee) const {
+  const TargetMachine &TM = getTLI()->getTargetMachine();
+
+  const FeatureBitset &CallerBits =
+      TM.getSubtargetImpl(*Caller)->getFeatureBits();
+  const FeatureBitset &CalleeBits =
+      TM.getSubtargetImpl(*Callee)->getFeatureBits();
+
+  // Inline a callee if its target-features are a subset of the callers
+  // target-features.
+  return (CallerBits & CalleeBits) == CalleeBits;
+}
+
diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h
@@ -68,6 +68,12 @@ class RISCVTTIImpl : public BasicTTIImplBase<RISCVTTIImpl> {
       : BaseT(TM, F.getDataLayout()), ST(TM->getSubtargetImpl(F)),
         TLI(ST->getTargetLowering()) {}
 
+  bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
+                        IntrinsicInst *II) const;
+  //This might be the problem
+  bool areInlineCompatible(const Function *Caller,
+                           const Function *Callee) const;
+
   /// Return the cost of materializing an immediate for a value operand of
   /// a store instruction.
   InstructionCost getStoreImmCost(Type *VecTy, TTI::OperandValueInfo OpInfo,