Skip to content

Commit e8ad87c

Browse files
yetingkyetingk
authored
[Asan] Provide TTI hook to provide memory reference infromation of target intrinsics. (#97070)
Previously asan considers target intrinsics as black boxes, so asan could not instrument accurate check. This patch provide TTI hooks to make targets describe their intrinsic informations to asan. Note, 1. this patch renames InterestingMemoryOperand to MemoryRefInfo. 2. this patch does not support RVV indexed/segment load/store.
1 parent d09be91 commit e8ad87c

File tree

10 files changed

+2463
-57
lines changed

10 files changed

+2463
-57
lines changed

‎llvm/include/llvm/Analysis/MemoryRefInfo.h

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
//===--------- Definition of the MemoryRefInfo class -*- C++ -*------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file defines MemoryRefInfo class that is used when getting
10+
// the information of a memory reference instruction.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
#ifndef LLVM_ANALYSIS_MEMORYREFINFO_H
14+
#define LLVM_ANALYSIS_MEMORYREFINFO_H
15+
16+
#include "llvm/IR/DataLayout.h"
17+
#include "llvm/IR/Instruction.h"
18+
#include "llvm/Support/TypeSize.h"
19+
20+
namespace llvm {
21+
class MemoryRefInfo {
22+
public:
23+
Use *PtrUse = nullptr;
24+
bool IsWrite;
25+
Type *OpType;
26+
TypeSize TypeStoreSize = TypeSize::getFixed(0);
27+
MaybeAlign Alignment;
28+
// The mask Value, if we're looking at a masked load/store.
29+
Value *MaybeMask;
30+
// The EVL Value, if we're looking at a vp intrinsic.
31+
Value *MaybeEVL;
32+
// The Stride Value, if we're looking at a strided load/store.
33+
Value *MaybeStride;
34+
35+
MemoryRefInfo() = default;
36+
MemoryRefInfo(Instruction *I, unsigned OperandNo, bool IsWrite,
37+
class Type *OpType, MaybeAlign Alignment,
38+
Value *MaybeMask = nullptr, Value *MaybeEVL = nullptr,
39+
Value *MaybeStride = nullptr)
40+
: IsWrite(IsWrite), OpType(OpType), Alignment(Alignment),
41+
MaybeMask(MaybeMask), MaybeEVL(MaybeEVL), MaybeStride(MaybeStride) {
42+
const DataLayout &DL = I->getDataLayout();
43+
TypeStoreSize = DL.getTypeStoreSizeInBits(OpType);
44+
PtrUse = &I->getOperandUse(OperandNo);
45+
}
46+
47+
Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
48+
Value *getPtr() { return PtrUse->get(); }
49+
operator bool() { return PtrUse != nullptr; }
50+
};
51+
52+
} // namespace llvm
53+
#endif

‎llvm/include/llvm/Analysis/TargetTransformInfo.h

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
#include "llvm/ADT/APInt.h"
2525
#include "llvm/ADT/SmallBitVector.h"
26+
#include "llvm/Analysis/MemoryRefInfo.h"
2627
#include "llvm/IR/FMF.h"
2728
#include "llvm/IR/InstrTypes.h"
2829
#include "llvm/IR/PassManager.h"
@@ -961,6 +962,10 @@ class TargetTransformInfo {
961962
MemCmpExpansionOptions enableMemCmpExpansion(bool OptSize,
962963
bool IsZeroCmp) const;
963964

965+
// Add MemoryRefInfo of Intrinsic \p II into array \p Interesting.
966+
bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
967+
IntrinsicInst *II) const;
968+
964969
/// Should the Select Optimization pass be enabled and ran.
965970
bool enableSelectOptimize() const;
966971

@@ -1951,6 +1956,8 @@ class TargetTransformInfo::Concept {
19511956
virtual bool enableAggressiveInterleaving(bool LoopHasReductions) = 0;
19521957
virtual MemCmpExpansionOptions
19531958
enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const = 0;
1959+
virtual bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
1960+
IntrinsicInst *II) const = 0;
19541961
virtual bool enableSelectOptimize() = 0;
19551962
virtual bool shouldTreatInstructionLikeSelect(const Instruction *I) = 0;
19561963
virtual bool enableInterleavedAccessVectorization() = 0;
@@ -2510,6 +2517,12 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
25102517
bool IsZeroCmp) const override {
25112518
return Impl.enableMemCmpExpansion(OptSize, IsZeroCmp);
25122519
}
2520+
2521+
bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
2522+
IntrinsicInst *II) const override {
2523+
return Impl.getMemoryRefInfo(Interesting, II);
2524+
}
2525+
25132526
bool enableSelectOptimize() override {
25142527
return Impl.enableSelectOptimize();
25152528
}

‎llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,11 @@ class TargetTransformInfoImplBase {
396396
return {};
397397
}
398398

