brintos

brintos / llvm-project-archived public Read only

0
0
Text · 25.2 KiB · 3427968 Raw
712 lines · cpp
1//===- DXILLegalizePass.cpp - Legalizes llvm IR for DXIL ------------------===//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 "DXILLegalizePass.h"10#include "DirectX.h"11#include "llvm/ADT/APInt.h"12#include "llvm/IR/Constants.h"13#include "llvm/IR/Function.h"14#include "llvm/IR/IRBuilder.h"15#include "llvm/IR/InstIterator.h"16#include "llvm/IR/Instruction.h"17#include "llvm/IR/Instructions.h"18#include "llvm/IR/Module.h"19#include "llvm/Pass.h"20#include "llvm/Transforms/Utils/BasicBlockUtils.h"21#include <functional>22 23#define DEBUG_TYPE "dxil-legalize"24 25using namespace llvm;26 27static bool legalizeFreeze(Instruction &I,28                           SmallVectorImpl<Instruction *> &ToRemove,29                           DenseMap<Value *, Value *>) {30  auto *FI = dyn_cast<FreezeInst>(&I);31  if (!FI)32    return false;33 34  FI->replaceAllUsesWith(FI->getOperand(0));35  ToRemove.push_back(FI);36  return true;37}38 39static bool fixI8UseChain(Instruction &I,40                          SmallVectorImpl<Instruction *> &ToRemove,41                          DenseMap<Value *, Value *> &ReplacedValues) {42 43  auto ProcessOperands = [&](SmallVector<Value *> &NewOperands) {44    Type *InstrType = IntegerType::get(I.getContext(), 32);45 46    for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {47      Value *Op = I.getOperand(OpIdx);48      if (ReplacedValues.count(Op) &&49          ReplacedValues[Op]->getType()->isIntegerTy())50        InstrType = ReplacedValues[Op]->getType();51    }52 53    for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {54      Value *Op = I.getOperand(OpIdx);55      if (ReplacedValues.count(Op))56        NewOperands.push_back(ReplacedValues[Op]);57      else if (auto *Imm = dyn_cast<ConstantInt>(Op)) {58        APInt Value = Imm->getValue();59        unsigned NewBitWidth = InstrType->getIntegerBitWidth();60        // Note: options here are sext or sextOrTrunc.61        // Since i8 isn't supported, we assume new values62        // will always have a higher bitness.63        assert(NewBitWidth > Value.getBitWidth() &&64               "Replacement's BitWidth should be larger than Current.");65        APInt NewValue = Value.sext(NewBitWidth);66        NewOperands.push_back(ConstantInt::get(InstrType, NewValue));67      } else {68        assert(!Op->getType()->isIntegerTy(8));69        NewOperands.push_back(Op);70      }71    }72  };73  IRBuilder<> Builder(&I);74  if (auto *Trunc = dyn_cast<TruncInst>(&I)) {75    if (Trunc->getDestTy()->isIntegerTy(8)) {76      ReplacedValues[Trunc] = Trunc->getOperand(0);77      ToRemove.push_back(Trunc);78      return true;79    }80  }81 82  if (auto *Store = dyn_cast<StoreInst>(&I)) {83    if (!Store->getValueOperand()->getType()->isIntegerTy(8))84      return false;85    SmallVector<Value *> NewOperands;86    ProcessOperands(NewOperands);87    Value *NewStore = Builder.CreateStore(NewOperands[0], NewOperands[1]);88    ReplacedValues[Store] = NewStore;89    ToRemove.push_back(Store);90    return true;91  }92 93  if (auto *Load = dyn_cast<LoadInst>(&I);94      Load && I.getType()->isIntegerTy(8)) {95    SmallVector<Value *> NewOperands;96    ProcessOperands(NewOperands);97    Type *ElementType = NewOperands[0]->getType();98    if (auto *AI = dyn_cast<AllocaInst>(NewOperands[0]))99      ElementType = AI->getAllocatedType();100    if (auto *GEP = dyn_cast<GetElementPtrInst>(NewOperands[0])) {101      ElementType = GEP->getSourceElementType();102    }103    if (ElementType->isArrayTy())104      ElementType = ElementType->getArrayElementType();105    LoadInst *NewLoad = Builder.CreateLoad(ElementType, NewOperands[0]);106    ReplacedValues[Load] = NewLoad;107    ToRemove.push_back(Load);108    return true;109  }110 111  if (auto *Load = dyn_cast<LoadInst>(&I);112      Load && isa<ConstantExpr>(Load->getPointerOperand())) {113    auto *CE = dyn_cast<ConstantExpr>(Load->getPointerOperand());114    if (!(CE->getOpcode() == Instruction::GetElementPtr))115      return false;116    auto *GEP = dyn_cast<GEPOperator>(CE);117    if (!GEP->getSourceElementType()->isIntegerTy(8))118      return false;119 120    Type *ElementType = Load->getType();121    ConstantInt *Offset = dyn_cast<ConstantInt>(GEP->getOperand(1));122    uint32_t ByteOffset = Offset->getZExtValue();123    uint32_t ElemSize = Load->getDataLayout().getTypeAllocSize(ElementType);124    uint32_t Index = ByteOffset / ElemSize;125 126    Value *PtrOperand = GEP->getPointerOperand();127    Type *GEPType = GEP->getPointerOperandType();128 129    if (auto *GV = dyn_cast<GlobalVariable>(PtrOperand))130      GEPType = GV->getValueType();131    if (auto *AI = dyn_cast<AllocaInst>(PtrOperand))132      GEPType = AI->getAllocatedType();133 134    if (auto *ArrTy = dyn_cast<ArrayType>(GEPType))135      GEPType = ArrTy;136    else137      GEPType = ArrayType::get(ElementType, 1); // its a scalar138 139    Value *NewGEP = Builder.CreateGEP(140        GEPType, PtrOperand, {Builder.getInt32(0), Builder.getInt32(Index)},141        GEP->getName(), GEP->getNoWrapFlags());142 143    LoadInst *NewLoad = Builder.CreateLoad(ElementType, NewGEP);144    ReplacedValues[Load] = NewLoad;145    Load->replaceAllUsesWith(NewLoad);146    ToRemove.push_back(Load);147    return true;148  }149 150  if (auto *BO = dyn_cast<BinaryOperator>(&I)) {151    if (!I.getType()->isIntegerTy(8))152      return false;153    SmallVector<Value *> NewOperands;154    ProcessOperands(NewOperands);155    Value *NewInst =156        Builder.CreateBinOp(BO->getOpcode(), NewOperands[0], NewOperands[1]);157    if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(&I)) {158      auto *NewBO = dyn_cast<BinaryOperator>(NewInst);159      if (NewBO && OBO->hasNoSignedWrap())160        NewBO->setHasNoSignedWrap();161      if (NewBO && OBO->hasNoUnsignedWrap())162        NewBO->setHasNoUnsignedWrap();163    }164    ReplacedValues[BO] = NewInst;165    ToRemove.push_back(BO);166    return true;167  }168 169  if (auto *Sel = dyn_cast<SelectInst>(&I)) {170    if (!I.getType()->isIntegerTy(8))171      return false;172    SmallVector<Value *> NewOperands;173    ProcessOperands(NewOperands);174    Value *NewInst = Builder.CreateSelect(Sel->getCondition(), NewOperands[1],175                                          NewOperands[2]);176    ReplacedValues[Sel] = NewInst;177    ToRemove.push_back(Sel);178    return true;179  }180 181  if (auto *Cmp = dyn_cast<CmpInst>(&I)) {182    if (!Cmp->getOperand(0)->getType()->isIntegerTy(8))183      return false;184    SmallVector<Value *> NewOperands;185    ProcessOperands(NewOperands);186    Value *NewInst =187        Builder.CreateCmp(Cmp->getPredicate(), NewOperands[0], NewOperands[1]);188    Cmp->replaceAllUsesWith(NewInst);189    ReplacedValues[Cmp] = NewInst;190    ToRemove.push_back(Cmp);191    return true;192  }193 194  if (auto *Cast = dyn_cast<CastInst>(&I)) {195    if (!Cast->getSrcTy()->isIntegerTy(8))196      return false;197 198    ToRemove.push_back(Cast);199    auto *Replacement = ReplacedValues[Cast->getOperand(0)];200    if (Cast->getType() == Replacement->getType()) {201      Cast->replaceAllUsesWith(Replacement);202      return true;203    }204 205    Value *AdjustedCast = nullptr;206    if (Cast->getOpcode() == Instruction::ZExt)207      AdjustedCast = Builder.CreateZExtOrTrunc(Replacement, Cast->getType());208    if (Cast->getOpcode() == Instruction::SExt)209      AdjustedCast = Builder.CreateSExtOrTrunc(Replacement, Cast->getType());210 211    if (AdjustedCast)212      Cast->replaceAllUsesWith(AdjustedCast);213  }214  if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {215    if (!GEP->getType()->isPointerTy() ||216        !GEP->getSourceElementType()->isIntegerTy(8))217      return false;218 219    Value *BasePtr = GEP->getPointerOperand();220    if (ReplacedValues.count(BasePtr))221      BasePtr = ReplacedValues[BasePtr];222 223    Type *ElementType = BasePtr->getType();224 225    if (auto *AI = dyn_cast<AllocaInst>(BasePtr))226      ElementType = AI->getAllocatedType();227    if (auto *GV = dyn_cast<GlobalVariable>(BasePtr))228      ElementType = GV->getValueType();229 230    Type *GEPType = ElementType;231    if (auto *ArrTy = dyn_cast<ArrayType>(ElementType))232      ElementType = ArrTy->getArrayElementType();233    else234      GEPType = ArrayType::get(ElementType, 1); // its a scalar235 236    ConstantInt *Offset = dyn_cast<ConstantInt>(GEP->getOperand(1));237    // Note: i8 to i32 offset conversion without emitting IR requires constant238    // ints. Since offset conversion is common, we can safely assume Offset is239    // always a ConstantInt, so no need to have a conditional bail out on240    // nullptr, instead assert this is the case.241    assert(Offset && "Offset is expected to be a ConstantInt");242    uint32_t ByteOffset = Offset->getZExtValue();243    uint32_t ElemSize = GEP->getDataLayout().getTypeAllocSize(ElementType);244    assert(ElemSize > 0 && "ElementSize must be set");245    uint32_t Index = ByteOffset / ElemSize;246    Value *NewGEP = Builder.CreateGEP(247        GEPType, BasePtr, {Builder.getInt32(0), Builder.getInt32(Index)},248        GEP->getName(), GEP->getNoWrapFlags());249    ReplacedValues[GEP] = NewGEP;250    GEP->replaceAllUsesWith(NewGEP);251    ToRemove.push_back(GEP);252    return true;253  }254  return false;255}256 257static bool upcastI8AllocasAndUses(Instruction &I,258                                   SmallVectorImpl<Instruction *> &ToRemove,259                                   DenseMap<Value *, Value *> &ReplacedValues) {260  auto *AI = dyn_cast<AllocaInst>(&I);261  if (!AI || !AI->getAllocatedType()->isIntegerTy(8))262    return false;263 264  Type *SmallestType = nullptr;265 266  auto ProcessLoad = [&](LoadInst *Load) {267    for (User *LU : Load->users()) {268      Type *Ty = nullptr;269      if (CastInst *Cast = dyn_cast<CastInst>(LU))270        Ty = Cast->getType();271      else if (CallInst *CI = dyn_cast<CallInst>(LU)) {272        if (CI->getIntrinsicID() == Intrinsic::memset)273          Ty = Type::getInt32Ty(CI->getContext());274      }275 276      if (!Ty)277        continue;278 279      if (!SmallestType ||280          Ty->getPrimitiveSizeInBits() < SmallestType->getPrimitiveSizeInBits())281        SmallestType = Ty;282    }283  };284 285  for (User *U : AI->users()) {286    if (auto *Load = dyn_cast<LoadInst>(U))287      ProcessLoad(Load);288    else if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {289      for (User *GU : GEP->users()) {290        if (auto *Load = dyn_cast<LoadInst>(GU))291          ProcessLoad(Load);292      }293    }294  }295 296  if (!SmallestType)297    return false; // no valid casts found298 299  // Replace alloca300  IRBuilder<> Builder(AI);301  auto *NewAlloca = Builder.CreateAlloca(SmallestType);302  ReplacedValues[AI] = NewAlloca;303  ToRemove.push_back(AI);304  return true;305}306 307static bool308downcastI64toI32InsertExtractElements(Instruction &I,309                                      SmallVectorImpl<Instruction *> &ToRemove,310                                      DenseMap<Value *, Value *> &) {311 312  if (auto *Extract = dyn_cast<ExtractElementInst>(&I)) {313    Value *Idx = Extract->getIndexOperand();314    auto *CI = dyn_cast<ConstantInt>(Idx);315    if (CI && CI->getBitWidth() == 64) {316      IRBuilder<> Builder(Extract);317      int64_t IndexValue = CI->getSExtValue();318      auto *Idx32 =319          ConstantInt::get(Type::getInt32Ty(I.getContext()), IndexValue);320      Value *NewExtract = Builder.CreateExtractElement(321          Extract->getVectorOperand(), Idx32, Extract->getName());322 323      Extract->replaceAllUsesWith(NewExtract);324      ToRemove.push_back(Extract);325      return true;326    }327  }328 329  if (auto *Insert = dyn_cast<InsertElementInst>(&I)) {330    Value *Idx = Insert->getOperand(2);331    auto *CI = dyn_cast<ConstantInt>(Idx);332    if (CI && CI->getBitWidth() == 64) {333      int64_t IndexValue = CI->getSExtValue();334      auto *Idx32 =335          ConstantInt::get(Type::getInt32Ty(I.getContext()), IndexValue);336      IRBuilder<> Builder(Insert);337      Value *Insert32Index = Builder.CreateInsertElement(338          Insert->getOperand(0), Insert->getOperand(1), Idx32,339          Insert->getName());340 341      Insert->replaceAllUsesWith(Insert32Index);342      ToRemove.push_back(Insert);343      return true;344    }345  }346  return false;347}348 349static void emitMemcpyExpansion(IRBuilder<> &Builder, Value *Dst, Value *Src,350                                ConstantInt *Length) {351 352  uint64_t ByteLength = Length->getZExtValue();353  // If length to copy is zero, no memcpy is needed.354  if (ByteLength == 0)355    return;356 357  const DataLayout &DL = Builder.GetInsertBlock()->getModule()->getDataLayout();358 359  auto GetArrTyFromVal = [](Value *Val) -> ArrayType * {360    assert(isa<AllocaInst>(Val) ||361           isa<GlobalVariable>(Val) &&362               "Expected Val to be an Alloca or Global Variable");363    if (auto *Alloca = dyn_cast<AllocaInst>(Val))364      return dyn_cast<ArrayType>(Alloca->getAllocatedType());365    if (auto *GlobalVar = dyn_cast<GlobalVariable>(Val))366      return dyn_cast<ArrayType>(GlobalVar->getValueType());367    return nullptr;368  };369 370  ArrayType *DstArrTy = GetArrTyFromVal(Dst);371  assert(DstArrTy && "Expected Dst of memcpy to be a Pointer to an Array Type");372  if (auto *DstGlobalVar = dyn_cast<GlobalVariable>(Dst))373    assert(!DstGlobalVar->isConstant() &&374           "The Dst of memcpy must not be a constant Global Variable");375  [[maybe_unused]] ArrayType *SrcArrTy = GetArrTyFromVal(Src);376  assert(SrcArrTy && "Expected Src of memcpy to be a Pointer to an Array Type");377 378  Type *DstElemTy = DstArrTy->getElementType();379  uint64_t DstElemByteSize = DL.getTypeStoreSize(DstElemTy);380  assert(DstElemByteSize > 0 && "Dst element type store size must be set");381  Type *SrcElemTy = SrcArrTy->getElementType();382  [[maybe_unused]] uint64_t SrcElemByteSize = DL.getTypeStoreSize(SrcElemTy);383  assert(SrcElemByteSize > 0 && "Src element type store size must be set");384 385  // This assumption simplifies implementation and covers currently-known386  // use-cases for DXIL. It may be relaxed in the future if required.387  assert(DstElemTy == SrcElemTy &&388         "The element types of Src and Dst arrays must match");389 390  [[maybe_unused]] uint64_t DstArrNumElems = DstArrTy->getArrayNumElements();391  assert(DstElemByteSize * DstArrNumElems >= ByteLength &&392         "Dst array size must be at least as large as the memcpy length");393  [[maybe_unused]] uint64_t SrcArrNumElems = SrcArrTy->getArrayNumElements();394  assert(SrcElemByteSize * SrcArrNumElems >= ByteLength &&395         "Src array size must be at least as large as the memcpy length");396 397  uint64_t NumElemsToCopy = ByteLength / DstElemByteSize;398  assert(ByteLength % DstElemByteSize == 0 &&399         "memcpy length must be divisible by array element type");400  for (uint64_t I = 0; I < NumElemsToCopy; ++I) {401    SmallVector<Value *, 2> Indices = {Builder.getInt32(0),402                                       Builder.getInt32(I)};403    Value *SrcPtr = Builder.CreateInBoundsGEP(SrcArrTy, Src, Indices, "gep");404    Value *SrcVal = Builder.CreateLoad(SrcElemTy, SrcPtr);405    Value *DstPtr = Builder.CreateInBoundsGEP(DstArrTy, Dst, Indices, "gep");406    Builder.CreateStore(SrcVal, DstPtr);407  }408}409 410static void emitMemsetExpansion(IRBuilder<> &Builder, Value *Dst, Value *Val,411                                ConstantInt *SizeCI,412                                DenseMap<Value *, Value *> &ReplacedValues) {413  [[maybe_unused]] const DataLayout &DL =414      Builder.GetInsertBlock()->getModule()->getDataLayout();415  [[maybe_unused]] uint64_t OrigSize = SizeCI->getZExtValue();416 417  AllocaInst *Alloca = dyn_cast<AllocaInst>(Dst);418 419  assert(Alloca && "Expected memset on an Alloca");420  assert(OrigSize == Alloca->getAllocationSize(DL)->getFixedValue() &&421         "Expected for memset size to match DataLayout size");422 423  Type *AllocatedTy = Alloca->getAllocatedType();424  ArrayType *ArrTy = dyn_cast<ArrayType>(AllocatedTy);425  assert(ArrTy && "Expected Alloca for an Array Type");426 427  Type *ElemTy = ArrTy->getElementType();428  uint64_t Size = ArrTy->getArrayNumElements();429 430  [[maybe_unused]] uint64_t ElemSize = DL.getTypeStoreSize(ElemTy);431 432  assert(ElemSize > 0 && "Size must be set");433  assert(OrigSize == ElemSize * Size && "Size in bytes must match");434 435  Value *TypedVal = Val;436 437  if (Val->getType() != ElemTy) {438    if (ReplacedValues[Val]) {439      // Note for i8 replacements if we know them we should use them.440      // Further if this is a constant ReplacedValues will return null441      // so we will stick to TypedVal = Val442      TypedVal = ReplacedValues[Val];443 444    } else {445      // This case Val is a ConstantInt so the cast folds away.446      // However if we don't do the cast the store below ends up being447      // an i8.448      TypedVal = Builder.CreateIntCast(Val, ElemTy, false);449    }450  }451 452  for (uint64_t I = 0; I < Size; ++I) {453    Value *Zero = Builder.getInt32(0);454    Value *Offset = Builder.getInt32(I);455    Value *Ptr = Builder.CreateGEP(ArrTy, Dst, {Zero, Offset}, "gep");456    Builder.CreateStore(TypedVal, Ptr);457  }458}459 460// Expands the instruction `I` into corresponding loads and stores if it is a461// memcpy call. In that case, the call instruction is added to the `ToRemove`462// vector. `ReplacedValues` is unused.463static bool legalizeMemCpy(Instruction &I,464                           SmallVectorImpl<Instruction *> &ToRemove,465                           DenseMap<Value *, Value *> &ReplacedValues) {466 467  CallInst *CI = dyn_cast<CallInst>(&I);468  if (!CI)469    return false;470 471  Intrinsic::ID ID = CI->getIntrinsicID();472  if (ID != Intrinsic::memcpy)473    return false;474 475  IRBuilder<> Builder(&I);476  Value *Dst = CI->getArgOperand(0);477  Value *Src = CI->getArgOperand(1);478  ConstantInt *Length = dyn_cast<ConstantInt>(CI->getArgOperand(2));479  assert(Length && "Expected Length to be a ConstantInt");480  [[maybe_unused]] ConstantInt *IsVolatile =481      dyn_cast<ConstantInt>(CI->getArgOperand(3));482  assert(IsVolatile && "Expected IsVolatile to be a ConstantInt");483  assert(IsVolatile->getZExtValue() == 0 && "Expected IsVolatile to be false");484  emitMemcpyExpansion(Builder, Dst, Src, Length);485  ToRemove.push_back(CI);486  return true;487}488 489static bool legalizeMemSet(Instruction &I,490                           SmallVectorImpl<Instruction *> &ToRemove,491                           DenseMap<Value *, Value *> &ReplacedValues) {492 493  CallInst *CI = dyn_cast<CallInst>(&I);494  if (!CI)495    return false;496 497  Intrinsic::ID ID = CI->getIntrinsicID();498  if (ID != Intrinsic::memset)499    return false;500 501  IRBuilder<> Builder(&I);502  Value *Dst = CI->getArgOperand(0);503  Value *Val = CI->getArgOperand(1);504  ConstantInt *Size = dyn_cast<ConstantInt>(CI->getArgOperand(2));505  assert(Size && "Expected Size to be a ConstantInt");506  emitMemsetExpansion(Builder, Dst, Val, Size, ReplacedValues);507  ToRemove.push_back(CI);508  return true;509}510 511static bool updateFnegToFsub(Instruction &I,512                             SmallVectorImpl<Instruction *> &ToRemove,513                             DenseMap<Value *, Value *> &) {514  const Intrinsic::ID ID = I.getOpcode();515  if (ID != Instruction::FNeg)516    return false;517 518  IRBuilder<> Builder(&I);519  Value *In = I.getOperand(0);520  Value *Zero = ConstantFP::get(In->getType(), -0.0);521  I.replaceAllUsesWith(Builder.CreateFSub(Zero, In));522  ToRemove.push_back(&I);523  return true;524}525 526static bool527legalizeGetHighLowi64Bytes(Instruction &I,528                           SmallVectorImpl<Instruction *> &ToRemove,529                           DenseMap<Value *, Value *> &ReplacedValues) {530  if (auto *BitCast = dyn_cast<BitCastInst>(&I)) {531    if (BitCast->getDestTy() ==532            FixedVectorType::get(Type::getInt32Ty(I.getContext()), 2) &&533        BitCast->getSrcTy()->isIntegerTy(64)) {534      ToRemove.push_back(BitCast);535      ReplacedValues[BitCast] = BitCast->getOperand(0);536      return true;537    }538  }539 540  if (auto *Extract = dyn_cast<ExtractElementInst>(&I)) {541    if (!dyn_cast<BitCastInst>(Extract->getVectorOperand()))542      return false;543    auto *VecTy = dyn_cast<FixedVectorType>(Extract->getVectorOperandType());544    if (VecTy && VecTy->getElementType()->isIntegerTy(32) &&545        VecTy->getNumElements() == 2) {546      if (auto *Index = dyn_cast<ConstantInt>(Extract->getIndexOperand())) {547        unsigned Idx = Index->getZExtValue();548        IRBuilder<> Builder(&I);549 550        auto *Replacement = ReplacedValues[Extract->getVectorOperand()];551        assert(Replacement && "The BitCast replacement should have been set "552                              "before working on ExtractElementInst.");553        if (Idx == 0) {554          Value *LowBytes = Builder.CreateTrunc(555              Replacement, Type::getInt32Ty(I.getContext()));556          ReplacedValues[Extract] = LowBytes;557        } else {558          assert(Idx == 1);559          Value *LogicalShiftRight = Builder.CreateLShr(560              Replacement,561              ConstantInt::get(562                  Replacement->getType(),563                  APInt(Replacement->getType()->getIntegerBitWidth(), 32)));564          Value *HighBytes = Builder.CreateTrunc(565              LogicalShiftRight, Type::getInt32Ty(I.getContext()));566          ReplacedValues[Extract] = HighBytes;567        }568        ToRemove.push_back(Extract);569        Extract->replaceAllUsesWith(ReplacedValues[Extract]);570        return true;571      }572    }573  }574  return false;575}576 577static bool578legalizeScalarLoadStoreOnArrays(Instruction &I,579                                SmallVectorImpl<Instruction *> &ToRemove,580                                DenseMap<Value *, Value *> &) {581 582  Value *PtrOp;583  unsigned PtrOpIndex;584  [[maybe_unused]] Type *LoadStoreTy;585  if (auto *LI = dyn_cast<LoadInst>(&I)) {586    PtrOp = LI->getPointerOperand();587    PtrOpIndex = LI->getPointerOperandIndex();588    LoadStoreTy = LI->getType();589  } else if (auto *SI = dyn_cast<StoreInst>(&I)) {590    PtrOp = SI->getPointerOperand();591    PtrOpIndex = SI->getPointerOperandIndex();592    LoadStoreTy = SI->getValueOperand()->getType();593  } else594    return false;595 596  // If the load/store is not of a single-value type (i.e., scalar or vector)597  // then we do not modify it. It shouldn't be a vector either because the598  // dxil-data-scalarization pass is expected to run before this, but it's not599  // incorrect to apply this transformation to vector load/stores.600  if (!LoadStoreTy->isSingleValueType())601    return false;602 603  Type *ArrayTy;604  if (auto *GlobalVarPtrOp = dyn_cast<GlobalVariable>(PtrOp))605    ArrayTy = GlobalVarPtrOp->getValueType();606  else if (auto *AllocaPtrOp = dyn_cast<AllocaInst>(PtrOp))607    ArrayTy = AllocaPtrOp->getAllocatedType();608  else609    return false;610 611  if (!isa<ArrayType>(ArrayTy))612    return false;613 614  assert(ArrayTy->getArrayElementType() == LoadStoreTy &&615         "Expected array element type to be the same as to the scalar load or "616         "store type");617 618  Value *Zero = ConstantInt::get(Type::getInt32Ty(I.getContext()), 0);619  Value *GEP = GetElementPtrInst::Create(620      ArrayTy, PtrOp, {Zero, Zero}, GEPNoWrapFlags::all(), "", I.getIterator());621  I.setOperand(PtrOpIndex, GEP);622  return true;623}624 625namespace {626class DXILLegalizationPipeline {627 628public:629  DXILLegalizationPipeline() { initializeLegalizationPipeline(); }630 631  bool runLegalizationPipeline(Function &F) {632    bool MadeChange = false;633    SmallVector<Instruction *> ToRemove;634    DenseMap<Value *, Value *> ReplacedValues;635    for (int Stage = 0; Stage < NumStages; ++Stage) {636      ToRemove.clear();637      ReplacedValues.clear();638      for (auto &I : instructions(F)) {639        for (auto &LegalizationFn : LegalizationPipeline[Stage])640          MadeChange |= LegalizationFn(I, ToRemove, ReplacedValues);641      }642 643      for (auto *Inst : reverse(ToRemove))644        Inst->eraseFromParent();645    }646    return MadeChange;647  }648 649private:650  enum LegalizationStage { Stage1 = 0, Stage2 = 1, NumStages };651 652  using LegalizationFnTy =653      std::function<bool(Instruction &, SmallVectorImpl<Instruction *> &,654                         DenseMap<Value *, Value *> &)>;655 656  SmallVector<LegalizationFnTy> LegalizationPipeline[NumStages];657 658  void initializeLegalizationPipeline() {659    LegalizationPipeline[Stage1].push_back(upcastI8AllocasAndUses);660    LegalizationPipeline[Stage1].push_back(fixI8UseChain);661    LegalizationPipeline[Stage1].push_back(legalizeGetHighLowi64Bytes);662    LegalizationPipeline[Stage1].push_back(legalizeFreeze);663    LegalizationPipeline[Stage1].push_back(legalizeMemCpy);664    LegalizationPipeline[Stage1].push_back(legalizeMemSet);665    LegalizationPipeline[Stage1].push_back(updateFnegToFsub);666    // Note: legalizeGetHighLowi64Bytes and667    // downcastI64toI32InsertExtractElements both modify extractelement, so they668    // must run staggered stages. legalizeGetHighLowi64Bytes runs first b\c it669    // removes extractelements, reducing the number that670    // downcastI64toI32InsertExtractElements needs to handle.671    LegalizationPipeline[Stage2].push_back(672        downcastI64toI32InsertExtractElements);673    LegalizationPipeline[Stage2].push_back(legalizeScalarLoadStoreOnArrays);674  }675};676 677class DXILLegalizeLegacy : public FunctionPass {678 679public:680  bool runOnFunction(Function &F) override;681  DXILLegalizeLegacy() : FunctionPass(ID) {}682 683  static char ID; // Pass identification.684};685} // namespace686 687PreservedAnalyses DXILLegalizePass::run(Function &F,688                                        FunctionAnalysisManager &FAM) {689  DXILLegalizationPipeline DXLegalize;690  bool MadeChanges = DXLegalize.runLegalizationPipeline(F);691  if (!MadeChanges)692    return PreservedAnalyses::all();693  PreservedAnalyses PA;694  return PA;695}696 697bool DXILLegalizeLegacy::runOnFunction(Function &F) {698  DXILLegalizationPipeline DXLegalize;699  return DXLegalize.runLegalizationPipeline(F);700}701 702char DXILLegalizeLegacy::ID = 0;703 704INITIALIZE_PASS_BEGIN(DXILLegalizeLegacy, DEBUG_TYPE, "DXIL Legalizer", false,705                      false)706INITIALIZE_PASS_END(DXILLegalizeLegacy, DEBUG_TYPE, "DXIL Legalizer", false,707                    false)708 709FunctionPass *llvm::createDXILLegalizeLegacyPass() {710  return new DXILLegalizeLegacy();711}712