DxilCondenseResources.cpp

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// DxilCondenseResources.cpp                                                 //
// Copyright (C) Microsoft Corporation. All rights reserved.                 //
// This file is distributed under the University of Illinois Open Source     //
// License. See LICENSE.TXT for details.                                     //
//                                                                           //
// Provides a pass to make resource IDs zero-based and dense.                //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#include "dxc/DXIL/DxilInstructions.h"
#include "dxc/DXIL/DxilMetadataHelper.h"
#include "dxc/DXIL/DxilModule.h"
#include "dxc/DXIL/DxilOperations.h"
#include "dxc/DXIL/DxilResourceBinding.h"
#include "dxc/DXIL/DxilSignatureElement.h"
#include "dxc/DXIL/DxilTypeSystem.h"
#include "dxc/DXIL/DxilUtil.h"
#include "dxc/DxcBindingTable/DxcBindingTable.h"
#include "dxc/HLSL/DxilGenerationPass.h"
#include "dxc/HLSL/DxilSpanAllocator.h"
#include "dxc/HLSL/HLMatrixType.h"
#include "dxc/HLSL/HLModule.h"
#include "dxc/Support/Global.h"
#include "llvm/Analysis/DxilValueCache.h"

#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Utils/Local.h"
#include <memory>
#include <unordered_set>

using namespace llvm;
using namespace hlsl;

// Resource rangeID remap.
namespace {
struct ResourceID {
  DXIL::ResourceClass Class; // Resource class.
  unsigned ID;               // Resource ID, as specified on entry.

  bool operator<(const ResourceID &other) const {
    if (Class < other.Class)
      return true;
    if (Class > other.Class)
      return false;
    if (ID < other.ID)
      return true;
    return false;
  }
};

struct RemapEntry {
  ResourceID ResID;           // Resource identity, as specified on entry.
  DxilResourceBase *Resource; // In-memory resource representation.
  unsigned Index; // Index in resource vector - new ID for the resource.
};

typedef std::map<ResourceID, RemapEntry> RemapEntryCollection;

template <typename TResource>
void BuildRewrites(const std::vector<std::unique_ptr<TResource>> &Rs,
                   RemapEntryCollection &C) {
  const unsigned s = (unsigned)Rs.size();
  for (unsigned i = 0; i < s; ++i) {
    const std::unique_ptr<TResource> &R = Rs[i];
    if (R->GetID() != i) {
      ResourceID RId = {R->GetClass(), R->GetID()};
      RemapEntry RE = {RId, R.get(), i};
      C[RId] = RE;
    }
  }
}

// Build m_rewrites, returns 'true' if any rewrites are needed.
bool BuildRewriteMap(RemapEntryCollection &rewrites, DxilModule &DM) {
  BuildRewrites(DM.GetCBuffers(), rewrites);
  BuildRewrites(DM.GetSRVs(), rewrites);
  BuildRewrites(DM.GetUAVs(), rewrites);
  BuildRewrites(DM.GetSamplers(), rewrites);

  return !rewrites.empty();
}

} // namespace

class DxilResourceRegisterAllocator {
private:
  SpacesAllocator<unsigned, hlsl::DxilCBuffer> m_reservedCBufferRegisters;
  SpacesAllocator<unsigned, hlsl::DxilSampler> m_reservedSamplerRegisters;
  SpacesAllocator<unsigned, hlsl::DxilResource> m_reservedUAVRegisters;
  SpacesAllocator<unsigned, hlsl::DxilResource> m_reservedSRVRegisters;

  template <typename T>
  static void
  GatherReservedRegisters(const std::vector<std::unique_ptr<T>> &ResourceList,
                          SpacesAllocator<unsigned, T> &SAlloc) {
    for (auto &res : ResourceList) {
      if (res->IsAllocated()) {
        typename SpacesAllocator<unsigned, T>::Allocator &Alloc =
            SAlloc.Get(res->GetSpaceID());
        Alloc.ForceInsertAndClobber(res.get(), res->GetLowerBound(),
                                    res->GetUpperBound());
        if (res->IsUnbounded())
          Alloc.SetUnbounded(res.get());
      }
    }
  }

  template <typename T>
  static bool
  AllocateRegisters(LLVMContext &Ctx,
                    const std::vector<std::unique_ptr<T>> &resourceList,
                    SpacesAllocator<unsigned, T> &ReservedRegisters,
                    unsigned AutoBindingSpace) {
    bool bChanged = false;
    SpacesAllocator<unsigned, T> SAlloc;

    // Reserve explicitly allocated resources
    for (auto &res : resourceList) {
      const unsigned space = res->GetSpaceID();
      typename SpacesAllocator<unsigned, T>::Allocator &alloc =
          SAlloc.Get(space);
      typename SpacesAllocator<unsigned, T>::Allocator &reservedAlloc =
          ReservedRegisters.Get(space);

      if (res->IsAllocated()) {
        const unsigned reg = res->GetLowerBound();
        const T *conflict = nullptr;
        if (res->IsUnbounded()) {
          const T *unbounded = alloc.GetUnbounded();
          if (unbounded) {
            dxilutil::EmitErrorOnGlobalVariable(
                Ctx, dyn_cast<GlobalVariable>(res->GetGlobalSymbol()),
                Twine("more than one unbounded resource (") +
                    unbounded->GetGlobalName() + (" and ") +
                    res->GetGlobalName() + (") in space ") + Twine(space));
          } else {
            conflict = alloc.Insert(res.get(), reg, res->GetUpperBound());
            if (!conflict) {
              alloc.SetUnbounded(res.get());
              reservedAlloc.SetUnbounded(res.get());
            }
          }
        } else {
          conflict = alloc.Insert(res.get(), reg, res->GetUpperBound());
        }
        if (conflict) {
          dxilutil::EmitErrorOnGlobalVariable(
              Ctx, dyn_cast<GlobalVariable>(res->GetGlobalSymbol()),
              ((res->IsUnbounded()) ? Twine("unbounded ") : Twine("")) +
                  Twine("resource ") + res->GetGlobalName() +
                  Twine(" at register ") + Twine(reg) +
                  Twine(" overlaps with resource ") +
                  conflict->GetGlobalName() + Twine(" at register ") +
                  Twine(conflict->GetLowerBound()) + Twine(", space ") +
                  Twine(space));
        } else {
          // Also add this to the reserved (unallocatable) range, if it wasn't
          // already there.
          reservedAlloc.ForceInsertAndClobber(res.get(), res->GetLowerBound(),
                                              res->GetUpperBound());
        }
      }
    }

    // Allocate unallocated resources
    for (auto &res : resourceList) {
      if (res->IsAllocated())
        continue;

      unsigned space = res->GetSpaceID();
      if (space == UINT_MAX)
        space = AutoBindingSpace;
      typename SpacesAllocator<unsigned, T>::Allocator &alloc =
          SAlloc.Get(space);
      typename SpacesAllocator<unsigned, T>::Allocator &reservedAlloc =
          ReservedRegisters.Get(space);

      unsigned reg = 0;
      unsigned end = 0;
      bool allocateSpaceFound = false;
      if (res->IsUnbounded()) {
        if (alloc.GetUnbounded() != nullptr) {
          const T *unbounded = alloc.GetUnbounded();
          dxilutil::EmitErrorOnGlobalVariable(
              Ctx, dyn_cast<GlobalVariable>(res->GetGlobalSymbol()),
              Twine("more than one unbounded resource (") +
                  unbounded->GetGlobalName() + Twine(" and ") +
                  res->GetGlobalName() + Twine(") in space ") + Twine(space));
          continue;
        }

        if (reservedAlloc.FindForUnbounded(reg)) {
          end = UINT_MAX;
          allocateSpaceFound = true;
        }
      } else if (reservedAlloc.Find(res->GetRangeSize(), reg)) {
        end = reg + res->GetRangeSize() - 1;
        allocateSpaceFound = true;
      }

      if (allocateSpaceFound) {
        bool success = reservedAlloc.Insert(res.get(), reg, end) == nullptr;
        DXASSERT_NOMSG(success);

        success = alloc.Insert(res.get(), reg, end) == nullptr;
        DXASSERT_NOMSG(success);

        if (res->IsUnbounded()) {
          alloc.SetUnbounded(res.get());
          reservedAlloc.SetUnbounded(res.get());
        }

        res->SetLowerBound(reg);
        res->SetSpaceID(space);
        bChanged = true;
      } else {
        dxilutil::EmitErrorOnGlobalVariable(
            Ctx, dyn_cast<GlobalVariable>(res->GetGlobalSymbol()),
            ((res->IsUnbounded()) ? Twine("unbounded ") : Twine("")) +
                Twine("resource ") + res->GetGlobalName() +
                Twine(" could not be allocated"));
      }
    }

    return bChanged;
  }

public:
  void GatherReservedRegisters(DxilModule &DM) {
    // For backcompat with FXC, shader models 5.0 and below will not
    // auto-allocate resources at a register explicitely assigned to even an
    // unused resource.
    if (DM.GetLegacyResourceReservation()) {
      GatherReservedRegisters(DM.GetCBuffers(), m_reservedCBufferRegisters);
      GatherReservedRegisters(DM.GetSamplers(), m_reservedSamplerRegisters);
      GatherReservedRegisters(DM.GetUAVs(), m_reservedUAVRegisters);
      GatherReservedRegisters(DM.GetSRVs(), m_reservedSRVRegisters);
    }
  }

  bool AllocateRegisters(DxilModule &DM) {
    uint32_t AutoBindingSpace = DM.GetAutoBindingSpace();
    if (AutoBindingSpace == UINT_MAX) {
      // For libraries, we don't allocate unless AutoBindingSpace is set.
      if (DM.GetShaderModel()->IsLib())
        return false;
      // For shaders, we allocate in space 0 by default.
      AutoBindingSpace = 0;
    }

    bool bChanged = false;
    bChanged |= AllocateRegisters(DM.GetCtx(), DM.GetCBuffers(),
                                  m_reservedCBufferRegisters, AutoBindingSpace);
    bChanged |= AllocateRegisters(DM.GetCtx(), DM.GetSamplers(),
                                  m_reservedSamplerRegisters, AutoBindingSpace);
    bChanged |= AllocateRegisters(DM.GetCtx(), DM.GetUAVs(),
                                  m_reservedUAVRegisters, AutoBindingSpace);
    bChanged |= AllocateRegisters(DM.GetCtx(), DM.GetSRVs(),
                                  m_reservedSRVRegisters, AutoBindingSpace);
    return bChanged;
  }
};

bool llvm::AreDxilResourcesDense(llvm::Module *M,
                                 hlsl::DxilResourceBase **ppNonDense) {
  DxilModule &DM = M->GetOrCreateDxilModule();
  RemapEntryCollection rewrites;
  if (BuildRewriteMap(rewrites, DM)) {
    *ppNonDense = rewrites.begin()->second.Resource;
    return false;
  } else {
    *ppNonDense = nullptr;
    return true;
  }
}

static bool GetConstantLegalGepForSplitAlloca(GetElementPtrInst *gep,
                                              DxilValueCache *DVC,
                                              int64_t *ret) {
  if (gep->getNumIndices() != 2) {
    return false;
  }

  if (ConstantInt *Index0 = dyn_cast<ConstantInt>(gep->getOperand(1))) {
    if (Index0->getLimitedValue() != 0) {
      return false;
    }
  } else {
    return false;
  }

  if (ConstantInt *C = DVC->GetConstInt(gep->getOperand(2))) {
    int64_t index = C->getSExtValue();
    *ret = index;
    return true;
  }

  return false;
}

static bool LegalizeResourceArrays(Module &M, DxilValueCache *DVC) {
  SmallVector<AllocaInst *, 16> Allocas;

  bool Changed = false;

  // Find all allocas
  for (Function &F : M) {
    if (F.empty())
      continue;

    BasicBlock &BB = F.getEntryBlock();
    for (Instruction &I : BB) {
      if (AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
        Type *ty = AI->getAllocatedType();
        // Only handle single dimentional array. Since this pass runs after
        // MultiDimArrayToOneDimArray, it should handle all arrays.
        if (ty->isArrayTy() &&
            hlsl::dxilutil::IsHLSLResourceType(ty->getArrayElementType()))
          Allocas.push_back(AI);
      }
    }
  }

  SmallVector<AllocaInst *, 16> ScalarAllocas;
  std::unordered_map<GetElementPtrInst *, int64_t> ConstIndices;

  for (AllocaInst *AI : Allocas) {
    Type *ty = AI->getAllocatedType();
    Type *resType = ty->getArrayElementType();

    ScalarAllocas.clear();
    ConstIndices.clear();

    bool SplitAlloca = true;

    for (User *U : AI->users()) {
      if (GetElementPtrInst *gep = dyn_cast<GetElementPtrInst>(U)) {
        int64_t index = 0;
        if (!GetConstantLegalGepForSplitAlloca(gep, DVC, &index)) {
          SplitAlloca = false;
          break;
        }

        // Out of bounds. Out of bounds GEP's will trigger and error later.
        if (index < 0 || index >= (int64_t)ty->getArrayNumElements()) {
          SplitAlloca = false;
          Changed = true;
          dxilutil::EmitErrorOnInstruction(
              gep, "Accessing resource array with out-out-bounds index.");
        }
        ConstIndices[gep] = index;
      } else {
        SplitAlloca = false;
        break;
      }
    }

    if (SplitAlloca) {

      IRBuilder<> B(AI);
      ScalarAllocas.resize(ty->getArrayNumElements());

      for (auto it = AI->user_begin(), end = AI->user_end(); it != end;) {
        GetElementPtrInst *gep = cast<GetElementPtrInst>(*(it++));
        assert(ConstIndices.count(gep));
        int64_t idx = ConstIndices[gep];

        AllocaInst *ScalarAI = ScalarAllocas[idx];
        if (!ScalarAI) {
          ScalarAI = B.CreateAlloca(resType);
          ScalarAllocas[idx] = ScalarAI;
        }

        gep->replaceAllUsesWith(ScalarAI);
        gep->eraseFromParent();
      }

      AI->eraseFromParent();

      Changed = true;
    }
  }

  return Changed;
}

typedef std::unordered_map<std::string, DxilResourceBase *> ResourceMap;
template <typename T>
static inline void GatherResources(const std::vector<std::unique_ptr<T>> &List,
                                   ResourceMap *Map) {
  for (const std::unique_ptr<T> &ptr : List) {
    (*Map)[ptr->GetGlobalName()] = ptr.get();
  }
}

static bool LegalizeResources(Module &M, DxilValueCache *DVC) {

  bool Changed = false;

  Changed |= LegalizeResourceArrays(M, DVC);

  // Simple pass to collect resource PHI's
  SmallVector<PHINode *, 8> PHIs;

  for (Function &F : M) {
    for (BasicBlock &BB : F) {
      for (Instruction &I : BB) {
        if (PHINode *PN = dyn_cast<PHINode>(&I)) {
          if (hlsl::dxilutil::IsHLSLResourceType(PN->getType())) {
            PHIs.push_back(PN);
          }
        } else {
          break;
        }
      }
    }
  }

  SmallVector<Instruction *, 8> DCEWorklist;

  // Try to simplify those PHI's with DVC and collect them in DCEWorklist
  for (unsigned Attempt = 0, MaxAttempt = PHIs.size(); Attempt < MaxAttempt;
       Attempt++) {
    bool LocalChanged = false;
    for (unsigned i = 0; i < PHIs.size(); i++) {
      PHINode *PN = PHIs[i];
      if (Value *V = DVC->GetValue(PN)) {
        PN->replaceAllUsesWith(V);
        LocalChanged = true;
        DCEWorklist.push_back(PN);
        PHIs.erase(PHIs.begin() + i);
      } else {
        i++;
      }
    }

    Changed |= LocalChanged;
    if (!LocalChanged)
      break;
  }

  // Collect Resource GV loads
  for (GlobalVariable &GV : M.globals()) {
    Type *Ty = GV.getType()->getPointerElementType();
    while (Ty->isArrayTy())
      Ty = Ty->getArrayElementType();
    if (!hlsl::dxilutil::IsHLSLResourceType(Ty))
      continue;

    SmallVector<User *, 4> WorkList(GV.user_begin(), GV.user_end());
    while (WorkList.size()) {
      User *U = WorkList.pop_back_val();
      if (LoadInst *Load = dyn_cast<LoadInst>(U)) {
        DCEWorklist.push_back(Load);
      } else if (GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
        for (User *GepU : GEP->users())
          WorkList.push_back(GepU);
      }
    }
  }

  // Simple DCE
  while (DCEWorklist.size()) {
    Instruction *I = DCEWorklist.back();
    DCEWorklist.pop_back();
    if (llvm::isInstructionTriviallyDead(I)) {
      for (Use &Op : I->operands())
        if (Instruction *OpI = dyn_cast<Instruction>(Op.get()))
          DCEWorklist.push_back(OpI);
      I->eraseFromParent();
      // Remove the instruction from the worklist if it still exists in it.
      DCEWorklist.erase(std::remove(DCEWorklist.begin(), DCEWorklist.end(), I),
                        DCEWorklist.end());
      Changed = true;
    }
  }

  return Changed;
}