399+
bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
400+
IntrinsicInst *II) const {
401+
return false;
402+
}
403+
399404
bool enableSelectOptimize() const { return true; }
400405

401406
bool shouldTreatInstructionLikeSelect(const Instruction *I) {

‎llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

Lines changed: 1 addition & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
1515

1616
#include "llvm/Analysis/CFG.h"
17+
#include "llvm/Analysis/MemoryRefInfo.h"
1718
#include "llvm/Analysis/PostDominators.h"
1819
#include "llvm/IR/Dominators.h"
1920
#include "llvm/IR/Instruction.h"
@@ -22,37 +23,6 @@
2223

2324
namespace llvm {
2425

25-
class InterestingMemoryOperand {
26-
public:
27-
Use *PtrUse;
28-
bool IsWrite;
29-
Type *OpType;
30-
TypeSize TypeStoreSize = TypeSize::getFixed(0);
31-
MaybeAlign Alignment;
32-
// The mask Value, if we're looking at a masked load/store.
33-
Value *MaybeMask;
34-
// The EVL Value, if we're looking at a vp intrinsic.
35-
Value *MaybeEVL;
36-
// The Stride Value, if we're looking at a strided load/store.
37-
Value *MaybeStride;
38-
39-
InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
40-
class Type *OpType, MaybeAlign Alignment,
41-
Value *MaybeMask = nullptr,
42-
Value *MaybeEVL = nullptr,
43-
Value *MaybeStride = nullptr)
44-
: IsWrite(IsWrite), OpType(OpType), Alignment(Alignment),
45-
MaybeMask(MaybeMask), MaybeEVL(MaybeEVL), MaybeStride(MaybeStride) {
46-
const DataLayout &DL = I->getDataLayout();
47-
TypeStoreSize = DL.getTypeStoreSizeInBits(OpType);
48-
PtrUse = &I->getOperandUse(OperandNo);
49-
}
50-
51-
Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
52-
53-
Value *getPtr() { return PtrUse->get(); }
54-
};
55-
5626
// Get AddressSanitizer parameters.
5727
void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
5828
bool IsKasan, uint64_t *ShadowBase,

‎llvm/lib/Analysis/TargetTransformInfo.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,6 +626,11 @@ TargetTransformInfo::enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const {
626626
return TTIImpl->enableMemCmpExpansion(OptSize, IsZeroCmp);
627627
}
628628

629+
bool TargetTransformInfo::getMemoryRefInfo(
630+
SmallVectorImpl<MemoryRefInfo> &Interesting, IntrinsicInst *II) const {
631+
return TTIImpl->getMemoryRefInfo(Interesting, II);
632+
}
633+
629634
bool TargetTransformInfo::enableSelectOptimize() const {
630635
return TTIImpl->enableSelectOptimize();
631636
}

‎llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include "llvm/CodeGen/CostTable.h"
1515
#include "llvm/CodeGen/TargetLowering.h"
1616
#include "llvm/IR/Instructions.h"
17+
#include "llvm/IR/IntrinsicsRISCV.h"
1718
#include "llvm/IR/PatternMatch.h"
1819
#include <cmath>
1920
#include <optional>
@@ -36,6 +37,83 @@ static cl::opt<unsigned> SLPMaxVF(
3637
"exclusively by SLP vectorizer."),
3738
cl::Hidden);
3839

