478 lines · cpp
1//===- DXILDataScalarization.cpp - Perform DXIL Data Legalization ---------===//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 "DXILDataScalarization.h"10#include "DirectX.h"11#include "llvm/ADT/PostOrderIterator.h"12#include "llvm/ADT/STLExtras.h"13#include "llvm/IR/DerivedTypes.h"14#include "llvm/IR/GlobalVariable.h"15#include "llvm/IR/IRBuilder.h"16#include "llvm/IR/InstVisitor.h"17#include "llvm/IR/Instructions.h"18#include "llvm/IR/Module.h"19#include "llvm/IR/Operator.h"20#include "llvm/IR/PassManager.h"21#include "llvm/IR/ReplaceConstant.h"22#include "llvm/IR/Type.h"23#include "llvm/Support/Casting.h"24#include "llvm/Transforms/Utils/Cloning.h"25#include "llvm/Transforms/Utils/Local.h"26 27#define DEBUG_TYPE "dxil-data-scalarization"28static const int MaxVecSize = 4;29 30using namespace llvm;31 32class DXILDataScalarizationLegacy : public ModulePass {33 34public:35 bool runOnModule(Module &M) override;36 DXILDataScalarizationLegacy() : ModulePass(ID) {}37 38 static char ID; // Pass identification.39};40 41static bool findAndReplaceVectors(Module &M);42 43class DataScalarizerVisitor : public InstVisitor<DataScalarizerVisitor, bool> {44public:45 DataScalarizerVisitor() : GlobalMap() {}46 bool visit(Function &F);47 // InstVisitor methods. They return true if the instruction was scalarized,48 // false if nothing changed.49 bool visitAllocaInst(AllocaInst &AI);50 bool visitInstruction(Instruction &I) { return false; }51 bool visitSelectInst(SelectInst &SI) { return false; }52 bool visitICmpInst(ICmpInst &ICI) { return false; }53 bool visitFCmpInst(FCmpInst &FCI) { return false; }54 bool visitUnaryOperator(UnaryOperator &UO) { return false; }55 bool visitBinaryOperator(BinaryOperator &BO) { return false; }56 bool visitGetElementPtrInst(GetElementPtrInst &GEPI);57 bool visitCastInst(CastInst &CI) { return false; }58 bool visitBitCastInst(BitCastInst &BCI) { return false; }59 bool visitInsertElementInst(InsertElementInst &IEI);60 bool visitExtractElementInst(ExtractElementInst &EEI);61 bool visitShuffleVectorInst(ShuffleVectorInst &SVI) { return false; }62 bool visitPHINode(PHINode &PHI) { return false; }63 bool visitLoadInst(LoadInst &LI);64 bool visitStoreInst(StoreInst &SI);65 bool visitCallInst(CallInst &ICI) { return false; }66 bool visitFreezeInst(FreezeInst &FI) { return false; }67 friend bool findAndReplaceVectors(llvm::Module &M);68 69private:70 typedef std::pair<AllocaInst *, SmallVector<Value *, 4>> AllocaAndGEPs;71 typedef SmallDenseMap<Value *, AllocaAndGEPs>72 VectorToArrayMap; // A map from a vector-typed Value to its corresponding73 // AllocaInst and GEPs to each element of an array74 VectorToArrayMap VectorAllocaMap;75 AllocaAndGEPs createArrayFromVector(IRBuilder<> &Builder, Value *Vec,76 const Twine &Name);77 bool replaceDynamicInsertElementInst(InsertElementInst &IEI);78 bool replaceDynamicExtractElementInst(ExtractElementInst &EEI);79 80 GlobalVariable *lookupReplacementGlobal(Value *CurrOperand);81 DenseMap<GlobalVariable *, GlobalVariable *> GlobalMap;82};83 84bool DataScalarizerVisitor::visit(Function &F) {85 bool MadeChange = false;86 ReversePostOrderTraversal<Function *> RPOT(&F);87 for (BasicBlock *BB : make_early_inc_range(RPOT)) {88 for (Instruction &I : make_early_inc_range(*BB))89 MadeChange |= InstVisitor::visit(I);90 }91 VectorAllocaMap.clear();92 return MadeChange;93}94 95GlobalVariable *96DataScalarizerVisitor::lookupReplacementGlobal(Value *CurrOperand) {97 if (GlobalVariable *OldGlobal = dyn_cast<GlobalVariable>(CurrOperand)) {98 auto It = GlobalMap.find(OldGlobal);99 if (It != GlobalMap.end()) {100 return It->second; // Found, return the new global101 }102 }103 return nullptr; // Not found104}105 106// Helper function to check if a type is a vector or an array of vectors107static bool isVectorOrArrayOfVectors(Type *T) {108 if (isa<VectorType>(T))109 return true;110 if (ArrayType *ArrayTy = dyn_cast<ArrayType>(T))111 return isVectorOrArrayOfVectors(ArrayTy->getElementType());112 return false;113}114 115// Recursively creates an array-like version of a given vector type.116static Type *equivalentArrayTypeFromVector(Type *T) {117 if (auto *VecTy = dyn_cast<VectorType>(T))118 return ArrayType::get(VecTy->getElementType(),119 dyn_cast<FixedVectorType>(VecTy)->getNumElements());120 if (auto *ArrayTy = dyn_cast<ArrayType>(T)) {121 Type *NewElementType =122 equivalentArrayTypeFromVector(ArrayTy->getElementType());123 return ArrayType::get(NewElementType, ArrayTy->getNumElements());124 }125 // If it's not a vector or array, return the original type.126 return T;127}128 129bool DataScalarizerVisitor::visitAllocaInst(AllocaInst &AI) {130 Type *AllocatedType = AI.getAllocatedType();131 if (!isVectorOrArrayOfVectors(AllocatedType))132 return false;133 134 IRBuilder<> Builder(&AI);135 Type *NewType = equivalentArrayTypeFromVector(AllocatedType);136 AllocaInst *ArrAlloca =137 Builder.CreateAlloca(NewType, nullptr, AI.getName() + ".scalarized");138 ArrAlloca->setAlignment(AI.getAlign());139 AI.replaceAllUsesWith(ArrAlloca);140 AI.eraseFromParent();141 return true;142}143 144bool DataScalarizerVisitor::visitLoadInst(LoadInst &LI) {145 Value *PtrOperand = LI.getPointerOperand();146 ConstantExpr *CE = dyn_cast<ConstantExpr>(PtrOperand);147 if (CE && CE->getOpcode() == Instruction::GetElementPtr) {148 GetElementPtrInst *OldGEP = cast<GetElementPtrInst>(CE->getAsInstruction());149 OldGEP->insertBefore(LI.getIterator());150 IRBuilder<> Builder(&LI);151 LoadInst *NewLoad = Builder.CreateLoad(LI.getType(), OldGEP, LI.getName());152 NewLoad->setAlignment(LI.getAlign());153 LI.replaceAllUsesWith(NewLoad);154 LI.eraseFromParent();155 visitGetElementPtrInst(*OldGEP);156 return true;157 }158 if (GlobalVariable *NewGlobal = lookupReplacementGlobal(PtrOperand))159 LI.setOperand(LI.getPointerOperandIndex(), NewGlobal);160 return false;161}162 163bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) {164 165 Value *PtrOperand = SI.getPointerOperand();166 ConstantExpr *CE = dyn_cast<ConstantExpr>(PtrOperand);167 if (CE && CE->getOpcode() == Instruction::GetElementPtr) {168 GetElementPtrInst *OldGEP = cast<GetElementPtrInst>(CE->getAsInstruction());169 OldGEP->insertBefore(SI.getIterator());170 IRBuilder<> Builder(&SI);171 StoreInst *NewStore = Builder.CreateStore(SI.getValueOperand(), OldGEP);172 NewStore->setAlignment(SI.getAlign());173 SI.replaceAllUsesWith(NewStore);174 SI.eraseFromParent();175 visitGetElementPtrInst(*OldGEP);176 return true;177 }178 if (GlobalVariable *NewGlobal = lookupReplacementGlobal(PtrOperand))179 SI.setOperand(SI.getPointerOperandIndex(), NewGlobal);180 181 return false;182}183 184DataScalarizerVisitor::AllocaAndGEPs185DataScalarizerVisitor::createArrayFromVector(IRBuilder<> &Builder, Value *Vec,186 const Twine &Name = "") {187 // If there is already an alloca for this vector, return it188 if (VectorAllocaMap.contains(Vec))189 return VectorAllocaMap[Vec];190 191 auto InsertPoint = Builder.GetInsertPoint();192 193 // Allocate the array to hold the vector elements194 Builder.SetInsertPointPastAllocas(Builder.GetInsertBlock()->getParent());195 Type *ArrTy = equivalentArrayTypeFromVector(Vec->getType());196 AllocaInst *ArrAlloca =197 Builder.CreateAlloca(ArrTy, nullptr, Name + ".alloca");198 const uint64_t ArrNumElems = ArrTy->getArrayNumElements();199 200 // Create loads and stores to populate the array immediately after the201 // original vector's defining instruction if available, else immediately after202 // the alloca203 if (auto *Instr = dyn_cast<Instruction>(Vec))204 Builder.SetInsertPoint(Instr->getNextNode());205 SmallVector<Value *, 4> GEPs(ArrNumElems);206 for (unsigned I = 0; I < ArrNumElems; ++I) {207 Value *EE = Builder.CreateExtractElement(Vec, I, Name + ".extract");208 GEPs[I] = Builder.CreateInBoundsGEP(209 ArrTy, ArrAlloca, {Builder.getInt32(0), Builder.getInt32(I)},210 Name + ".index");211 Builder.CreateStore(EE, GEPs[I]);212 }213 214 VectorAllocaMap.insert({Vec, {ArrAlloca, GEPs}});215 Builder.SetInsertPoint(InsertPoint);216 return {ArrAlloca, GEPs};217}218 219/// Returns a pair of Value* with the first being a GEP into ArrAlloca using220/// indices {0, Index}, and the second Value* being a Load of the GEP221static std::pair<Value *, Value *>222dynamicallyLoadArray(IRBuilder<> &Builder, AllocaInst *ArrAlloca, Value *Index,223 const Twine &Name = "") {224 Type *ArrTy = ArrAlloca->getAllocatedType();225 Value *GEP = Builder.CreateInBoundsGEP(226 ArrTy, ArrAlloca, {Builder.getInt32(0), Index}, Name + ".index");227 Value *Load =228 Builder.CreateLoad(ArrTy->getArrayElementType(), GEP, Name + ".load");229 return std::make_pair(GEP, Load);230}231 232bool DataScalarizerVisitor::replaceDynamicInsertElementInst(233 InsertElementInst &IEI) {234 IRBuilder<> Builder(&IEI);235 236 Value *Vec = IEI.getOperand(0);237 Value *Val = IEI.getOperand(1);238 Value *Index = IEI.getOperand(2);239 240 AllocaAndGEPs ArrAllocaAndGEPs =241 createArrayFromVector(Builder, Vec, IEI.getName());242 AllocaInst *ArrAlloca = ArrAllocaAndGEPs.first;243 Type *ArrTy = ArrAlloca->getAllocatedType();244 SmallVector<Value *, 4> &ArrGEPs = ArrAllocaAndGEPs.second;245 246 auto GEPAndLoad =247 dynamicallyLoadArray(Builder, ArrAlloca, Index, IEI.getName());248 Value *GEP = GEPAndLoad.first;249 Value *Load = GEPAndLoad.second;250 251 Builder.CreateStore(Val, GEP);252 Value *NewIEI = PoisonValue::get(Vec->getType());253 for (unsigned I = 0; I < ArrTy->getArrayNumElements(); ++I) {254 Value *Load = Builder.CreateLoad(ArrTy->getArrayElementType(), ArrGEPs[I],255 IEI.getName() + ".load");256 NewIEI = Builder.CreateInsertElement(NewIEI, Load, Builder.getInt32(I),257 IEI.getName() + ".insert");258 }259 260 // Store back the original value so the Alloca can be reused for subsequent261 // insertelement instructions on the same vector262 Builder.CreateStore(Load, GEP);263 264 IEI.replaceAllUsesWith(NewIEI);265 IEI.eraseFromParent();266 return true;267}268 269bool DataScalarizerVisitor::visitInsertElementInst(InsertElementInst &IEI) {270 // If the index is a constant then we don't need to scalarize it271 Value *Index = IEI.getOperand(2);272 if (isa<ConstantInt>(Index))273 return false;274 return replaceDynamicInsertElementInst(IEI);275}276 277bool DataScalarizerVisitor::replaceDynamicExtractElementInst(278 ExtractElementInst &EEI) {279 IRBuilder<> Builder(&EEI);280 281 AllocaAndGEPs ArrAllocaAndGEPs =282 createArrayFromVector(Builder, EEI.getVectorOperand(), EEI.getName());283 AllocaInst *ArrAlloca = ArrAllocaAndGEPs.first;284 285 auto GEPAndLoad = dynamicallyLoadArray(Builder, ArrAlloca,286 EEI.getIndexOperand(), EEI.getName());287 Value *Load = GEPAndLoad.second;288 289 EEI.replaceAllUsesWith(Load);290 EEI.eraseFromParent();291 return true;292}293 294bool DataScalarizerVisitor::visitExtractElementInst(ExtractElementInst &EEI) {295 // If the index is a constant then we don't need to scalarize it296 Value *Index = EEI.getIndexOperand();297 if (isa<ConstantInt>(Index))298 return false;299 return replaceDynamicExtractElementInst(EEI);300}301 302bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {303 GEPOperator *GOp = cast<GEPOperator>(&GEPI);304 Value *PtrOperand = GOp->getPointerOperand();305 Type *GEPType = GOp->getSourceElementType();306 307 // Replace a GEP ConstantExpr pointer operand with a GEP instruction so that308 // it can be visited309 if (auto *PtrOpGEPCE = dyn_cast<ConstantExpr>(PtrOperand);310 PtrOpGEPCE && PtrOpGEPCE->getOpcode() == Instruction::GetElementPtr) {311 GetElementPtrInst *OldGEPI =312 cast<GetElementPtrInst>(PtrOpGEPCE->getAsInstruction());313 OldGEPI->insertBefore(GEPI.getIterator());314 315 IRBuilder<> Builder(&GEPI);316 SmallVector<Value *> Indices(GEPI.indices());317 Value *NewGEP =318 Builder.CreateGEP(GEPI.getSourceElementType(), OldGEPI, Indices,319 GEPI.getName(), GEPI.getNoWrapFlags());320 assert(isa<GetElementPtrInst>(NewGEP) &&321 "Expected newly-created GEP to be an instruction");322 GetElementPtrInst *NewGEPI = cast<GetElementPtrInst>(NewGEP);323 324 GEPI.replaceAllUsesWith(NewGEPI);325 GEPI.eraseFromParent();326 visitGetElementPtrInst(*OldGEPI);327 visitGetElementPtrInst(*NewGEPI);328 return true;329 }330 331 Type *NewGEPType = equivalentArrayTypeFromVector(GEPType);332 Value *NewPtrOperand = PtrOperand;333 if (GlobalVariable *NewGlobal = lookupReplacementGlobal(PtrOperand))334 NewPtrOperand = NewGlobal;335 336 bool NeedsTransform = NewPtrOperand != PtrOperand || NewGEPType != GEPType;337 if (!NeedsTransform)338 return false;339 340 IRBuilder<> Builder(&GEPI);341 SmallVector<Value *, MaxVecSize> Indices(GOp->idx_begin(), GOp->idx_end());342 Value *NewGEP = Builder.CreateGEP(NewGEPType, NewPtrOperand, Indices,343 GOp->getName(), GOp->getNoWrapFlags());344 345 GOp->replaceAllUsesWith(NewGEP);346 347 if (auto *OldGEPI = dyn_cast<GetElementPtrInst>(GOp))348 OldGEPI->eraseFromParent();349 350 return true;351}352 353static Constant *transformInitializer(Constant *Init, Type *OrigType,354 Type *NewType, LLVMContext &Ctx) {355 // Handle ConstantAggregateZero (zero-initialized constants)356 if (isa<ConstantAggregateZero>(Init)) {357 return ConstantAggregateZero::get(NewType);358 }359 360 // Handle UndefValue (undefined constants)361 if (isa<UndefValue>(Init)) {362 return UndefValue::get(NewType);363 }364 365 // Handle vector to array transformation366 if (isa<VectorType>(OrigType) && isa<ArrayType>(NewType)) {367 // Convert vector initializer to array initializer368 SmallVector<Constant *, MaxVecSize> ArrayElements;369 if (ConstantVector *ConstVecInit = dyn_cast<ConstantVector>(Init)) {370 for (unsigned I = 0; I < ConstVecInit->getNumOperands(); ++I)371 ArrayElements.push_back(ConstVecInit->getOperand(I));372 } else if (ConstantDataVector *ConstDataVecInit =373 llvm::dyn_cast<llvm::ConstantDataVector>(Init)) {374 for (unsigned I = 0; I < ConstDataVecInit->getNumElements(); ++I)375 ArrayElements.push_back(ConstDataVecInit->getElementAsConstant(I));376 } else {377 assert(false && "Expected a ConstantVector or ConstantDataVector for "378 "vector initializer!");379 }380 381 return ConstantArray::get(cast<ArrayType>(NewType), ArrayElements);382 }383 384 // Handle array of vectors transformation385 if (auto *ArrayTy = dyn_cast<ArrayType>(OrigType)) {386 auto *ArrayInit = dyn_cast<ConstantArray>(Init);387 assert(ArrayInit && "Expected a ConstantArray for array initializer!");388 389 SmallVector<Constant *, MaxVecSize> NewArrayElements;390 for (unsigned I = 0; I < ArrayTy->getNumElements(); ++I) {391 // Recursively transform array elements392 Constant *NewElemInit = transformInitializer(393 ArrayInit->getOperand(I), ArrayTy->getElementType(),394 cast<ArrayType>(NewType)->getElementType(), Ctx);395 NewArrayElements.push_back(NewElemInit);396 }397 398 return ConstantArray::get(cast<ArrayType>(NewType), NewArrayElements);399 }400 401 // If not a vector or array, return the original initializer402 return Init;403}404 405static bool findAndReplaceVectors(Module &M) {406 bool MadeChange = false;407 LLVMContext &Ctx = M.getContext();408 IRBuilder<> Builder(Ctx);409 DataScalarizerVisitor Impl;410 for (GlobalVariable &G : M.globals()) {411 Type *OrigType = G.getValueType();412 413 Type *NewType = equivalentArrayTypeFromVector(OrigType);414 if (OrigType != NewType) {415 // Create a new global variable with the updated type416 // Note: Initializer is set via transformInitializer417 GlobalVariable *NewGlobal = new GlobalVariable(418 M, NewType, G.isConstant(), G.getLinkage(),419 /*Initializer=*/nullptr, G.getName() + ".scalarized", &G,420 G.getThreadLocalMode(), G.getAddressSpace(),421 G.isExternallyInitialized());422 423 // Copy relevant attributes424 NewGlobal->setUnnamedAddr(G.getUnnamedAddr());425 if (G.getAlignment() > 0) {426 NewGlobal->setAlignment(G.getAlign());427 }428 429 if (G.hasInitializer()) {430 Constant *Init = G.getInitializer();431 Constant *NewInit = transformInitializer(Init, OrigType, NewType, Ctx);432 NewGlobal->setInitializer(NewInit);433 }434 435 // Note: we want to do G.replaceAllUsesWith(NewGlobal);, but it assumes436 // type equality. Instead we will use the visitor pattern.437 Impl.GlobalMap[&G] = NewGlobal;438 }439 }440 441 for (auto &F : make_early_inc_range(M.functions())) {442 if (F.isDeclaration())443 continue;444 MadeChange |= Impl.visit(F);445 }446 447 // Remove the old globals after the iteration448 for (auto &[Old, New] : Impl.GlobalMap) {449 Old->eraseFromParent();450 MadeChange = true;451 }452 return MadeChange;453}454 455PreservedAnalyses DXILDataScalarization::run(Module &M,456 ModuleAnalysisManager &) {457 bool MadeChanges = findAndReplaceVectors(M);458 if (!MadeChanges)459 return PreservedAnalyses::all();460 PreservedAnalyses PA;461 return PA;462}463 464bool DXILDataScalarizationLegacy::runOnModule(Module &M) {465 return findAndReplaceVectors(M);466}467 468char DXILDataScalarizationLegacy::ID = 0;469 470INITIALIZE_PASS_BEGIN(DXILDataScalarizationLegacy, DEBUG_TYPE,471 "DXIL Data Scalarization", false, false)472INITIALIZE_PASS_END(DXILDataScalarizationLegacy, DEBUG_TYPE,473 "DXIL Data Scalarization", false, false)474 475ModulePass *llvm::createDXILDataScalarizationLegacyPass() {476 return new DXILDataScalarizationLegacy();477}478