213 lines · cpp
1//===- SeedCollector.cpp -------------------------------------------------===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h"10#include "llvm/Analysis/LoopAccessAnalysis.h"11#include "llvm/Analysis/ValueTracking.h"12#include "llvm/IR/Type.h"13#include "llvm/SandboxIR/Instruction.h"14#include "llvm/SandboxIR/Utils.h"15#include "llvm/Support/Compiler.h"16#include "llvm/Support/Debug.h"17 18using namespace llvm;19namespace llvm::sandboxir {20 21static cl::opt<unsigned> SeedBundleSizeLimit(22 "sbvec-seed-bundle-size-limit", cl::init(32), cl::Hidden,23 cl::desc("Limit the size of the seed bundle to cap compilation time."));24 25static cl::opt<unsigned> SeedGroupsLimit(26 "sbvec-seed-groups-limit", cl::init(256), cl::Hidden,27 cl::desc("Limit the number of collected seeds groups in a BB to "28 "cap compilation time."));29 30ArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,31 unsigned MaxVecRegBits,32 bool ForcePowerOf2) {33 // Use uint32_t here for compatibility with IsPowerOf2_3234 35 // BitCount tracks the size of the working slice. From that we can tell36 // when the working slice's size is a power-of-two and when it exceeds37 // the legal size in MaxVecBits.38 uint32_t BitCount = 0;39 uint32_t NumElements = 0;40 // Tracks the most recent slice where NumElements gave a power-of-2 BitCount41 uint32_t NumElementsPowerOfTwo = 0;42 uint32_t BitCountPowerOfTwo = 0;43 // Can't start a slice with a used instruction.44 assert(!isUsed(StartIdx) && "Expected unused at StartIdx");45 for (Instruction *S : drop_begin(Seeds, StartIdx)) {46 // Stop if this instruction is used. This needs to be done before47 // getNumBits() because a "used" instruction may have been erased.48 if (isUsed(StartIdx + NumElements))49 break;50 uint32_t InstBits = Utils::getNumBits(S);51 // Stop if adding it puts the slice over the limit.52 if (BitCount + InstBits > MaxVecRegBits)53 break;54 NumElements++;55 BitCount += InstBits;56 if (ForcePowerOf2 && isPowerOf2_32(BitCount)) {57 NumElementsPowerOfTwo = NumElements;58 BitCountPowerOfTwo = BitCount;59 }60 }61 if (ForcePowerOf2) {62 NumElements = NumElementsPowerOfTwo;63 BitCount = BitCountPowerOfTwo;64 }65 66 // Return any non-empty slice67 if (NumElements > 1) {68 assert((!ForcePowerOf2 || isPowerOf2_32(BitCount)) &&69 "Must be a power of two");70 return ArrayRef<Instruction *>(&Seeds[StartIdx], NumElements);71 }72 return {};73}74 75template <typename LoadOrStoreT>76SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI,77 bool AllowDiffTypes) const {78 assert((isa<LoadInst>(LSI) || isa<StoreInst>(LSI)) &&79 "Expected Load or Store!");80 Value *Ptr = Utils::getMemInstructionBase(LSI);81 Instruction::Opcode Op = LSI->getOpcode();82 Type *Ty;83 if (AllowDiffTypes) {84 Ty = nullptr;85 } else {86 Ty = Utils::getExpectedType(LSI);87 if (auto *VTy = dyn_cast<VectorType>(Ty))88 Ty = VTy->getElementType();89 }90 return {Ptr, Ty, Op};91}92 93// Explicit instantiations94template SeedContainer::KeyT95SeedContainer::getKey<LoadInst>(LoadInst *LSI, bool AllowDiffTypes) const;96template SeedContainer::KeyT97SeedContainer::getKey<StoreInst>(StoreInst *LSI, bool AllowDiffTypes) const;98 99bool SeedContainer::erase(Instruction *I) {100 assert((isa<LoadInst>(I) || isa<StoreInst>(I)) && "Expected Load or Store!");101 auto It = SeedLookupMap.find(I);102 if (It == SeedLookupMap.end())103 return false;104 SeedBundle *Bndl = It->second;105 Bndl->setUsed(I);106 return true;107}108 109template <typename LoadOrStoreT>110void SeedContainer::insert(LoadOrStoreT *LSI, bool AllowDiffTypes) {111 // Find the bundle containing seeds for this symbol and type-of-access.112 auto &BundleVec = Bundles[getKey(LSI, AllowDiffTypes)];113 // Fill this vector of bundles front to back so that only the last bundle in114 // the vector may have available space. This avoids iteration to find one with115 // space.116 if (BundleVec.empty() || BundleVec.back()->size() == SeedBundleSizeLimit)117 BundleVec.emplace_back(std::make_unique<MemSeedBundle<LoadOrStoreT>>(LSI));118 else119 BundleVec.back()->insert(LSI, SE);120 121 SeedLookupMap[LSI] = BundleVec.back().get();122}123 124// Explicit instantiations125template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *,126 bool);127template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<StoreInst>(StoreInst *,128 bool);129 130#ifndef NDEBUG131void SeedContainer::print(raw_ostream &OS) const {132 for (const auto &Pair : Bundles) {133 auto [I, Ty, Opc] = Pair.first;134 const auto &SeedsVec = Pair.second;135 std::string RefType = dyn_cast<LoadInst>(I) ? "Load"136 : dyn_cast<StoreInst>(I) ? "Store"137 : "Other";138 OS << "[Inst=" << *I << " Ty=" << Ty << " " << RefType << "]\n";139 for (const auto &SeedPtr : SeedsVec) {140 SeedPtr->dump(OS);141 OS << "\n";142 }143 }144 OS << "\n";145}146 147LLVM_DUMP_METHOD void SeedContainer::dump() const { print(dbgs()); }148#endif // NDEBUG149 150template <typename LoadOrStoreT> static bool isValidMemSeed(LoadOrStoreT *LSI) {151 if (!LSI->isSimple())152 return false;153 auto *Ty = Utils::getExpectedType(LSI);154 // Omit types that are architecturally unvectorizable155 if (Ty->isX86_FP80Ty() || Ty->isPPC_FP128Ty())156 return false;157 // Omit vector types without compile-time-known lane counts158 if (isa<ScalableVectorType>(Ty))159 return false;160 if (auto *VTy = dyn_cast<FixedVectorType>(Ty))161 return VectorType::isValidElementType(VTy->getElementType());162 return VectorType::isValidElementType(Ty);163}164 165template bool isValidMemSeed<LoadInst>(LoadInst *LSI);166template bool isValidMemSeed<StoreInst>(StoreInst *LSI);167 168SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,169 bool CollectStores, bool CollectLoads,170 bool AllowDiffTypes)171 : StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) {172 173 if (!CollectStores && !CollectLoads)174 return;175 176 EraseCallbackID = Ctx.registerEraseInstrCallback([this](Instruction *I) {177 if (auto SI = dyn_cast<StoreInst>(I))178 StoreSeeds.erase(SI);179 else if (auto LI = dyn_cast<LoadInst>(I))180 LoadSeeds.erase(LI);181 });182 183 // Actually collect the seeds.184 for (auto &I : *BB) {185 if (StoreInst *SI = dyn_cast<StoreInst>(&I))186 if (CollectStores && isValidMemSeed(SI))187 StoreSeeds.insert(SI, AllowDiffTypes);188 if (LoadInst *LI = dyn_cast<LoadInst>(&I))189 if (CollectLoads && isValidMemSeed(LI))190 LoadSeeds.insert(LI, AllowDiffTypes);191 // Cap compilation time.192 if (totalNumSeedGroups() > SeedGroupsLimit)193 break;194 }195}196 197SeedCollector::~SeedCollector() {198 Ctx.unregisterEraseInstrCallback(EraseCallbackID);199}200 201#ifndef NDEBUG202void SeedCollector::print(raw_ostream &OS) const {203 OS << "=== StoreSeeds ===\n";204 StoreSeeds.print(OS);205 OS << "=== LoadSeeds ===\n";206 LoadSeeds.print(OS);207}208 209void SeedCollector::dump() const { print(dbgs()); }210#endif211 212} // namespace llvm::sandboxir213