namespace {
class DxilLowerCreateHandleForLib : public ModulePass {
private:
  RemapEntryCollection m_rewrites;
  DxilModule *m_DM;
  bool m_HasDbgInfo;
  bool m_bIsLib;
  bool m_bLegalizationFailed;

public:
  static char ID; // Pass identification, replacement for typeid
  explicit DxilLowerCreateHandleForLib() : ModulePass(ID) {}

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<DxilValueCache>();
  }

  StringRef getPassName() const override {
    return "DXIL Lower createHandleForLib";
  }

  bool runOnModule(Module &M) override {
    DxilModule &DM = M.GetOrCreateDxilModule();
    m_DM = &DM;
    // Clear llvm used to remove unused resource.
    m_DM->ClearLLVMUsed();
    m_bIsLib = DM.GetShaderModel()->IsLib();
    m_bLegalizationFailed = false;

    FailOnPoisonResources();

    bool bChanged = false;
    if (DM.GetShaderModel()->IsSM66Plus()) {
      bChanged = PatchDynamicTBuffers(DM);
      SetNonUniformIndexForDynamicResource(DM);
    }

    unsigned numResources = DM.GetCBuffers().size() + DM.GetUAVs().size() +
                            DM.GetSRVs().size() + DM.GetSamplers().size();

    if (!numResources) {
      // Remove createHandleFromHandle when not a lib
      if (!m_bIsLib)
        RemoveCreateHandleFromHandle(DM);
      return false;
    }
    // Switch tbuffers to SRVs, as they have been treated as cbuffers up to this
    // point.
    if (DM.GetCBuffers().size())
      bChanged |= PatchTBuffers(DM);

    // Assign resource binding overrides.
    hlsl::ApplyBindingTableFromMetadata(DM);

    // Gather reserved resource registers while we still have
    // unused resources that might have explicit register assignments.
    DxilResourceRegisterAllocator ResourceRegisterAllocator;
    ResourceRegisterAllocator.GatherReservedRegisters(DM);

    // Remove unused resources.
    if (DM.GetUnusedResourceBinding() == UnusedResourceBinding::Strip)
      bChanged |= DM.RemoveResourcesWithUnusedSymbols();

    unsigned newResources = DM.GetCBuffers().size() + DM.GetUAVs().size() +
                            DM.GetSRVs().size() + DM.GetSamplers().size();

    if (0 == newResources)
      return bChanged;

    {
      DxilValueCache *DVC = &getAnalysis<DxilValueCache>();
      bool bLocalChanged = LegalizeResources(M, DVC);
      if (bLocalChanged &&
          DM.GetUnusedResourceBinding() == UnusedResourceBinding::Strip) {
        // Remove unused resources.
        bChanged |= DM.RemoveResourcesWithUnusedSymbols();
      }
      bChanged |= bLocalChanged;
    }

    bChanged |= ResourceRegisterAllocator.AllocateRegisters(DM);

    if (DM.GetUnusedResourceBinding() == UnusedResourceBinding::ReserveAll)
      bChanged |= DM.RemoveResourcesWithUnusedSymbols();

    // Fill in top-level CBuffer variable usage bit
    UpdateCBufferUsage();

    if (m_bIsLib &&
        DM.GetShaderModel()->GetMinor() == ShaderModel::kOfflineMinor)
      return bChanged;

    // Make sure no select on resource.
    bChanged |= RemovePhiOnResource();

    if (m_bLegalizationFailed)
      return bChanged;

    if (m_bIsLib) {
      if (DM.GetOP()->UseMinPrecision())
        bChanged |= UpdateStructTypeForLegacyLayout();

      return bChanged;
    }

    bChanged = true;

    // Load up debug information, to cross-reference values and the instructions
    // used to load them.
    m_HasDbgInfo = llvm::getDebugMetadataVersionFromModule(M) != 0;

    GenerateDxilResourceHandles();

    if (DM.GetOP()->UseMinPrecision())
      UpdateStructTypeForLegacyLayout();

    // Change resource symbol into undef.
    UpdateResourceSymbols();

    // Remove createHandleFromHandle when not a lib.
    RemoveCreateHandleFromHandle(DM);

    // Remove unused createHandleForLib functions.
    dxilutil::RemoveUnusedFunctions(M, DM.GetEntryFunction(),
                                    DM.GetPatchConstantFunction(), m_bIsLib);

    // Erase type annotations for structures no longer used
    DM.GetTypeSystem().EraseUnusedStructAnnotations();

    return bChanged;
  }

private:
  void FailOnPoisonResources();
  bool RemovePhiOnResource();
  void UpdateResourceSymbols();
  void ReplaceResourceUserWithHandle(DxilResource &res, LoadInst *load,
                                     Instruction *handle);
  void TranslateDxilResourceUses(DxilResourceBase &res);
  void GenerateDxilResourceHandles();
  bool UpdateStructTypeForLegacyLayout();
  // Switch CBuffer for SRV for TBuffers.
  bool PatchDynamicTBuffers(DxilModule &DM);
  bool PatchTBuffers(DxilModule &DM);
  void PatchTBufferUse(Value *V, DxilModule &DM, DenseSet<Value *> &patchedSet);
  void UpdateCBufferUsage();
  void SetNonUniformIndexForDynamicResource(DxilModule &DM);
  void RemoveCreateHandleFromHandle(DxilModule &DM);
};

} // namespace

// Phi on resource.
namespace {

typedef std::unordered_map<Value *, Value *> ValueToValueMap;
typedef llvm::SetVector<Value *> ValueSetVector;
typedef llvm::SmallVector<Value *, 4> IndexVector;
typedef std::unordered_map<Value *, IndexVector> ValueToIdxMap;

//#define SUPPORT_SELECT_ON_ALLOCA

// Errors:
class ResourceUseErrors {
  bool m_bErrorsReported;

public:
  ResourceUseErrors() : m_bErrorsReported(false) {}

  enum ErrorCode : unsigned int {
    // Collision between use of one resource GV and another.
    // All uses must be guaranteed to resolve to only one GV.
    // Additionally, when writing resource to alloca, all uses
    // of that alloca are considered resolving to a single GV.
    GVConflicts,

    // static global resources are disallowed for libraries at this time.
    // for non-library targets, they should have been eliminated already.
    StaticGVUsed,

    // user function calls with resource params or return type are
    // are currently disallowed for libraries.
    UserCallsWithResources,

    // When searching up from store pointer looking for alloca,
    // we encountered an unexpted value type
    UnexpectedValuesFromStorePointer,

    // Without SUPPORT_SELECT_ON_ALLOCA, phi/select on alloca based
    // pointer is disallowed, since this scenario is still untested.
    // This error also covers any other unknown alloca pointer uses.
    // Supported:
    // alloca (-> gep)? -> load -> ...
    // alloca (-> gep)? -> store.
    // Unsupported without SUPPORT_SELECT_ON_ALLOCA:
    // alloca (-> gep)? -> phi/select -> ...
    AllocaUserDisallowed,
    MismatchHandleAnnotation,
    MixDynamicResourceWithBindingResource,
    MismatchIsSampler,
#ifdef SUPPORT_SELECT_ON_ALLOCA
    // Conflict in select/phi between GV pointer and alloca pointer.  This
    // algorithm can't handle this case.
    AllocaSelectConflict,
#endif

    ErrorCodeCount
  };

  const StringRef ErrorText[ErrorCodeCount] = {
      "local resource not guaranteed to map to unique global resource.",
      "static global resource use is disallowed for library functions.",
      "exported library functions cannot have resource parameters or return "
      "value.",
      "internal error: unexpected instruction type when looking for alloca "
      "from store.",
      "phi/select disallowed on pointers to local resources.",
      "mismatch handle annotation",
      "possible mixing dynamic resource and binding resource",
      "merging sampler handle and resource handle",
#ifdef SUPPORT_SELECT_ON_ALLOCA
      ,
      "unable to resolve merge of global and local resource pointers."
#endif
  };

  ValueSetVector ErrorSets[ErrorCodeCount];

  // Ulitimately, the goal of ErrorUsers is to mark all create handles
  // so we don't try to report errors on them again later.
  std::unordered_set<Value *> ErrorUsers; // users of error values
  bool AddErrorUsers(Value *V) {
    auto it = ErrorUsers.insert(V);
    if (!it.second)
      return false; // already there
    if (isa<GEPOperator>(V) || isa<LoadInst>(V) || isa<PHINode>(V) ||
        isa<SelectInst>(V) || isa<AllocaInst>(V)) {
      for (auto U : V->users()) {
        AddErrorUsers(U);
      }
    } else if (isa<StoreInst>(V)) {
      AddErrorUsers(cast<StoreInst>(V)->getPointerOperand());
    }
    // create handle will be marked, but users not followed
    return true;
  }
  void ReportError(ErrorCode ec, Value *V) {
    DXASSERT_NOMSG(ec < ErrorCodeCount);
    if (!ErrorSets[ec].insert(V))
      return; // Error already reported
    AddErrorUsers(V);
    m_bErrorsReported = true;
    if (Instruction *I = dyn_cast<Instruction>(V)) {
      dxilutil::EmitErrorOnInstruction(I, ErrorText[ec]);
    } else {
      StringRef Name = V->getName();
      std::string escName;
      if (isa<Function>(V)) {
        llvm::raw_string_ostream os(escName);
        dxilutil::PrintEscapedString(Name, os);
        os.flush();
        Name = escName;
      }
      V->getContext().emitError(Twine(ErrorText[ec]) + " Value: " + Name);
    }
  }

  bool ErrorsReported() { return m_bErrorsReported; }
};

unsigned CountArrayDimensions(Type *Ty,
                              // Optionally collect dimensions
                              SmallVector<unsigned, 4> *dims = nullptr) {
  if (Ty->isPointerTy())
    Ty = Ty->getPointerElementType();
  unsigned dim = 0;
  if (dims)
    dims->clear();
  while (Ty->isArrayTy()) {
    if (dims)
      dims->push_back(Ty->getArrayNumElements());
    dim++;
    Ty = Ty->getArrayElementType();
  }
  return dim;
}

// Delete unused CleanupInsts, restarting when changed
// Return true if something was deleted
bool CleanupUnusedValues(std::unordered_set<Instruction *> &CleanupInsts) {
  //  - delete unused CleanupInsts, restarting when changed
  bool bAnyChanges = false;
  bool bChanged = false;
  do {
    bChanged = false;
    for (auto it = CleanupInsts.begin(); it != CleanupInsts.end();) {
      Instruction *I = *(it++);
      if (I->user_empty()) {
        // Add instructions operands CleanupInsts
        for (unsigned iOp = 0; iOp < I->getNumOperands(); iOp++) {
          if (Instruction *opI = dyn_cast<Instruction>(I->getOperand(iOp)))
            CleanupInsts.insert(opI);
        }
        I->eraseFromParent();
        CleanupInsts.erase(I);
        bChanged = true;
      }
    }
    if (bChanged)
      bAnyChanges = true;
  } while (bChanged);
  return bAnyChanges;
}

// Helper class for legalizing resource use
// Convert select/phi on resources to select/phi on index to GEP on GV.
// Convert resource alloca to index alloca.
// Assumes createHandleForLib has no select/phi
class LegalizeResourceUseHelper {
  // Change:
  //  gep1 = GEP gRes, i1
  //  res1 = load gep1
  //  gep2 = GEP gRes, i2
  //  gep3 = GEP gRes, i3
  //  gep4 = phi gep2, gep3           <-- handle select/phi on GEP
  //  res4 = load gep4
  //  res5 = phi res1, res4
  //  res6 = load GEP gRes, 23        <-- handle constant GepExpression
  //  res = select cnd2, res5, res6
  //  handle = createHandleForLib(res)
  // To:
  //  i4 = phi i2, i3
  //  i5 = phi i1, i4
  //  i6 = select cnd, i5, 23
  //  gep = GEP gRes, i6
  //  res = load gep
  //  handle = createHandleForLib(res)

  // Also handles alloca
  //  resArray = alloca [2 x Resource]
  //  gep1 = GEP gRes, i1
  //  res1 = load gep1
  //  gep2 = GEP gRes, i2
  //  gep3 = GEP gRes, i3
  //  phi4 = phi gep2, gep3
  //  res4 = load phi4
  //  gep5 = GEP resArray, 0
  //  gep6 = GEP resArray, 1
  //  store gep5, res1
  //  store gep6, res4
  //  gep7 = GEP resArray, i7   <-- dynamically index array
  //  res = load gep7
  //  handle = createHandleForLib(res)
  // Desired result:
  //  idxArray = alloca [2 x i32]
  //  phi4 = phi i2, i3
  //  gep5 = GEP idxArray, 0
  //  gep6 = GEP idxArray, 1
  //  store gep5, i1
  //  store gep6, phi4
  //  gep7 = GEP idxArray, i7
  //  gep8 = GEP gRes, gep7
  //  res = load gep8
  //  handle = createHandleForLib(res)

  // Also handles multi-dim resource index and multi-dim resource array allocas

  // Basic algorithm:
  // - recursively mark each GV user with GV (ValueToResourceGV)
  //  - verify only one GV used for any given value
  // - handle allocas by searching up from store for alloca
  //  - then recursively mark alloca users
  // - ResToIdxReplacement keeps track of vector of indices that
  //   will be used to replace a given resource value or pointer
  // - Next, create selects/phis for indices corresponding to
  //   selects/phis on resource pointers or values.
  //  - leave incoming index values undef for now
  // - Create index allocas to replace resource allocas
  // - Create GEPs on index allocas to replace GEPs on resource allocas
  // - Create index loads on index allocas to replace loads on resource alloca
  // GEP
  // - Fill in replacements for GEPs on resource GVs
  //  - copy replacement index vectors to corresponding loads
  // - Create index stores to replace resource stores to alloca/GEPs
  // - Update selects/phis incoming index values
  // - SimplifyMerges: replace index phis/selects on same value with that value
  //  - RemappedValues[phi/select] set to replacement value
  //  - use LookupValue from now on when reading from ResToIdxReplacement
  // - Update handles by replacing load/GEP chains that go through select/phi
  //   with direct GV GEP + load, with select/phi on GEP indices instead.

public:
  ResourceUseErrors m_Errors;

  ValueToValueMap ValueToResourceGV;
  ValueToIdxMap ResToIdxReplacement;
  // Value sets we can use to iterate
  ValueSetVector Selects, GEPs, Stores, Handles;
  ValueSetVector Allocas, AllocaGEPs, AllocaLoads;
#ifdef SUPPORT_SELECT_ON_ALLOCA
  ValueSetVector AllocaSelects;
#endif

  std::unordered_set<Value *> NonUniformSet;

  // New index selects created by pass, so we can try simplifying later
  ValueSetVector NewSelects;

  // Values that have been replaced with other values need remapping
  ValueToValueMap RemappedValues;

  // Things to clean up if no users:
  std::unordered_set<Instruction *> CleanupInsts;

  GlobalVariable *LookupResourceGV(Value *V) {
    auto itGV = ValueToResourceGV.find(V);
    if (itGV == ValueToResourceGV.end())
      return nullptr;
    return cast<GlobalVariable>(itGV->second);
  }

  // Follow RemappedValues, return input if not remapped
  Value *LookupValue(Value *V) {
    auto it = RemappedValues.find(V);
    SmallPtrSet<Value *, 4> visited;
    while (it != RemappedValues.end()) {
      // Cycles should not happen, but are bad if they do.
      if (visited.count(it->second)) {
        // When remapping values to be replaced, we add them to RemappedValues
        // so we don't use dead values stored in other sets/maps. Circular
        // remaps that should not happen
        DXASSERT(false, "otherwise, circular remapping");
        llvm_unreachable("cycles detected in value remapping");
        break;
      }
      V = it->second;
      it = RemappedValues.find(V);
      if (it != RemappedValues.end())
        visited.insert(V);
    }
    return V;
  }

  bool AreLoadUsersTrivial(LoadInst *LI) {
    for (auto U : LI->users()) {
      if (CallInst *CI = dyn_cast<CallInst>(U)) {
        Function *F = CI->getCalledFunction();
        DxilModule &DM = F->getParent()->GetDxilModule();
        hlsl::OP *hlslOP = DM.GetOP();
        if (hlslOP->IsDxilOpFunc(F)) {
          hlsl::OP::OpCodeClass opClass;
          if (hlslOP->GetOpCodeClass(F, opClass) &&
              opClass == DXIL::OpCodeClass::CreateHandleForLib) {
            continue;
          }
        }
      }
      return false;
    }
    return true;
  }

