Lower llvm.dx.rawBufferLoad to dxil ops

llvm/include/llvm/Analysis/DXILResource.h

@@ -264,6 +264,8 @@ class ResourceInfo {
 class DXILResourceMap {
   SmallVector<dxil::ResourceInfo> Resources;
   DenseMap<CallInst *, unsigned> CallMap;
+  // Mapping from Resource use to Resource Handle
+  DenseMap<CallInst *, CallInst *> ResUseToHandleMap;
   unsigned FirstUAV = 0;
   unsigned FirstCBuffer = 0;
   unsigned FirstSampler = 0;
@@ -335,6 +337,36 @@ class DXILResourceMap {
   }
 
   void print(raw_ostream &OS) const;
+
+  void updateResourceMap(CallInst *origCallInst, CallInst *newCallInst);
+
+  // Update ResUseMap with multiple new resource uses
+  void updateResUseMap(CallInst *origResUse,
+                       std::vector<Value *> &multiNewResUse);
+
+  // Update ResUseMap with single new resource use
+  void updateResUseMap(CallInst *origResUse, CallInst *newResUse) {
+    assert((origResUse != nullptr) && (newResUse != nullptr) &&
+           (origResUse != newResUse) && "Wrong Inputs");
+
+    updateResUseMapCommon(origResUse, newResUse, /*keepOrigResUseInMap=*/false);
+  }
+
+  CallInst *findResHandleByUse(CallInst *resUse) {
+    auto Pos = ResUseToHandleMap.find(resUse);
+    assert((Pos != ResUseToHandleMap.end()) &&
+           "Can't find the resource handle");
+
+    return Pos->second;
+  }
+
+private:
+  void updateResUseMapCommon(CallInst *origResUse, CallInst *newResUse,
+                             bool keepOrigResUseInMap) {
+    ResUseToHandleMap.try_emplace(newResUse, findResHandleByUse(origResUse));
+    if (!keepOrigResUseInMap)
+      ResUseToHandleMap.erase(origResUse);
+  }
 };
 
 class DXILResourceAnalysis : public AnalysisInfoMixin<DXILResourceAnalysis> {

llvm/include/llvm/IR/IntrinsicsDirectX.td

@@ -27,6 +27,9 @@ def int_dx_handle_fromBinding
           [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i1_ty],
           [IntrNoMem]>;
 
+def int_dx_rawBufferLoad
+    : DefaultAttrsIntrinsic<[llvm_any_ty], [llvm_any_ty, llvm_i32_ty, llvm_i32_ty]>;
+
 def int_dx_typedBufferLoad
     : DefaultAttrsIntrinsic<[llvm_any_ty], [llvm_any_ty, llvm_i32_ty],
                             [IntrReadMem]>;

llvm/lib/Analysis/DXILResource.cpp

@@ -719,6 +719,12 @@ DXILResourceMap::DXILResourceMap(
     if (Resources.empty() || RI != Resources.back())
       Resources.push_back(RI);
     CallMap[CI] = Resources.size() - 1;
+
+    // Build ResUseToHandleMap
+    for (auto it = CI->users().begin(); it != CI->users().end(); ++it) {
+      CallInst *CI_Use = dyn_cast<CallInst>(*it);
+      ResUseToHandleMap[CI_Use] = CI;
+    }
   }
 
   unsigned Size = Resources.size();
@@ -744,6 +750,61 @@ DXILResourceMap::DXILResourceMap(
   }
 }
 
+// Parameter origCallInst: original Resource Handle
+// Parameter newCallInst:  new Resource Handle
+//
+// This function is needed when origCallInst's lowered to newCallInst.
+//
+// Because origCallInst and its uses will be replaced by newCallInst and new def
+// instructions after lowering. The [origCallInst, resource info] entry in
+// CallMap and [origCallInst's use, origCallInst] entries in ResUseToHandleMap
+// have to be updated per the changes in lowering.
+//
+// What this function does are:
+//   1. Add [newCallInst, resource info] entry in CallMap
+//   2. Remove [origCallInst, resource info] entry in CallMap
+//   3. Remap [origCallInst's use, origCallInst] entries to
+//      [origCallInst's use, newCallInst] entries in ResUseToHandleMap
+//
+// Remove those entries related to origCallInst in maps is necessary since
+// origCallInst's no longer existing after lowering. Moreover, keeping those
+// entries in maps will crash DXILResourceMap::print function
+//
+// FYI:
+// Make sure to invoke this function before origCallInst->replaceAllUsesWith()
+// and origCallInst->eraseFromParent() since this function needs to visit
+// origCallInst and its uses.
+//
+void DXILResourceMap::updateResourceMap(CallInst *origCallInst,
+                                        CallInst *newCallInst) {
+  assert((origCallInst != nullptr) && (newCallInst != nullptr) &&
+         (origCallInst != newCallInst));
+
+  CallMap.try_emplace(newCallInst, CallMap[origCallInst]);
+  CallMap.erase(origCallInst);
+
+  // Update ResUseToHandleMap since Resource Handle changed
+  for (auto it = origCallInst->users().begin();
+       it != origCallInst->users().end(); ++it) {
+    CallInst *CI_Use = dyn_cast<CallInst>(*it);
+    ResUseToHandleMap[CI_Use] = newCallInst;
+  }
+}
+
+  void DXILResourceMap::updateResUseMap(CallInst *origResUse,
+                                      std::vector<Value *> &multiNewResUse) {
+    assert((origResUse != nullptr) && "Wrong Inputs");
+
+    for (int i = 0; i < multiNewResUse.size(); ++i) {
+      CallInst *newResUse = dyn_cast<CallInst>(multiNewResUse[i]);
+      assert(newResUse != nullptr);
+
+      bool keepOrigResUseInMap =
+          i == (multiNewResUse.size() - 1) ? false : true;
+      updateResUseMapCommon(origResUse, newResUse, keepOrigResUseInMap);
+    }
+  }
+
 void DXILResourceMap::print(raw_ostream &OS) const {
   for (unsigned I = 0, E = Resources.size(); I != E; ++I) {
     OS << "Binding " << I << ":\n";
@@ -756,6 +817,14 @@ void DXILResourceMap::print(raw_ostream &OS) const {
     CI->print(OS);
     OS << "\n";
   }
+
+  for (const auto &[ResUse, ResHandle] : ResUseToHandleMap) {
+    OS << "\n";
+    OS << "Resource " << CallMap.find(ResHandle)->second;
+    OS << " is used by ";
+    ResUse->print(OS);
+    OS << "\n";
+  }
 }
 
 //===----------------------------------------------------------------------===//

llvm/lib/Target/DirectX/DXIL.td

@@ -854,6 +854,17 @@ def AnnotateHandle : DXILOp<216, annotateHandle> {
   let stages = [Stages<DXIL1_6, [all_stages]>];
 }
 
+def RawBufferLoad : DXILOp<139, rawBufferLoad> {
+  let Doc = "reads from a ByteAddressBuffer or StructuredBuffer";
+  // Handle, Coord0, Coord1, mask, alignment
+  let arguments = [HandleTy, Int32Ty, Int32Ty, Int8Ty, Int32Ty];
+  let result = OverloadTy;
+  let overloads =
+      [Overloads<DXIL1_0,
+                 [ResRetHalfTy, ResRetFloatTy, ResRetInt16Ty, ResRetInt32Ty]>];
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+}
+
 def CreateHandleFromBinding : DXILOp<217, createHandleFromBinding> {
   let Doc = "create resource handle from binding";
   let arguments = [ResBindTy, Int32Ty, Int1Ty];