| 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-exception |
| 6 | // |
| 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" |
| 28 | static const int MaxVecSize = 4; |
| 29 | |
| 30 | using namespace llvm; |
| 31 | |
| 32 | class DXILDataScalarizationLegacy : public ModulePass { |
| 33 | |
| 34 | public: |
| 35 | bool runOnModule(Module &M) override; |
| 36 | DXILDataScalarizationLegacy() : ModulePass(ID) {} |
| 37 | |
| 38 | static char ID; // Pass identification. |
| 39 | }; |
| 40 | |
| 41 | static bool findAndReplaceVectors(Module &M); |
| 42 | |
| 43 | class DataScalarizerVisitor : public InstVisitor<DataScalarizerVisitor, bool> { |
| 44 | public: |
| 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 | |
| 69 | private: |
| 70 | typedef std::tuple<AllocaInst *, Type *, SmallVector<Value *, 4>> |
| 71 | AllocaAndGEPs; |
| 72 | typedef SmallDenseMap<Value *, AllocaAndGEPs> |
| 73 | VectorToArrayMap; // A map from a vector-typed Value to its corresponding |
| 74 | // AllocaInst and GEPs to each element of an array |
| 75 | VectorToArrayMap VectorAllocaMap; |
| 76 | AllocaAndGEPs createArrayFromVector(IRBuilder<> &Builder, Value *Vec, |
| 77 | const Twine &Name); |
| 78 | bool replaceDynamicInsertElementInst(InsertElementInst &IEI); |
| 79 | bool replaceDynamicExtractElementInst(ExtractElementInst &EEI); |
| 80 | |
| 81 | GlobalVariable *lookupReplacementGlobal(Value *CurrOperand); |
| 82 | DenseMap<GlobalVariable *, GlobalVariable *> GlobalMap; |
| 83 | }; |
| 84 | |
| 85 | bool DataScalarizerVisitor::visit(Function &F) { |
| 86 | bool MadeChange = false; |
| 87 | ReversePostOrderTraversal<Function *> RPOT(&F); |
| 88 | for (BasicBlock *BB : make_early_inc_range(Range&: RPOT)) { |
| 89 | for (Instruction &I : make_early_inc_range(Range&: *BB)) |
| 90 | MadeChange |= InstVisitor::visit(I); |
| 91 | } |
| 92 | VectorAllocaMap.clear(); |
| 93 | return MadeChange; |
| 94 | } |
| 95 | |
| 96 | GlobalVariable * |
| 97 | DataScalarizerVisitor::lookupReplacementGlobal(Value *CurrOperand) { |
| 98 | if (GlobalVariable *OldGlobal = dyn_cast<GlobalVariable>(Val: CurrOperand)) { |
| 99 | auto It = GlobalMap.find(Val: OldGlobal); |
| 100 | if (It != GlobalMap.end()) { |
| 101 | return It->second; // Found, return the new global |
| 102 | } |
| 103 | } |
| 104 | return nullptr; // Not found |
| 105 | } |
| 106 | |
| 107 | // Helper function to check if a type is a vector or an array of vectors |
| 108 | static bool isVectorOrArrayOfVectors(Type *T) { |
| 109 | if (isa<VectorType>(Val: T)) |
| 110 | return true; |
| 111 | if (ArrayType *ArrayTy = dyn_cast<ArrayType>(Val: T)) |
| 112 | return isVectorOrArrayOfVectors(T: ArrayTy->getElementType()); |
| 113 | return false; |
| 114 | } |
| 115 | |
| 116 | // Recursively creates an array-like version of a given vector type. |
| 117 | static Type *equivalentArrayTypeFromVector(Type *T) { |
| 118 | if (auto *VecTy = dyn_cast<VectorType>(Val: T)) |
| 119 | return ArrayType::get(ElementType: VecTy->getElementType(), |
| 120 | NumElements: dyn_cast<FixedVectorType>(Val: VecTy)->getNumElements()); |
| 121 | if (auto *ArrayTy = dyn_cast<ArrayType>(Val: T)) { |
| 122 | Type *NewElementType = |
| 123 | equivalentArrayTypeFromVector(T: ArrayTy->getElementType()); |
| 124 | return ArrayType::get(ElementType: NewElementType, NumElements: ArrayTy->getNumElements()); |
| 125 | } |
| 126 | // If it's not a vector or array, return the original type. |
| 127 | return T; |
| 128 | } |
| 129 | |
| 130 | bool DataScalarizerVisitor::visitAllocaInst(AllocaInst &AI) { |
| 131 | Type *AllocatedType = AI.getAllocatedType(); |
| 132 | if (!isVectorOrArrayOfVectors(T: AllocatedType)) |
| 133 | return false; |
| 134 | |
| 135 | IRBuilder<> Builder(&AI); |
| 136 | Type *NewType = equivalentArrayTypeFromVector(T: AllocatedType); |
| 137 | AllocaInst *ArrAlloca = |
| 138 | Builder.CreateAlloca(Ty: NewType, ArraySize: nullptr, Name: AI.getName() + ".scalarized" ); |
| 139 | ArrAlloca->setAlignment(AI.getAlign()); |
| 140 | AI.replaceAllUsesWith(V: ArrAlloca); |
| 141 | AI.eraseFromParent(); |
| 142 | return true; |
| 143 | } |
| 144 | |
| 145 | bool DataScalarizerVisitor::visitLoadInst(LoadInst &LI) { |
| 146 | Value *PtrOperand = LI.getPointerOperand(); |
| 147 | ConstantExpr *CE = dyn_cast<ConstantExpr>(Val: PtrOperand); |
| 148 | if (CE && CE->getOpcode() == Instruction::GetElementPtr) { |
| 149 | GetElementPtrInst *OldGEP = cast<GetElementPtrInst>(Val: CE->getAsInstruction()); |
| 150 | OldGEP->insertBefore(InsertPos: LI.getIterator()); |
| 151 | IRBuilder<> Builder(&LI); |
| 152 | LoadInst *NewLoad = Builder.CreateLoad(Ty: LI.getType(), Ptr: OldGEP, Name: LI.getName()); |
| 153 | NewLoad->setAlignment(LI.getAlign()); |
| 154 | LI.replaceAllUsesWith(V: NewLoad); |
| 155 | LI.eraseFromParent(); |
| 156 | visitGetElementPtrInst(GEPI&: *OldGEP); |
| 157 | return true; |
| 158 | } |
| 159 | if (GlobalVariable *NewGlobal = lookupReplacementGlobal(CurrOperand: PtrOperand)) |
| 160 | LI.setOperand(i_nocapture: LI.getPointerOperandIndex(), Val_nocapture: NewGlobal); |
| 161 | return false; |
| 162 | } |
| 163 | |
| 164 | bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) { |
| 165 | |
| 166 | Value *PtrOperand = SI.getPointerOperand(); |
| 167 | ConstantExpr *CE = dyn_cast<ConstantExpr>(Val: PtrOperand); |
| 168 | if (CE && CE->getOpcode() == Instruction::GetElementPtr) { |
| 169 | GetElementPtrInst *OldGEP = cast<GetElementPtrInst>(Val: CE->getAsInstruction()); |
| 170 | OldGEP->insertBefore(InsertPos: SI.getIterator()); |
| 171 | IRBuilder<> Builder(&SI); |
| 172 | StoreInst *NewStore = Builder.CreateStore(Val: SI.getValueOperand(), Ptr: OldGEP); |
| 173 | NewStore->setAlignment(SI.getAlign()); |
| 174 | SI.replaceAllUsesWith(V: NewStore); |
| 175 | SI.eraseFromParent(); |
| 176 | visitGetElementPtrInst(GEPI&: *OldGEP); |
| 177 | return true; |
| 178 | } |
| 179 | if (GlobalVariable *NewGlobal = lookupReplacementGlobal(CurrOperand: PtrOperand)) |
| 180 | SI.setOperand(i_nocapture: SI.getPointerOperandIndex(), Val_nocapture: NewGlobal); |
| 181 | |
| 182 | return false; |
| 183 | } |
| 184 | |
| 185 | DataScalarizerVisitor::AllocaAndGEPs |
| 186 | DataScalarizerVisitor::createArrayFromVector(IRBuilder<> &Builder, Value *Vec, |
| 187 | const Twine &Name = "" ) { |
| 188 | // If there is already an alloca for this vector, return it |
| 189 | if (VectorAllocaMap.contains(Val: Vec)) |
| 190 | return VectorAllocaMap[Vec]; |
| 191 | |
| 192 | auto InsertPoint = Builder.GetInsertPoint(); |
| 193 | |
| 194 | // Allocate the array to hold the vector elements |
| 195 | Builder.SetInsertPointPastAllocas(Builder.GetInsertBlock()->getParent()); |
| 196 | Type *ArrTy = equivalentArrayTypeFromVector(T: Vec->getType()); |
| 197 | // DXIL indexable temps cannot hold i1 elements; booleans occupy 32 bits in |
| 198 | // memory. Widen i1 element arrays to i32. |
| 199 | Type *ArrElemTy = ArrTy->getArrayElementType(); |
| 200 | bool WidenBool = ArrElemTy->isIntegerTy(BitWidth: 1); |
| 201 | if (WidenBool) { |
| 202 | ArrElemTy = Builder.getInt32Ty(); |
| 203 | ArrTy = ArrayType::get(ElementType: ArrElemTy, NumElements: ArrTy->getArrayNumElements()); |
| 204 | } |
| 205 | AllocaInst *ArrAlloca = |
| 206 | Builder.CreateAlloca(Ty: ArrTy, ArraySize: nullptr, Name: Name + ".alloca" ); |
| 207 | const uint64_t ArrNumElems = ArrTy->getArrayNumElements(); |
| 208 | |
| 209 | // Create loads and stores to populate the array immediately after the |
| 210 | // original vector's defining instruction if available, else immediately after |
| 211 | // the alloca |
| 212 | if (auto *Instr = dyn_cast<Instruction>(Val: Vec)) |
| 213 | Builder.SetInsertPoint(Instr->getNextNode()); |
| 214 | SmallVector<Value *, 4> GEPs(ArrNumElems); |
| 215 | for (unsigned I = 0; I < ArrNumElems; ++I) { |
| 216 | Value *EE = Builder.CreateExtractElement(Vec, Idx: I, Name: Name + ".extract" ); |
| 217 | if (WidenBool) |
| 218 | EE = Builder.CreateZExt(V: EE, DestTy: ArrElemTy, Name: Name + ".zext" ); |
| 219 | GEPs[I] = Builder.CreateInBoundsGEP( |
| 220 | Ty: ArrTy, Ptr: ArrAlloca, IdxList: {Builder.getInt32(C: 0), Builder.getInt32(C: I)}, |
| 221 | Name: Name + ".index" ); |
| 222 | Builder.CreateStore(Val: EE, Ptr: GEPs[I]); |
| 223 | } |
| 224 | |
| 225 | VectorAllocaMap.insert(KV: {Vec, {ArrAlloca, ArrTy, GEPs}}); |
| 226 | Builder.SetInsertPoint(InsertPoint); |
| 227 | return {ArrAlloca, ArrTy, GEPs}; |
| 228 | } |
| 229 | |
| 230 | /// Returns a pair of Value* with the first being a GEP into ArrAlloca using |
| 231 | /// indices {0, Index}, and the second Value* being a Load of the GEP |
| 232 | static std::pair<Value *, Value *> |
| 233 | dynamicallyLoadArray(IRBuilder<> &Builder, AllocaInst *ArrAlloca, Type *ArrTy, |
| 234 | Value *Index, const Twine &Name = "" ) { |
| 235 | Value *GEP = Builder.CreateInBoundsGEP( |
| 236 | Ty: ArrTy, Ptr: ArrAlloca, IdxList: {Builder.getInt32(C: 0), Index}, Name: Name + ".index" ); |
| 237 | Value *Load = |
| 238 | Builder.CreateLoad(Ty: ArrTy->getArrayElementType(), Ptr: GEP, Name: Name + ".load" ); |
| 239 | return std::make_pair(x&: GEP, y&: Load); |
| 240 | } |
| 241 | |
| 242 | bool DataScalarizerVisitor::replaceDynamicInsertElementInst( |
| 243 | InsertElementInst &IEI) { |
| 244 | IRBuilder<> Builder(&IEI); |
| 245 | |
| 246 | Value *Vec = IEI.getOperand(i_nocapture: 0); |
| 247 | Value *Val = IEI.getOperand(i_nocapture: 1); |
| 248 | Value *Index = IEI.getOperand(i_nocapture: 2); |
| 249 | |
| 250 | AllocaAndGEPs ArrAllocaAndGEPs = |
| 251 | createArrayFromVector(Builder, Vec, Name: IEI.getName()); |
| 252 | AllocaInst *ArrAlloca = std::get<0>(t&: ArrAllocaAndGEPs); |
| 253 | Type *ArrTy = std::get<1>(t&: ArrAllocaAndGEPs); |
| 254 | SmallVector<Value *, 4> &ArrGEPs = std::get<2>(t&: ArrAllocaAndGEPs); |
| 255 | |
| 256 | // The array element type may have been widened (e.g. i1 -> i32) so that the |
| 257 | // indexable temp uses a legal DXIL memory type. Convert between the vector |
| 258 | // element type and the (possibly wider) array element type as needed. |
| 259 | Type *ArrElemTy = ArrTy->getArrayElementType(); |
| 260 | Type *VecElemTy = cast<VectorType>(Val: Vec->getType())->getElementType(); |
| 261 | bool WidenBool = ArrElemTy != VecElemTy && VecElemTy->isIntegerTy(BitWidth: 1); |
| 262 | |
| 263 | auto GEPAndLoad = |
| 264 | dynamicallyLoadArray(Builder, ArrAlloca, ArrTy, Index, Name: IEI.getName()); |
| 265 | Value *GEP = GEPAndLoad.first; |
| 266 | Value *Load = GEPAndLoad.second; |
| 267 | |
| 268 | Value *StoreVal = Val; |
| 269 | if (WidenBool) |
| 270 | StoreVal = Builder.CreateZExt(V: Val, DestTy: ArrElemTy, Name: IEI.getName() + ".zext" ); |
| 271 | Builder.CreateStore(Val: StoreVal, Ptr: GEP); |
| 272 | Value *NewIEI = PoisonValue::get(T: Vec->getType()); |
| 273 | for (unsigned I = 0; I < ArrTy->getArrayNumElements(); ++I) { |
| 274 | Value *EltLoad = |
| 275 | Builder.CreateLoad(Ty: ArrElemTy, Ptr: ArrGEPs[I], Name: IEI.getName() + ".load" ); |
| 276 | if (WidenBool) |
| 277 | EltLoad = |
| 278 | Builder.CreateTrunc(V: EltLoad, DestTy: VecElemTy, Name: IEI.getName() + ".trunc" ); |
| 279 | NewIEI = Builder.CreateInsertElement(Vec: NewIEI, NewElt: EltLoad, Idx: Builder.getInt32(C: I), |
| 280 | Name: IEI.getName() + ".insert" ); |
| 281 | } |
| 282 | |
| 283 | // Store back the original value so the Alloca can be reused for subsequent |
| 284 | // insertelement instructions on the same vector |
| 285 | Builder.CreateStore(Val: Load, Ptr: GEP); |
| 286 | |
| 287 | IEI.replaceAllUsesWith(V: NewIEI); |
| 288 | IEI.eraseFromParent(); |
| 289 | return true; |
| 290 | } |
| 291 | |
| 292 | bool DataScalarizerVisitor::visitInsertElementInst(InsertElementInst &IEI) { |
| 293 | // If the index is a constant then we don't need to scalarize it |
| 294 | Value *Index = IEI.getOperand(i_nocapture: 2); |
| 295 | if (isa<ConstantInt>(Val: Index)) |
| 296 | return false; |
| 297 | return replaceDynamicInsertElementInst(IEI); |
| 298 | } |
| 299 | |
| 300 | bool DataScalarizerVisitor::( |
| 301 | ExtractElementInst &EEI) { |
| 302 | IRBuilder<> Builder(&EEI); |
| 303 | |
| 304 | AllocaAndGEPs ArrAllocaAndGEPs = |
| 305 | createArrayFromVector(Builder, Vec: EEI.getVectorOperand(), Name: EEI.getName()); |
| 306 | AllocaInst *ArrAlloca = std::get<0>(t&: ArrAllocaAndGEPs); |
| 307 | Type *ArrTy = std::get<1>(t&: ArrAllocaAndGEPs); |
| 308 | |
| 309 | auto GEPAndLoad = dynamicallyLoadArray(Builder, ArrAlloca, ArrTy, |
| 310 | Index: EEI.getIndexOperand(), Name: EEI.getName()); |
| 311 | Value *Load = GEPAndLoad.second; |
| 312 | |
| 313 | // The array element type may have been widened (e.g. i1 -> i32) so that the |
| 314 | // indexable temp uses a legal DXIL memory type. Truncate back to the original |
| 315 | // element type of the extractelement if necessary. |
| 316 | if (Load->getType() != EEI.getType()) { |
| 317 | assert(Load->getType()->isIntegerTy(32) && EEI.getType()->isIntegerTy(1) && |
| 318 | "Unexpected type mismatch: only i32 -> i1 widening is supported" ); |
| 319 | Load = Builder.CreateTrunc(V: Load, DestTy: EEI.getType(), Name: EEI.getName() + ".trunc" ); |
| 320 | } |
| 321 | |
| 322 | EEI.replaceAllUsesWith(V: Load); |
| 323 | EEI.eraseFromParent(); |
| 324 | return true; |
| 325 | } |
| 326 | |
| 327 | bool DataScalarizerVisitor::(ExtractElementInst &EEI) { |
| 328 | // If the index is a constant then we don't need to scalarize it |
| 329 | Value *Index = EEI.getIndexOperand(); |
| 330 | if (isa<ConstantInt>(Val: Index)) |
| 331 | return false; |
| 332 | return replaceDynamicExtractElementInst(EEI); |
| 333 | } |
| 334 | |
| 335 | bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) { |
| 336 | GEPOperator *GOp = cast<GEPOperator>(Val: &GEPI); |
| 337 | Value *PtrOperand = GOp->getPointerOperand(); |
| 338 | Type *GEPType = GOp->getSourceElementType(); |
| 339 | |
| 340 | // Replace a GEP ConstantExpr pointer operand with a GEP instruction so that |
| 341 | // it can be visited |
| 342 | if (auto *PtrOpGEPCE = dyn_cast<ConstantExpr>(Val: PtrOperand); |
| 343 | PtrOpGEPCE && PtrOpGEPCE->getOpcode() == Instruction::GetElementPtr) { |
| 344 | GetElementPtrInst *OldGEPI = |
| 345 | cast<GetElementPtrInst>(Val: PtrOpGEPCE->getAsInstruction()); |
| 346 | OldGEPI->insertBefore(InsertPos: GEPI.getIterator()); |
| 347 | |
| 348 | IRBuilder<> Builder(&GEPI); |
| 349 | SmallVector<Value *> Indices(GEPI.indices()); |
| 350 | Value *NewGEP = |
| 351 | Builder.CreateGEP(Ty: GEPI.getSourceElementType(), Ptr: OldGEPI, IdxList: Indices, |
| 352 | Name: GEPI.getName(), NW: GEPI.getNoWrapFlags()); |
| 353 | assert(isa<GetElementPtrInst>(NewGEP) && |
| 354 | "Expected newly-created GEP to be an instruction" ); |
| 355 | GetElementPtrInst *NewGEPI = cast<GetElementPtrInst>(Val: NewGEP); |
| 356 | |
| 357 | GEPI.replaceAllUsesWith(V: NewGEPI); |
| 358 | GEPI.eraseFromParent(); |
| 359 | visitGetElementPtrInst(GEPI&: *OldGEPI); |
| 360 | visitGetElementPtrInst(GEPI&: *NewGEPI); |
| 361 | return true; |
| 362 | } |
| 363 | |
| 364 | Type *NewGEPType = equivalentArrayTypeFromVector(T: GEPType); |
| 365 | Value *NewPtrOperand = PtrOperand; |
| 366 | if (GlobalVariable *NewGlobal = lookupReplacementGlobal(CurrOperand: PtrOperand)) |
| 367 | NewPtrOperand = NewGlobal; |
| 368 | |
| 369 | bool NeedsTransform = NewPtrOperand != PtrOperand || NewGEPType != GEPType; |
| 370 | if (!NeedsTransform) |
| 371 | return false; |
| 372 | |
| 373 | IRBuilder<> Builder(&GEPI); |
| 374 | SmallVector<Value *, MaxVecSize> Indices(GOp->idx_begin(), GOp->idx_end()); |
| 375 | Value *NewGEP = Builder.CreateGEP(Ty: NewGEPType, Ptr: NewPtrOperand, IdxList: Indices, |
| 376 | Name: GOp->getName(), NW: GOp->getNoWrapFlags()); |
| 377 | |
| 378 | GOp->replaceAllUsesWith(V: NewGEP); |
| 379 | |
| 380 | if (auto *OldGEPI = dyn_cast<GetElementPtrInst>(Val: GOp)) |
| 381 | OldGEPI->eraseFromParent(); |
| 382 | |
| 383 | return true; |
| 384 | } |
| 385 | |
| 386 | static Constant *transformInitializer(Constant *Init, Type *OrigType, |
| 387 | Type *NewType, LLVMContext &Ctx) { |
| 388 | // Handle ConstantAggregateZero (zero-initialized constants) |
| 389 | if (isa<ConstantAggregateZero>(Val: Init)) { |
| 390 | return ConstantAggregateZero::get(Ty: NewType); |
| 391 | } |
| 392 | |
| 393 | // Handle UndefValue (undefined constants) |
| 394 | if (isa<UndefValue>(Val: Init)) { |
| 395 | return UndefValue::get(T: NewType); |
| 396 | } |
| 397 | |
| 398 | // Handle vector to array transformation |
| 399 | if (isa<VectorType>(Val: OrigType) && isa<ArrayType>(Val: NewType)) { |
| 400 | // Convert vector initializer to array initializer |
| 401 | SmallVector<Constant *, MaxVecSize> ArrayElements; |
| 402 | |
| 403 | unsigned E = cast<FixedVectorType>(Val: OrigType)->getNumElements(); |
| 404 | for (unsigned I = 0; I != E; ++I) |
| 405 | if (Constant *Elt = Init->getAggregateElement(Elt: I)) |
| 406 | ArrayElements.push_back(Elt); |
| 407 | |
| 408 | assert(ArrayElements.size() == E && |
| 409 | "Expected fixed length constant aggregate for vector initializer!" ); |
| 410 | return ConstantArray::get(T: cast<ArrayType>(Val: NewType), V: ArrayElements); |
| 411 | } |
| 412 | |
| 413 | // Handle array of vectors transformation |
| 414 | if (auto *ArrayTy = dyn_cast<ArrayType>(Val: OrigType)) { |
| 415 | auto *ArrayInit = dyn_cast<ConstantArray>(Val: Init); |
| 416 | assert(ArrayInit && "Expected a ConstantArray for array initializer!" ); |
| 417 | |
| 418 | SmallVector<Constant *, MaxVecSize> NewArrayElements; |
| 419 | for (unsigned I = 0; I < ArrayTy->getNumElements(); ++I) { |
| 420 | // Recursively transform array elements |
| 421 | Constant *NewElemInit = transformInitializer( |
| 422 | Init: ArrayInit->getOperand(i_nocapture: I), OrigType: ArrayTy->getElementType(), |
| 423 | NewType: cast<ArrayType>(Val: NewType)->getElementType(), Ctx); |
| 424 | NewArrayElements.push_back(Elt: NewElemInit); |
| 425 | } |
| 426 | |
| 427 | return ConstantArray::get(T: cast<ArrayType>(Val: NewType), V: NewArrayElements); |
| 428 | } |
| 429 | |
| 430 | // If not a vector or array, return the original initializer |
| 431 | return Init; |
| 432 | } |
| 433 | |
| 434 | static bool findAndReplaceVectors(Module &M) { |
| 435 | bool MadeChange = false; |
| 436 | LLVMContext &Ctx = M.getContext(); |
| 437 | IRBuilder<> Builder(Ctx); |
| 438 | DataScalarizerVisitor Impl; |
| 439 | for (GlobalVariable &G : M.globals()) { |
| 440 | Type *OrigType = G.getValueType(); |
| 441 | |
| 442 | Type *NewType = equivalentArrayTypeFromVector(T: OrigType); |
| 443 | if (OrigType != NewType) { |
| 444 | // Create a new global variable with the updated type |
| 445 | // Note: Initializer is set via transformInitializer |
| 446 | GlobalVariable *NewGlobal = new GlobalVariable( |
| 447 | M, NewType, G.isConstant(), G.getLinkage(), |
| 448 | /*Initializer=*/nullptr, G.getName() + ".scalarized" , &G, |
| 449 | G.getThreadLocalMode(), G.getAddressSpace(), |
| 450 | G.isExternallyInitialized()); |
| 451 | |
| 452 | // Copy relevant attributes |
| 453 | NewGlobal->setUnnamedAddr(G.getUnnamedAddr()); |
| 454 | if (G.getAlign()) { |
| 455 | NewGlobal->setAlignment(G.getAlign()); |
| 456 | } |
| 457 | |
| 458 | if (G.hasInitializer()) { |
| 459 | Constant *Init = G.getInitializer(); |
| 460 | Constant *NewInit = transformInitializer(Init, OrigType, NewType, Ctx); |
| 461 | NewGlobal->setInitializer(NewInit); |
| 462 | } |
| 463 | |
| 464 | // Note: we want to do G.replaceAllUsesWith(NewGlobal);, but it assumes |
| 465 | // type equality. Instead we will use the visitor pattern. |
| 466 | Impl.GlobalMap[&G] = NewGlobal; |
| 467 | } |
| 468 | } |
| 469 | |
| 470 | for (auto &F : make_early_inc_range(Range: M.functions())) { |
| 471 | if (F.isDeclaration()) |
| 472 | continue; |
| 473 | MadeChange |= Impl.visit(F); |
| 474 | } |
| 475 | |
| 476 | // Remove the old globals after the iteration |
| 477 | for (auto &[Old, New] : Impl.GlobalMap) { |
| 478 | Old->eraseFromParent(); |
| 479 | MadeChange = true; |
| 480 | } |
| 481 | return MadeChange; |
| 482 | } |
| 483 | |
| 484 | PreservedAnalyses DXILDataScalarization::run(Module &M, |
| 485 | ModuleAnalysisManager &) { |
| 486 | bool MadeChanges = findAndReplaceVectors(M); |
| 487 | if (!MadeChanges) |
| 488 | return PreservedAnalyses::all(); |
| 489 | PreservedAnalyses PA; |
| 490 | return PA; |
| 491 | } |
| 492 | |
| 493 | bool DXILDataScalarizationLegacy::runOnModule(Module &M) { |
| 494 | return findAndReplaceVectors(M); |
| 495 | } |
| 496 | |
| 497 | char DXILDataScalarizationLegacy::ID = 0; |
| 498 | |
| 499 | INITIALIZE_PASS_BEGIN(DXILDataScalarizationLegacy, DEBUG_TYPE, |
| 500 | "DXIL Data Scalarization" , false, false) |
| 501 | INITIALIZE_PASS_END(DXILDataScalarizationLegacy, DEBUG_TYPE, |
| 502 | "DXIL Data Scalarization" , false, false) |
| 503 | |
| 504 | ModulePass *llvm::createDXILDataScalarizationLegacyPass() { |
| 505 | return new DXILDataScalarizationLegacy(); |
| 506 | } |
| 507 | |