//===- DXILResourceAccess.cpp - Resource access via load/store ------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "DXILResourceAccess.h"
#include "DirectX.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/DXILResource.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/Analysis/VectorUtils.h"
#include "llvm/Frontend/HLSL/HLSLResource.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/IntrinsicsDirectX.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/IR/User.h"
#include "llvm/IR/ValueHandle.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/DXILABI.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/KnownBits.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include <optional>

#define DEBUG_TYPE "dxil-resource-access"

using namespace llvm;

[[noreturn]] static void
diagnoseNonUniqueResourceAccess(Instruction *I,
                                ArrayRef<IntrinsicInst *> Handles) {
  LLVMContext &Context = I->getContext();
  std::string InstStr;
  raw_string_ostream InstOS(InstStr);
  I->print(InstOS);
  Context.diagnose(
      DiagnosticInfoGeneric("At resource access:" + Twine(InstStr), DS_Note));

  for (auto *Handle : Handles) {
    std::string HandleStr;
    raw_string_ostream HandleOS(HandleStr);
    Handle->print(HandleOS);
    Context.diagnose(DiagnosticInfoGeneric(
        "Uses resource handle:" + Twine(HandleStr), DS_Note));
  }
  report_fatal_error(
      "Resource access is not guaranteed to map to a unique global resource",
      /*gen_crash_diag=*/false);
}

static Value *traverseGEPOffsets(const DataLayout &DL, IRBuilder<> &Builder,
                                 Value *Ptr, uint64_t AccessSize) {
  Value *Offset = nullptr;

  while (Ptr) {
    if ([[maybe_unused]] auto *II = dyn_cast<IntrinsicInst>(Ptr)) {
      assert((II->getIntrinsicID() == Intrinsic::dx_resource_getpointer ||
              II->getIntrinsicID() == Intrinsic::dx_resource_getbasepointer) &&
             "Resource access through unexpected intrinsic");
      return Offset ? Offset : ConstantInt::get(Builder.getInt32Ty(), 0);
    }

    auto *GEP = dyn_cast<GetElementPtrInst>(Ptr);
    assert(GEP && "Resource access through unexpected instruction");

    unsigned NumIndices = GEP->getNumIndices();
    uint64_t IndexScale = DL.getTypeAllocSize(GEP->getSourceElementType());
    APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
    Value *GEPOffset;
    if (GEP->accumulateConstantOffset(DL, ConstantOffset)) {
      // We have a constant offset (in bytes).
      GEPOffset =
          ConstantInt::get(DL.getIndexType(GEP->getType()), ConstantOffset);
      IndexScale = 1;
    } else if (NumIndices == 1) {
      // If we have a single index we're indexing into a top level array. This
      // generally only happens with cbuffers.
      GEPOffset = *GEP->idx_begin();
    } else if (NumIndices == 2) {
      // If we have two indices, this should be an access through a pointer.
      auto *IndexIt = GEP->idx_begin();
      assert(cast<ConstantInt>(IndexIt)->getZExtValue() == 0 &&
             "GEP is not indexing through pointer");
      GEPOffset = *(++IndexIt);
    } else
      llvm_unreachable("Unhandled GEP structure for resource access");

    uint64_t ElemSize = AccessSize;
    if (!(IndexScale % ElemSize)) {
      // If our scale is an exact multiple of the access size, adjust the
      // scaling to avoid an unnecessary division.
      IndexScale /= ElemSize;
      ElemSize = 1;
    }
    if (IndexScale != 1)
      GEPOffset = Builder.CreateMul(
          GEPOffset, ConstantInt::get(Builder.getInt32Ty(), IndexScale));
    if (ElemSize != 1)
      GEPOffset = Builder.CreateUDiv(
          GEPOffset, ConstantInt::get(Builder.getInt32Ty(), ElemSize));

    Offset = Offset ? Builder.CreateAdd(Offset, GEPOffset) : GEPOffset;
    Ptr = GEP->getPointerOperand();
  }

  llvm_unreachable("GEP of null pointer?");
}

static void createTypedBufferStore(IntrinsicInst *II, StoreInst *SI,
                                   dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = SI->getDataLayout();
  IRBuilder<> Builder(SI);
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
  Type *ScalarType = ContainedType->getScalarType();
  Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());

  Value *V = SI->getValueOperand();
  if (V->getType() == ContainedType) {
    // V is already the right type.
    assert(SI->getPointerOperand() == II &&
           "Store of whole element has mismatched address to store to");
  } else if (V->getType() == ScalarType) {
    // We're storing a scalar, so we need to load the current value and only
    // replace the relevant part.
    auto *Load = Builder.CreateIntrinsic(
        LoadType, Intrinsic::dx_resource_load_typedbuffer,
        {II->getOperand(0), II->getOperand(1)});
    auto *Struct = Builder.CreateExtractValue(Load, {0});

    uint64_t AccessSize = DL.getTypeSizeInBits(ScalarType) / 8;
    Value *Offset =
        traverseGEPOffsets(DL, Builder, SI->getPointerOperand(), AccessSize);
    V = Builder.CreateInsertElement(Struct, V, Offset);
  } else {
    llvm_unreachable("Store to typed resource has invalid type");
  }

  auto *Inst = Builder.CreateIntrinsic(
      Builder.getVoidTy(), Intrinsic::dx_resource_store_typedbuffer,
      {II->getOperand(0), II->getOperand(1), V});
  SI->replaceAllUsesWith(Inst);
}

/// Build a zero-initialized offset operand matching the shape of the given
/// coordinate operand. Accesses through `operator[]` never have offsets.
static Value *getNullOffsetsFor(IRBuilder<> &Builder, Value *Coords) {
  Type *CoordTy = Coords->getType();
  Type *OffsetTy;
  if (auto *VecTy = dyn_cast<FixedVectorType>(CoordTy))
    OffsetTy =
        FixedVectorType::get(Builder.getInt32Ty(), VecTy->getNumElements());
  else
    OffsetTy = Builder.getInt32Ty();
  return Constant::getNullValue(OffsetTy);
}

static void createTextureStore(IntrinsicInst *II, StoreInst *SI,
                               dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = SI->getDataLayout();
  IRBuilder<> Builder(SI);
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
  Type *ScalarType = ContainedType->getScalarType();

  Value *Handle = II->getOperand(0);
  Value *Coords = II->getOperand(1);

  Value *V = SI->getValueOperand();
  if (V->getType() == ContainedType) {
    // V is already the right type.
    assert(SI->getPointerOperand() == II &&
           "Store of whole element has mismatched address to store to");
  } else if (V->getType() == ScalarType) {
    // We're storing a scalar, so we need to load the current value and only
    // replace the relevant part. For operator[] the mip level and the offsets
    // are always zero; DXILOpLowering drops the mip level for UAVs.
    Value *MipLevel = Builder.getInt32(0);
    Value *Offsets = getNullOffsetsFor(Builder, Coords);
    auto *Load = Builder.CreateIntrinsic(ContainedType,
                                         Intrinsic::dx_resource_load_level,
                                         {Handle, Coords, MipLevel, Offsets});

    uint64_t AccessSize = DL.getTypeSizeInBits(ScalarType) / 8;
    Value *Offset =
        traverseGEPOffsets(DL, Builder, SI->getPointerOperand(), AccessSize);
    V = Builder.CreateInsertElement(Load, V, Offset);
  } else {
    llvm_unreachable("Store to texture resource has invalid type");
  }

  auto *Inst = Builder.CreateIntrinsic(Builder.getVoidTy(),
                                       Intrinsic::dx_resource_store_texture,
                                       {Handle, Coords, V});
  SI->replaceAllUsesWith(Inst);
}

static void emitRawStore(IRBuilder<> &Builder, Value *Buffer, Value *Index,
                         Value *Offset, Value *V, dxil::ResourceTypeInfo &RTI) {
  // For raw buffer (ie, HLSL's ByteAddressBuffer), we need to fold the access
  // entirely into the index.
  if (!RTI.isStruct()) {
    auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
    if (!ConstantOffset || !ConstantOffset->isZero())
      Index = Builder.CreateAdd(Index, Offset);
    Offset = llvm::PoisonValue::get(Builder.getInt32Ty());
  }

  Builder.CreateIntrinsic(Builder.getVoidTy(),
                          Intrinsic::dx_resource_store_rawbuffer,
                          {Buffer, Index, Offset, V});
}