  // This is used to quickly skip the common case where no work is needed
  bool AreGEPUsersTrivial(GEPOperator *GEP) {
    if (GlobalVariable *GV = LookupResourceGV(GEP)) {
      if (GEP->getPointerOperand() != LookupResourceGV(GEP))
        return false;
    }
    for (auto U : GEP->users()) {
      if (LoadInst *LI = dyn_cast<LoadInst>(U)) {
        if (AreLoadUsersTrivial(LI))
          continue;
      }
      return false;
    }
    return true;
  }

  // AssignResourceGVFromStore is used on pointer being stored to.
  // Follow GEP/Phi/Select up to Alloca, then CollectResourceGVUsers on Alloca
  void AssignResourceGVFromStore(GlobalVariable *GV, Value *V,
                                 SmallPtrSet<Value *, 4> &visited,
                                 bool bNonUniform) {
    // Prevent cycles as we search up
    if (visited.count(V) != 0)
      return;
    // Verify and skip if already processed
    auto it = ValueToResourceGV.find(V);
    if (it != ValueToResourceGV.end()) {
      if (it->second != GV) {
        m_Errors.ReportError(ResourceUseErrors::GVConflicts, V);
      }
      return;
    }
    if (AllocaInst *AI = dyn_cast<AllocaInst>(V)) {
      CollectResourceGVUsers(GV, AI, /*bAlloca*/ true, bNonUniform);
      return;
    } else if (GEPOperator *GEP = dyn_cast<GEPOperator>(V)) {
      // follow the pointer up
      AssignResourceGVFromStore(GV, GEP->getPointerOperand(), visited,
                                bNonUniform);
      return;
    } else if (PHINode *Phi = dyn_cast<PHINode>(V)) {
#ifdef SUPPORT_SELECT_ON_ALLOCA
      // follow all incoming values
      for (auto it : Phi->operand_values())
        AssignResourceGVFromStore(GV, it, visited, bNonUniform);
#else
      m_Errors.ReportError(ResourceUseErrors::AllocaUserDisallowed, V);
#endif
      return;
    } else if (SelectInst *Sel = dyn_cast<SelectInst>(V)) {
#ifdef SUPPORT_SELECT_ON_ALLOCA
      // follow all incoming values
      AssignResourceGVFromStore(GV, Sel->getTrueValue(), visited, bNonUniform);
      AssignResourceGVFromStore(GV, Sel->getFalseValue(), visited, bNonUniform);
#else
      m_Errors.ReportError(ResourceUseErrors::AllocaUserDisallowed, V);
#endif
      return;
    } else if (isa<GlobalVariable>(V) &&
               cast<GlobalVariable>(V)->getLinkage() ==
                   GlobalVariable::LinkageTypes::InternalLinkage) {
      // this is writing to global static, which is disallowed at this point.
      m_Errors.ReportError(ResourceUseErrors::StaticGVUsed, V);
      return;
    } else {
      // Most likely storing to output parameter
      m_Errors.ReportError(ResourceUseErrors::UserCallsWithResources, V);
      return;
    }
    return;
  }

  // Recursively mark values with GV, following users.
  // Starting value V should be GV itself.
  // Returns true if value/uses reference no other GV in map.
  void CollectResourceGVUsers(GlobalVariable *GV, Value *V,
                              bool bAlloca = false, bool bNonUniform = false) {
    // Recursively tag value V and its users as using GV.
    auto it = ValueToResourceGV.find(V);
    if (it != ValueToResourceGV.end()) {
      if (it->second != GV) {
        m_Errors.ReportError(ResourceUseErrors::GVConflicts, V);
#ifdef SUPPORT_SELECT_ON_ALLOCA
      } else {
        // if select/phi, make sure bAlloca is consistent
        if (isa<PHINode>(V) || isa<SelectInst>(V))
          if ((bAlloca && AllocaSelects.count(V) == 0) ||
              (!bAlloca && Selects.count(V) == 0))
            m_Errors.ReportError(ResourceUseErrors::AllocaSelectConflict, V);
#endif
      }
      return;
    }
    ValueToResourceGV[V] = GV;
    if (GV == V) {
      // Just add and recurse users
      // make sure bAlloca is clear for users
      bAlloca = false;
    } else if (GEPOperator *GEP = dyn_cast<GEPOperator>(V)) {
      if (bAlloca)
        AllocaGEPs.insert(GEP);
      else if (!AreGEPUsersTrivial(GEP))
        GEPs.insert(GEP);
      else
        return; // Optimization: skip trivial GV->GEP->load->createHandle
      if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP)) {
        if (DxilMDHelper::IsMarkedNonUniform(GEPInst))
          bNonUniform = true;
      }
    } else if (LoadInst *LI = dyn_cast<LoadInst>(V)) {
      if (bAlloca)
        AllocaLoads.insert(LI);
      // clear bAlloca for users
      bAlloca = false;
      if (bNonUniform)
        NonUniformSet.insert(LI);
    } else if (StoreInst *SI = dyn_cast<StoreInst>(V)) {
      Stores.insert(SI);
      if (!bAlloca) {
        // Find and mark allocas this store could be storing to
        SmallPtrSet<Value *, 4> visited;
        AssignResourceGVFromStore(GV, SI->getPointerOperand(), visited,
                                  bNonUniform);
      }
      return;
    } else if (PHINode *Phi = dyn_cast<PHINode>(V)) {
      if (bAlloca) {
#ifdef SUPPORT_SELECT_ON_ALLOCA
        AllocaSelects.insert(Phi);
#else
        m_Errors.ReportError(ResourceUseErrors::AllocaUserDisallowed, V);
#endif
      } else {
        Selects.insert(Phi);
      }
    } else if (SelectInst *Sel = dyn_cast<SelectInst>(V)) {
      if (bAlloca) {
#ifdef SUPPORT_SELECT_ON_ALLOCA
        AllocaSelects.insert(Sel);
#else
        m_Errors.ReportError(ResourceUseErrors::AllocaUserDisallowed, V);
#endif
      } else {
        Selects.insert(Sel);
      }
    } else if (AllocaInst *AI = dyn_cast<AllocaInst>(V)) {
      Allocas.insert(AI);
      // set bAlloca for users
      bAlloca = true;
    } else if (Constant *C = dyn_cast<Constant>(V)) {
      // skip @llvm.used entry
      return;
    } else if (BitCastInst *BCI = dyn_cast<BitCastInst>(V)) {
      DXASSERT(onlyUsedByLifetimeMarkers(BCI),
               "expected bitcast to only be used by lifetime intrinsics");
      return;
    } else if (bAlloca) {
      m_Errors.ReportError(ResourceUseErrors::AllocaUserDisallowed, V);
    } else {
      // Must be createHandleForLib or user function call.
      CallInst *CI = cast<CallInst>(V);
      Function *F = CI->getCalledFunction();
      DxilModule &DM = GV->getParent()->GetDxilModule();
      hlsl::OP *hlslOP = DM.GetOP();
      if (hlslOP->IsDxilOpFunc(F)) {
        hlsl::OP::OpCodeClass opClass;
        if (hlslOP->GetOpCodeClass(F, opClass) &&
            (opClass == DXIL::OpCodeClass::CreateHandleForLib)) {
          Handles.insert(CI);
          if (bNonUniform)
            NonUniformSet.insert(CI);
          return;
        }
      }
      // This could be user call with resource param, which is disallowed for
      // lib_6_3
      m_Errors.ReportError(ResourceUseErrors::UserCallsWithResources, V);
      return;
    }

