Changes in DXILResourceMap for lowering llvm.dx.rawbufferload to dxil ops

Author: lizhengxing

llvm/include/llvm/Analysis/DXILResource.h

@@ -340,6 +340,11 @@ class DXILResourceMap {
 
   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");

llvm/lib/Analysis/DXILResource.cpp

@@ -791,6 +791,20 @@ void DXILResourceMap::updateResourceMap(CallInst *origCallInst,
   }
 }
 
+  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";

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -704,6 +704,7 @@ class OpLowerer {
       Value *EVI = IRB.CreateExtractValue(Op, 0);
 
       Intrin->replaceAllUsesWith(EVI);
+      DRM.updateResUseMap(Intrin, Op);
       Intrin->eraseFromParent();
 
       return Error::success();
@@ -761,6 +762,7 @@ class OpLowerer {
     // TODO:
     // Remove the dx.op.rawbufferload without any uses now?
 
+    DRM.updateResUseMap(Intrin, bufLds);
     Intrin->eraseFromParent();
 
     return Error::success();
@@ -771,16 +773,14 @@ class OpLowerer {
 
     return replaceFunction(F, [&](CallInst *CI) -> Error {
       IRB.SetInsertPoint(CI);
-#if 0
-      auto *It = DRM.find(dyn_cast<CallInst>(CI->getArgOperand(0)));
+
+      auto *It = DRM.find(DRM.findResHandleByUse(CI));
       assert(It != DRM.end() && "Resource not in map?");
       dxil::ResourceInfo &RI = *It;
 
-      assert((RI.getResourceKind() == dxil::ResourceKind::StructuredBuffer) ||
-             (RI.getResourceKind() == dxil::ResourceKind::RawBuffer));
-#else
-      ResourceKind RCKind = dxil::ResourceKind::StructuredBuffer;
-#endif
+      ResourceKind RCKind =  RI.getResourceKind();
+      assert((RCKind == dxil::ResourceKind::StructuredBuffer) ||
+             (RCKind == dxil::ResourceKind::RawBuffer));
 
       Type *Ty = CI->getType();
       std::vector<Value *> bufLds;