static void createRawStores(IntrinsicInst *II, StoreInst *SI,
                            dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = SI->getDataLayout();
  IRBuilder<> Builder(SI);

  Value *V = SI->getValueOperand();
  assert(!V->getType()->isAggregateType() &&
         "Resource store should be scalar or vector type");

  Value *Index = II->getOperand(1);
  // The offset for the rawbuffer load and store ops is always in bytes.
  uint64_t AccessSize = 1;
  Value *Offset =
      traverseGEPOffsets(DL, Builder, SI->getPointerOperand(), AccessSize);

  auto *VT = dyn_cast<FixedVectorType>(V->getType());
  if (VT && VT->getNumElements() > 4) {
    // Split into stores of at most 4 elements.
    Type *EltTy = VT->getElementType();
    Value *Stride = ConstantInt::get(Builder.getInt32Ty(),
                                     4 * (DL.getTypeSizeInBits(EltTy) / 8));

    SmallVector<int, 4> Indices;
    for (unsigned int I = 0, N = VT->getNumElements(); I < N; I += 4) {
      if (I > 0)
        Offset = Builder.CreateAdd(Offset, Stride);

      for (unsigned int J = I, E = std::min(N, J + 4); J < E; ++J)
        Indices.push_back(J);
      Value *Part = Builder.CreateShuffleVector(V, Indices);
      emitRawStore(Builder, II->getOperand(0), Index, Offset, Part, RTI);

      Indices.clear();
    }
  } else
    emitRawStore(Builder, II->getOperand(0), Index, Offset, V, RTI);
}

static void createStoreIntrinsic(IntrinsicInst *II, StoreInst *SI,
                                 dxil::ResourceTypeInfo &RTI) {
  switch (RTI.getResourceKind()) {
  case dxil::ResourceKind::TypedBuffer:
    return createTypedBufferStore(II, SI, RTI);
  case dxil::ResourceKind::RawBuffer:
  case dxil::ResourceKind::StructuredBuffer:
    return createRawStores(II, SI, RTI);
  case dxil::ResourceKind::Texture1D:
  case dxil::ResourceKind::Texture2D:
  case dxil::ResourceKind::Texture3D:
  case dxil::ResourceKind::Texture1DArray:
  case dxil::ResourceKind::Texture2DArray:
    return createTextureStore(II, SI, RTI);
  case dxil::ResourceKind::Texture2DMS:
  case dxil::ResourceKind::Texture2DMSArray:
  case dxil::ResourceKind::TextureCube:
  case dxil::ResourceKind::TextureCubeArray:
  case dxil::ResourceKind::FeedbackTexture2D:
  case dxil::ResourceKind::FeedbackTexture2DArray:
    reportFatalUsageError(
        "DXIL Store not implemented for this texture resource kind");
    return;
  case dxil::ResourceKind::CBuffer:
  case dxil::ResourceKind::Sampler:
  case dxil::ResourceKind::TBuffer:
  case dxil::ResourceKind::RTAccelerationStructure:
  case dxil::ResourceKind::Invalid:
  case dxil::ResourceKind::NumEntries:
    llvm_unreachable("Invalid resource kind for store");
  }
  llvm_unreachable("Unhandled case in switch");
}

static std::optional<dxil::AtomicBinOpCode>
getAtomicBinOpCode(AtomicRMWInst::BinOp BinOp) {
  switch (BinOp) {
  case AtomicRMWInst::Add:
    return dxil::AtomicBinOpCode::Add;
  case AtomicRMWInst::And:
    return dxil::AtomicBinOpCode::And;
  case AtomicRMWInst::Or:
    return dxil::AtomicBinOpCode::Or;
  case AtomicRMWInst::Xor:
    return dxil::AtomicBinOpCode::Xor;
  case AtomicRMWInst::Min:
    return dxil::AtomicBinOpCode::IMin;
  case AtomicRMWInst::Max:
    return dxil::AtomicBinOpCode::IMax;
  case AtomicRMWInst::UMin:
    return dxil::AtomicBinOpCode::UMin;
  case AtomicRMWInst::UMax:
    return dxil::AtomicBinOpCode::UMax;
  case AtomicRMWInst::Xchg:
    return dxil::AtomicBinOpCode::Exchange;
  case AtomicRMWInst::Sub:
  case AtomicRMWInst::Nand:
  case AtomicRMWInst::FAdd:
  case AtomicRMWInst::FSub:
  case AtomicRMWInst::FMax:
  case AtomicRMWInst::FMin:
  case AtomicRMWInst::FMaximum:
  case AtomicRMWInst::FMinimum:
  case AtomicRMWInst::FMaximumNum:
  case AtomicRMWInst::FMinimumNum:
  case AtomicRMWInst::UIncWrap:
  case AtomicRMWInst::UDecWrap:
  case AtomicRMWInst::USubCond:
  case AtomicRMWInst::USubSat:
  case AtomicRMWInst::BAD_BINOP:
    return std::nullopt;
  }
  llvm_unreachable("Unhandled atomicrmw operation");
}

// Compute the (coord0, coord1) pair for a buffer resource atomic operation.
// Non-struct buffers (RawBuffer or TypedBuffer) fold the byte offset into the
// index and leave coord1 poison. Only StructuredBuffer atomics use both a
// struct index and a byte offset.
static std::pair<Value *, Value *>
getAtomicResourceCoords(IntrinsicInst *II, Value *PointerOperand,
                        dxil::ResourceTypeInfo &RTI, IRBuilder<> &Builder,
                        const DataLayout &DL) {
  Value *Index = II->getOperand(1);

  // The offset for the rawbuffer load/store/atomic ops is always in bytes.
  uint64_t AccessSize = 1;
  Value *Offset = traverseGEPOffsets(DL, Builder, PointerOperand, AccessSize);

  if (!RTI.isStruct()) {
    auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
    if (!ConstantOffset || !ConstantOffset->isZero())
      Index = Builder.CreateAdd(Index, Offset);
    Offset = llvm::PoisonValue::get(Builder.getInt32Ty());
  }

  return {Index, Offset};
}

// The coordinates of a texture access are a scalar or a vector with one element
// per texture dimension, including the array slice if there is one. These map
// directly onto the coordinate operands of the atomic ops.
static SmallVector<Value *, 3> getTextureAtomicCoords(IntrinsicInst *II,
                                                      IRBuilder<> &Builder) {
  Value *Coords = II->getOperand(1);
  SmallVector<Value *, 3> CoordArgs;
  if (auto *VecTy = dyn_cast<FixedVectorType>(Coords->getType())) {
    assert(VecTy->getNumElements() <= 3 && "Too many texture coordinates");
    for (unsigned I = 0, E = VecTy->getNumElements(); I != E; ++I)
      CoordArgs.push_back(Builder.CreateExtractElement(Coords, I));
  } else {
    CoordArgs.push_back(Coords);
  }
  return CoordArgs;
}

static void emitAtomicBinOp(IRBuilder<> &Builder, AtomicRMWInst *AI,
                            Value *Handle, ArrayRef<Value *> Coords) {
  assert(!Coords.empty() && Coords.size() <= 3 &&
         "Atomic operations take between one and three coordinates");

  std::optional<dxil::AtomicBinOpCode> BinOpCode =
      getAtomicBinOpCode(AI->getOperation());
  if (!BinOpCode) {
    reportFatalUsageError("DXIL resource atomicrmw operation not implemented");
    return;
  }

  // DXIL has no floating-point atomic op. A float exchange only moves the bit
  // pattern, so cast the value to an integer of the same width, exchange, and
  // cast the result back. This matches what DXC emits.
  Value *Val = AI->getValOperand();
  Type *ValTy = Val->getType();
  Type *OpTy = ValTy;
  if (ValTy->isFloatingPointTy()) {
    OpTy = Builder.getIntNTy(ValTy->getPrimitiveSizeInBits());
    Val = Builder.CreateBitCast(Val, OpTy);
  }

  SmallVector<Value *, 6> Args{
      Handle, Builder.getInt32(static_cast<uint32_t>(*BinOpCode))};
  append_range(Args, Coords);
  Args.append(3 - Coords.size(), PoisonValue::get(Builder.getInt32Ty()));
  Args.push_back(Val);

  // Emit the target-independent intrinsic; DXILOpLowering lowers it to the
  // DXIL `AtomicBinOp` op and handles the target-ext-typed handle cast via
  // its `createTmpHandleCast` bookkeeping.
  Value *Result =
      Builder.CreateIntrinsic(OpTy, Intrinsic::dx_resource_atomic_binop, Args);

  if (OpTy != ValTy)
    Result = Builder.CreateBitCast(Result, ValTy);

  AI->replaceAllUsesWith(Result);
}