    // Recurse users
    for (auto U : V->users())
      CollectResourceGVUsers(GV, U, bAlloca, bNonUniform);
    return;
  }

  // Remove conflicting values from sets before
  // transforming the remainder.
  void RemoveConflictingValue(Value *V) {
    bool bRemoved = false;
    if (isa<GEPOperator>(V)) {
      bRemoved = GEPs.remove(V) || AllocaGEPs.remove(V);
    } else if (isa<LoadInst>(V)) {
      bRemoved = AllocaLoads.remove(V);
    } else if (isa<StoreInst>(V)) {
      bRemoved = Stores.remove(V);
    } else if (isa<PHINode>(V) || isa<SelectInst>(V)) {
      bRemoved = Selects.remove(V);
#ifdef SUPPORT_SELECT_ON_ALLOCA
      bRemoved |= AllocaSelects.remove(V);
#endif
    } else if (isa<AllocaInst>(V)) {
      bRemoved = Allocas.remove(V);
    } else if (isa<CallInst>(V)) {
      bRemoved = Handles.remove(V);
      return; // don't recurse
    }
    if (bRemoved) {
      // Recurse users
      for (auto U : V->users())
        RemoveConflictingValue(U);
    }
  }
  void RemoveConflicts() {
    for (auto V : m_Errors.ErrorSets[ResourceUseErrors::GVConflicts]) {
      RemoveConflictingValue(V);
      ValueToResourceGV.erase(V);
    }
  }

  void CreateSelects() {
    if (Selects.empty()
#ifdef SUPPORT_SELECT_ON_ALLOCA
        && AllocaSelects.empty()
#endif
    )
      return;
    LLVMContext &Ctx =
#ifdef SUPPORT_SELECT_ON_ALLOCA
        Selects.empty() ? AllocaSelects[0]->getContext() :
#endif
                        Selects[0]->getContext();
    Type *i32Ty = IntegerType::getInt32Ty(Ctx);
#ifdef SUPPORT_SELECT_ON_ALLOCA
    for (auto &SelectSet : {Selects, AllocaSelects}) {
      bool bAlloca = !(&SelectSet == &Selects);
#else
    for (auto &SelectSet : {Selects}) {
#endif
      for (auto pValue : SelectSet) {
        Type *SelectTy = i32Ty;
#ifdef SUPPORT_SELECT_ON_ALLOCA
        // For alloca case, type needs to match dimensionality of incoming value
        if (bAlloca) {
          // TODO: Not sure if this case will actually work
          //      (or whether it can even be generated from HLSL)
          Type *Ty = pValue->getType();
          SmallVector<unsigned, 4> dims;
          unsigned dim = CountArrayDimensions(Ty, &dims);
          for (unsigned i = 0; i < dim; i++)
            SelectTy = ArrayType::get(SelectTy, (uint64_t)dims[dim - i - 1]);
          if (Ty->isPointerTy())
            SelectTy = PointerType::get(SelectTy, 0);
        }
#endif
        Value *UndefValue = UndefValue::get(SelectTy);
        if (PHINode *Phi = dyn_cast<PHINode>(pValue)) {
          GlobalVariable *GV = LookupResourceGV(Phi);
          if (!GV)
            continue; // skip value removed due to conflict
          IRBuilder<> PhiBuilder(Phi);
          unsigned gvDim = CountArrayDimensions(GV->getType());
          IndexVector &idxVector = ResToIdxReplacement[Phi];
          idxVector.resize(gvDim, nullptr);
          unsigned numIncoming = Phi->getNumIncomingValues();
          for (unsigned i = 0; i < gvDim; i++) {
            PHINode *newPhi = PhiBuilder.CreatePHI(SelectTy, numIncoming);
            NewSelects.insert(newPhi);
            idxVector[i] = newPhi;
            for (unsigned j = 0; j < numIncoming; j++) {
              // Set incoming values to undef until next pass
              newPhi->addIncoming(UndefValue, Phi->getIncomingBlock(j));
            }
          }
        } else if (SelectInst *Sel = dyn_cast<SelectInst>(pValue)) {
          GlobalVariable *GV = LookupResourceGV(Sel);
          if (!GV)
            continue; // skip value removed due to conflict
          IRBuilder<> Builder(Sel);
          unsigned gvDim = CountArrayDimensions(GV->getType());
          IndexVector &idxVector = ResToIdxReplacement[Sel];
          idxVector.resize(gvDim, nullptr);
          for (unsigned i = 0; i < gvDim; i++) {
            Value *newSel = Builder.CreateSelect(Sel->getCondition(),
                                                 UndefValue, UndefValue);
            NewSelects.insert(newSel);
            idxVector[i] = newSel;
          }
        } else {
          DXASSERT(false, "otherwise, non-select/phi in Selects set");
        }
      }
    }
  }

  // Create index allocas to replace resource allocas
  void CreateIndexAllocas() {
    if (Allocas.empty())
      return;
    Type *i32Ty = IntegerType::getInt32Ty(Allocas[0]->getContext());
    for (auto pValue : Allocas) {
      AllocaInst *pAlloca = cast<AllocaInst>(pValue);
      GlobalVariable *GV = LookupResourceGV(pAlloca);
      if (!GV)
        continue; // skip value removed due to conflict
      IRBuilder<> AllocaBuilder(pAlloca);
      unsigned gvDim = CountArrayDimensions(GV->getType());
      SmallVector<unsigned, 4> dimVector;
      unsigned allocaTyDim =
          CountArrayDimensions(pAlloca->getType(), &dimVector);
      Type *pIndexType = i32Ty;
      for (unsigned i = 0; i < allocaTyDim; i++) {
        pIndexType = ArrayType::get(pIndexType, dimVector[allocaTyDim - i - 1]);
      }
      Value *arraySize = pAlloca->getArraySize();
      IndexVector &idxVector = ResToIdxReplacement[pAlloca];
      idxVector.resize(gvDim, nullptr);
      for (unsigned i = 0; i < gvDim; i++) {
        AllocaInst *pAlloca = AllocaBuilder.CreateAlloca(pIndexType, arraySize);
        pAlloca->setAlignment(4);
        idxVector[i] = pAlloca;
      }
    }
  }

  // Add corresponding GEPs for index allocas
  IndexVector &ReplaceAllocaGEP(GetElementPtrInst *GEP) {
    IndexVector &idxVector = ResToIdxReplacement[GEP];
    if (!idxVector.empty())
      return idxVector;

    Value *Ptr = GEP->getPointerOperand();

    // Recurse for partial GEPs
    IndexVector &ptrIndices =
        isa<GetElementPtrInst>(Ptr)
            ? ReplaceAllocaGEP(cast<GetElementPtrInst>(Ptr))
            : ResToIdxReplacement[Ptr];

    IRBuilder<> Builder(GEP);
    SmallVector<Value *, 4> gepIndices;
    for (auto it = GEP->idx_begin(), idxEnd = GEP->idx_end(); it != idxEnd;
         it++)
      gepIndices.push_back(*it);
    idxVector.resize(ptrIndices.size(), nullptr);
    for (unsigned i = 0; i < ptrIndices.size(); i++) {
      idxVector[i] = Builder.CreateInBoundsGEP(ptrIndices[i], gepIndices);
    }
    return idxVector;
  }

  void ReplaceAllocaGEPs() {
    for (auto V : AllocaGEPs) {
      ReplaceAllocaGEP(cast<GetElementPtrInst>(V));
    }
  }

  void ReplaceAllocaLoads() {
    for (auto V : AllocaLoads) {
      LoadInst *LI = cast<LoadInst>(V);
      Value *Ptr = LI->getPointerOperand();
      IRBuilder<> Builder(LI);
      IndexVector &idxVector = ResToIdxReplacement[V];
      IndexVector &ptrIndices = ResToIdxReplacement[Ptr];
      idxVector.resize(ptrIndices.size(), nullptr);
      for (unsigned i = 0; i < ptrIndices.size(); i++) {
        idxVector[i] = Builder.CreateLoad(ptrIndices[i]);
      }
    }
  }

  // Add GEP to ResToIdxReplacement with indices from incoming + GEP
  IndexVector &ReplaceGVGEPs(GEPOperator *GEP) {
    IndexVector &idxVector = ResToIdxReplacement[GEP];
    // Skip if already done
    // (we recurse into partial GEP and iterate all GEPs)
    if (!idxVector.empty())
      return idxVector;

    Type *i32Ty = IntegerType::getInt32Ty(GEP->getContext());
    Constant *Zero = Constant::getIntegerValue(i32Ty, APInt(32, 0));

    Value *Ptr = GEP->getPointerOperand();

    unsigned idx = 0;
    if (GlobalVariable *GV = dyn_cast<GlobalVariable>(Ptr)) {
      unsigned gvDim = CountArrayDimensions(GV->getType());
      idxVector.resize(gvDim, Zero);
    } else if (isa<GEPOperator>(Ptr) || isa<PHINode>(Ptr) ||
               isa<SelectInst>(Ptr)) {
      // Recurse for partial GEPs
      IndexVector &ptrIndices = isa<GEPOperator>(Ptr)
                                    ? ReplaceGVGEPs(cast<GEPOperator>(Ptr))
                                    : ResToIdxReplacement[Ptr];
      unsigned ptrDim = CountArrayDimensions(Ptr->getType());
      unsigned gvDim = ptrIndices.size();
      DXASSERT(
          ptrDim <= gvDim,
          "otherwise incoming pointer has more dimensions than associated GV");
      unsigned gepStart = gvDim - ptrDim;
      // Copy indices and add ours
      idxVector.resize(ptrIndices.size(), Zero);
      for (; idx < gepStart; idx++)
        idxVector[idx] = ptrIndices[idx];
    }
    if (GEP->hasIndices()) {
      auto itIdx = GEP->idx_begin();
      ++itIdx; // Always skip leading zero (we don't support GV+n pointer arith)
      while (itIdx != GEP->idx_end())
        idxVector[idx++] = *itIdx++;
    }
    return idxVector;
  }

  // Add GEPs to ResToIdxReplacement and update loads
  void ReplaceGVGEPs() {
    if (GEPs.empty())
      return;
    for (auto V : GEPs) {
      GEPOperator *GEP = cast<GEPOperator>(V);
      IndexVector &gepVector = ReplaceGVGEPs(GEP);
      for (auto U : GEP->users()) {
        if (LoadInst *LI = dyn_cast<LoadInst>(U)) {
          // Just copy incoming indices
          ResToIdxReplacement[LI] = gepVector;
        }
      }
    }
  }

  // Create new index stores for incoming indices
  void ReplaceStores() {
    // generate stores of incoming indices to corresponding index pointers
    if (Stores.empty())
      return;
    Type *i32Ty = IntegerType::getInt32Ty(Stores[0]->getContext());
    for (auto V : Stores) {
      StoreInst *SI = cast<StoreInst>(V);
      IRBuilder<> Builder(SI);
      IndexVector &idxVector = ResToIdxReplacement[SI];
      Value *Ptr = SI->getPointerOperand();
      Value *Val = SI->getValueOperand();
      IndexVector &ptrIndices = ResToIdxReplacement[Ptr];
      IndexVector &valIndices = ResToIdxReplacement[Val];
      // If Val is not found, it is treated as an undef value that will
      // translate to an undef index, which may still be valid if it's never
      // used.
      Value *UndefIndex =
          valIndices.size() > 0 ? nullptr : UndefValue::get(i32Ty);
      DXASSERT_NOMSG(valIndices.size() == 0 ||
                     ptrIndices.size() == valIndices.size());
      idxVector.resize(ptrIndices.size(), nullptr);
      for (unsigned i = 0; i < idxVector.size(); i++) {
        idxVector[i] = Builder.CreateStore(
            UndefIndex ? UndefIndex : valIndices[i], ptrIndices[i]);
      }
    }
  }

  // For each Phi/Select: update matching incoming values for new phis
  void UpdateSelects() {
    if (Selects.empty())
      return;
    Type *i32Ty = IntegerType::getInt32Ty(Selects[0]->getContext());
    for (auto V : Selects) {
      // update incoming index values corresponding to incoming resource values
      IndexVector &idxVector = ResToIdxReplacement[V];
      Instruction *I = cast<Instruction>(V);
      unsigned numOperands = I->getNumOperands();
      unsigned startOp = isa<PHINode>(V) ? 0 : 1;
      for (unsigned iOp = startOp; iOp < numOperands; iOp++) {
        Value *Val = I->getOperand(iOp);
        IndexVector &incomingIndices = ResToIdxReplacement[Val];
        // If Val is not found, it is treated as an undef value that will
        // translate to an undef index, which may still be valid if it's never
        // used.
        Value *UndefIndex =
            incomingIndices.size() > 0 ? nullptr : UndefValue::get(i32Ty);
        DXASSERT_NOMSG(incomingIndices.size() == 0 ||
                       idxVector.size() == incomingIndices.size());
        for (unsigned i = 0; i < idxVector.size(); i++) {
          // must be instruction (phi/select)
          Instruction *indexI = cast<Instruction>(idxVector[i]);
          indexI->setOperand(iOp, UndefIndex ? UndefIndex : incomingIndices[i]);
        }

        // Now clear incoming operand (adding to cleanup) to break cycles
        if (Instruction *OpI = dyn_cast<Instruction>(I->getOperand(iOp)))
          CleanupInsts.insert(OpI);
        I->setOperand(iOp, UndefValue::get(I->getType()));
      }
    }
  }

  // ReplaceHandles
  //  - iterate handles
  //    - insert GEP using new indices associated with resource value
  //    - load resource from new GEP
  //    - replace resource use in createHandleForLib with new load
  // Assumes: no users of handle are phi/select or store
  void ReplaceHandles() {
    if (Handles.empty())
      return;
    Type *i32Ty = IntegerType::getInt32Ty(Handles[0]->getContext());
    Constant *Zero = Constant::getIntegerValue(i32Ty, APInt(32, 0));
    for (auto V : Handles) {
      CallInst *CI = cast<CallInst>(V);
      DxilInst_CreateHandleForLib createHandle(CI);
      Value *res = createHandle.get_Resource();
      // Skip extra work if nothing between load and create handle
      if (LoadInst *LI = dyn_cast<LoadInst>(res)) {
        Value *Ptr = LI->getPointerOperand();
        if (GEPOperator *GEP = dyn_cast<GEPOperator>(Ptr))
          Ptr = GEP->getPointerOperand();
        if (isa<GlobalVariable>(Ptr))
          continue;
      }
      GlobalVariable *GV = LookupResourceGV(res);
      if (!GV)
        continue; // skip value removed due to conflict
      IRBuilder<> Builder(CI);
      IndexVector &idxVector = ResToIdxReplacement[res];
      DXASSERT(idxVector.size() == CountArrayDimensions(GV->getType()),
               "replacements empty or invalid");
      SmallVector<Value *, 4> gepIndices;
      gepIndices.push_back(Zero);
      for (auto idxVal : idxVector)
        gepIndices.push_back(LookupValue(idxVal));
      Value *GEP = Builder.CreateInBoundsGEP(GV, gepIndices);
      // Mark new GEP instruction non-uniform if necessary
      if (NonUniformSet.count(res) != 0 || NonUniformSet.count(CI) != 0)
        if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP))
          DxilMDHelper::MarkNonUniform(GEPInst);
      LoadInst *LI = Builder.CreateLoad(GEP);
      createHandle.set_Resource(LI);
      if (Instruction *resI = dyn_cast<Instruction>(res))
        CleanupInsts.insert(resI);
    }
  }

  void SimplifyMerges() {
    // Loop if changed
    bool bChanged = false;
    do {
      bChanged = false;
      for (auto V : NewSelects) {
        if (LookupValue(V) != V)
          continue;
        Instruction *I = cast<Instruction>(V);
        unsigned startOp = isa<PHINode>(I) ? 0 : 1;
        Value *newV = dxilutil::MergeSelectOnSameValue(
            cast<Instruction>(V), startOp, I->getNumOperands());
        if (newV) {
          RemappedValues[V] = newV;
          bChanged = true;
        }
      }
    } while (bChanged);
  }

  void CleanupDeadInsts() {
    // Assuming everything was successful:
    // delete stores to allocas to remove cycles
    for (auto V : Stores) {
      StoreInst *SI = cast<StoreInst>(V);
      if (Instruction *I = dyn_cast<Instruction>(SI->getValueOperand()))
        CleanupInsts.insert(I);
      if (Instruction *I = dyn_cast<Instruction>(SI->getPointerOperand()))
        CleanupInsts.insert(I);
      SI->eraseFromParent();
    }
    CleanupUnusedValues(CleanupInsts);
  }

  void VerifyComplete(DxilModule &DM) {
    // Check that all handles now resolve to a global variable, otherwise,
    // they are likely loading from resource function parameter, which
    // is disallowed.
    hlsl::OP *hlslOP = DM.GetOP();
    for (Function &F : DM.GetModule()->functions()) {
      if (hlslOP->IsDxilOpFunc(&F)) {
        hlsl::OP::OpCodeClass opClass;
        if (hlslOP->GetOpCodeClass(&F, opClass) &&
            opClass == DXIL::OpCodeClass::CreateHandleForLib) {
          for (auto U : F.users()) {
            CallInst *CI = cast<CallInst>(U);
            if (m_Errors.ErrorUsers.count(CI))
              continue; // Error already reported
            DxilInst_CreateHandleForLib createHandle(CI);
            Value *res = createHandle.get_Resource();
            LoadInst *LI = dyn_cast<LoadInst>(res);
            if (LI) {
              Value *Ptr = LI->getPointerOperand();
              if (GEPOperator *GEP = dyn_cast<GEPOperator>(Ptr))
                Ptr = GEP->getPointerOperand();
              if (isa<GlobalVariable>(Ptr))
                continue;
            }
            // handle wasn't processed
            // Right now, the most likely cause is user call with resources, but
            // this should be updated if there are other reasons for this to
            // happen.
            m_Errors.ReportError(ResourceUseErrors::UserCallsWithResources, U);
          }
        }
      }
    }
  }

  // Fix resource global variable properties to external constant
  bool SetExternalConstant(GlobalVariable *GV) {
    if (GV->hasInitializer() || !GV->isConstant() ||
        GV->getLinkage() != GlobalVariable::LinkageTypes::ExternalLinkage) {
      GV->setInitializer(nullptr);
      GV->setConstant(true);
      GV->setLinkage(GlobalVariable::LinkageTypes::ExternalLinkage);
      return true;
    }
    return false;
  }

  bool CollectResources(DxilModule &DM) {
    bool bChanged = false;
    for (const auto &res : DM.GetCBuffers()) {
      if (GlobalVariable *GV =
              dyn_cast<GlobalVariable>(res->GetGlobalSymbol())) {
        bChanged |= SetExternalConstant(GV);
        CollectResourceGVUsers(GV, GV);
      }
    }
    for (const auto &res : DM.GetSRVs()) {
      if (GlobalVariable *GV =
              dyn_cast<GlobalVariable>(res->GetGlobalSymbol())) {
        bChanged |= SetExternalConstant(GV);
        CollectResourceGVUsers(GV, GV);
      }
    }
    for (const auto &res : DM.GetUAVs()) {
      if (GlobalVariable *GV =
              dyn_cast<GlobalVariable>(res->GetGlobalSymbol())) {
        bChanged |= SetExternalConstant(GV);
        CollectResourceGVUsers(GV, GV);
      }
    }
    for (const auto &res : DM.GetSamplers()) {
      if (GlobalVariable *GV =
              dyn_cast<GlobalVariable>(res->GetGlobalSymbol())) {
        bChanged |= SetExternalConstant(GV);
        CollectResourceGVUsers(GV, GV);
      }
    }
    return bChanged;
  }

  void DoTransform(hlsl::OP *hlslOP) {
    RemoveConflicts();
    CreateSelects();
    CreateIndexAllocas();
    ReplaceAllocaGEPs();
    ReplaceAllocaLoads();
    ReplaceGVGEPs();
    ReplaceStores();
    UpdateSelects();
    SimplifyMerges();
    ReplaceHandles();
    if (!m_Errors.ErrorsReported())
      CleanupDeadInsts();
  }

  bool ErrorsReported() { return m_Errors.ErrorsReported(); }

  bool runOnModule(llvm::Module &M) {
    DxilModule &DM = M.GetOrCreateDxilModule();
    hlsl::OP *hlslOP = DM.GetOP();

    bool bChanged = CollectResources(DM);

    // If no selects or allocas are involved, there isn't anything to do
    if (Selects.empty() && Allocas.empty())
      return bChanged;

    DoTransform(hlslOP);
    VerifyComplete(DM);

    return true;
  }
};

class DxilLegalizeResources : public ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid
  explicit DxilLegalizeResources() : ModulePass(ID) {}

  StringRef getPassName() const override {
    return "DXIL Legalize Resource Use";
  }

  bool runOnModule(Module &M) override {
    LegalizeResourceUseHelper helper;
    return helper.runOnModule(M);
  }

private:
};

} // namespace

char DxilLegalizeResources::ID = 0;

ModulePass *llvm::createDxilLegalizeResources() {
  return new DxilLegalizeResources();
}

INITIALIZE_PASS(DxilLegalizeResources, "hlsl-dxil-legalize-resources",
                "DXIL legalize resource use", false, false)

bool DxilLowerCreateHandleForLib::RemovePhiOnResource() {
  LegalizeResourceUseHelper helper;
  bool bChanged = helper.runOnModule(*m_DM->GetModule());
  if (helper.ErrorsReported())
    m_bLegalizationFailed = true;
  return bChanged;
}

