1 files changed, 192 insertions, 0 deletions
diff --git a/llvm/lib/Target/DirectX/DXILResourceAccess.cpp b/llvm/lib/Target/DirectX/DXILResourceAccess.cpp
new file mode 100644
index 000000000000..1ff8f09f066d
--- /dev/null
+++ b/llvm/lib/Target/DirectX/DXILResourceAccess.cpp
@@ -0,0 +1,192 @@
+//===- 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/Analysis/DXILResource.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/InitializePasses.h"
+
+#define DEBUG_TYPE "dxil-resource-access"
+
+using namespace llvm;
+
+static void replaceTypedBufferAccess(IntrinsicInst *II,
+                                     dxil::ResourceTypeInfo &RTI) {
+  const DataLayout &DL = II->getDataLayout();
+
+  auto *HandleType = cast<TargetExtType>(II->getOperand(0)->getType());
+  assert(HandleType->getName() == "dx.TypedBuffer" &&
+         "Unexpected typed buffer type");
+  Type *ContainedType = HandleType->getTypeParameter(0);
+
+  // We need the size of an element in bytes so that we can calculate the offset
+  // in elements given a total offset in bytes later.
+  Type *ScalarType = ContainedType->getScalarType();
+  uint64_t ScalarSize = DL.getTypeSizeInBits(ScalarType) / 8;
+
+  // Process users keeping track of indexing accumulated from GEPs.
+  struct AccessAndIndex {
+    User *Access;
+    Value *Index;
+  };
+  SmallVector<AccessAndIndex> Worklist;
+  for (User *U : II->users())
+    Worklist.push_back({U, nullptr});
+
+  SmallVector<Instruction *> DeadInsts;
+  while (!Worklist.empty()) {
+    AccessAndIndex Current = Worklist.back();
+    Worklist.pop_back();
+
+    if (auto *GEP = dyn_cast<GetElementPtrInst>(Current.Access)) {
+      IRBuilder<> Builder(GEP);
+
+      Value *Index;
+      APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+      if (GEP->accumulateConstantOffset(DL, ConstantOffset)) {
+        APInt Scaled = ConstantOffset.udiv(ScalarSize);
+        Index = ConstantInt::get(Builder.getInt32Ty(), Scaled);
+      } else {
+        auto IndexIt = GEP->idx_begin();
+        assert(cast<ConstantInt>(IndexIt)->getZExtValue() == 0 &&
+               "GEP is not indexing through pointer");
+        ++IndexIt;
+        Index = *IndexIt;
+        assert(++IndexIt == GEP->idx_end() && "Too many indices in GEP");
+      }
+
+      for (User *U : GEP->users())
+        Worklist.push_back({U, Index});
+      DeadInsts.push_back(GEP);
+
+    } else if (auto *SI = dyn_cast<StoreInst>(Current.Access)) {
+      assert(SI->getValueOperand() != II && "Pointer escaped!");
+      IRBuilder<> Builder(SI);
+
+      Value *V = SI->getValueOperand();
+      if (V->getType() == ContainedType) {
+        // V is already the right type.
+      } 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(
+            ContainedType, Intrinsic::dx_resource_load_typedbuffer,
+            {II->getOperand(0), II->getOperand(1)});
+        // If we have an offset from seeing a GEP earlier, use it.
+        Value *IndexOp = Current.Index
+                             ? Current.Index
+                             : ConstantInt::get(Builder.getInt32Ty(), 0);
+        V = Builder.CreateInsertElement(Load, V, IndexOp);
+      } 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);
+      DeadInsts.push_back(SI);
+
+    } else if (auto *LI = dyn_cast<LoadInst>(Current.Access)) {
+      IRBuilder<> Builder(LI);
+      Value *V = Builder.CreateIntrinsic(
+          ContainedType, Intrinsic::dx_resource_load_typedbuffer,
+          {II->getOperand(0), II->getOperand(1)});
+      if (Current.Index)
+        V = Builder.CreateExtractElement(V, Current.Index);
+
+      LI->replaceAllUsesWith(V);
+      DeadInsts.push_back(LI);
+
+    } 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) {
+  bool Changed = false;
+  SmallVector<std::pair<IntrinsicInst *, dxil::ResourceTypeInfo>> Resources;
+  for (BasicBlock &BB : F)
+    for (Instruction &I : BB)
+      if (auto *II = dyn_cast<IntrinsicInst>(&I))
+        if (II->getIntrinsicID() == Intrinsic::dx_resource_getpointer) {
+          auto *HandleTy = cast<TargetExtType>(II->getArgOperand(0)->getType());
+          Resources.emplace_back(II, DRTM[HandleTy]);
+        }
+
+  for (auto &[II, RI] : Resources) {
+    if (RI.isTyped()) {
+      Changed = true;
+      replaceTypedBufferAccess(II, RI);
+    }
+
+    // TODO: handle other resource types. We should probably have an
+    // `unreachable` here once we've added support for all of them.
+  }
+
+  return Changed;
+}
+
+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 MadeChanges = transformResourcePointers(F, *DRTM);
+  if (!MadeChanges)
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  PA.preserve<DXILResourceTypeAnalysis>();
+  PA.preserve<DominatorTreeAnalysis>();
+  return PA;
+}
+
+namespace {
+class DXILResourceAccessLegacy : public FunctionPass {
+public:
+  bool runOnFunction(Function &F) override {
+    DXILResourceTypeMap &DRTM =
+        getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
+
+    return transformResourcePointers(F, DRTM);
+  }
+  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>();
+    AU.addPreserved<DominatorTreeWrapperPass>();
+  }
+};
+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();
+}