static void createBufferAtomicBinOp(IntrinsicInst *II, AtomicRMWInst *AI,
                                    dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = AI->getDataLayout();
  IRBuilder<> Builder(AI);
  auto [Index, Offset] =
      getAtomicResourceCoords(II, AI->getPointerOperand(), RTI, Builder, DL);

  emitAtomicBinOp(Builder, AI, II->getOperand(0), {Index, Offset});
}

static void createTextureAtomicBinOp(IntrinsicInst *II, AtomicRMWInst *AI,
                                     dxil::ResourceTypeInfo &RTI) {
  // A texture atomic operates on a whole texel, so a multi-component texel has
  // no single addressable component. A scalar float texel is allowed, because
  // emitAtomicBinOp exchanges its bit pattern as an integer.
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
  if (!ContainedType->isIntegerTy() && !ContainedType->isFloatingPointTy()) {
    reportFatalUsageError("DXIL atomicrmw requires a texture resource with a "
                          "scalar element type");
    return;
  }

  IRBuilder<> Builder(AI);

  emitAtomicBinOp(Builder, AI, II->getOperand(0),
                  getTextureAtomicCoords(II, Builder));
}

static void emitAtomicCompareExchange(IRBuilder<> &Builder,
                                      AtomicCmpXchgInst *AI, Value *Handle,
                                      ArrayRef<Value *> Coords) {
  assert(!Coords.empty() && Coords.size() <= 3 &&
         "Atomic operations take between one and three coordinates");

  Value *Compare = AI->getCompareOperand();
  Value *NewValue = AI->getNewValOperand();

  SmallVector<Value *, 6> Args{Handle};
  append_range(Args, Coords);
  Args.append(3 - Coords.size(), PoisonValue::get(Builder.getInt32Ty()));
  Args.push_back(Compare);
  Args.push_back(NewValue);

  Value *Original = Builder.CreateIntrinsic(
      NewValue->getType(), Intrinsic::dx_resource_atomic_compare_exchange,
      Args);

  // `cmpxchg` yields a { original, success } pair, but the DXIL op returns
  // only the original value. DXIL has no way to express the success flag, and
  // no HLSL builtin reads it, so replace the users of the pair directly
  // instead of building it again. No pass after this one removes dead code.
  SmallVector<ExtractValueInst *> Extracts;
  for (User *U : AI->users()) {
    auto *EV = dyn_cast<ExtractValueInst>(U);
    if (!EV || EV->getIndices()[0] != 0)
      reportFatalUsageError("DXIL cmpxchg provides only the original value");
    Extracts.push_back(EV);
  }

  for (ExtractValueInst *EV : Extracts) {
    EV->replaceAllUsesWith(Original);
    EV->eraseFromParent();
  }
}

static void createBufferAtomicCompareExchange(IntrinsicInst *II,
                                              AtomicCmpXchgInst *AI,
                                              dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = AI->getDataLayout();
  IRBuilder<> Builder(AI);
  auto [Index, Offset] =
      getAtomicResourceCoords(II, AI->getPointerOperand(), RTI, Builder, DL);

  emitAtomicCompareExchange(Builder, AI, II->getOperand(0), {Index, Offset});
}

// `cmpxchg` operands are always integers, so unlike atomicrmw there is no
// float element type to convert here. The element type must be scalar.
static void createTextureAtomicCompareExchange(IntrinsicInst *II,
                                               AtomicCmpXchgInst *AI,
                                               dxil::ResourceTypeInfo &RTI) {
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
  if (!ContainedType->isIntegerTy() && !ContainedType->isFloatingPointTy()) {
    reportFatalUsageError("DXIL cmpxchg requires a texture resource with a "
                          "scalar element type");
    return;
  }

  IRBuilder<> Builder(AI);

  emitAtomicCompareExchange(Builder, AI, II->getOperand(0),
                            getTextureAtomicCoords(II, Builder));
}

static void createAtomicBinOpIntrinsic(IntrinsicInst *II, AtomicRMWInst *AI,
                                       dxil::ResourceTypeInfo &RTI) {
  switch (RTI.getResourceKind()) {
  case dxil::ResourceKind::TypedBuffer:
  case dxil::ResourceKind::RawBuffer:
  case dxil::ResourceKind::StructuredBuffer:
    return createBufferAtomicBinOp(II, AI, RTI);
  case dxil::ResourceKind::Texture1D:
  case dxil::ResourceKind::Texture2D:
  case dxil::ResourceKind::Texture3D:
  case dxil::ResourceKind::Texture1DArray:
  case dxil::ResourceKind::Texture2DArray:
    return createTextureAtomicBinOp(II, AI, RTI);
  case dxil::ResourceKind::Texture2DMS:
  case dxil::ResourceKind::Texture2DMSArray:
  case dxil::ResourceKind::TextureCube:
  case dxil::ResourceKind::TextureCubeArray:
  case dxil::ResourceKind::FeedbackTexture2D:
  case dxil::ResourceKind::FeedbackTexture2DArray:
    reportFatalUsageError(
        "DXIL atomicrmw not implemented for this texture resource kind");
    return;
  case dxil::ResourceKind::CBuffer:
  case dxil::ResourceKind::Sampler:
  case dxil::ResourceKind::TBuffer:
    reportFatalUsageError(
        "DXIL atomicrmw not implemented for this resource type");
    return;
  case dxil::ResourceKind::RTAccelerationStructure:
  case dxil::ResourceKind::Invalid:
  case dxil::ResourceKind::NumEntries:
    llvm_unreachable("Invalid resource kind for atomicrmw");
  }
  llvm_unreachable("Unhandled case in switch");
}

static void createAtomicCompareExchangeIntrinsic(IntrinsicInst *II,
                                                 AtomicCmpXchgInst *AI,
                                                 dxil::ResourceTypeInfo &RTI) {
  switch (RTI.getResourceKind()) {
  case dxil::ResourceKind::TypedBuffer:
  case dxil::ResourceKind::RawBuffer:
  case dxil::ResourceKind::StructuredBuffer:
    return createBufferAtomicCompareExchange(II, AI, RTI);
  case dxil::ResourceKind::Texture1D:
  case dxil::ResourceKind::Texture2D:
  case dxil::ResourceKind::Texture3D:
  case dxil::ResourceKind::Texture1DArray:
  case dxil::ResourceKind::Texture2DArray:
    return createTextureAtomicCompareExchange(II, AI, RTI);
  case dxil::ResourceKind::Texture2DMS:
  case dxil::ResourceKind::Texture2DMSArray:
  case dxil::ResourceKind::TextureCube:
  case dxil::ResourceKind::TextureCubeArray:
  case dxil::ResourceKind::FeedbackTexture2D:
  case dxil::ResourceKind::FeedbackTexture2DArray:
    reportFatalUsageError(
        "DXIL cmpxchg not implemented for this texture resource kind");
    return;
  case dxil::ResourceKind::CBuffer:
  case dxil::ResourceKind::Sampler:
  case dxil::ResourceKind::TBuffer:
    reportFatalUsageError(
        "DXIL cmpxchg not implemented for this resource type");
    return;
  case dxil::ResourceKind::RTAccelerationStructure:
  case dxil::ResourceKind::Invalid:
  case dxil::ResourceKind::NumEntries:
    llvm_unreachable("Invalid resource kind for cmpxchg");
  }
  llvm_unreachable("Unhandled case in switch");
}

static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
                                  dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = LI->getDataLayout();
  IRBuilder<> Builder(LI);
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
  Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());

  Value *V =
      Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_typedbuffer,
                              {II->getOperand(0), II->getOperand(1)});
  V = Builder.CreateExtractValue(V, {0});

  Type *ScalarType = ContainedType->getScalarType();
  uint64_t AccessSize = DL.getTypeSizeInBits(ScalarType) / 8;
  Value *Offset =
      traverseGEPOffsets(DL, Builder, LI->getPointerOperand(), AccessSize);
  auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
  if (!ConstantOffset || !ConstantOffset->isZero())
    V = Builder.CreateExtractElement(V, Offset);

  // If we loaded a <1 x ...> instead of a scalar (presumably to feed a
  // shufflevector), then make sure we're maintaining the resulting type.
  if (auto *VT = dyn_cast<FixedVectorType>(LI->getType()))
    if (VT->getNumElements() == 1 && !isa<FixedVectorType>(V->getType()))
      V = Builder.CreateInsertElement(PoisonValue::get(VT), V,
                                      Builder.getInt32(0));

  LI->replaceAllUsesWith(V);
}