// LegacyLayout.
namespace {

StructType *
UpdateStructTypeForLegacyLayout(StructType *ST, DxilTypeSystem &TypeSys,
                                Module &M,
                                bool includeTopLevelResource = false);

Type *UpdateFieldTypeForLegacyLayout(Type *Ty, DxilFieldAnnotation &annotation,
                                     DxilTypeSystem &TypeSys, Module &M) {
  DXASSERT(!Ty->isPointerTy(), "struct field should not be a pointer");

  if (Ty->isArrayTy()) {
    Type *EltTy = Ty->getArrayElementType();
    Type *UpdatedTy =
        UpdateFieldTypeForLegacyLayout(EltTy, annotation, TypeSys, M);
    if (EltTy == UpdatedTy)
      return Ty;
    else if (UpdatedTy)
      return ArrayType::get(UpdatedTy, Ty->getArrayNumElements());
    else
      return nullptr;
  } else if (hlsl::HLMatrixType::isa(Ty)) {
    DXASSERT(annotation.HasMatrixAnnotation(), "must a matrix");
    HLMatrixType MatTy = HLMatrixType::cast(Ty);
    unsigned rows = MatTy.getNumRows();
    unsigned cols = MatTy.getNumColumns();
    Type *EltTy = MatTy.getElementTypeForReg();

    // Get cols and rows from annotation.
    const DxilMatrixAnnotation &matrix = annotation.GetMatrixAnnotation();
    if (matrix.Orientation == MatrixOrientation::RowMajor) {
      rows = matrix.Rows;
      cols = matrix.Cols;
    } else {
      DXASSERT_NOMSG(matrix.Orientation == MatrixOrientation::ColumnMajor);
      cols = matrix.Rows;
      rows = matrix.Cols;
    }

    EltTy = UpdateFieldTypeForLegacyLayout(EltTy, annotation, TypeSys, M);
    Type *rowTy = VectorType::get(EltTy, cols);

    // Matrix should be aligned like array if rows > 1,
    // otherwise, it's just like a vector.
    if (rows > 1)
      return ArrayType::get(rowTy, rows);
    else
      return rowTy;
  } else if (StructType *ST = dyn_cast<StructType>(Ty)) {
    return UpdateStructTypeForLegacyLayout(ST, TypeSys, M);
  } else if (FixedVectorType *VT = dyn_cast<FixedVectorType>(Ty)) {
    Type *EltTy = VT->getElementType();
    Type *UpdatedTy =
        UpdateFieldTypeForLegacyLayout(EltTy, annotation, TypeSys, M);
    if (EltTy == UpdatedTy)
      return Ty;
    else
      return VectorType::get(UpdatedTy, VT->getNumElements());
  } else {
    Type *i32Ty = Type::getInt32Ty(Ty->getContext());
    // Basic types.
    if (Ty->isHalfTy()) {
      return Type::getFloatTy(Ty->getContext());
    } else if (IntegerType *ITy = dyn_cast<IntegerType>(Ty)) {
      if (ITy->getBitWidth() < 32)
        return i32Ty;
      else
        return Ty;
    } else
      return Ty;
  }
}

StructType *UpdateStructTypeForLegacyLayout(StructType *ST,
                                            DxilTypeSystem &TypeSys, Module &M,
                                            bool includeTopLevelResource) {
  bool bUpdated = false;
  unsigned fieldsCount = ST->getNumElements();
  std::vector<Type *> fieldTypes;
  fieldTypes.reserve(fieldsCount);
  DxilStructAnnotation *SA = TypeSys.GetStructAnnotation(ST);

  if (!includeTopLevelResource && dxilutil::IsHLSLResourceType(ST))
    return nullptr;

  // After reflection is stripped from library, this will be null if no update
  // is required.
  if (!SA) {
    return ST;
  }

  if (SA->IsEmptyStruct()) {
    return ST;
  }

  // Resource fields must be deleted, since they don't actually
  // show up in the structure layout.
  // fieldMap maps from new field index to old field index for porting
  // annotations
  std::vector<unsigned> fieldMap;
  fieldMap.reserve(fieldsCount);

  for (unsigned i = 0; i < fieldsCount; i++) {
    Type *EltTy = ST->getElementType(i);
    Type *UpdatedTy = UpdateFieldTypeForLegacyLayout(
        EltTy, SA->GetFieldAnnotation(i), TypeSys, M);
    if (UpdatedTy != nullptr) {
      fieldMap.push_back(i);
      fieldTypes.push_back(UpdatedTy);
    }
    if (EltTy != UpdatedTy)
      bUpdated = true;
  }

  if (!bUpdated) {
    return ST;
  } else {
    std::string legacyName =
        std::string(DXIL::kHostLayoutTypePrefix) + ST->getName().str();
    if (StructType *legacyST = M.getTypeByName(legacyName))
      return legacyST;

    StructType *NewST =
        StructType::create(ST->getContext(), fieldTypes, legacyName);

    // Only add annotation if struct is not empty.
    if (NewST->getNumElements() > 0) {
      DxilStructAnnotation *NewSA = TypeSys.AddStructAnnotation(NewST);

      // Clone annotation.
      NewSA->SetCBufferSize(SA->GetCBufferSize());
      NewSA->SetNumTemplateArgs(SA->GetNumTemplateArgs());
      for (unsigned i = 0; i < SA->GetNumTemplateArgs(); i++) {
        NewSA->GetTemplateArgAnnotation(i) = SA->GetTemplateArgAnnotation(i);
      }
      // Remap with deleted resource fields
      for (unsigned i = 0; i < NewSA->GetNumFields(); i++) {
        NewSA->GetFieldAnnotation(i) = SA->GetFieldAnnotation(fieldMap[i]);
      }
      TypeSys.FinishStructAnnotation(*NewSA);
    }

    return NewST;
  }
}

bool UpdateStructTypeForLegacyLayout(DxilResourceBase &Res,
                                     DxilTypeSystem &TypeSys, DxilModule &DM) {
  Module &M = *DM.GetModule();
  Constant *Symbol = Res.GetGlobalSymbol();
  Type *ElemTy = Res.GetHLSLType()->getPointerElementType();
  // Support Array of ConstantBuffer/StructuredBuffer.
  llvm::SmallVector<unsigned, 4> arrayDims;
  ElemTy = dxilutil::StripArrayTypes(ElemTy, &arrayDims);
  StructType *ST = cast<StructType>(ElemTy);
  if (ST->isOpaque()) {
    DXASSERT(Res.GetClass() == DxilResourceBase::Class::CBuffer,
             "Only cbuffer can have opaque struct.");
    return false;
  }

  Type *UpdatedST = UpdateStructTypeForLegacyLayout(
      ST, TypeSys, M, Res.GetKind() == DXIL::ResourceKind::StructuredBuffer);
  if (ST != UpdatedST) {
    // Support Array of ConstantBuffer/StructuredBuffer.
    Type *UpdatedTy = dxilutil::WrapInArrayTypes(UpdatedST, arrayDims);
    GlobalVariable *NewGV = cast<GlobalVariable>(
        M.getOrInsertGlobal(Symbol->getName().str() + "_legacy", UpdatedTy));
    Res.SetGlobalSymbol(NewGV);
    Res.SetHLSLType(NewGV->getType());
    OP *hlslOP = DM.GetOP();

    if (DM.GetShaderModel()->IsLib()) {
      TypeSys.EraseStructAnnotation(ST);
      // If it's a library, we need to replace the GV which involves a few
      // replacements
      Function *NF =
          hlslOP->GetOpFunc(hlsl::OP::OpCode::CreateHandleForLib, UpdatedST);
      Value *opArg =
          hlslOP->GetI32Const((unsigned)hlsl::OP::OpCode::CreateHandleForLib);
      auto replaceResLd = [&NF, &opArg](LoadInst *ldInst, Value *NewPtr) {
        if (!ldInst->user_empty()) {
          IRBuilder<> Builder = IRBuilder<>(ldInst);
          LoadInst *newLoad = Builder.CreateLoad(NewPtr);
          Value *args[] = {opArg, newLoad};

          for (auto user = ldInst->user_begin(), E = ldInst->user_end();
               user != E;) {
            CallInst *CI = cast<CallInst>(*(user++));
            CallInst *newCI = CallInst::Create(NF, args, "", CI);
            CI->replaceAllUsesWith(newCI);
            CI->eraseFromParent();
          }
        }
        ldInst->eraseFromParent();
      };
      // Merge GEP to simplify replace old GV.
      if (!arrayDims.empty())
        dxilutil::MergeGepUse(Symbol);
      // Replace old GV.
      for (auto UserIt = Symbol->user_begin(), userEnd = Symbol->user_end();
           UserIt != userEnd;) {
        Value *User = *(UserIt++);

        if (LoadInst *ldInst = dyn_cast<LoadInst>(User)) {
          replaceResLd(ldInst, NewGV);
        } else if (GEPOperator *GEP = dyn_cast<GEPOperator>(User)) {
          IRBuilder<> Builder(GEP->getContext());
          StringRef Name = "";
          if (Instruction *I = dyn_cast<Instruction>(GEP)) {
            Builder.SetInsertPoint(I);
            Name = GEP->getName();
          }
          SmallVector<Value *, 8> Indices(GEP->idx_begin(), GEP->idx_end());
          Value *NewPtr = Builder.CreateGEP(NewGV, Indices);
          for (auto GEPUserIt = GEP->user_begin(), GEPuserEnd = GEP->user_end();
               GEPUserIt != GEPuserEnd;) {
            Value *User = *(GEPUserIt++);
            if (LoadInst *ldInst = dyn_cast<LoadInst>(User)) {
              replaceResLd(ldInst, NewPtr);
            } else {
              User->dump();
              DXASSERT(0, "unsupported user when update resouce type");
            }
          }
          if (Instruction *I = dyn_cast<Instruction>(GEP))
            I->eraseFromParent();
        } else {
          User->dump();
          DXASSERT(0, "unsupported user when update resouce type");
        }
      }
    } else {
      // If not a library, the GV should be deleted
      for (auto UserIt = Symbol->user_begin(); UserIt != Symbol->user_end();) {
        Value *User = *(UserIt++);

        if (Instruction *I = dyn_cast<Instruction>(User)) {
          if (!User->user_empty())
            I->replaceAllUsesWith(UndefValue::get(I->getType()));

          I->eraseFromParent();
        } else {
          ConstantExpr *CE = cast<ConstantExpr>(User);
          if (!CE->user_empty())
            CE->replaceAllUsesWith(UndefValue::get(CE->getType()));
        }
      }
    }
    Symbol->removeDeadConstantUsers();

    if (GlobalVariable *GV = dyn_cast<GlobalVariable>(Symbol))
      GV->eraseFromParent();

    return true;
  }

  return false;
}

bool UpdateStructTypeForLegacyLayoutOnDM(DxilModule &DM) {
  DxilTypeSystem &TypeSys = DM.GetTypeSystem();
  bool bChanged = false;
  for (auto &CBuf : DM.GetCBuffers()) {
    bChanged |= UpdateStructTypeForLegacyLayout(*CBuf.get(), TypeSys, DM);
  }

  for (auto &UAV : DM.GetUAVs()) {
    if (DXIL::IsStructuredBuffer(UAV->GetKind()))
      bChanged |= UpdateStructTypeForLegacyLayout(*UAV.get(), TypeSys, DM);
  }

  for (auto &SRV : DM.GetSRVs()) {
    if (SRV->IsStructuredBuffer() || SRV->IsTBuffer())
      bChanged |= UpdateStructTypeForLegacyLayout(*SRV.get(), TypeSys, DM);
  }

  return bChanged;
}

} // namespace

void DxilLowerCreateHandleForLib::FailOnPoisonResources() {
  // A previous pass replaced all undef resources with constant zero resources.
  // If those made it here, the program is malformed.
  for (Function &Func : this->m_DM->GetModule()->functions()) {
    hlsl::OP::OpCodeClass OpcodeClass;
    if (m_DM->GetOP()->GetOpCodeClass(&Func, OpcodeClass) &&
        OpcodeClass == OP::OpCodeClass::CreateHandleForLib) {
      Type *ResTy = Func.getFunctionType()->getParamType(
          DXIL::OperandIndex::kCreateHandleForLibResOpIdx);
      Constant *PoisonRes = ConstantAggregateZero::get(ResTy);
      for (User *PoisonUser : PoisonRes->users())
        if (Instruction *PoisonUserInst = dyn_cast<Instruction>(PoisonUser))
          dxilutil::EmitResMappingError(PoisonUserInst);
    }
  }
}

bool DxilLowerCreateHandleForLib::UpdateStructTypeForLegacyLayout() {
  return UpdateStructTypeForLegacyLayoutOnDM(*m_DM);
}

// Change ResourceSymbol to undef if don't need.
void DxilLowerCreateHandleForLib::UpdateResourceSymbols() {
  auto UpdateResourceSymbol = [](DxilResourceBase *res) {
    if (GlobalVariable *GV = dyn_cast<GlobalVariable>(res->GetGlobalSymbol())) {
      GV->removeDeadConstantUsers();
      DXASSERT(GV->user_empty(), "else resource not lowered");
      res->SetGlobalSymbol(UndefValue::get(GV->getType()));
      if (GV->user_empty())
        GV->eraseFromParent();
    }
  };

  for (auto &&C : m_DM->GetCBuffers()) {
    UpdateResourceSymbol(C.get());
  }
  for (auto &&Srv : m_DM->GetSRVs()) {
    UpdateResourceSymbol(Srv.get());
  }
  for (auto &&Uav : m_DM->GetUAVs()) {
    UpdateResourceSymbol(Uav.get());
  }
  for (auto &&S : m_DM->GetSamplers()) {
    UpdateResourceSymbol(S.get());
  }
}

// Lower createHandleForLib
namespace {

Value *flattenGepIdx(GEPOperator *GEP) {
  Value *idx = nullptr;
  if (GEP->getNumIndices() == 2) {
    // one dim array of resource
    idx = (GEP->idx_begin() + 1)->get();
  } else {
    gep_type_iterator GEPIt = gep_type_begin(GEP), E = gep_type_end(GEP);
    // Must be instruction for multi dim array.
    std::unique_ptr<IRBuilder<>> Builder;
    if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP)) {
      Builder = llvm::make_unique<IRBuilder<>>(GEPInst);
    } else {
      Builder = llvm::make_unique<IRBuilder<>>(GEP->getContext());
    }
    for (; GEPIt != E; ++GEPIt) {
      if (GEPIt->isArrayTy()) {
        unsigned arraySize = GEPIt->getArrayNumElements();
        Value *tmpIdx = GEPIt.getOperand();
        if (idx == nullptr)
          idx = tmpIdx;
        else {
          idx = Builder->CreateMul(idx, Builder->getInt32(arraySize));
          idx = Builder->CreateAdd(idx, tmpIdx);
        }
      }
    }
  }
  return idx;
}

} // namespace

void DxilLowerCreateHandleForLib::ReplaceResourceUserWithHandle(
    DxilResource &res, LoadInst *load, Instruction *handle) {
  for (auto resUser = load->user_begin(), E = load->user_end(); resUser != E;) {
    Value *V = *(resUser++);
    CallInst *CI = dyn_cast<CallInst>(V);
    DxilInst_CreateHandleForLib createHandle(CI);
    DXASSERT(createHandle, "must be createHandle");
    CI->replaceAllUsesWith(handle);
    CI->eraseFromParent();
  }

  if (res.GetClass() == DXIL::ResourceClass::UAV) {
    // Before this pass, the global resources might not have been mapped with
    // all the uses. Now we're 100% sure who uses what resources (otherwise the
    // compilation would have failed), so we do a round on marking UAV's as
    // having counter.
    static auto IsDxilOp = [](Value *V, hlsl::OP::OpCode Op) -> bool {
      Instruction *I = dyn_cast<Instruction>(V);
      if (!I)
        return false;
      return hlsl::OP::IsDxilOpFuncCallInst(I, Op);
    };

    // Search all users for update counter
    bool updateAnnotateHandle =
        res.IsGloballyCoherent() || res.IsReorderCoherent();
    if (!res.HasCounter()) {
      for (User *U : handle->users()) {
        if (IsDxilOp(U, hlsl::OP::OpCode::BufferUpdateCounter)) {
          res.SetHasCounter(true);
          break;
        } else if (IsDxilOp(U, hlsl::OP::OpCode::AnnotateHandle)) {
          for (User *UU : U->users()) {
            if (IsDxilOp(UU, hlsl::OP::OpCode::BufferUpdateCounter)) {
              res.SetHasCounter(true);
              updateAnnotateHandle = true;
              break;
            }
          }
          if (updateAnnotateHandle)
            break;
        }
      }
    }
    if (updateAnnotateHandle) {
      // Update resource props with counter flag
      DxilResourceProperties RP =
          resource_helper::loadPropsFromResourceBase(&res);
      // Require ShaderModule to reconstruct resource property constant
      const ShaderModel *pSM = m_DM->GetShaderModel();

      SmallVector<Instruction *, 4> annotHandles;
      for (User *U : handle->users()) {
        DxilInst_AnnotateHandle annotateHandle(cast<Instruction>(U));
        if (annotateHandle) {
          annotHandles.emplace_back(cast<Instruction>(U));
        }
      }
      if (!annotHandles.empty()) {
        Instruction *firstAnnot = annotHandles.pop_back_val();
        DxilInst_AnnotateHandle annotateHandle(firstAnnot);
        // Update props.
        Constant *propsConst = resource_helper::getAsConstant(
            RP, annotateHandle.get_props()->getType(), *pSM);
        annotateHandle.set_props(propsConst);
        if (!annotHandles.empty()) {
          // Move firstAnnot after handle.
          firstAnnot->removeFromParent();
          firstAnnot->insertAfter(handle);
          // Remove redundant annotate handles.
          for (auto *annotHdl : annotHandles) {
            annotHdl->replaceAllUsesWith(firstAnnot);
            annotHdl->eraseFromParent();
          }
        }
      }
    }
  }

  load->eraseFromParent();
}