40+
bool RISCVTTIImpl::getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
41+
IntrinsicInst *II) const {
42+
const DataLayout &DL = getDataLayout();
43+
Intrinsic::ID IID = II->getIntrinsicID();
44+
LLVMContext &C = II->getContext();
45+
bool HasMask = false;
46+
47+
switch (IID) {
48+
case Intrinsic::riscv_vle_mask:
49+
case Intrinsic::riscv_vse_mask:
50+
HasMask = true;
51+
[[fallthrough]];
52+
case Intrinsic::riscv_vle:
53+
case Intrinsic::riscv_vse: {
54+
// Intrinsic interface:
55+
// riscv_vle(merge, ptr, vl)
56+
// riscv_vle_mask(merge, ptr, mask, vl, policy)
57+
// riscv_vse(val, ptr, vl)
58+
// riscv_vse_mask(val, ptr, mask, vl, policy)
59+
bool IsWrite = II->getType()->isVoidTy();
60+
Type *Ty = IsWrite ? II->getArgOperand(0)->getType() : II->getType();
61+
const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
62+
unsigned VLIndex = RVVIInfo->VLOperand;
63+
unsigned PtrOperandNo = VLIndex - 1 - HasMask;
64+
MaybeAlign Alignment =
65+
II->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
66+
Type *MaskType = Ty->getWithNewType(Type::getInt1Ty(C));
67+
Value *Mask = ConstantInt::getTrue(MaskType);
68+
if (HasMask)
69+
Mask = II->getArgOperand(VLIndex - 1);
70+
Value *EVL = II->getArgOperand(VLIndex);
71+
Interesting.emplace_back(II, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
72+
EVL);
73+
return true;
74+
}
75+
case Intrinsic::riscv_vlse_mask:
76+
case Intrinsic::riscv_vsse_mask:
77+
HasMask = true;
78+
[[fallthrough]];
79+
case Intrinsic::riscv_vlse:
80+
case Intrinsic::riscv_vsse: {
81+
// Intrinsic interface:
82+
// riscv_vlse(merge, ptr, stride, vl)
83+
// riscv_vlse_mask(merge, ptr, stride, mask, vl, policy)
84+
// riscv_vsse(val, ptr, stride, vl)
85+
// riscv_vsse_mask(val, ptr, stride, mask, vl, policy)
86+
bool IsWrite = II->getType()->isVoidTy();
87+
Type *Ty = IsWrite ? II->getArgOperand(0)->getType() : II->getType();
88+
const auto *RVVIInfo = RISCVVIntrinsicsTable::getRISCVVIntrinsicInfo(IID);
89+
unsigned VLIndex = RVVIInfo->VLOperand;
90+
unsigned PtrOperandNo = VLIndex - 2 - HasMask;
91+
MaybeAlign Alignment =
92+
II->getArgOperand(PtrOperandNo)->getPointerAlignment(DL);
93+
94+
Value *Stride = II->getArgOperand(PtrOperandNo + 1);
95+
// Use the pointer alignment as the element alignment if the stride is a
96+
// multiple of the pointer alignment. Otherwise, the element alignment
97+
// should be the greatest common divisor of pointer alignment and stride.
98+
// For simplicity, just consider unalignment for elements.
99+
unsigned PointerAlign = Alignment.valueOrOne().value();
100+
if (!isa<ConstantInt>(Stride) ||
101+
cast<ConstantInt>(Stride)->getZExtValue() % PointerAlign != 0)
102+
Alignment = Align(1);
103+
104+
Type *MaskType = Ty->getWithNewType(Type::getInt1Ty(C));
105+
Value *Mask = ConstantInt::getTrue(MaskType);
106+
if (HasMask)
107+
Mask = II->getArgOperand(VLIndex - 1);
108+
Value *EVL = II->getArgOperand(VLIndex);
109+
Interesting.emplace_back(II, PtrOperandNo, IsWrite, Ty, Alignment, Mask,
110+
EVL, Stride);
111+
return true;
112+
}
113+
}
114+
return false;
115+
}
116+
39117
InstructionCost
40118
RISCVTTIImpl::getRISCVInstructionCost(ArrayRef<unsigned> OpCodes, MVT VT,
41119
TTI::TargetCostKind CostKind) {

‎llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,9 @@ class RISCVTTIImpl : public BasicTTIImplBase<RISCVTTIImpl> {
6060
: BaseT(TM, F.getDataLayout()), ST(TM->getSubtargetImpl(F)),
6161
TLI(ST->getTargetLowering()) {}
6262

63+
bool getMemoryRefInfo(SmallVectorImpl<MemoryRefInfo> &Interesting,
64+
IntrinsicInst *II) const;
65+
6366
bool areInlineCompatible(const Function *Caller,
6467
const Function *Callee) const;
6568

‎llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "llvm/Analysis/MemoryBuiltins.h"
3030
#include "llvm/Analysis/StackSafetyAnalysis.h"
3131
#include "llvm/Analysis/TargetLibraryInfo.h"
32+
#include "llvm/Analysis/TargetTransformInfo.h"
3233
#include "llvm/Analysis/ValueTracking.h"
3334
#include "llvm/BinaryFormat/MachO.h"
3435
#include "llvm/Demangle/Demangle.h"
@@ -754,12 +755,13 @@ struct AddressSanitizer {
754755
bool isInterestingAlloca(const AllocaInst &AI);
755756

756757
bool ignoreAccess(Instruction *Inst, Value *Ptr);
757-
void getInterestingMemoryOperands(
758-
Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
758+
void getMemoryRefInfos(Instruction *I,
759+
SmallVectorImpl<MemoryRefInfo> &Interesting,
760+
const TargetTransformInfo *TTI);
759761

760-
void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
761-
InterestingMemoryOperand &O, bool UseCalls,
762-
const DataLayout &DL, RuntimeCallInserter &RTCI);
762+
void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis, MemoryRefInfo &O,
763+
bool UseCalls, const DataLayout &DL,
764+
RuntimeCallInserter &RTCI);
763765
void instrumentPointerComparisonOrSubtraction(Instruction *I,
764766
RuntimeCallInserter &RTCI);
765767
void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
@@ -795,7 +797,8 @@ struct AddressSanitizer {
795797
void instrumentMemIntrinsic(MemIntrinsic *MI, RuntimeCallInserter &RTCI);
796798
Value *memToShadow(Value *Shadow, IRBuilder<> &IRB);
797799
bool suppressInstrumentationSiteForDebug(int &Instrumented);
798-
bool instrumentFunction(Function &F, const TargetLibraryInfo *TLI);
800+
bool instrumentFunction(Function &F, const TargetLibraryInfo *TLI,
801+
const TargetTransformInfo *TTI);
799802
bool maybeInsertAsanInitAtFunctionEntry(Function &F);
800803
bool maybeInsertDynamicShadowAtFunctionEntry(Function &F);
801804
void markEscapedLocalAllocas(Function &F);
@@ -1264,7 +1267,8 @@ PreservedAnalyses AddressSanitizerPass::run(Module &M,
12641267
Options.MaxInlinePoisoningSize, Options.CompileKernel, Options.Recover,
12651268
Options.UseAfterScope, Options.UseAfterReturn);
12661269
const TargetLibraryInfo &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
1267-
Modified |= FunctionSanitizer.instrumentFunction(F, &TLI);
1270+
const TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
1271+
Modified |= FunctionSanitizer.instrumentFunction(F, &TLI, &TTI);
12681272
}
12691273
Modified |= ModuleSanitizer.instrumentModule(M);
12701274
if (!Modified)
@@ -1401,8 +1405,9 @@ bool AddressSanitizer::ignoreAccess(Instruction *Inst, Value *Ptr) {
14011405
return false;
14021406
}
14031407

1404-
void AddressSanitizer::getInterestingMemoryOperands(
1405-
Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting) {
1408+
void AddressSanitizer::getMemoryRefInfos(
1409+
Instruction *I, SmallVectorImpl<MemoryRefInfo> &Interesting,
1410+
const TargetTransformInfo *TTI) {
14061411
// Do not instrument the load fetching the dynamic shadow address.
14071412
if (LocalDynamicShadow == I)
14081413
return;
@@ -1520,6 +1525,9 @@ void AddressSanitizer::getInterestingMemoryOperands(
15201525
break;
15211526
}
15221527
default:
1528+
if (auto *II = dyn_cast<IntrinsicInst>(I))
1529+
if (TTI->getMemoryRefInfo(Interesting, II))
1530+
return;
15231531
for (unsigned ArgNo = 0; ArgNo < CI->arg_size(); ArgNo++) {
15241532
if (!ClInstrumentByval || !CI->isByValArgument(ArgNo) ||
15251533
ignoreAccess(I, CI->getArgOperand(ArgNo)))
@@ -1682,7 +1690,7 @@ void AddressSanitizer::instrumentMaskedLoadOrStore(
16821690
}
16831691

16841692
void AddressSanitizer::instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
1685-
InterestingMemoryOperand &O, bool UseCalls,
1693+
MemoryRefInfo &O, bool UseCalls,
16861694
const DataLayout &DL,
16871695
RuntimeCallInserter &RTCI) {
16881696
Value *Addr = O.getPtr();
@@ -2940,7 +2948,8 @@ bool AddressSanitizer::suppressInstrumentationSiteForDebug(int &Instrumented) {
29402948
}
29412949

29422950
bool AddressSanitizer::instrumentFunction(Function &F,
2943-
const TargetLibraryInfo *TLI) {
2951+
const TargetLibraryInfo *TLI,
2952+
const TargetTransformInfo *TTI) {
29442953
if (F.empty())
29452954
return false;
29462955
if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage) return false;
@@ -2980,7 +2989,7 @@ bool AddressSanitizer::instrumentFunction(Function &F,
29802989
// We want to instrument every address only once per basic block (unless there
29812990
// are calls between uses).
29822991
SmallPtrSet<Value *, 16> TempsToInstrument;
2983-
SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
2992+
SmallVector<MemoryRefInfo, 16> OperandsToInstrument;
29842993
SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
29852994
SmallVector<Instruction *, 8> NoReturnCalls;
29862995
SmallVector<BasicBlock *, 16> AllBlocks;
@@ -2996,8 +3005,8 @@ bool AddressSanitizer::instrumentFunction(Function &F,
29963005
// Skip instructions inserted by another instrumentation.
29973006
if (Inst.hasMetadata(LLVMContext::MD_nosanitize))
29983007
continue;
2999-
SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
3000-
getInterestingMemoryOperands(&Inst, InterestingOperands);
3008+
SmallVector<MemoryRefInfo, 1> InterestingOperands;
3009+
getMemoryRefInfos(&Inst, InterestingOperands, TTI);
30013010

30023011
if (!InterestingOperands.empty()) {
30033012
for (auto &Operand : InterestingOperands) {

0 commit comments

Comments
 (0)