static void createTextureLoad(IntrinsicInst *II, LoadInst *LI,
                              dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = LI->getDataLayout();
  IRBuilder<> Builder(LI);
  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);

  Value *Handle = II->getOperand(0);
  Value *Coords = II->getOperand(1);

  // For operator[], mip level is 0.
  Value *MipLevel = Builder.getInt32(0);

  // For operator[], offsets are zero.
  Value *Offsets = getNullOffsetsFor(Builder, Coords);

  Value *V =
      Builder.CreateIntrinsic(ContainedType, Intrinsic::dx_resource_load_level,
                              {Handle, Coords, MipLevel, Offsets});

  Type *ScalarType = ContainedType->getScalarType();
  uint64_t AccessSize = DL.getTypeSizeInBits(ScalarType) / 8;
  Value *Offset =
      traverseGEPOffsets(DL, Builder, LI->getPointerOperand(), AccessSize);
  auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
  if (!ConstantOffset || !ConstantOffset->isZero())
    V = Builder.CreateExtractElement(V, Offset);

  // If we loaded a <1 x ...> instead of a scalar (presumably to feed a
  // shufflevector), then make sure we're maintaining the resulting type.
  if (auto *VT = dyn_cast<FixedVectorType>(LI->getType()))
    if (VT->getNumElements() == 1 && !isa<FixedVectorType>(V->getType()))
      V = Builder.CreateInsertElement(PoisonValue::get(VT), V,
                                      Builder.getInt32(0));

  LI->replaceAllUsesWith(V);
}

static Value *emitRawLoad(IRBuilder<> &Builder, Type *Ty, Value *Buffer,
                          Value *Index, Value *Offset,
                          dxil::ResourceTypeInfo &RTI) {
  // For raw buffer (ie, HLSL's ByteAddressBuffer), we need to fold the access
  // entirely into the index.
  if (!RTI.isStruct()) {
    auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
    if (!ConstantOffset || !ConstantOffset->isZero())
      Index = Builder.CreateAdd(Index, Offset);
    Offset = llvm::PoisonValue::get(Builder.getInt32Ty());
  }

  // The load intrinsic includes the bit for CheckAccessFullyMapped, so we need
  // to add that to the return type.
  Type *TypeWithCheck = StructType::get(Ty, Builder.getInt1Ty());
  Value *V = Builder.CreateIntrinsic(TypeWithCheck,
                                     Intrinsic::dx_resource_load_rawbuffer,
                                     {Buffer, Index, Offset});
  return Builder.CreateExtractValue(V, {0});
}

static void createRawLoads(IntrinsicInst *II, LoadInst *LI,
                           dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = LI->getDataLayout();
  IRBuilder<> Builder(LI);

  Value *Index = II->getOperand(1);
  // The offset for the rawbuffer load and store ops is always in bytes.
  uint64_t AccessSize = 1;
  Value *Offset =
      traverseGEPOffsets(DL, Builder, LI->getPointerOperand(), AccessSize);

  // TODO: We could make this handle aggregates by walking the structure and
  // handling each field individually, but we don't ever generate code that
  // would hit that so it seems superfluous.
  assert(!LI->getType()->isAggregateType() &&
         "Resource load should be scalar or vector type");

  Value *V;
  if (auto *VT = dyn_cast<FixedVectorType>(LI->getType())) {
    // Split into loads of at most 4 elements.
    Type *EltTy = VT->getElementType();
    Value *Stride = ConstantInt::get(Builder.getInt32Ty(),
                                     4 * (DL.getTypeSizeInBits(EltTy) / 8));

    SmallVector<Value *> Parts;
    for (unsigned int I = 0, N = VT->getNumElements(); I < N; I += 4) {
      Type *Ty = FixedVectorType::get(EltTy, N - I < 4 ? N - I : 4);
      if (I > 0)
        Offset = Builder.CreateAdd(Offset, Stride);
      Parts.push_back(
          emitRawLoad(Builder, Ty, II->getOperand(0), Index, Offset, RTI));
    }

    V = Parts.size() > 1 ? concatenateVectors(Builder, Parts) : Parts[0];
  } else
    V = emitRawLoad(Builder, LI->getType(), II->getOperand(0), Index, Offset,
                    RTI);

  LI->replaceAllUsesWith(V);
}

namespace {
/// Helper for building a `load.cbufferrow` intrinsic given a simple type.
struct CBufferRowIntrin {
  Intrinsic::ID IID;
  Type *RetTy;
  unsigned int EltSize;
  unsigned int NumElts;

  CBufferRowIntrin(const DataLayout &DL, Type *Ty) {
    assert(Ty == Ty->getScalarType() && "Expected scalar type");

    switch (DL.getTypeSizeInBits(Ty)) {
    case 16:
      IID = Intrinsic::dx_resource_load_cbufferrow_8;
      RetTy = StructType::get(Ty, Ty, Ty, Ty, Ty, Ty, Ty, Ty);
      EltSize = 2;
      NumElts = 8;
      break;
    case 32:
      IID = Intrinsic::dx_resource_load_cbufferrow_4;
      RetTy = StructType::get(Ty, Ty, Ty, Ty);
      EltSize = 4;
      NumElts = 4;
      break;
    case 64:
      IID = Intrinsic::dx_resource_load_cbufferrow_2;
      RetTy = StructType::get(Ty, Ty);
      EltSize = 8;
      NumElts = 2;
      break;
    default:
      llvm_unreachable("Only 16, 32, and 64 bit types supported");
    }
  }
};
} // namespace

static Value *loadCBufferRowIntoWindow(IRBuilder<> &Builder,
                                       const CBufferRowIntrin &Intrin,
                                       Value *Handle, Value *FirstRow,
                                       Value *Window, unsigned RowNumber,
                                       StringRef Name) {
  Value *RowIndex =
      RowNumber == 0 ? FirstRow
                     : Builder.CreateAdd(FirstRow, Builder.getInt32(RowNumber));
  Value *RowLoad = Builder.CreateIntrinsic(
      Intrin.RetTy, Intrin.IID, {Handle, RowIndex}, nullptr, Name + ".load");
  for (unsigned J = 0; J != Intrin.NumElts; ++J) {
    Value *Element = Builder.CreateExtractValue(RowLoad, {J});
    Window = Builder.CreateInsertElement(
        Window, Element, Builder.getInt32(RowNumber * Intrin.NumElts + J));
  }
  return Window;
}

static Value *createDynamicCBufferLoad(IntrinsicInst *II, LoadInst *LI,
                                       const CBufferRowIntrin &Intrin,
                                       Value *BaseRow, Value *ByteOffset,
                                       const KnownBits &WithinRow) {
  IRBuilder<> Builder(LI);
  Type *Ty = LI->getType();
  auto *VT = dyn_cast<FixedVectorType>(Ty);
  unsigned NumElements = VT ? VT->getNumElements() : 1;
  StringRef Name = LI->getName();
  Value *Handle = II->getOperand(0);
  Value *RowSize = Builder.getInt32(hlsl::CBufferRowSizeInBytes);
  Value *FirstRow =
      Builder.CreateAdd(BaseRow, Builder.CreateUDiv(ByteOffset, RowSize));
  Value *Component = Builder.CreateUDiv(Builder.CreateURem(ByteOffset, RowSize),
                                        Builder.getInt32(Intrin.EltSize));
  unsigned MinComponent =
      WithinRow.getMinValue().getZExtValue() / Intrin.EltSize;
  unsigned MaxComponent =
      WithinRow.getMaxValue().getZExtValue() / Intrin.EltSize;
  unsigned MinRows = 1 + (MinComponent + NumElements - 1) / Intrin.NumElts;
  unsigned MaxRows = 1 + (MaxComponent + NumElements - 1) / Intrin.NumElts;
  auto *WindowTy =
      FixedVectorType::get(Ty->getScalarType(), MaxRows * Intrin.NumElts);
  Value *Window = PoisonValue::get(WindowTy);
  for (unsigned I = 0; I != MinRows; ++I)
    Window = loadCBufferRowIntoWindow(Builder, Intrin, Handle, FirstRow, Window,
                                      I, Name);

  // The unknown starting component can require at most one extra row.
  // Load it only when needed, rather than accessing beyond the load.
  if (MaxRows != MinRows) {
    BasicBlock *Head = LI->getParent();
    Value *LastComponent =
        Builder.CreateAdd(Component, Builder.getInt32(NumElements - 1));
    Value *CrossesRow = Builder.CreateICmpUGE(
        LastComponent, Builder.getInt32(MinRows * Intrin.NumElts));
    Instruction *ThenTerm = SplitBlockAndInsertIfThen(CrossesRow, LI, false);
    setExplicitlyUnknownBranchWeightsIfProfiled(*Head->getTerminator(),
                                                DEBUG_TYPE);
    Builder.SetInsertPoint(ThenTerm);
    Value *NextWindow = loadCBufferRowIntoWindow(
        Builder, Intrin, Handle, FirstRow, Window, MinRows, Name);
    Builder.SetInsertPoint(LI);
    auto *Rows = Builder.CreatePHI(WindowTy, 2, Name + ".rows");
    Rows->addIncoming(Window, Head);
    Rows->addIncoming(NextWindow, ThenTerm->getParent());
    Window = Rows;
  }

  Value *Result = PoisonValue::get(Ty);
  for (unsigned I = 0; I != NumElements; ++I) {
    Value *Index =
        I == 0 ? Component : Builder.CreateAdd(Component, Builder.getInt32(I));
    Value *Element =
        Builder.CreateExtractElement(Window, Index, Name + ".extract");
    Result =
        VT ? Builder.CreateInsertElement(Result, Element, Builder.getInt32(I))
           : Element;
  }
  return Result;
}