void DxilLowerCreateHandleForLib::TranslateDxilResourceUses(
    DxilResourceBase &res) {
  OP *hlslOP = m_DM->GetOP();
  // Generate createHandleFromBinding for sm66 and later.
  bool bCreateFromBinding = m_DM->GetShaderModel()->IsSM66Plus();
  OP::OpCode createOp = bCreateFromBinding ? OP::OpCode::CreateHandleFromBinding
                                           : OP::OpCode::CreateHandle;
  Function *createHandle =
      hlslOP->GetOpFunc(createOp, llvm::Type::getVoidTy(m_DM->GetCtx()));
  Value *opArg = hlslOP->GetU32Const((unsigned)createOp);

  bool isViewResource = res.GetClass() == DXIL::ResourceClass::SRV ||
                        res.GetClass() == DXIL::ResourceClass::UAV;
  bool isROV = isViewResource && static_cast<DxilResource &>(res).IsROV();
  std::string handleName =
      (res.GetGlobalName() + Twine("_") + Twine(res.GetResClassName())).str();
  if (isViewResource)
    handleName += (Twine("_") + Twine(res.GetResDimName())).str();
  if (isROV)
    handleName += "_ROV";

  Value *resClassArg = hlslOP->GetU8Const(
      static_cast<std::underlying_type<DxilResourceBase::Class>::type>(
          res.GetClass()));
  Value *resIDArg = hlslOP->GetU32Const(res.GetID());
  // resLowerBound will be added after allocation in DxilCondenseResources.
  Value *resLowerBound = hlslOP->GetU32Const(res.GetLowerBound());

  Value *isUniformRes = hlslOP->GetI1Const(0);

  Value *GV = res.GetGlobalSymbol();
  DXASSERT(isa<GlobalValue>(GV),
           "DxilLowerCreateHandleForLib cannot deal with unused resources.");

  Module *pM = m_DM->GetModule();
  // TODO: add debug info to create handle.
  DIVariable *DIV = nullptr;
  DILocation *DL = nullptr;
  if (m_HasDbgInfo) {
    DebugInfoFinder &Finder = m_DM->GetOrCreateDebugInfoFinder();
    DIV =
        dxilutil::FindGlobalVariableDebugInfo(cast<GlobalVariable>(GV), Finder);
    if (DIV)
      // TODO: how to get col?
      DL =
          DILocation::get(pM->getContext(), DIV->getLine(), 1, DIV->getScope());
  }

  bool isResArray = res.GetRangeSize() > 1;
  std::unordered_map<Function *, Instruction *> handleMapOnFunction;

  Value *createHandleArgs[] = {opArg, resClassArg, resIDArg, resLowerBound,
                               isUniformRes};

  DxilResourceBinding binding =
      resource_helper::loadBindingFromResourceBase(&res);
  Value *bindingV = resource_helper::getAsConstant(
      binding, hlslOP->GetResourceBindingType(), *m_DM->GetShaderModel());

  Value *createHandleFromBindingArgs[] = {opArg, bindingV, resLowerBound,
                                          isUniformRes};

  MutableArrayRef<Value *> Args(bCreateFromBinding ? createHandleFromBindingArgs
                                                   : createHandleArgs,
                                bCreateFromBinding ? 4 : 5);

  const unsigned resIdxOpIdx = bCreateFromBinding
                                   ? DxilInst_CreateHandleFromBinding::arg_index
                                   : DxilInst_CreateHandle::arg_index;
  const unsigned nonUniformOpIdx =
      bCreateFromBinding ? DxilInst_CreateHandleFromBinding::arg_nonUniformIndex
                         : DxilInst_CreateHandle::arg_nonUniformIndex;

  for (iplist<Function>::iterator F : pM->getFunctionList()) {
    if (!F->isDeclaration()) {
      if (!isResArray) {
        IRBuilder<> Builder(dxilutil::FindAllocaInsertionPt(F));
        if (m_HasDbgInfo) {
          // TODO: set debug info.
          // Builder.SetCurrentDebugLocation(DL);
        }
        handleMapOnFunction[F] =
            Builder.CreateCall(createHandle, Args, handleName);
      }
    }
  }

  for (auto U = GV->user_begin(), E = GV->user_end(); U != E;) {
    User *user = *(U++);
    // Skip unused user.
    if (user->user_empty())
      continue;

    if (LoadInst *ldInst = dyn_cast<LoadInst>(user)) {
      Function *userF = ldInst->getParent()->getParent();
      DXASSERT(handleMapOnFunction.count(userF), "must exist");
      Instruction *handle = handleMapOnFunction[userF];
      ReplaceResourceUserWithHandle(static_cast<DxilResource &>(res), ldInst,
                                    handle);
    } else if (GEPOperator *GEP = dyn_cast<GEPOperator>(user)) {
      Value *idx = flattenGepIdx(GEP);

      Args[resIdxOpIdx] = idx;

      Args[nonUniformOpIdx] = isUniformRes;

      Instruction *handle = nullptr;
      if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP)) {
        IRBuilder<> Builder = IRBuilder<>(GEPInst);
        if (DxilMDHelper::IsMarkedNonUniform(GEPInst)) {
          // Mark nonUniform.
          Args[nonUniformOpIdx] = hlslOP->GetI1Const(1);
          // Clear nonUniform on GEP.
          GEPInst->setMetadata(DxilMDHelper::kDxilNonUniformAttributeMDName,
                               nullptr);
        }
        Args[resIdxOpIdx] = Builder.CreateAdd(idx, resLowerBound);
        handle = Builder.CreateCall(createHandle, Args, handleName);
      }

      for (auto GEPU = GEP->user_begin(), GEPE = GEP->user_end();
           GEPU != GEPE;) {
        // Must be load inst.
        LoadInst *ldInst = cast<LoadInst>(*(GEPU++));
        if (handle) {
          ReplaceResourceUserWithHandle(static_cast<DxilResource &>(res),
                                        ldInst, handle);
        } else {
          IRBuilder<> Builder = IRBuilder<>(ldInst);
          Args[resIdxOpIdx] = Builder.CreateAdd(idx, resLowerBound);
          Instruction *localHandle =
              Builder.CreateCall(createHandle, Args, handleName);
          ReplaceResourceUserWithHandle(static_cast<DxilResource &>(res),
                                        ldInst, localHandle);
        }
      }

      if (Instruction *I = dyn_cast<Instruction>(GEP)) {
        I->eraseFromParent();
      }
    } else if (BitCastInst *BCI = dyn_cast<BitCastInst>(user)) {
      DXASSERT(onlyUsedByLifetimeMarkers(BCI),
               "expected bitcast to only be used by lifetime intrinsics");
      for (auto BCIU = BCI->user_begin(), BCIE = BCI->user_end();
           BCIU != BCIE;) {
        IntrinsicInst *II = cast<IntrinsicInst>(*(BCIU++));
        II->eraseFromParent();
      }
      BCI->eraseFromParent();
    } else if (ConstantExpr *CE = dyn_cast<ConstantExpr>(user)) {
      // A GEPOperator can also be a ConstantExpr, so it must be checked before
      // this code.
      DXASSERT(CE->getOpcode() == Instruction::BitCast, "expected bitcast");
      DXASSERT(onlyUsedByLifetimeMarkers(CE),
               "expected ConstantExpr to only be used by lifetime intrinsics");
      for (auto CEU = CE->user_begin(), CEE = CE->user_end(); CEU != CEE;) {
        IntrinsicInst *II = cast<IntrinsicInst>(*(CEU++));
        II->eraseFromParent();
      }
    } else {
      DXASSERT(false,
               "AddOpcodeParamForIntrinsic in CodeGen did not patch uses "
               "to only have ld/st refer to temp object");
    }
  }
  // Erase unused handle.
  for (auto It : handleMapOnFunction) {
    Instruction *I = It.second;
    if (I->user_empty())
      I->eraseFromParent();
  }
}

void DxilLowerCreateHandleForLib::GenerateDxilResourceHandles() {
  for (size_t i = 0; i < m_DM->GetCBuffers().size(); i++) {
    DxilCBuffer &C = m_DM->GetCBuffer(i);
    TranslateDxilResourceUses(C);
  }
  // Create sampler handle first, may be used by SRV operations.
  for (size_t i = 0; i < m_DM->GetSamplers().size(); i++) {
    DxilSampler &S = m_DM->GetSampler(i);
    TranslateDxilResourceUses(S);
  }

  for (size_t i = 0; i < m_DM->GetSRVs().size(); i++) {
    DxilResource &SRV = m_DM->GetSRV(i);
    TranslateDxilResourceUses(SRV);
  }

  for (size_t i = 0; i < m_DM->GetUAVs().size(); i++) {
    DxilResource &UAV = m_DM->GetUAV(i);
    TranslateDxilResourceUses(UAV);
  }
}

// TBuffer.
namespace {
void InitTBuffer(const DxilCBuffer *pSource, DxilResource *pDest) {
  pDest->SetKind(pSource->GetKind());
  pDest->SetCompType(DXIL::ComponentType::U32);
  pDest->SetSampleCount(0);
  pDest->SetElementStride(0);
  pDest->SetGloballyCoherent(false);
  pDest->SetReorderCoherent(false);
  pDest->SetHasCounter(false);
  pDest->SetRW(false);
  pDest->SetROV(false);
  pDest->SetID(pSource->GetID());
  pDest->SetSpaceID(pSource->GetSpaceID());
  pDest->SetLowerBound(pSource->GetLowerBound());
  pDest->SetRangeSize(pSource->GetRangeSize());
  pDest->SetGlobalSymbol(pSource->GetGlobalSymbol());
  pDest->SetGlobalName(pSource->GetGlobalName());
  pDest->SetHandle(pSource->GetHandle());
  pDest->SetHLSLType(pSource->GetHLSLType());
}

void PatchTBufferLoad(CallInst *handle, DxilModule &DM,
                      DenseSet<Value *> &patchedSet) {
  if (patchedSet.count(handle))
    return;
  patchedSet.insert(handle);
  hlsl::OP *hlslOP = DM.GetOP();
  llvm::LLVMContext &Ctx = DM.GetCtx();
  Type *doubleTy = Type::getDoubleTy(Ctx);
  Type *i64Ty = Type::getInt64Ty(Ctx);

  // Replace corresponding cbuffer loads with typed buffer loads
  for (auto U = handle->user_begin(); U != handle->user_end();) {
    User *user = *(U++);
    CallInst *I = dyn_cast<CallInst>(user);
    // Could also be store for out arg in lib.
    DXASSERT(isa<StoreInst>(user) || (I && OP::IsDxilOpFuncCallInst(I)),
             "otherwise unexpected user of CreateHandle value");
    if (!I)
      continue;
    DXIL::OpCode opcode = OP::GetDxilOpFuncCallInst(I);
    if (opcode == DXIL::OpCode::CBufferLoadLegacy) {
      DxilInst_CBufferLoadLegacy cbLoad(I);

      StructType *cbRetTy = cast<StructType>(I->getType());
      // elements will be 4, or 8 for native 16-bit types, which require special
      // handling.
      bool cbRet8Elt = cbRetTy->getNumElements() > 4;

      // Replace with appropriate buffer load instruction
      IRBuilder<> Builder(I);
      opcode = OP::OpCode::BufferLoad;
      Type *Ty = Type::getInt32Ty(Ctx);
      Function *BufLoad = hlslOP->GetOpFunc(opcode, Ty);
      Constant *opArg = hlslOP->GetU32Const((unsigned)opcode);
      Value *undefI = UndefValue::get(Type::getInt32Ty(Ctx));
      Value *offset = cbLoad.get_regIndex();
      CallInst *load =
          Builder.CreateCall(BufLoad, {opArg, handle, offset, undefI});

      // Find extractelement uses of cbuffer load and replace + generate bitcast
      // as necessary
      for (auto LU = I->user_begin(); LU != I->user_end();) {
        ExtractValueInst *evInst = dyn_cast<ExtractValueInst>(*(LU++));
        DXASSERT(evInst && evInst->getNumIndices() == 1,
                 "user of cbuffer load result should be extractvalue");
        uint64_t idx = evInst->getIndices()[0];
        Type *EltTy = evInst->getType();
        IRBuilder<> EEBuilder(evInst);
        Value *result = nullptr;
        if (EltTy != Ty) {
          // extract two values and DXIL::OpCode::MakeDouble or construct i64
          if ((EltTy == doubleTy) || (EltTy == i64Ty)) {
            DXASSERT(idx < 2, "64-bit component index out of range");

            // This assumes big endian order in tbuffer elements (is this
            // correct?)
            Value *low = EEBuilder.CreateExtractValue(load, idx * 2);
            Value *high = EEBuilder.CreateExtractValue(load, idx * 2 + 1);
            if (EltTy == doubleTy) {
              opcode = OP::OpCode::MakeDouble;
              Function *MakeDouble = hlslOP->GetOpFunc(opcode, doubleTy);
              Constant *opArg = hlslOP->GetU32Const((unsigned)opcode);
              result = EEBuilder.CreateCall(MakeDouble, {opArg, low, high});
            } else {
              high = EEBuilder.CreateZExt(high, i64Ty);
              low = EEBuilder.CreateZExt(low, i64Ty);
              high = EEBuilder.CreateShl(high, hlslOP->GetU64Const(32));
              result = EEBuilder.CreateOr(high, low);
            }
          } else {
            if (cbRet8Elt) {
              DXASSERT_NOMSG(cbRetTy->getNumElements() == 8);
              DXASSERT_NOMSG(EltTy->getScalarSizeInBits() == 16);
              // Translate extract from 16bit x 8 to extract and translate from
              // i32 by 4
              result = EEBuilder.CreateExtractValue(load, idx >> 1);
              if (idx & 1)
                result = EEBuilder.CreateLShr(result, 16);
              result = EEBuilder.CreateTrunc(result, Type::getInt16Ty(Ctx));
              if (EltTy->isHalfTy())
                result = EEBuilder.CreateBitCast(result, EltTy);
            } else {
              result = EEBuilder.CreateExtractValue(load, idx);
              if (Ty->getScalarSizeInBits() > EltTy->getScalarSizeInBits()) {
                if (EltTy->isIntegerTy()) {
                  result = EEBuilder.CreateTrunc(result, EltTy);
                } else {
                  result =
                      EEBuilder.CreateBitCast(result, Type::getFloatTy(Ctx));
                  result = EEBuilder.CreateFPTrunc(result, EltTy);
                }
              } else {
                result = EEBuilder.CreateBitCast(result, EltTy);
              }
            }
          }
        } else {
          result = EEBuilder.CreateExtractValue(load, idx);
        }

        evInst->replaceAllUsesWith(result);
        evInst->eraseFromParent();
      }
    } else if (opcode == DXIL::OpCode::CBufferLoad) {
      // TODO: Handle this, or prevent this for tbuffer
      DXASSERT(false, "otherwise CBufferLoad used for tbuffer rather than "
                      "CBufferLoadLegacy");
    } else if (opcode == DXIL::OpCode::AnnotateHandle) {
      PatchTBufferLoad(cast<CallInst>(I), DM, patchedSet);
      continue;
    } else if (opcode == DXIL::OpCode::BufferLoad) {
      // Already translated, skip.
      continue;
    } else {
      DXASSERT(false, "otherwise unexpected user of CreateHandle value");
    }
    I->eraseFromParent();
  }
}

} // namespace

void DxilLowerCreateHandleForLib::PatchTBufferUse(
    Value *V, DxilModule &DM, DenseSet<Value *> &patchedSet) {
  for (User *U : V->users()) {
    if (CallInst *CI = dyn_cast<CallInst>(U)) {
      // Patch dxil call.
      if (hlsl::OP::IsDxilOpFuncCallInst(CI))
        PatchTBufferLoad(CI, DM, patchedSet);
    } else {
      PatchTBufferUse(U, DM, patchedSet);
    }
  }
}

bool DxilLowerCreateHandleForLib::PatchDynamicTBuffers(DxilModule &DM) {
  hlsl::OP *hlslOP = DM.GetOP();
  Function *AnnotHandleFn = hlslOP->GetOpFunc(DXIL::OpCode::AnnotateHandle,
                                              Type::getVoidTy(DM.GetCtx()));
  if (AnnotHandleFn->user_empty()) {
    AnnotHandleFn->eraseFromParent();
    return false;
  }
  bool bUpdated = false;
  for (User *U : AnnotHandleFn->users()) {
    CallInst *CI = cast<CallInst>(U);
    DxilInst_AnnotateHandle annot(CI);
    DxilResourceProperties RP = resource_helper::loadPropsFromAnnotateHandle(
        annot, *DM.GetShaderModel());

    if (RP.getResourceKind() != DXIL::ResourceKind::TBuffer)
      continue;
    // Skip handle from createHandleForLib which take care in PatchTBuffers.
    if (CallInst *HdlCI = dyn_cast<CallInst>(annot.get_res())) {
      if (hlslOP->IsDxilOpFuncCallInst(HdlCI)) {
        if (hlslOP->GetDxilOpFuncCallInst(HdlCI) ==
            DXIL::OpCode::CreateHandleForLib)
          continue;
      }
    }

    DenseSet<Value *> patchedSet;
    PatchTBufferLoad(CI, DM, patchedSet);
    bUpdated = true;
  }
  return bUpdated;
}

bool DxilLowerCreateHandleForLib::PatchTBuffers(DxilModule &DM) {
  bool bChanged = false;
  // move tbuffer resources to SRVs
  Module &M = *DM.GetModule();
  const ShaderModel &SM = *DM.GetShaderModel();
  DenseSet<Value *> patchedSet;

  // First, patch users of AnnotateHandle calls if we have them.
  // This will pick up uses in lib_6_x functions that otherwise
  // would be missed.
  if (SM.IsSM66Plus()) {
    for (auto it : DM.GetOP()->GetOpFuncList(DXIL::OpCode::AnnotateHandle)) {
      Function *F = it.second;
      for (auto U = F->user_begin(); U != F->user_end();) {
        User *user = *(U++);
        if (CallInst *CI = dyn_cast<CallInst>(user)) {
          DxilInst_AnnotateHandle AH(CI);
          if (AH) {
            DxilResourceProperties RP =
                resource_helper::loadPropsFromAnnotateHandle(AH, SM);
            if (RP.getResourceKind() == DXIL::ResourceKind::TBuffer)
              PatchTBufferLoad(CI, DM, patchedSet);
          }
        }
      }
    }
  }

  unsigned offset = DM.GetSRVs().size();
  for (auto it = DM.GetCBuffers().begin(); it != DM.GetCBuffers().end(); it++) {
    DxilCBuffer *CB = it->get();
    if (CB->GetKind() == DXIL::ResourceKind::TBuffer) {
      auto srv = make_unique<DxilResource>();
      InitTBuffer(CB, srv.get());
      srv->SetID(offset++);
      DM.AddSRV(std::move(srv));
      GlobalVariable *GV = dyn_cast<GlobalVariable>(CB->GetGlobalSymbol());
      if (GV == nullptr)
        continue;
      PatchTBufferUse(GV, DM, patchedSet);
      // Set global symbol for cbuffer to an unused value so it can be removed
      // in RemoveUnusedResourceSymbols.
      Type *Ty = GV->getType()->getElementType();
      GlobalVariable *NewGV = new GlobalVariable(
          M, Ty, GV->isConstant(), GV->getLinkage(), /*Initializer*/ nullptr,
          GV->getName(),
          /*InsertBefore*/ nullptr, GV->getThreadLocalMode(),
          GV->getType()->getAddressSpace(), GV->isExternallyInitialized());
      CB->SetGlobalSymbol(NewGV);
      bChanged = true;
    }
  }
  return bChanged;
}

