[llvm] [DirectX] Add Resource uses to Resource Handle map in DXILResourceMap (PR #112798)

via llvm-commits llvm-commits at lists.llvm.org
Thu Oct 17 16:29:55 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-directx

Author: Zhengxing li (lizhengxing)

<details>
<summary>Changes</summary>

When lowering some resource use intrisics to DXIL operations, it needs to know the information of the resource that the intrisics are using.

This PR adds Resource uses to Resource Handle map in DXILResourceMap. It helps the resource uses to find the resource information.

This PR is also useful to #<!-- -->106188

---
Full diff: https://github.com/llvm/llvm-project/pull/112798.diff


4 Files Affected:

- (modified) llvm/include/llvm/Analysis/DXILResource.h (+25) 
- (modified) llvm/lib/Analysis/DXILResource.cpp (+31) 
- (modified) llvm/lib/Target/DirectX/DXILOpLowering.cpp (+9) 
- (added) llvm/test/Analysis/DXILResource/dx-resource-map.ll (+48) 


``````````diff
diff --git a/llvm/include/llvm/Analysis/DXILResource.h b/llvm/include/llvm/Analysis/DXILResource.h
index 6b577c02f05450..abd0e8e918aa05 100644
--- a/llvm/include/llvm/Analysis/DXILResource.h
+++ b/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,29 @@ class DXILResourceMap {
   }
 
   void print(raw_ostream &OS) const;
+
+  void updateResourceMap(CallInst *origCallInst, CallInst *newCallInst);
+
+  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> {
diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index 2802480481690d..801cb0aa9876d5 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/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,23 @@ DXILResourceMap::DXILResourceMap(
   }
 }
 
+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::print(raw_ostream &OS) const {
   for (unsigned I = 0, E = Resources.size(); I != E; ++I) {
     OS << "Binding " << I << ":\n";
@@ -756,6 +779,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";
+  }
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Target/DirectX/DXILOpLowering.cpp b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
index 99df4850872078..8e800da9968303 100644
--- a/llvm/lib/Target/DirectX/DXILOpLowering.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
@@ -204,6 +204,8 @@ class OpLowerer {
 
       Value *Cast = createTmpHandleCast(*OpCall, CI->getType());
 
+      DRM.updateResourceMap(CI, *OpCall);
+
       CI->replaceAllUsesWith(Cast);
       CI->eraseFromParent();
       return Error::success();
@@ -248,6 +250,8 @@ class OpLowerer {
 
       Value *Cast = createTmpHandleCast(*OpAnnotate, CI->getType());
 
+      DRM.updateResourceMap(CI, *OpBind);
+
       CI->replaceAllUsesWith(Cast);
       CI->eraseFromParent();
 
@@ -412,6 +416,9 @@ class OpLowerer {
           OpCode::BufferLoad, Args, CI->getName(), NewRetTy);
       if (Error E = OpCall.takeError())
         return E;
+
+      DRM.updateResUseMap(CI, *OpCall);
+
       if (Error E = replaceResRetUses(CI, *OpCall, HasCheckBit))
         return E;
 
@@ -456,6 +463,8 @@ class OpLowerer {
       if (Error E = OpCall.takeError())
         return E;
 
+      DRM.updateResUseMap(CI, *OpCall);
+
       CI->eraseFromParent();
       return Error::success();
     });
diff --git a/llvm/test/Analysis/DXILResource/dx-resource-map.ll b/llvm/test/Analysis/DXILResource/dx-resource-map.ll
new file mode 100644
index 00000000000000..5f196e3ba7e825
--- /dev/null
+++ b/llvm/test/Analysis/DXILResource/dx-resource-map.ll
@@ -0,0 +1,48 @@
+; RUN: opt -S -disable-output -disable-output -passes="print<dxil-resource>,dxil-op-lower,print<dxil-resource>" -mtriple=dxil-pc-shadermodel6.6-compute < %s 2>&1 | FileCheck %s -check-prefixes=CHECK,CHECK_SM66
+; RUN: opt -S -disable-output -disable-output -passes="print<dxil-resource>,dxil-op-lower,print<dxil-resource>" -mtriple=dxil-pc-shadermodel6.2-compute < %s 2>&1 | FileCheck %s -check-prefixes=CHECK,CHECK_SM62
+
+define void @test_typedbuffer() {
+  ; RWBuffer<float4> Buf : register(u5, space3)
+  %uav1 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
+              @llvm.dx.handle.fromBinding.tdx.TypedBuffer_f32_1_0(
+                  i32 3, i32 5, i32 1, i32 0, i1 false)
+  ; CHECK: Binding [[UAV1:[0-9]+]]:
+  ; CHECK:   Symbol: ptr undef
+  ; CHECK:   Name: ""
+  ; CHECK:   Binding:
+  ; CHECK:     Record ID: 0
+  ; CHECK:     Space: 3
+  ; CHECK:     Lower Bound: 5
+  ; CHECK:     Size: 1
+  ; CHECK:   Class: UAV
+  ; CHECK:   Kind: TypedBuffer
+  ; CHECK:   Globally Coherent: 0
+  ; CHECK:   HasCounter: 0
+  ; CHECK:   IsROV: 0
+  ; CHECK:   Element Type: f32
+  ; CHECK:   Element Count: 4
+
+  ; CHECK:     Call bound to [[UAV1]]:  %uav1 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0) @llvm.dx.handle.fromBinding.tdx.TypedBuffer_v4f32_1_0_0t(i32 3, i32 5, i32 1, i32 0, i1 false)
+  ; CHECK-DAG: Resource [[UAV1]] is used by   %data0 = call <4 x float> @llvm.dx.typedBufferLoad.v4f32.tdx.TypedBuffer_v4f32_1_0_0t(target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav1, i32 0)
+  ; CHECK-DAG: Resource [[UAV1]] is used by   call void @llvm.dx.typedBufferStore.tdx.TypedBuffer_v4f32_1_0_0t.v4f32(target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav1, i32 2, <4 x float> %data0)
+
+  %data0 = call <4 x float> @llvm.dx.typedBufferLoad(
+      target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav1, i32 0)
+  call void @llvm.dx.typedBufferStore(
+      target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav1,
+      i32 2, <4 x float> %data0)
+
+  ;
+  ;;; After dxil-op-lower, the DXILResourceMap info should be updated.
+  ;
+  ; CHECK_SM66:     Call bound to [[UAV1]]:  %uav11 = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 218, %dx.types.ResBind { i32 5, i32 5, i32 3, i8 1 }, i32 0, i1 false)
+  ; CHECK_SM66-DAG: Resource [[UAV1]] is used by   %data02 = call %dx.types.ResRet.f32 @dx.op.bufferLoad.f32(i32 68, %dx.types.Handle %uav1_annot, i32 0, i32 undef)
+  ; CHECK_SM66-DAG: Resource [[UAV1]] is used by   call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %uav1_annot, i32 2, i32 undef, float %9, float %10, float %11, float %12, i8 15)
+  ;
+  ; CHECK_SM62:     Call bound to [[UAV1]]:  %uav11 = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 1, i32 0, i32 0, i1 false)
+  ; CHECK_SM62-DAG: Resource [[UAV1]] is used by   %data02 = call %dx.types.ResRet.f32 @dx.op.bufferLoad.f32(i32 68, %dx.types.Handle %uav11, i32 0, i32 undef)
+  ; CHECK_SM62-DAG: Resource [[UAV1]] is used by   call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %uav11, i32 2, i32 undef, float %9, float %10, float %11, float %12, i8 15)
+
+  ret void
+}
+

``````````

</details>


https://github.com/llvm/llvm-project/pull/112798


More information about the llvm-commits mailing list