static KnownBits computeCBufferOffsetKnownBits(Value *Offset,
                                               const DataLayout &DL,
                                               const Instruction *CxtI) {
  KnownBits Known = computeKnownBits(Offset, DL, nullptr, CxtI);
  auto *Phi = dyn_cast<PHINode>(Offset);
  if (!Phi || Phi->getNumIncomingValues() == 0)
    return Known;

  // PHI queries limit recursion into each input, which can hide row alignment.
  KnownBits IncomingKnown(Known.getBitWidth());
  IncomingKnown.setAllConflict();
  for (unsigned I = 0; I != Phi->getNumIncomingValues(); ++I)
    IncomingKnown = IncomingKnown.intersectWith(
        computeKnownBits(Phi->getIncomingValue(I), DL, nullptr,
                         Phi->getIncomingBlock(I)->getTerminator()));
  return Known.unionWith(IncomingKnown);
}

static void createCBufferLoad(IntrinsicInst *II, LoadInst *LI,
                              dxil::ResourceTypeInfo &RTI) {
  const DataLayout &DL = LI->getDataLayout();

  Type *Ty = LI->getType();
  assert(!isa<StructType>(Ty) && "Structs not handled yet");
  CBufferRowIntrin Intrin(DL, Ty->getScalarType());

  StringRef Name = LI->getName();
  Value *Handle = II->getOperand(0);

  IRBuilder<> Builder(LI);

  ConstantInt *GlobalOffset =
      II->getIntrinsicID() == Intrinsic::dx_resource_getbasepointer
          ? ConstantInt::get(Builder.getInt32Ty(), 0)
          : dyn_cast<ConstantInt>(II->getOperand(1));
  assert(GlobalOffset && "CBuffer getpointer index must be constant");

  uint64_t GlobalOffsetVal = GlobalOffset->getZExtValue();
  Value *CurrentRow = ConstantInt::get(
      Builder.getInt32Ty(), GlobalOffsetVal / hlsl::CBufferRowSizeInBytes);
  unsigned int CurrentIndex =
      (GlobalOffsetVal % hlsl::CBufferRowSizeInBytes) / Intrin.EltSize;

  // Every object in a cbuffer either fits in a row or is aligned to a row. This
  // means that only the very last pointer access can point into a row.
  auto *LastGEP = dyn_cast<GEPOperator>(LI->getPointerOperand());
  if (!LastGEP) {
    // If we don't have a GEP at all we're just accessing the resource through
    // the result of getpointer directly.
    assert(LI->getPointerOperand() == II &&
           "Unexpected indirect access to resource without GEP");
  } else {
    Value *GEPOffset = traverseGEPOffsets(
        DL, Builder, LastGEP->getPointerOperand(), hlsl::CBufferRowSizeInBytes);
    CurrentRow = Builder.CreateAdd(GEPOffset, CurrentRow);

    APInt ConstantOffset(DL.getIndexTypeSizeInBits(LastGEP->getType()), 0);
    if (LastGEP->accumulateConstantOffset(DL, ConstantOffset)) {
      APInt Remainder(DL.getIndexTypeSizeInBits(LastGEP->getType()),
                      hlsl::CBufferRowSizeInBytes);
      APInt::udivrem(ConstantOffset, Remainder, ConstantOffset, Remainder);
      CurrentRow = Builder.CreateAdd(
          CurrentRow, ConstantInt::get(Builder.getInt32Ty(), ConstantOffset));
      CurrentIndex += Remainder.udiv(Intrin.EltSize).getZExtValue();
    } else {
      assert(LastGEP->getNumIndices() == 1 &&
             "Last GEP of cbuffer access is not array or struct access");
      // Aggregate strides are row-aligned, but merged byte offsets can select
      // different components within a row.
      Value *RowIndex = Builder.CreateSExtOrTrunc(*LastGEP->idx_begin(),
                                                  Builder.getInt32Ty());
      uint64_t IndexScale =
          DL.getTypeAllocSize(LastGEP->getSourceElementType());
      if (IndexScale < hlsl::CBufferRowSizeInBytes) {
        assert(hlsl::CBufferRowSizeInBytes % IndexScale == 0 &&
               "CBuffer GEP index is not row aligned");
        Value *ByteOffset =
            IndexScale == 1
                ? RowIndex
                : Builder.CreateMul(RowIndex, Builder.getInt32(IndexScale));
        uint64_t BaseWithinRow = GlobalOffsetVal % hlsl::CBufferRowSizeInBytes;
        if (BaseWithinRow != 0)
          ByteOffset =
              Builder.CreateAdd(ByteOffset, Builder.getInt32(BaseWithinRow));
        Value *RowSize = Builder.getInt32(hlsl::CBufferRowSizeInBytes);
        KnownBits WithinRow =
            computeCBufferOffsetKnownBits(ByteOffset, DL, LI) &
            KnownBits::makeConstant(APInt(32, hlsl::CBufferRowSizeInBytes - 1));
        if (!WithinRow.isConstant()) {
          LI->replaceAllUsesWith(createDynamicCBufferLoad(
              II, LI, Intrin, CurrentRow, ByteOffset, WithinRow));
          return;
        }
        RowIndex = Builder.CreateUDiv(ByteOffset, RowSize, "",
                                      WithinRow.getConstant().isZero());
        CurrentIndex = WithinRow.getConstant().getZExtValue() / Intrin.EltSize;
      } else if (IndexScale > hlsl::CBufferRowSizeInBytes) {
        assert(IndexScale % hlsl::CBufferRowSizeInBytes == 0 &&
               "CBuffer GEP index is not row aligned");
        RowIndex = Builder.CreateMul(
            RowIndex,
            ConstantInt::get(Builder.getInt32Ty(),
                             IndexScale / hlsl::CBufferRowSizeInBytes));
      }
      if (auto *BaseRow = dyn_cast<ConstantInt>(CurrentRow);
          BaseRow && BaseRow->isZero())
        CurrentRow = RowIndex;
      else
        CurrentRow = Builder.CreateAdd(CurrentRow, RowIndex);
    }
  }

  auto *CBufLoad = Builder.CreateIntrinsic(
      Intrin.RetTy, Intrin.IID, {Handle, CurrentRow}, nullptr, Name + ".load");
  auto *Elt =
      Builder.CreateExtractValue(CBufLoad, {CurrentIndex++}, Name + ".extract");

  // At this point we've loaded the first scalar of our result, but our original
  // type may have been a vector.
  unsigned int Remaining =
      ((DL.getTypeSizeInBits(Ty) / 8) / Intrin.EltSize) - 1;
  if (Remaining == 0) {
    // We only have a single element, so we're done.
    Value *Result = Elt;

    // However, if we loaded a <1 x T>, then we need to adjust the type.
    if (auto *VT = dyn_cast<FixedVectorType>(Ty)) {
      assert(VT->getNumElements() == 1 && "Can't have multiple elements here");
      Result = Builder.CreateInsertElement(PoisonValue::get(VT), Result,
                                           Builder.getInt32(0), Name);
    }
    LI->replaceAllUsesWith(Result);
    return;
  }

  // Walk each element and extract it, wrapping to new rows as needed.
  SmallVector<Value *> Extracts{Elt};
  while (Remaining--) {
    CurrentIndex %= Intrin.NumElts;

    if (CurrentIndex == 0) {
      CurrentRow = Builder.CreateAdd(CurrentRow,
                                     ConstantInt::get(Builder.getInt32Ty(), 1));
      CBufLoad = Builder.CreateIntrinsic(Intrin.RetTy, Intrin.IID,
                                         {Handle, CurrentRow}, nullptr,
                                         Name + ".load");
    }

    Extracts.push_back(Builder.CreateExtractValue(CBufLoad, {CurrentIndex++},
                                                  Name + ".extract"));
  }

  // Finally, we build up the original loaded value.
  Value *Result = PoisonValue::get(Ty);
  for (int I = 0, E = Extracts.size(); I < E; ++I)
    Result = Builder.CreateInsertElement(
        Result, Extracts[I], Builder.getInt32(I), Name + formatv(".upto{}", I));
  LI->replaceAllUsesWith(Result);
}