typedef DenseMap<Value *, unsigned> OffsetForValueMap;

// Find the imm offset part from a value.
// It must exist unless offset is 0.
static unsigned GetCBOffset(Value *V, OffsetForValueMap &visited) {
  auto it = visited.find(V);
  if (it != visited.end())
    return it->second;
  visited[V] = 0;
  unsigned result = 0;
  if (ConstantInt *Imm = dyn_cast<ConstantInt>(V)) {
    result = Imm->getLimitedValue();
  } else if (BinaryOperator *BO = dyn_cast<BinaryOperator>(V)) {
    switch (BO->getOpcode()) {
    case Instruction::Add: {
      unsigned left = GetCBOffset(BO->getOperand(0), visited);
      unsigned right = GetCBOffset(BO->getOperand(1), visited);
      result = left + right;
    } break;
    case Instruction::Or: {
      unsigned left = GetCBOffset(BO->getOperand(0), visited);
      unsigned right = GetCBOffset(BO->getOperand(1), visited);
      result = left | right;
    } break;
    default:
      break;
    }
  } else if (SelectInst *SI = dyn_cast<SelectInst>(V)) {
    result = std::min(GetCBOffset(SI->getOperand(1), visited),
                      GetCBOffset(SI->getOperand(2), visited));
  } else if (PHINode *PN = dyn_cast<PHINode>(V)) {
    result = UINT_MAX;
    for (unsigned i = 0, ops = PN->getNumIncomingValues(); i < ops; ++i) {
      result = std::min(result, GetCBOffset(PN->getIncomingValue(i), visited));
    }
  }
  visited[V] = result;
  return result;
}

typedef std::map<unsigned, DxilFieldAnnotation *> FieldAnnotationByOffsetMap;

// Returns size in bits of the field if it's a basic type, otherwise 0.
static unsigned MarkCBUse(unsigned offset,
                          FieldAnnotationByOffsetMap &fieldMap) {
  auto it = fieldMap.upper_bound(offset);
  it--;
  if (it != fieldMap.end()) {
    it->second->SetCBVarUsed(true);
    return it->second->GetCompType().GetSizeInBits();
  }
  return 0;
}

// Detect patterns of lshr v,16 or trunc to 16-bits and return low and high
// word usage.
static const unsigned kLowWordUsed = 1;
static const unsigned kHighWordUsed = 2;
static const unsigned kLowHighWordMask = kLowWordUsed | kHighWordUsed;
static unsigned DetectLowAndHighWordUsage(ExtractValueInst *EV) {
  unsigned result = 0;
  if (EV->getType()->getScalarSizeInBits() == 32) {
    for (auto U : EV->users()) {
      Instruction *I = cast<Instruction>(U);
      if (I->getOpcode() == Instruction::LShr) {
        ConstantInt *CShift = dyn_cast<ConstantInt>(I->getOperand(1));
        if (CShift && CShift->getLimitedValue() == 16)
          result |= kHighWordUsed;
      } else if (I->getOpcode() == Instruction::Trunc &&
                 I->getType()->getPrimitiveSizeInBits() == 16) {
        result |= kLowWordUsed;
      } else {
        // Assume whole dword is used, return 0
        return 0;
      }
      if ((result & kLowHighWordMask) == kLowHighWordMask)
        break;
    }
  }
  return result;
}

static unsigned GetOffsetForCBExtractValue(ExtractValueInst *EV,
                                           bool bMinPrecision,
                                           unsigned &lowHighWordUsage) {
  DXASSERT(EV->getNumIndices() == 1,
           "otherwise, unexpected indices/type for extractvalue");
  unsigned typeSize = 4;
  unsigned bits = EV->getType()->getScalarSizeInBits();
  if (bits == 64)
    typeSize = 8;
  else if (bits == 16 && !bMinPrecision)
    typeSize = 2;
  lowHighWordUsage = DetectLowAndHighWordUsage(EV);
  return (EV->getIndices().front() * typeSize);
}

// Marks up to two CB uses for the case where only 16-bit type(s)
// are being used from lower or upper word of a tbuffer load,
// which is always 4 x 32 instead of 8 x 16, like cbuffer.
static void MarkCBUsesForExtractElement(unsigned offset,
                                        FieldAnnotationByOffsetMap &fieldMap,
                                        ExtractValueInst *EV,
                                        bool bMinPrecision) {
  unsigned lowHighWordUsage = 0;
  unsigned evOffset =
      GetOffsetForCBExtractValue(EV, bMinPrecision, lowHighWordUsage);

  // For tbuffer, where value extracted is always 32-bits:
  // If lowHighWordUsage is 0, it means 32-bits used.
  // If field marked is < 32 bits, we still need to mark the high 16-bits as
  // used, in case there is another 16-bit field.
  // Since MarkCBUse could return 0 on non-basic type field, look for 16
  // when determining whether we still need to mark high word as used.
  bool highUnmarked = EV->getType()->getScalarSizeInBits() == 32;
  if (!lowHighWordUsage || 0 != (lowHighWordUsage & kLowWordUsed))
    highUnmarked &= MarkCBUse(offset + evOffset, fieldMap) == 16;
  if (highUnmarked &&
      (!lowHighWordUsage || 0 != (lowHighWordUsage & kHighWordUsed)))
    MarkCBUse(offset + evOffset + 2, fieldMap);
}

static void CollectInPhiChain(PHINode *cbUser, unsigned offset,
                              std::unordered_set<Value *> &userSet,
                              FieldAnnotationByOffsetMap &fieldMap,
                              bool bMinPrecision) {
  if (userSet.count(cbUser) > 0)
    return;

  userSet.insert(cbUser);
  for (User *cbU : cbUser->users()) {
    if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(cbU)) {
      MarkCBUsesForExtractElement(offset, fieldMap, EV, bMinPrecision);
    } else {
      PHINode *phi = cast<PHINode>(cbU);
      CollectInPhiChain(phi, offset, userSet, fieldMap, bMinPrecision);
    }
  }
}

static void
CollectCBufferMemberUsage(Value *V, FieldAnnotationByOffsetMap &legacyFieldMap,
                          FieldAnnotationByOffsetMap &newFieldMap,
                          hlsl::OP *hlslOP, bool bMinPrecision,
                          OffsetForValueMap &visited) {
  for (auto U : V->users()) {
    if (Constant *C = dyn_cast<Constant>(U)) {
      CollectCBufferMemberUsage(C, legacyFieldMap, newFieldMap, hlslOP,
                                bMinPrecision, visited);
    } else if (LoadInst *LI = dyn_cast<LoadInst>(U)) {
      CollectCBufferMemberUsage(U, legacyFieldMap, newFieldMap, hlslOP,
                                bMinPrecision, visited);
    } else if (CallInst *CI = dyn_cast<CallInst>(U)) {
      if (hlslOP->IsDxilOpFuncCallInst(CI)) {
        hlsl::OP::OpCode op = hlslOP->GetDxilOpFuncCallInst(CI);
        if (op == DXIL::OpCode::CreateHandleForLib) {
          CollectCBufferMemberUsage(U, legacyFieldMap, newFieldMap, hlslOP,
                                    bMinPrecision, visited);
        } else if (op == DXIL::OpCode::AnnotateHandle) {
          CollectCBufferMemberUsage(U, legacyFieldMap, newFieldMap, hlslOP,
                                    bMinPrecision, visited);
        } else if (op == DXIL::OpCode::CBufferLoadLegacy ||
                   op == DXIL::OpCode::BufferLoad) {
          Value *resIndex = (op == DXIL::OpCode::CBufferLoadLegacy)
                                ? DxilInst_CBufferLoadLegacy(CI).get_regIndex()
                                : DxilInst_BufferLoad(CI).get_index();
          unsigned offset = GetCBOffset(resIndex, visited) << 4;
          for (User *cbU : U->users()) {
            if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(cbU)) {
              MarkCBUsesForExtractElement(offset, legacyFieldMap, EV,
                                          bMinPrecision);
            } else {
              PHINode *phi = cast<PHINode>(cbU);
              std::unordered_set<Value *> userSet;
              CollectInPhiChain(phi, offset, userSet, legacyFieldMap,
                                bMinPrecision);
            }
          }
        } else if (op == DXIL::OpCode::CBufferLoad) {
          DxilInst_CBufferLoad cbload(CI);
          Value *byteOffset = cbload.get_byteOffset();
          unsigned offset = GetCBOffset(byteOffset, visited);
          MarkCBUse(offset, newFieldMap);
        }
      }
    }
  }
}

void DxilLowerCreateHandleForLib::UpdateCBufferUsage() {
  DxilTypeSystem &TypeSys = m_DM->GetTypeSystem();
  hlsl::OP *hlslOP = m_DM->GetOP();
  const DataLayout &DL = m_DM->GetModule()->getDataLayout();
  const auto &CBuffers = m_DM->GetCBuffers();
  OffsetForValueMap visited;

  SmallVector<std::pair<GlobalVariable *, Type *>, 4> CBufferVars;

  // Collect cbuffers
  for (auto it = CBuffers.begin(); it != CBuffers.end(); it++) {
    DxilCBuffer *CB = it->get();
    GlobalVariable *GV = dyn_cast<GlobalVariable>(CB->GetGlobalSymbol());
    if (GV == nullptr)
      continue;
    CBufferVars.emplace_back(GV, CB->GetHLSLType());
  }

  // Collect tbuffers
  for (auto &it : m_DM->GetSRVs()) {
    if (it->GetKind() != DXIL::ResourceKind::TBuffer)
      continue;
    GlobalVariable *GV = dyn_cast<GlobalVariable>(it->GetGlobalSymbol());
    if (GV == nullptr)
      continue;
    CBufferVars.emplace_back(GV, it->GetHLSLType());
  }

  for (auto GV_Ty : CBufferVars) {
    auto GV = GV_Ty.first;
    Type *ElemTy = GV_Ty.second->getPointerElementType();
    ElemTy = dxilutil::StripArrayTypes(ElemTy, nullptr);
    StructType *ST = cast<StructType>(ElemTy);
    DxilStructAnnotation *SA = TypeSys.GetStructAnnotation(ST);
    if (SA == nullptr)
      continue;
    // If elements < 2, it's used if it exists.
    // Only old-style cbuffer { ... } will have more than one member, and
    // old-style cbuffers are the only ones that report usage per member.
    if (ST->getStructNumElements() < 2) {
      continue;
    }

    // Create offset maps for legacy layout and new compact layout, while
    // resetting usage flags
    const StructLayout *SL = DL.getStructLayout(ST);
    FieldAnnotationByOffsetMap legacyFieldMap, newFieldMap;
    for (unsigned i = 0; i < SA->GetNumFields(); ++i) {
      DxilFieldAnnotation &FA = SA->GetFieldAnnotation(i);
      FA.SetCBVarUsed(false);
      legacyFieldMap[FA.GetCBufferOffset()] = &FA;
      newFieldMap[(unsigned)SL->getElementOffset(i)] = &FA;
    }
    CollectCBufferMemberUsage(GV, legacyFieldMap, newFieldMap, hlslOP,
                              m_DM->GetUseMinPrecision(), visited);
  }
}

void DxilLowerCreateHandleForLib::SetNonUniformIndexForDynamicResource(
    DxilModule &DM) {
  hlsl::OP *hlslOP = DM.GetOP();
  Value *TrueVal = hlslOP->GetI1Const(true);
  for (auto it : hlslOP->GetOpFuncList(DXIL::OpCode::CreateHandleFromHeap)) {
    Function *F = it.second;
    if (!F)
      continue;
    for (User *U : F->users()) {
      CallInst *CI = cast<CallInst>(U);
      if (!DxilMDHelper::IsMarkedNonUniform(CI))
        continue;
      // Set NonUniform to be true.
      CI->setOperand(DxilInst_CreateHandleFromHeap::arg_nonUniformIndex,
                     TrueVal);
      // Clear nonUniform metadata.
      CI->setMetadata(DxilMDHelper::kDxilNonUniformAttributeMDName, nullptr);
    }
  }
}

// Remove createHandleFromHandle when not a lib
void DxilLowerCreateHandleForLib::RemoveCreateHandleFromHandle(DxilModule &DM) {
  hlsl::OP *hlslOP = DM.GetOP();
  Type *HdlTy = hlslOP->GetHandleType();
  for (auto it : hlslOP->GetOpFuncList(DXIL::OpCode::CreateHandleForLib)) {
    Function *F = it.second;
    if (!F)
      continue;
    if (it.first != HdlTy)
      continue;
    for (auto it = F->users().begin(); it != F->users().end();) {
      User *U = *(it++);
      CallInst *CI = cast<CallInst>(U);
      DxilInst_CreateHandleForLib Hdl(CI);
      Value *Res = Hdl.get_Resource();
      CI->replaceAllUsesWith(Res);
      CI->eraseFromParent();
    }
    break;
  }
}

char DxilLowerCreateHandleForLib::ID = 0;

ModulePass *llvm::createDxilLowerCreateHandleForLibPass() {
  return new DxilLowerCreateHandleForLib();
}

INITIALIZE_PASS_BEGIN(DxilLowerCreateHandleForLib,
                      "hlsl-dxil-lower-handle-for-lib",
                      "DXIL Lower createHandleForLib", false, false)
INITIALIZE_PASS_DEPENDENCY(DxilValueCache)
INITIALIZE_PASS_END(DxilLowerCreateHandleForLib,
                    "hlsl-dxil-lower-handle-for-lib",
                    "DXIL Lower createHandleForLib", false, false)

class DxilAllocateResourcesForLib : public ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid
  explicit DxilAllocateResourcesForLib()
      : ModulePass(ID), m_AutoBindingSpace(UINT_MAX) {}

  void applyOptions(PassOptions O) override {
    GetPassOptionUInt32(O, "auto-binding-space", &m_AutoBindingSpace, UINT_MAX);
  }
  StringRef getPassName() const override {
    return "DXIL Allocate Resources For Library";
  }

  bool runOnModule(Module &M) override {
    DxilModule &DM = M.GetOrCreateDxilModule();
    // Must specify a default space, and must apply to library.
    // Use DxilCondenseResources instead for shaders.
    if ((m_AutoBindingSpace == UINT_MAX) || !DM.GetShaderModel()->IsLib())
      return false;

    bool hasResource = DM.GetCBuffers().size() || DM.GetUAVs().size() ||
                       DM.GetSRVs().size() || DM.GetSamplers().size();

    if (hasResource) {
      DM.SetAutoBindingSpace(m_AutoBindingSpace);

      DxilResourceRegisterAllocator ResourceRegisterAllocator;
      ResourceRegisterAllocator.AllocateRegisters(DM);
    }
    return true;
  }

private:
  uint32_t m_AutoBindingSpace;
};

char DxilAllocateResourcesForLib::ID = 0;

ModulePass *llvm::createDxilAllocateResourcesForLibPass() {
  return new DxilAllocateResourcesForLib();
}

INITIALIZE_PASS(DxilAllocateResourcesForLib,
                "hlsl-dxil-allocate-resources-for-lib",
                "DXIL Allocate Resources For Library", false, false)

namespace {
struct CreateHandleFromHeapArgs {
  Value *Index;
  bool isSampler;
  bool isNonUniform;
  // All incoming handle args are confirmed.
  // If not resolved, some of the incoming handle is not from
  // createHandleFromHeap. Might be resolved after linking for lib.
  bool isResolved;
  void merge(CreateHandleFromHeapArgs &args, ResourceUseErrors &Errors,
             Value *mergeHdl) {
    if (args.isSampler != isSampler) {
      // Report error.
      Errors.ReportError(ResourceUseErrors::ErrorCode::MismatchIsSampler,
                         mergeHdl);
    }
    args.isNonUniform |= isNonUniform;
  }
};

} // namespace

// Helper class for legalizing dynamic resource use
// Convert select/phi on resources to select/phi on index.
// TODO: support case when save dynamic resource as local array element.
// TODO: share code with LegalizeResourceUseHelper.
class LegalizeDynamicResourceUseHelper {

public:
  ResourceUseErrors m_Errors;
  DenseMap<Value *, CreateHandleFromHeapArgs> HandleToArgs;
  // Value sets we can use to iterate
  ValueSetVector HandleSelects;
  ResourceUseErrors Errors;
  std::unordered_set<Instruction *> CleanupInsts;

