[PIX] Fix shared instrumentation resource handling

lib/DxilPIXPasses/DxilDebugBreakInstrumentation.cpp

@@ -120,10 +120,8 @@ bool DxilDebugBreakInstrumentation::runOnModule(Module &M) {
     CI->eraseFromParent();
   }
 
-  // Clean up the now-unused declaration. Not strictly required for
-  // correctness, but keeps the module free of dead references.
-  if (DebugBreakFunc->use_empty())
-    DebugBreakFunc->eraseFromParent();
+  PIXPassHelpers::eraseIfUnused(DM, DebugBreakFunc);
+  PIXPassHelpers::eraseIfUnused(DM, AtomicOpFunc);
 
   const bool modified = (PixUAVResource != nullptr);
 

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -1436,23 +1436,6 @@ bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
 
   auto &values = m_FunctionToValues[BC.Builder.GetInsertBlock()->getParent()];
 
-  // PIX binds two UAVs when running this instrumentation: one for raygen
-  // shaders and another for the hitgroups and miss shaders. Since PIX invokes
-  // this pass at the library level, which may contain examples of both types,
-  // PIX can't really specify which UAV index to use per-shader. This pass
-  // therefore just has to know this:
-  constexpr unsigned int RayGenUAVRegister = 0;
-  constexpr unsigned int HitGroupAndMissUAVRegister = 1;
-  unsigned int UAVRegisterId = RayGenUAVRegister;
-  switch (shaderKind) {
-  case DXIL::ShaderKind::ClosestHit:
-  case DXIL::ShaderKind::Intersection:
-  case DXIL::ShaderKind::AnyHit:
-  case DXIL::ShaderKind::Miss:
-    UAVRegisterId = HitGroupAndMissUAVRegister;
-    break;
-  }
-
   values.UAVHandle = PIXPassHelpers::CreateHandleForResource(
       DM, Builder, uav, "PIX_DebugUAV_Handle");
 

lib/DxilPIXPasses/DxilNonUniformResourceIndexInstrumentation.cpp

@@ -149,6 +149,9 @@ bool DxilNonUniformResourceIndexInstrumentation::runOnModule(Module &M) {
 
   const bool modified = (PixUAVResource != nullptr);
 
+  PIXPassHelpers::eraseIfUnused(DM, WaveActiveAllEqualFunc);
+  PIXPassHelpers::eraseIfUnused(DM, AtomicOpFunc);
+
   if (modified) {
     DM.ReEmitDxilResources();
 

lib/DxilPIXPasses/DxilOutputColorBecomesConstant.cpp

@@ -125,6 +125,8 @@ bool DxilOutputColorBecomesConstant::runOnModule(Module &M) {
       [&hasIntOutputs](CallInst *) { hasIntOutputs = true; });
 
   if (!hasFloatOutputs && !hasIntOutputs) {
+    PIXPassHelpers::eraseIfUnused(DM, FloatOutputFunction);
+    PIXPassHelpers::eraseIfUnused(DM, IntOutputFunction);
     return false;
   }
 
@@ -251,6 +253,9 @@ bool DxilOutputColorBecomesConstant::runOnModule(Module &M) {
         });
   }
 
+  PIXPassHelpers::eraseIfUnused(DM, FloatOutputFunction);
+  PIXPassHelpers::eraseIfUnused(DM, IntOutputFunction);
+
   return Modified;
 }
 

lib/DxilPIXPasses/DxilPIXAddTidToAmplificationShaderPayload.cpp

@@ -183,8 +183,11 @@ bool DxilPIXAddTidToAmplificationShaderPayload::runOnModule(Module &M) {
                    {DispatchMeshOpcode, DispatchMesh.get_threadGroupCountX(),
                     DispatchMesh.get_threadGroupCountY(),
                     DispatchMesh.get_threadGroupCountZ(), NewStructAlloca});
+      llvm::Function *OriginalDispatchMeshFn =
+          cast<CallInst>(&*I)->getCalledFunction();
       I->removeFromParent();
       delete &*I;
+      PIXPassHelpers::eraseIfUnused(DM, OriginalDispatchMeshFn);
       // Validation requires exactly one DispatchMesh in an AS, so we can exit
       // after the first one:
       DM.ReEmitDxilResources();

lib/DxilPIXPasses/DxilPIXMeshShaderOutputInstrumentation.cpp

@@ -310,6 +310,8 @@ bool DxilPIXMeshShaderOutputInstrumentation::runOnModule(Module &M) {
     }
 
     if (getMeshPayloadInstructions != nullptr) {
+      llvm::Function *OriginalGetMeshPayloadFunction =
+          cast<CallInst>(getMeshPayloadInstructions)->getCalledFunction();
 
       Function *DxilFunc = HlslOP->GetOpFunc(
           OP::OpCode::GetMeshPayload, expanded.ExpandedPayloadStructPtrType);
@@ -326,6 +328,7 @@ bool DxilPIXMeshShaderOutputInstrumentation::runOnModule(Module &M) {
       ReplaceAllUsesOfInstructionWithNewValueAndDeleteInstruction(
           getMeshPayloadInstructions, payload,
           expanded.ExpandedPayloadStructType);
+      PIXPassHelpers::eraseIfUnused(DM, OriginalGetMeshPayloadFunction);
     }
   }
 
@@ -378,9 +381,11 @@ bool DxilPIXMeshShaderOutputInstrumentation::runOnModule(Module &M) {
       {Type::getInt16Ty(Ctx), int16ValueIndicator},
       {Type::getFloatTy(Ctx), floatValueIndicator},
       {Type::getHalfTy(Ctx), float16ValueIndicator}};
+  SmallVector<Function *, 4> StoreVertexOutputFunctions;
 
   for (auto const &Overload : StoreVertexOutputOverloads) {
     F = HlslOP->GetOpFunc(DXIL::OpCode::StoreVertexOutput, Overload.type);
+    StoreVertexOutputFunctions.push_back(F);
     FunctionUses = F->uses();
     for (auto FI = FunctionUses.begin(); FI != FunctionUses.end();) {
       auto &FunctionUse = *FI++;
@@ -419,6 +424,10 @@ bool DxilPIXMeshShaderOutputInstrumentation::runOnModule(Module &M) {
     }
   }
 
+  for (Function *StoreVertexOutputFunction : StoreVertexOutputFunctions) {
+    PIXPassHelpers::eraseIfUnused(DM, StoreVertexOutputFunction);
+  }
+
   DM.ReEmitDxilResources();
 
   return true;

lib/DxilPIXPasses/DxilRemoveDiscards.cpp

@@ -17,6 +17,8 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 
+#include "PixPassHelpers.h"
+
 using namespace llvm;
 using namespace hlsl;
 
@@ -53,6 +55,8 @@ bool DxilRemoveDiscards::runOnModule(Module &M) {
     Modified = true;
   }
 
+  PIXPassHelpers::eraseIfUnused(DM, DiscardFunction);
+
   return Modified;
 }