static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI,
                                dxil::ResourceTypeInfo &RTI) {
  switch (RTI.getResourceKind()) {
  case dxil::ResourceKind::TypedBuffer:
    return createTypedBufferLoad(II, LI, RTI);
  case dxil::ResourceKind::RawBuffer:
  case dxil::ResourceKind::StructuredBuffer:
    return createRawLoads(II, LI, RTI);
  case dxil::ResourceKind::CBuffer:
    return createCBufferLoad(II, LI, RTI);
  case dxil::ResourceKind::Texture1D:
  case dxil::ResourceKind::Texture2D:
  case dxil::ResourceKind::Texture2DMS:
  case dxil::ResourceKind::Texture3D:
  case dxil::ResourceKind::TextureCube:
  case dxil::ResourceKind::Texture1DArray:
  case dxil::ResourceKind::Texture2DArray:
  case dxil::ResourceKind::Texture2DMSArray:
  case dxil::ResourceKind::TextureCubeArray:
    return createTextureLoad(II, LI, RTI);
  case dxil::ResourceKind::FeedbackTexture2D:
  case dxil::ResourceKind::FeedbackTexture2DArray:
  case dxil::ResourceKind::TBuffer:
    reportFatalUsageError("Load not yet implemented for resource type");
    return;
  case dxil::ResourceKind::Sampler:
  case dxil::ResourceKind::RTAccelerationStructure:
  case dxil::ResourceKind::Invalid:
  case dxil::ResourceKind::NumEntries:
    llvm_unreachable("Invalid resource kind for load");
  }
  llvm_unreachable("Unhandled case in switch");
}

static Instruction *getHandleOperand(Instruction *AI) {
  if (auto *LI = dyn_cast<LoadInst>(AI))
    return dyn_cast<Instruction>(LI->getPointerOperand());
  if (auto *SI = dyn_cast<StoreInst>(AI))
    return dyn_cast<Instruction>(SI->getPointerOperand());
  if (auto *RMWI = dyn_cast<AtomicRMWInst>(AI))
    return dyn_cast<Instruction>(RMWI->getPointerOperand());
  if (auto *CXI = dyn_cast<AtomicCmpXchgInst>(AI))
    return dyn_cast<Instruction>(CXI->getPointerOperand());
  if (auto *II = dyn_cast<IntrinsicInst>(AI))
    if (II->getIntrinsicID() == Intrinsic::dx_resource_updatecounter)
      return dyn_cast<Instruction>(II->getArgOperand(0));

  return nullptr;
}

static Instruction *getHandleRoot(Instruction *I) {
  while (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
    auto *PointerI = dyn_cast<Instruction>(GEP->getPointerOperand());
    if (!PointerI)
      break;
    I = PointerI;
  }
  return I;
}

static const std::array<Intrinsic::ID, 2> HandleIntrins = {
    Intrinsic::dx_resource_handlefrombinding,
    Intrinsic::dx_resource_handlefromimplicitbinding,
};

static SmallVector<IntrinsicInst *> collectUsedHandles(Value *Ptr) {
  SmallVector<Value *> Worklist = {Ptr};
  SmallVector<IntrinsicInst *> Handles;
  SmallPtrSet<Value *, 4> VisitedPhis;

  while (!Worklist.empty()) {
    Value *X = Worklist.pop_back_val();

    if (!X->getType()->isPointerTy() && !X->getType()->isTargetExtTy())
      return {}; // Early exit on store/load into non-resource

    if (auto *Phi = dyn_cast<PHINode>(X)) {
      if (VisitedPhis.contains(X))
        continue;
      for (Use &V : Phi->incoming_values())
        Worklist.push_back(V.get());
      VisitedPhis.insert(Phi);
    } else if (auto *Select = dyn_cast<SelectInst>(X)) {
      Worklist.push_back(Select->getTrueValue());
      Worklist.push_back(Select->getFalseValue());
    } else if (auto *GEP = dyn_cast<GetElementPtrInst>(X)) {
      Worklist.push_back(GEP->getPointerOperand());
    } else if (auto *II = dyn_cast<IntrinsicInst>(X)) {
      Intrinsic::ID IID = II->getIntrinsicID();

      if (IID == Intrinsic::dx_resource_getpointer ||
          IID == Intrinsic::dx_resource_getbasepointer)
        Worklist.push_back(II->getArgOperand(/*Handle=*/0));

      if (llvm::is_contained(HandleIntrins, IID))
        Handles.push_back(II);
    }
  }

  return Handles;
}

static hlsl::Binding getHandleIntrinsicBinding(IntrinsicInst *Handle,
                                               DXILResourceTypeMap &DRTM) {
  assert(llvm::is_contained(HandleIntrins, Handle->getIntrinsicID()) &&
         "Only expects a Handle as determined from collectUsedHandles.");

  auto *HandleTy = cast<TargetExtType>(Handle->getType());
  dxil::ResourceClass Class = DRTM[HandleTy].getResourceClass();
  uint32_t Space = cast<ConstantInt>(Handle->getArgOperand(0))->getZExtValue();
  uint32_t LowerBound =
      cast<ConstantInt>(Handle->getArgOperand(1))->getZExtValue();
  uint32_t Size = cast<ConstantInt>(Handle->getArgOperand(2))->getZExtValue();
  uint32_t UpperBound = Size == UINT32_MAX ? UINT32_MAX : LowerBound + Size - 1;

  return hlsl::Binding(Class, Space, LowerBound, UpperBound, nullptr);
}

namespace {
/// Helper for propagating the current handle and ptr indices.
struct AccessIndices {
  Value *GetPtrIdx;
  Value *HandleIdx;
  Value *OffsetIdx;

  bool hasGetPtrIdx() { return GetPtrIdx != nullptr; }
  bool hasHandleIdx() { return HandleIdx != nullptr; }
  bool hasOffsetIdx() { return OffsetIdx != nullptr; }
};

/// AccessIndices cached per phi. The values may be placeholder phis that are
/// later folded away (via replaceAllUsesWith) and freed, so track them with
/// value handles rather than raw pointers.
struct TrackedAccessIndices {
  WeakTrackingVH GetPtrIdx;
  WeakTrackingVH HandleIdx;
  WeakTrackingVH OffsetIdx;

  TrackedAccessIndices() = default;
  TrackedAccessIndices(const AccessIndices &AI)
      : GetPtrIdx(AI.GetPtrIdx), HandleIdx(AI.HandleIdx),
        OffsetIdx(AI.OffsetIdx) {}
  operator AccessIndices() const { return {GetPtrIdx, HandleIdx, OffsetIdx}; }
};

using PhiAccessIndicesMap = SmallDenseMap<PHINode *, TrackedAccessIndices>;
} // namespace

// Compute the total byte offset described by a GEP as a single i32 value,
// handling both constant and variable indices.
static Value *accumulateGEPOffset(GetElementPtrInst *GEP,
                                  IRBuilder<> &Builder) {
  const DataLayout &DL = GEP->getDataLayout();
  unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType());
  SmallMapVector<Value *, APInt, 4> VariableOffsets;
  APInt ConstantOffset(BitWidth, 0);
  bool Success =
      GEP->collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset);
  assert(Success && "Unhandled GEP structure for resource access");
  (void)Success;

  Type *I32 = Builder.getInt32Ty();
  Value *Offset = ConstantInt::get(I32, ConstantOffset.getSExtValue());
  for (auto &[V, Scale] : VariableOffsets) {
    Value *Index = Builder.CreateSExtOrTrunc(V, I32);
    if (!Scale.isOne())
      Index =
          Builder.CreateMul(Index, ConstantInt::get(I32, Scale.getSExtValue()));
    Offset = Builder.CreateAdd(Offset, Index);
  }
  return Offset;
}