  void mergeHeapArgs(Value *SelHdl, Value *SelIdx, User::op_range Hdls) {
    CreateHandleFromHeapArgs args = {nullptr, false, false, true};

    for (Value *V : Hdls) {
      auto it = HandleToArgs.find(V);
      // keep invalid when V is not createHandleFromHeap.
      if (it == HandleToArgs.end()) {
        args.isResolved = false;
        continue;
      }
      CreateHandleFromHeapArgs &itArgs = it->second;
      if (!itArgs.isResolved) {
        args.isResolved = false;
        continue;
      }

      if (args.Index != nullptr) {
        args.merge(itArgs, Errors, SelHdl);
      } else {
        args.Index = SelIdx;
        args.isNonUniform = itArgs.isNonUniform;
        args.isSampler = itArgs.isSampler;
      }
    }
    // set Index when all incoming Hdls cannot be resolved.
    if (args.Index == nullptr)
      args.Index = SelIdx;
    HandleToArgs[SelHdl] = args;
  }

  void CreateSelectsForHandleSelects() {
    if (HandleSelects.empty())
      return;

    LLVMContext &Ctx = HandleSelects[0]->getContext();
    Type *i32Ty = Type::getInt32Ty(Ctx);
    Value *UndefValue = UndefValue::get(i32Ty);
    // Create select for each HandleSelects.
    for (auto &Select : HandleSelects) {
      if (PHINode *Phi = dyn_cast<PHINode>(Select)) {
        IRBuilder<> B(Phi);
        unsigned numIncoming = Phi->getNumIncomingValues();
        PHINode *newPhi = B.CreatePHI(i32Ty, numIncoming);
        for (unsigned j = 0; j < numIncoming; j++) {
          // Set incoming values to undef until next pass
          newPhi->addIncoming(UndefValue, Phi->getIncomingBlock(j));
        }
        mergeHeapArgs(Phi, newPhi, Phi->incoming_values());
      } else if (SelectInst *Sel = dyn_cast<SelectInst>(Select)) {
        IRBuilder<> B(Sel);
        Value *newSel =
            B.CreateSelect(Sel->getCondition(), UndefValue, UndefValue);
        User::op_range range = User::op_range(Sel->getOperandList() + 1,
                                              Sel->getOperandList() + 3);
        mergeHeapArgs(Sel, newSel, range);
      } else {
        DXASSERT(false, "otherwise, non-select/phi in Selects set");
      }
    }
  }
  // propagate CreateHandleFromHeapArgs for HandleSel which all operands are
  // other HandleSel.
  void PropagateHeapArgs() {
    SmallVector<Value *, 4> Candidates;
    for (auto &Select : HandleSelects) {
      CreateHandleFromHeapArgs &args = HandleToArgs[Select];
      if (args.isResolved)
        continue;
      Candidates.emplace_back(Select);
    }

    while (1) {
      SmallVector<Value *, 4> NextPass;
      for (auto &Select : Candidates) {
        CreateHandleFromHeapArgs &args = HandleToArgs[Select];
        if (PHINode *Phi = dyn_cast<PHINode>(Select)) {
          mergeHeapArgs(Phi, args.Index, Phi->incoming_values());
        } else if (SelectInst *Sel = dyn_cast<SelectInst>(Select)) {
          User::op_range range = User::op_range(Sel->getOperandList() + 1,
                                                Sel->getOperandList() + 3);
          mergeHeapArgs(Sel, args.Index, range);
        } else {
          DXASSERT(false, "otherwise, non-select/phi in Selects set");
        }

        if (args.isResolved)
          continue;
        NextPass.emplace_back(Select);
      }
      // Some node cannot be reached.
      if (NextPass.size() == Candidates.size())
        return;
      Candidates = NextPass;
    }
  }

  void UpdateSelectsForHandleSelect(hlsl::OP *hlslOP) {
    if (HandleSelects.empty())
      return;
    LLVMContext &Ctx = HandleSelects[0]->getContext();
    Type *pVoidTy = Type::getVoidTy(Ctx);
    // NOTE: phi of createHandleFromHeap and createHandleFromBinding
    // is not supported.
    Function *createHdlFromHeap =
        hlslOP->GetOpFunc(DXIL::OpCode::CreateHandleFromHeap, pVoidTy);
    Value *hdlFromHeapOP = hlslOP->GetI32Const(
        static_cast<unsigned>(DXIL::OpCode::CreateHandleFromHeap));
    for (auto &Select : HandleSelects) {
      if (PHINode *Phi = dyn_cast<PHINode>(Select)) {
        unsigned numIncoming = Phi->getNumIncomingValues();
        CreateHandleFromHeapArgs &args = HandleToArgs[Phi];
        PHINode *newPhi = cast<PHINode>(args.Index);
        if (args.isResolved) {
          for (unsigned j = 0; j < numIncoming; j++) {
            Value *V = Phi->getIncomingValue(j);
            auto it = HandleToArgs.find(V);
            DXASSERT(it != HandleToArgs.end(),
                     "args.isResolved should be false");
            CreateHandleFromHeapArgs &itArgs = it->second;
            newPhi->setIncomingValue(j, itArgs.Index);
          }

          IRBuilder<> B(Phi->getParent()->getFirstNonPHI());
          B.SetCurrentDebugLocation(Phi->getDebugLoc());
          Value *isSampler = hlslOP->GetI1Const(args.isSampler);
          // TODO: or args.IsNonUniform with !isUniform(Phi) with uniform
          // analysis.
          Value *isNonUniform = hlslOP->GetI1Const(args.isNonUniform);
          CallInst *newCI =
              B.CreateCall(createHdlFromHeap,
                           {hdlFromHeapOP, newPhi, isSampler, isNonUniform});
          Phi->replaceAllUsesWith(newCI);
          CleanupInsts.insert(Phi);
          // put newCI in HandleToArgs.
          HandleToArgs[newCI] = args;
        } else {
          newPhi->eraseFromParent();
        }
      } else if (SelectInst *Sel = dyn_cast<SelectInst>(Select)) {
        CreateHandleFromHeapArgs &args = HandleToArgs[Sel];
        SelectInst *newSel = cast<SelectInst>(args.Index);
        if (args.isResolved) {
          for (unsigned j = 1; j < 3; ++j) {
            Value *V = Sel->getOperand(j);
            auto it = HandleToArgs.find(V);
            DXASSERT(it != HandleToArgs.end(),
                     "args.isResolved should be false");
            CreateHandleFromHeapArgs &itArgs = it->second;
            newSel->setOperand(j, itArgs.Index);
          }

          IRBuilder<> B(newSel->getNextNode());
          B.SetCurrentDebugLocation(newSel->getDebugLoc());
          Value *isSampler = hlslOP->GetI1Const(args.isSampler);
          // TODO: or args.IsNonUniform with !isUniform(Phi).
          Value *isNonUniform = hlslOP->GetI1Const(args.isNonUniform);
          CallInst *newCI =
              B.CreateCall(createHdlFromHeap,
                           {hdlFromHeapOP, newSel, isSampler, isNonUniform});
          Sel->replaceAllUsesWith(newCI);
          CleanupInsts.insert(Sel);
          // put newCI in HandleToArgs.
          HandleToArgs[newCI] = args;
        } else {
          newSel->eraseFromParent();
        }
      } else {
        DXASSERT(false, "otherwise, non-select/phi in HandleSelects set");
      }
    }
  }

  void CollectResources(DxilModule &DM) {
    ValueSetVector tmpHandleSelects;
    hlsl::OP *hlslOP = DM.GetOP();
    if (hlslOP->IsDxilOpUsed(DXIL::OpCode::CreateHandleFromHeap)) {
      Function *F = hlslOP->GetOpFunc(DXIL::OpCode::CreateHandleFromHeap,
                                      Type::getVoidTy(DM.GetCtx()));
      for (User *U : F->users()) {
        DxilInst_CreateHandleFromHeap Hdl(cast<CallInst>(U));
        HandleToArgs[U] = {Hdl.get_index(), Hdl.get_samplerHeap_val(),
                           Hdl.get_nonUniformIndex_val(), true};
        for (User *HdlU : U->users()) {
          if (isa<PHINode>(HdlU) || isa<SelectInst>(HdlU)) {
            tmpHandleSelects.insert(HdlU);
          }
        }
      }
    }
    // Collect phi/sel of other phi/sel selected handles.
    while (!tmpHandleSelects.empty()) {
      HandleSelects.insert(tmpHandleSelects.begin(), tmpHandleSelects.end());
      ValueSetVector newHandleSelects;
      for (Value *Hdl : tmpHandleSelects) {
        for (User *HdlU : Hdl->users()) {
          if (HandleSelects.count(HdlU))
            continue;
          if (isa<PHINode>(HdlU) || isa<SelectInst>(HdlU)) {
            newHandleSelects.insert(HdlU);
          }
        }
      }
      tmpHandleSelects = newHandleSelects;
    }
  }

  void DoTransform(hlsl::OP *hlslOP) {
    CreateSelectsForHandleSelects();
    PropagateHeapArgs();
    UpdateSelectsForHandleSelect(hlslOP);
    CleanupUnusedValues(CleanupInsts);
  }

  bool runOnModule(llvm::Module &M) {
    DxilModule &DM = M.GetOrCreateDxilModule();
    hlsl::OP *hlslOP = DM.GetOP();

    CollectResources(DM);

    // If no selects or allocas are involved, there isn't anything to do
    if (HandleSelects.empty())
      return false;

    DoTransform(hlslOP);

    return true;
  }
};

namespace {
// Make sure no phi/sel on annotateHandle.
bool sinkAnnotateHandleAfterSelect(DxilModule &DM, Module &M) {
  // Collect AnnotateHandle calls.
  SmallVector<CallInst *, 4> annotHdls;
  hlsl::OP *op = DM.GetOP();
  LLVMContext &Ctx = M.getContext();
  Type *pVoidTy = Type::getVoidTy(Ctx);
  Function *annotHdlFn = op->GetOpFunc(DXIL::OpCode::AnnotateHandle, pVoidTy);
  for (auto it : op->GetOpFuncList(OP::OpCode::AnnotateHandle)) {
    Function *F = it.second;
    if (F == nullptr)
      continue;
    for (auto U = F->user_begin(); U != F->user_end();) {
      CallInst *CI = dyn_cast<CallInst>(*(U++));
      annotHdls.emplace_back(CI);
    }
  }
  if (annotHdls.empty())
    return false;

  SetVector<Instruction *> selectAnnotHdls;
  for (CallInst *CI : annotHdls) {
    for (User *U : CI->users()) {
      if (isa<PHINode>(U) || isa<SelectInst>(U))
        selectAnnotHdls.insert(cast<Instruction>(U));
    }
  }
  const ShaderModel *pSM = DM.GetShaderModel();
  Type *propsTy = op->GetResourcePropertiesType();
  Value *OpArg =
      op->GetI32Const(static_cast<unsigned>(DXIL::OpCode::AnnotateHandle));
  ResourceUseErrors Errors;
  Value *undefHdl = UndefValue::get(op->GetHandleType());
  // Sink annotateHandle after phi.
  for (Instruction *Hdl : selectAnnotHdls) {
    if (PHINode *phi = dyn_cast<PHINode>(Hdl)) {
      Value *props = nullptr;
      for (unsigned i = 0; i < phi->getNumIncomingValues(); ++i) {
        Value *V = phi->getIncomingValue(i);
        if (CallInst *CI = dyn_cast<CallInst>(V)) {
          DxilInst_AnnotateHandle annot(CI);
          if (annot) {
            if (props == nullptr) {
              props = annot.get_props();
            } else if (props != annot.get_props()) {
              props = resource_helper::tryMergeProps(
                  cast<Constant>(props), cast<Constant>(annot.get_props()),
                  propsTy, *pSM);
              if (props == nullptr) {
                Errors.ReportError(
                    ResourceUseErrors::ErrorCode::MismatchHandleAnnotation,
                    phi);
                props = annot.get_props();
              }
            }

            Value *res = annot.get_res();
            phi->setIncomingValue(i, res);
          }
        }
      }
      // Insert after phi.
      IRBuilder<> B(phi->getParent()->getFirstNonPHI());
      CallInst *annotCI = B.CreateCall(annotHdlFn, {OpArg, undefHdl, props});
      phi->replaceAllUsesWith(annotCI);
      annotCI->setArgOperand(DxilInst_AnnotateHandle::arg_res, phi);
    } else {
      SelectInst *sel = dyn_cast<SelectInst>(Hdl);
      Value *TVal = sel->getTrueValue();
      Value *FVal = sel->getFalseValue();
      Value *props = nullptr;
      if (CallInst *CI = dyn_cast<CallInst>(TVal)) {
        DxilInst_AnnotateHandle annot(CI);
        if (annot) {
          props = annot.get_props();
          Value *res = annot.get_res();
          sel->setOperand(1, res);
        }
      }
      if (CallInst *CI = dyn_cast<CallInst>(FVal)) {
        DxilInst_AnnotateHandle annot(CI);
        if (annot) {
          if (props == nullptr) {
            props = annot.get_props();
          } else if (props != annot.get_props()) {
            props = resource_helper::tryMergeProps(
                cast<Constant>(props), cast<Constant>(annot.get_props()),
                propsTy, *pSM);
            if (props == nullptr) {
              Errors.ReportError(
                  ResourceUseErrors::ErrorCode::MismatchHandleAnnotation, sel);
              props = annot.get_props();
            }
          }

          Value *res = annot.get_res();
          sel->setOperand(2, res);
        }
      }

      // Insert after sel.
      IRBuilder<> B(sel->getNextNode());
      CallInst *annotCI = B.CreateCall(annotHdlFn, {OpArg, undefHdl, props});
      sel->replaceAllUsesWith(annotCI);
      annotCI->setArgOperand(DxilInst_AnnotateHandle::arg_res, sel);
    }
  }
  return true;
}
} // namespace

// Remove redudant annotateHandle.
// Legalize phi on createHandleFromHeap.
class DxilCleanupDynamicResourceHandle : public ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid
  explicit DxilCleanupDynamicResourceHandle() : ModulePass(ID) {}

  StringRef getPassName() const override {
    return "DXIL Cleanup dynamic resource handle calls";
  }

  bool runOnModule(Module &M) override {
    DxilModule &DM = M.GetOrCreateDxilModule();

    // Nothing to do if Dxil ver < 1.6
    unsigned dxilMajor, dxilMinor;
    DM.GetShaderModel()->GetDxilVersion(dxilMajor, dxilMinor);
    if (DXIL::CompareVersions(dxilMajor, dxilMinor, 1, 6) < 0)
      return false;

    bool bChanged = sinkAnnotateHandleAfterSelect(DM, M);

    // Legalize phi on createHandleFromHeap.
    LegalizeDynamicResourceUseHelper helper;
    bChanged |= helper.runOnModule(M);

    hlsl::OP *op = DM.GetOP();
    const ShaderModel *pSM = DM.GetShaderModel();
    Type *propsTy = op->GetResourcePropertiesType();
    // Iterate AnnotateHandle calls and eliminate redundant annotate handle call
    // chains.
    for (auto it : op->GetOpFuncList(OP::OpCode::AnnotateHandle)) {
      Function *F = it.second;
      if (F == nullptr)
        continue;
      for (auto U = F->user_begin(); U != F->user_end();) {
        CallInst *CI = dyn_cast<CallInst>(*(U++));
        if (CI) {
          DxilInst_AnnotateHandle AH(CI);
          if (AH) {
            Value *Res = AH.get_res();
            // Skip handle from load global res.
            if (isa<LoadInst>(Res))
              continue;
            CallInst *CRes = dyn_cast<CallInst>(Res);
            if (!CRes)
              continue;
            DxilInst_AnnotateHandle PrevAH(CRes);
            if (PrevAH) {
              Value *mergedProps = resource_helper::tryMergeProps(
                  cast<Constant>(AH.get_props()),
                  cast<Constant>(PrevAH.get_props()), propsTy, *pSM);
              if (mergedProps == nullptr) {
                ResourceUseErrors Errors;
                Errors.ReportError(
                    ResourceUseErrors::ErrorCode::MismatchHandleAnnotation, CI);
              } else if (mergedProps != PrevAH.get_props()) {
                PrevAH.set_props(mergedProps);
              }
              CI->replaceAllUsesWith(Res);
              CI->eraseFromParent();
              bChanged = true;
            }
          }
        }
      }
    }

    return bChanged;
  }

private:
};

char DxilCleanupDynamicResourceHandle::ID = 0;

ModulePass *llvm::createDxilCleanupDynamicResourceHandlePass() {
  return new DxilCleanupDynamicResourceHandle();
}

INITIALIZE_PASS(DxilCleanupDynamicResourceHandle,
                "hlsl-dxil-cleanup-dynamic-resource-handle",
                "DXIL Cleanup dynamic resource handle calls", false, false)