// getAccessIndices traverses up the control flow that a ptr came from and
// propagates back the indicies used to access the resource (AccessIndices):
//
//  - GetPtrIdx is the index of dx.resource.getpointer
//  - HandleIdx is the index of dx.resource.handlefrom.*
//  - OffsetIdx is the accumulated byte offset of any GEPs in the ptr chain
static AccessIndices
getAccessIndices(Instruction *I, SmallSetVector<Instruction *, 16> &DeadInsts,
                 PhiAccessIndicesMap &VisitedPhis) {
  if (auto *II = dyn_cast<IntrinsicInst>(I)) {
    if (llvm::is_contained(HandleIntrins, II->getIntrinsicID())) {
      DeadInsts.insert(II);
      return {nullptr, II->getArgOperand(/*Index=*/3), nullptr};
    }

    Intrinsic::ID IID = II->getIntrinsicID();
    if (IID == Intrinsic::dx_resource_getpointer ||
        IID == Intrinsic::dx_resource_getbasepointer) {
      auto *V = dyn_cast<Instruction>(II->getArgOperand(/*Handle=*/0));
      auto AccessIdx = getAccessIndices(V, DeadInsts, VisitedPhis);
      assert(!AccessIdx.hasGetPtrIdx() &&
             "Encountered multiple dx.resource.getpointers in ptr chain?");
      IRBuilder<> Builder(II);
      AccessIdx.GetPtrIdx = ConstantInt::get(Builder.getInt32Ty(), 0);
      if (IID == Intrinsic::dx_resource_getpointer)
        AccessIdx.GetPtrIdx = II->getArgOperand(1);

      DeadInsts.insert(II);
      return AccessIdx;
    }
  }

  if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
    auto *V = dyn_cast<Instruction>(GEP->getPointerOperand());
    auto AccessIdx = getAccessIndices(V, DeadInsts, VisitedPhis);

    IRBuilder<> Builder(GEP);
    Value *GEPOffset = accumulateGEPOffset(GEP, Builder);
    AccessIdx.OffsetIdx =
        AccessIdx.hasOffsetIdx()
            ? Builder.CreateAdd(AccessIdx.OffsetIdx, GEPOffset)
            : GEPOffset;

    DeadInsts.insert(GEP);
    return AccessIdx;
  }

  if (auto *Phi = dyn_cast<PHINode>(I)) {
    // If we're already building indices for this phi, return a ref to the phi
    if (auto It = VisitedPhis.find(Phi); It != VisitedPhis.end())
      return It->second;

    unsigned NumEdges = Phi->getNumIncomingValues();
    assert(NumEdges != 0 && "Malformed Phi Node");

    IRBuilder<> Builder(Phi);
    std::unique_ptr<PHINode> GetPtrPhi(
        PHINode::Create(Builder.getInt32Ty(), NumEdges));
    std::unique_ptr<PHINode> HandlePhi(
        PHINode::Create(Builder.getInt32Ty(), NumEdges));
    std::unique_ptr<PHINode> OffsetPhi(
        PHINode::Create(Builder.getInt32Ty(), NumEdges));
    bool HasOffset = false;

    // Register a ref to this phi for a recursive phi. If a placeholder phi
    // below folds to a single value, replaceAllUsesWith updates any cached
    // entries that captured it.
    if (Phi->getType()->isTargetExtTy())
      VisitedPhis[Phi] = AccessIndices{/*GetPtrIdx=*/nullptr, HandlePhi.get(),
                                       /*OffsetIdx=*/nullptr};
    else if (Phi->getType()->isPointerTy())
      VisitedPhis[Phi] =
          AccessIndices{GetPtrPhi.get(), HandlePhi.get(), OffsetPhi.get()};

    for (unsigned Idx = 0; Idx < NumEdges; Idx++) {
      auto *BB = Phi->getIncomingBlock(Idx);
      auto *V = dyn_cast<Instruction>(Phi->getIncomingValue(Idx));
      auto AccessIdx = getAccessIndices(V, DeadInsts, VisitedPhis);
      if (AccessIdx.hasGetPtrIdx())
        GetPtrPhi->addIncoming(AccessIdx.GetPtrIdx, BB);
      HandlePhi->addIncoming(AccessIdx.HandleIdx, BB);
      if (AccessIdx.hasOffsetIdx()) {
        OffsetPhi->addIncoming(AccessIdx.OffsetIdx, BB);
        HasOffset = true;
      } else {
        OffsetPhi->addIncoming(ConstantInt::get(Builder.getInt32Ty(), 0), BB);
      }
    }

    Value *GetPtrIdx;
    if (GetPtrPhi->getNumIncomingValues() == 0) {
      GetPtrIdx = nullptr;
    } else if (Value *ConstantGetPtr = GetPtrPhi->hasConstantValue()) {
      GetPtrIdx = ConstantGetPtr;
      GetPtrPhi->replaceAllUsesWith(GetPtrIdx);
    } else {
      GetPtrIdx = GetPtrPhi.release();
      Builder.Insert(GetPtrIdx);
    }

    Value *HandleIdx;
    if (Value *ConstantHandle = HandlePhi->hasConstantValue()) {
      HandleIdx = ConstantHandle;
      HandlePhi->replaceAllUsesWith(HandleIdx);
    } else {
      HandleIdx = HandlePhi.release();
      Builder.Insert(HandleIdx);
    }

    Value *OffsetIdx = nullptr;
    if (HasOffset) {
      if (Value *ConstantIdx = OffsetPhi->hasConstantValue()) {
        OffsetIdx = ConstantIdx;
        OffsetPhi->replaceAllUsesWith(OffsetIdx);
      } else {
        OffsetIdx = OffsetPhi.release();
        Builder.Insert(OffsetIdx);
      }
    }

    // Replace the placeholder entry with the final indices.
    AccessIndices Result = {GetPtrIdx, HandleIdx, OffsetIdx};
    VisitedPhis[Phi] = Result;
    DeadInsts.insert(Phi);
    return Result;
  }

  if (auto *Select = dyn_cast<SelectInst>(I)) {
    auto *TrueV = dyn_cast<Instruction>(Select->getTrueValue());
    auto TrueAccessIdx = getAccessIndices(TrueV, DeadInsts, VisitedPhis);

    auto *FalseV = dyn_cast<Instruction>(Select->getFalseValue());
    auto FalseAccessIdx = getAccessIndices(FalseV, DeadInsts, VisitedPhis);

    IRBuilder<> Builder(Select);
    Value *GetPtrSelect = nullptr;
    Value *OffsetSelect = nullptr;

    if (TrueAccessIdx.hasGetPtrIdx() && FalseAccessIdx.hasGetPtrIdx())
      GetPtrSelect =
          Builder.CreateSelect(Select->getCondition(), TrueAccessIdx.GetPtrIdx,
                               FalseAccessIdx.GetPtrIdx, "", Select);

    auto *HandleSelect =
        Builder.CreateSelect(Select->getCondition(), TrueAccessIdx.HandleIdx,
                             FalseAccessIdx.HandleIdx, "", Select);

    if (TrueAccessIdx.hasOffsetIdx() || FalseAccessIdx.hasOffsetIdx()) {
      Value *Zero = ConstantInt::get(Builder.getInt32Ty(), 0);
      Value *TrueOffset =
          TrueAccessIdx.hasOffsetIdx() ? TrueAccessIdx.OffsetIdx : Zero;
      Value *FalseOffset =
          FalseAccessIdx.hasOffsetIdx() ? FalseAccessIdx.OffsetIdx : Zero;
      OffsetSelect = Builder.CreateSelect(Select->getCondition(), TrueOffset,
                                          FalseOffset, "", Select);
    }

    DeadInsts.insert(Select);
    return {GetPtrSelect, HandleSelect, OffsetSelect};
  }

  llvm_unreachable("collectUsedHandles should assure this does not occur");
}

static void
replaceHandleWithIndices(Instruction *Ptr, IntrinsicInst *OldHandle,
                         SmallSetVector<Instruction *, 16> &DeadInsts,
                         PhiAccessIndicesMap &VisitedPhis) {
  auto AccessIdx = getAccessIndices(Ptr, DeadInsts, VisitedPhis);
  assert(AccessIdx.hasHandleIdx() &&
         "Couldn't retrieve handle index. This is guaranteed by "
         "getAccessIndices");

  IRBuilder<> Builder(Ptr);
  if (isa<PHINode>(Ptr))
    Builder.SetInsertPoint(Ptr->getParent()->getFirstNonPHIIt());
  IntrinsicInst *Handle = cast<IntrinsicInst>(OldHandle->clone());
  Handle->setArgOperand(/*Index=*/3, AccessIdx.HandleIdx);
  Builder.Insert(Handle);

  if (Ptr->getType()->isPointerTy()) {
    assert(AccessIdx.hasGetPtrIdx() &&
           "Couldn't retrieve getpointer index. This is guaranteed by "
           "getAccessIndices");
    auto *GetPtr = Builder.CreateIntrinsic(Ptr->getType(),
                                           Intrinsic::dx_resource_getpointer,
                                           {Handle, AccessIdx.GetPtrIdx});

    // Reapply any byte offset that came from GEPs in the original ptr chain so
    // that later access lowering still sees it.
    Value *Result = GetPtr;
    if (AccessIdx.hasOffsetIdx())
      Result = Builder.CreatePtrAdd(GetPtr, AccessIdx.OffsetIdx);
    Ptr->replaceAllUsesWith(Result);
  } else {
    assert(Ptr->getType()->isTargetExtTy() && !AccessIdx.hasGetPtrIdx() &&
           "Unexpected resource access operand type");
    Ptr->replaceAllUsesWith(Handle);
  }

  DeadInsts.insert(Ptr);
}

// Try to legalize dx.resource.handlefrom.*.binding and dx.resource.getpointer
// calls with their respective index values and propagate the index values to
// be used at resource access.
//
// If it can't be transformed to be legal then:
//
// Reports an error if a resource access is not guaranteed into a unique global
// resource.
//
// Returns true if any changes are made.
static bool legalizeResourceHandles(Function &F, DXILResourceTypeMap &DRTM) {
  SmallSetVector<Instruction *, 16> DeadInsts;
  PhiAccessIndicesMap VisitedPhis;

  for (BasicBlock &BB : make_early_inc_range(F)) {
    for (Instruction &I : BB) {
      if (auto *HandleOp = getHandleOperand(&I)) {
        SmallVector<IntrinsicInst *> Handles = collectUsedHandles(HandleOp);
        unsigned NumHandles = Handles.size();
        if (NumHandles <= 1)
          continue; // Legal, no-replacement required

        bool SameGlobalBinding = true;
        hlsl::Binding B = getHandleIntrinsicBinding(Handles[0], DRTM);
        for (unsigned Idx = 1; Idx < NumHandles; Idx++)
          SameGlobalBinding &=
              (B == getHandleIntrinsicBinding(Handles[Idx], DRTM));

        if (!SameGlobalBinding)
          diagnoseNonUniqueResourceAccess(&I, Handles);

        replaceHandleWithIndices(getHandleRoot(HandleOp), Handles[0], DeadInsts,
                                 VisitedPhis);
      }
    }
  }

  bool MadeChanges = false;

  // Set up the phis to track if they are erased below
  SmallVector<WeakTrackingVH> ResourcePhis;
  for (const auto &HandleToIndex : VisitedPhis)
    ResourcePhis.push_back(HandleToIndex.first);

  for (auto *I : llvm::reverse(DeadInsts))
    if (I->hasNUses(0)) { // Handle can still be used outside of replaced path
      I->eraseFromParent();
      MadeChanges = true;
    }

  // Any remaining phi nodes are now looped with another phi node and have no
  // other uses
  for (WeakTrackingVH &VH : ResourcePhis)
    if (VH) // True if not removed above or already in this loop
      MadeChanges |= RecursivelyDeleteDeadPHINode(cast<PHINode>(VH));

  return MadeChanges;
}

static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
  SmallVector<User *> Worklist;
  for (User *U : II->users())
    Worklist.push_back(U);

  SmallVector<Instruction *> DeadInsts;
  while (!Worklist.empty()) {
    User *U = Worklist.back();
    Worklist.pop_back();

    if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
      for (User *U : GEP->users())
        Worklist.push_back(U);
      DeadInsts.push_back(GEP);

    } else if (auto *SI = dyn_cast<StoreInst>(U)) {
      assert(SI->getValueOperand() != II && "Pointer escaped!");
      createStoreIntrinsic(II, SI, RTI);
      DeadInsts.push_back(SI);

    } else if (auto *LI = dyn_cast<LoadInst>(U)) {
      createLoadIntrinsic(II, LI, RTI);
      DeadInsts.push_back(LI);
    } else if (auto *AI = dyn_cast<AtomicRMWInst>(U)) {
      createAtomicBinOpIntrinsic(II, AI, RTI);
      DeadInsts.push_back(AI);
    } else if (auto *CXI = dyn_cast<AtomicCmpXchgInst>(U)) {
      createAtomicCompareExchangeIntrinsic(II, CXI, RTI);
      DeadInsts.push_back(CXI);
    } else
      llvm_unreachable("Unhandled instruction - pointer escaped?");
  }

  // Traverse the now-dead instructions in RPO and remove them.
  for (Instruction *Dead : llvm::reverse(DeadInsts))
    Dead->eraseFromParent();
  II->eraseFromParent();
}

static bool transformResourcePointers(Function &F, DXILResourceTypeMap &DRTM) {
  SmallVector<std::pair<IntrinsicInst *, dxil::ResourceTypeInfo>> Resources;
  for (BasicBlock &BB : make_early_inc_range(F))
    for (Instruction &I : BB)
      if (auto *II = dyn_cast<IntrinsicInst>(&I))
        if (II->getIntrinsicID() == Intrinsic::dx_resource_getpointer ||
            II->getIntrinsicID() == Intrinsic::dx_resource_getbasepointer) {
          auto *HandleTy = cast<TargetExtType>(II->getArgOperand(0)->getType());
          assert(
              (DRTM[HandleTy].isCBuffer() ||
               II->getIntrinsicID() != Intrinsic::dx_resource_getbasepointer) &&
              "dx_resource_getbasepointer should only be used by cbuffers");
          Resources.emplace_back(II, DRTM[HandleTy]);
        }

  for (auto &[II, RI] : Resources)
    replaceAccess(II, RI);

  return !Resources.empty();
}

PreservedAnalyses DXILResourceAccess::run(Function &F,
                                          FunctionAnalysisManager &FAM) {
  auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
  DXILResourceTypeMap *DRTM =
      MAMProxy.getCachedResult<DXILResourceTypeAnalysis>(*F.getParent());
  assert(DRTM && "DXILResourceTypeAnalysis must be available");

  bool MadeHandleChanges = legalizeResourceHandles(F, *DRTM);
  bool MadeResourceChanges = transformResourcePointers(F, *DRTM);
  if (!(MadeHandleChanges || MadeResourceChanges))
    return PreservedAnalyses::all();

  PreservedAnalyses PA;
  PA.preserve<DXILResourceTypeAnalysis>();
  return PA;
}

namespace {
class DXILResourceAccessLegacy : public FunctionPass {
public:
  bool runOnFunction(Function &F) override {
    DXILResourceTypeMap &DRTM =
        getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
    bool MadeHandleChanges = legalizeResourceHandles(F, DRTM);
    bool MadeResourceChanges = transformResourcePointers(F, DRTM);
    return MadeHandleChanges || MadeResourceChanges;
  }
  StringRef getPassName() const override { return "DXIL Resource Access"; }
  DXILResourceAccessLegacy() : FunctionPass(ID) {}

  static char ID; // Pass identification.
  void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
    AU.addRequired<DXILResourceTypeWrapperPass>();
  }
};
char DXILResourceAccessLegacy::ID = 0;
} // end anonymous namespace

INITIALIZE_PASS_BEGIN(DXILResourceAccessLegacy, DEBUG_TYPE,
                      "DXIL Resource Access", false, false)
INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
INITIALIZE_PASS_END(DXILResourceAccessLegacy, DEBUG_TYPE,
                    "DXIL Resource Access", false, false)

FunctionPass *llvm::createDXILResourceAccessLegacyPass() {
  return new DXILResourceAccessLegacy();
}