PIX: Rationalize UAV generation (#6883)

The impetus for these changes was unexplained crashes in a display
driver while attempting to create a debug-instrumented shader for PIX.
The heart of it is the new test in pixtest.cpp: use the compiler to
generate a raw UAV, and then compare the generated DXIL with what PIX
generates for the same purpose.
Some of the PIX passes need only one UAV for a module, but some need two
or more. In the latter case, the previous code was a bit loose about
what it was doing with respect to adding the UAV resource, and creating
its handles for each interested function. Most of the actual changes
herein are to do with that.
Lastly, the PIX UAV is raw at the D3D API level, but the instrumentation
had been doing non-raw writes. No driver seemed to care, but I've fixed
it anyway.

(cherry picked from commit bf24b7a)

lib/DxilPIXPasses/DxilAddPixelHitInstrumentation.cpp

@@ -78,8 +78,8 @@ bool DxilAddPixelHitInstrumentation::runOnModule(Module &M) {
     IRBuilder<> Builder(dxilutil::FirstNonAllocaInsertionPt(
         PIXPassHelpers::GetEntryFunction(DM)));
 
-    HandleForUAV =
-        PIXPassHelpers::CreateUAV(DM, Builder, 0, "PIX_CountUAV_Handle");
+    HandleForUAV = PIXPassHelpers::CreateUAVOnceForModule(
+        DM, Builder, 0, "PIX_CountUAV_Handle");
 
     DM.ReEmitDxilResources();
   }

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -322,7 +322,8 @@ class DxilDebugInstrumentation : public ModulePass {
   void applyOptions(PassOptions O) override;
   bool runOnModule(Module &M) override;
 
-  bool RunOnFunction(Module &M, DxilModule &DM, llvm::Function *function);
+  bool RunOnFunction(Module &M, DxilModule &DM, hlsl::DxilResource *uav,
+                     llvm::Function *function);
 
 private:
   SystemValueIndices addRequiredSystemValues(BuilderContext &BC,
@@ -447,6 +448,7 @@ DxilDebugInstrumentation::addRequiredSystemValues(BuilderContext &BC,
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::ClosestHit:
   case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Node:
     // Dispatch* thread Id is not in the input signature
     break;
   case DXIL::ShaderKind::Vertex: {
@@ -703,6 +705,9 @@ void DxilDebugInstrumentation::addInvocationSelectionProlog(
   case DXIL::ShaderKind::Miss:
     ParameterTestResult = addRaygenShaderProlog(BC);
     break;
+  case DXIL::ShaderKind::Node:
+    ParameterTestResult = BC.HlslOP->GetI1Const(1);
+    break;
   case DXIL::ShaderKind::Compute:
   case DXIL::ShaderKind::Amplification:
   case DXIL::ShaderKind::Mesh:
@@ -895,10 +900,10 @@ uint32_t DxilDebugInstrumentation::addDebugEntryValue(BuilderContext &BC,
     BytesToBeEmitted += addDebugEntryValue(BC, AsFloat);
   } else {
     Function *StoreValue =
-        BC.HlslOP->GetOpFunc(OP::OpCode::BufferStore,
+        BC.HlslOP->GetOpFunc(OP::OpCode::RawBufferStore,
                              TheValue->getType()); // Type::getInt32Ty(BC.Ctx));
     Constant *StoreValueOpcode =
-        BC.HlslOP->GetU32Const((unsigned)DXIL::OpCode::BufferStore);
+        BC.HlslOP->GetU32Const((unsigned)DXIL::OpCode::RawBufferStore);
     UndefValue *Undef32Arg = UndefValue::get(Type::getInt32Ty(BC.Ctx));
     UndefValue *UndefArg = nullptr;
     if (TheValueTypeID == Type::TypeID::IntegerTyID) {
@@ -913,6 +918,7 @@ uint32_t DxilDebugInstrumentation::addDebugEntryValue(BuilderContext &BC,
     Constant *WriteMask_X = BC.HlslOP->GetI8Const(1);
 
     auto &values = m_FunctionToValues[BC.Builder.GetInsertBlock()->getParent()];
+    Constant *RawBufferStoreAlignment = BC.HlslOP->GetU32Const(4);
 
     (void)BC.Builder.CreateCall(
         StoreValue, {StoreValueOpcode,    // i32 opcode
@@ -923,7 +929,7 @@ uint32_t DxilDebugInstrumentation::addDebugEntryValue(BuilderContext &BC,
                      UndefArg, // unused values
                      UndefArg, // unused values
                      UndefArg, // unused values
-                     WriteMask_X});
+                     WriteMask_X, RawBufferStoreAlignment});
 
     assert(m_RemainingReservedSpaceInBytes >= 4); // check for underflow
     m_RemainingReservedSpaceInBytes -= 4;
@@ -1211,19 +1217,20 @@ bool DxilDebugInstrumentation::runOnModule(Module &M) {
 
   auto ShaderModel = DM.GetShaderModel();
   auto shaderKind = ShaderModel->GetKind();
-
+  auto HLSLBindId = 0;
+  auto *uav = PIXPassHelpers::CreateGlobalUAVResource(DM, HLSLBindId, "PIXUAV");
   bool modified = false;
   if (shaderKind == DXIL::ShaderKind::Library) {
     auto instrumentableFunctions =
         PIXPassHelpers::GetAllInstrumentableFunctions(DM);
     for (auto *F : instrumentableFunctions) {
-      if (RunOnFunction(M, DM, F)) {
+      if (RunOnFunction(M, DM, uav, F)) {
         modified = true;
       }
     }
   } else {
     llvm::Function *entryFunction = PIXPassHelpers::GetEntryFunction(DM);
-    modified = RunOnFunction(M, DM, entryFunction);
+    modified = RunOnFunction(M, DM, uav, entryFunction);
   }
   return modified;
 }
@@ -1379,6 +1386,7 @@ DxilDebugInstrumentation::FindInstrumentableInstructionsInBlock(
 }
 
 bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
+                                             hlsl::DxilResource *uav,
                                              llvm::Function *function) {
   DXIL::ShaderKind shaderKind =
       PIXPassHelpers::GetFunctionShaderKind(DM, function);
@@ -1397,6 +1405,7 @@ bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::ClosestHit:
   case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Node:
     break;
   default:
     return false;
@@ -1431,8 +1440,8 @@ bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
     break;
   }
 
-  values.UAVHandle = PIXPassHelpers::CreateUAV(DM, Builder, UAVRegisterId,
-                                               "PIX_DebugUAV_Handle");
+  values.UAVHandle = PIXPassHelpers::CreateHandleForResource(
+      DM, Builder, uav, "PIX_DebugUAV_Handle");
 
   auto SystemValues = addRequiredSystemValues(BC, shaderKind);
   addInvocationSelectionProlog(BC, SystemValues, shaderKind);

lib/DxilPIXPasses/DxilOutputColorBecomesConstant.cpp

@@ -161,7 +161,7 @@ bool DxilOutputColorBecomesConstant::runOnModule(Module &M) {
     std::unique_ptr<DxilCBuffer> pCBuf = llvm::make_unique<DxilCBuffer>();
     pCBuf->SetGlobalName("PIX_ConstantColorCBName");
     pCBuf->SetGlobalSymbol(UndefValue::get(CBStructTy));
-    pCBuf->SetID(0);
+    pCBuf->SetID(static_cast<unsigned int>(DM.GetCBuffers().size()));
     pCBuf->SetSpaceID(
         (unsigned int)-2); // This is the reserved-for-tools register space
     pCBuf->SetLowerBound(0);

lib/DxilPIXPasses/DxilPIXDXRInvocationsLog.cpp

@@ -88,9 +88,9 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
     IRBuilder<> Builder(dxilutil::FirstNonAllocaInsertionPt(entryFunction));
 
     // Add the UAVs that we're going to write to
-    CallInst *HandleForCountUAV = PIXPassHelpers::CreateUAV(
+    CallInst *HandleForCountUAV = PIXPassHelpers::CreateUAVOnceForModule(
         DM, Builder, /* registerID */ 0, "PIX_CountUAV_Handle");
-    CallInst *HandleForUAV = PIXPassHelpers::CreateUAV(
+    CallInst *HandleForUAV = PIXPassHelpers::CreateUAVOnceForModule(
         DM, Builder, /* registerID */ 1, "PIX_UAV_Handle");
 
     DM.ReEmitDxilResources();

lib/DxilPIXPasses/DxilPIXMeshShaderOutputInstrumentation.cpp

@@ -337,7 +337,7 @@ bool DxilPIXMeshShaderOutputInstrumentation::runOnModule(Module &M) {
 
   m_OffsetMask = BC.HlslOP->GetU32Const(UAVDumpingGroundOffset() - 1);
 
-  m_OutputUAV = CreateUAV(DM, Builder, 0, "PIX_DebugUAV_Handle");
+  m_OutputUAV = CreateUAVOnceForModule(DM, Builder, 0, "PIX_DebugUAV_Handle");
 
   if (FirstNewStructGetMeshPayload == nullptr) {
     Instruction *firstInsertionPt = dxilutil::FirstNonAllocaInsertionPt(

lib/DxilPIXPasses/DxilShaderAccessTracking.cpp

@@ -907,6 +907,9 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
       }
     }
 
+    auto uav =
+        PIXPassHelpers::CreateGlobalUAVResource(DM, 0u, "PIX_ShaderAccessUAV");
+
     for (auto *F : instrumentableFunctions) {
       DXIL::ShaderKind shaderKind = DXIL::ShaderKind::Invalid;
       if (!DM.HasDxilFunctionProps(F)) {
@@ -922,8 +925,8 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
 
       IRBuilder<> Builder(F->getEntryBlock().getFirstInsertionPt());
 
-      m_FunctionToUAVHandle[F] =
-          PIXPassHelpers::CreateUAV(DM, Builder, 0u, "PIX_CountUAV_Handle");
+      m_FunctionToUAVHandle[F] = PIXPassHelpers::CreateHandleForResource(
+          DM, Builder, uav, "PIX_ShaderAccessUAV_Handle");
       OP *HlslOP = DM.GetOP();
       for (int accessStyle = static_cast<int>(ResourceAccessStyle::None);
            accessStyle < static_cast<int>(ResourceAccessStyle::EndOfEnum);
@@ -998,6 +1001,10 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
           }
 
           for (unsigned iParam : handleParams) {
+            // Don't instrument the accesses to the UAV that we just added
+            if (Call->getArgOperand(iParam) ==
+                m_FunctionToUAVHandle[CallerParent->getParent()])
+              continue;
             auto res = GetResourceFromHandle(Call->getArgOperand(iParam), DM);
             if (res.accessStyle == AccessStyle::None) {
               continue;

lib/DxilPIXPasses/PixPassHelpers.cpp

@@ -66,20 +66,6 @@ void FindRayQueryHandlesForFunction(llvm::Function *F,
   }
 }
 
-static unsigned int
-GetNextRegisterIdForClass(hlsl::DxilModule &DM,
-                          DXIL::ResourceClass resourceClass) {
-  switch (resourceClass) {
-  case DXIL::ResourceClass::CBuffer:
-    return static_cast<unsigned int>(DM.GetCBuffers().size());
-  case DXIL::ResourceClass::UAV:
-    return static_cast<unsigned int>(DM.GetUAVs().size());
-  default:
-    DXASSERT(false, "Unexpected resource class");
-    return 0;
-  }
-}
-
 static bool IsDynamicResourceShaderModel(DxilModule &DM) {
   return DM.GetShaderModel()->IsSMAtLeast(6, 6);
 }
@@ -88,6 +74,16 @@ static bool ShaderModelRequiresAnnotateHandle(DxilModule &DM) {
   return DM.GetShaderModel()->IsSMAtLeast(6, 6);
 }
 
+static char const *RawUAVType() { return "struct.RWByteAddressBuffer"; }
+static char const *ShaderModelHandleTypeName(DxilModule &DM) {
+  // Prior to sm6.6, lib handles were typed after the resource they denote.
+  // In 6.6 and after, and in all non-lib shader models,
+  // all handles are dx.types.Handle.
+  if (!DM.GetShaderModel()->IsLib() || DM.GetShaderModel()->IsSM66Plus())
+    return "dx.types.Handle";
+  return RawUAVType();
+}
+
 llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,
                                         llvm::IRBuilder<> &Builder,
                                         hlsl::DxilResourceBase *resource,
@@ -98,16 +94,16 @@ llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,
 
   DXIL::ResourceClass resourceClass = resource->GetClass();
 
-  unsigned int resourceMetaDataId =
-      GetNextRegisterIdForClass(DM, resourceClass);
-
   auto const *shaderModel = DM.GetShaderModel();
+  Type *resourceHandleType =
+      DM.GetModule()->getTypeByName(ShaderModelHandleTypeName(DM));
   if (shaderModel->IsLib()) {
     llvm::Constant *object = resource->GetGlobalSymbol();
-    auto *load = Builder.CreateLoad(object);
+    Value *load = Builder.CreateLoad(object, resourceHandleType);
+    llvm::cast<LoadInst>(load)->setAlignment(4);
+    llvm::cast<LoadInst>(load)->setVolatile(false);
     Function *CreateHandleForLibOpFunc =
-        HlslOP->GetOpFunc(DXIL::OpCode::CreateHandleForLib,
-                          resource->GetHLSLType()->getVectorElementType());
+        HlslOP->GetOpFunc(DXIL::OpCode::CreateHandleForLib, load->getType());
     Constant *CreateHandleForLibOpcodeArg =
         HlslOP->GetU32Const((unsigned)DXIL::OpCode::CreateHandleForLib);
     auto *handle = Builder.CreateCall(CreateHandleForLibOpFunc,
@@ -170,10 +166,8 @@ llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,
     Constant *ClassArg = HlslOP->GetI8Const(
         static_cast<std::underlying_type<DxilResourceBase::Class>::type>(
             resourceClass));
-    Constant *MetaDataArg = HlslOP->GetU32Const(
-        resourceMetaDataId); // position of the metadata record in the
-                             // corresponding metadata list
-    Constant *IndexArg = HlslOP->GetU32Const(0); //
+    Constant *MetaDataArg = HlslOP->GetU32Const(resource->GetID());
+    Constant *IndexArg = HlslOP->GetU32Const(0);
     Constant *FalseArg =
         HlslOP->GetI1Const(0); // non-uniform resource index: false
     return Builder.CreateCall(
@@ -275,29 +269,34 @@ static void AddUAVToDxilDefinedGlobalRootSignatures(DxilModule &DM) {
 }
 
 // Set up a UAV with structure of a single int
-llvm::CallInst *CreateUAV(DxilModule &DM, IRBuilder<> &Builder,
-                          unsigned int registerId, const char *name) {
+hlsl::DxilResource *CreateGlobalUAVResource(hlsl::DxilModule &DM,
+                                            unsigned int hlslBindIndex,
+                                            const char *name) {
   LLVMContext &Ctx = DM.GetModule()->getContext();
 
-  const char *PIXStructTypeName = "struct.RWByteAddressBuffer";
+  const char *PIXStructTypeName = ShaderModelHandleTypeName(DM);
   llvm::StructType *UAVStructTy =
       DM.GetModule()->getTypeByName(PIXStructTypeName);
+
   if (UAVStructTy == nullptr) {
     SmallVector<llvm::Type *, 1> Elements{Type::getInt32Ty(Ctx)};
     UAVStructTy = llvm::StructType::create(Elements, PIXStructTypeName);
-
-    // Since we only have to do this once per module, we can do it now when
-    // we're adding the singular UAV structure type to the module:
-    AddUAVToDxilDefinedGlobalRootSignatures(DM);
-    AddUAVToShaderAttributeRootSignature(DM);
   }
 
+  // Since this function should only be called once per module,
+  // we can modify the root sig at the same time:
+  AddUAVToDxilDefinedGlobalRootSignatures(DM);
+  AddUAVToShaderAttributeRootSignature(DM);
+
+  unsigned int Id = static_cast<unsigned int>(DM.GetUAVs().size());
   std::unique_ptr<DxilResource> pUAV = llvm::make_unique<DxilResource>();
+  pUAV->SetID(Id);
 
   auto const *shaderModel = DM.GetShaderModel();
+  std::string PixUavName = "PIXUAV" + std::to_string(hlslBindIndex);
   if (shaderModel->IsLib()) {
-    auto *Global = DM.GetModule()->getOrInsertGlobal(
-        ("PIXUAV" + std::to_string(registerId)).c_str(), UAVStructTy);
+    auto *Global =
+        DM.GetModule()->getOrInsertGlobal(PixUavName.c_str(), UAVStructTy);
     GlobalVariable *NewGV = cast<GlobalVariable>(Global);
     NewGV->setConstant(true);
     NewGV->setLinkage(GlobalValue::ExternalLinkage);
@@ -308,17 +307,24 @@ llvm::CallInst *CreateUAV(DxilModule &DM, IRBuilder<> &Builder,
     pUAV->SetGlobalSymbol(UndefValue::get(UAVStructTy->getPointerTo()));
   }
   pUAV->SetGlobalName(name);
-  pUAV->SetID(GetNextRegisterIdForClass(DM, DXIL::ResourceClass::UAV));
   pUAV->SetRW(true); // sets UAV class
   pUAV->SetSpaceID(
-      (unsigned int)-2); // This is the reserved-for-tools register space
-  pUAV->SetSampleCount(1);
+      (unsigned int)-2);   // This is the reserved-for-tools register space
+  pUAV->SetSampleCount(0); // This is what compiler generates for a raw UAV
   pUAV->SetGloballyCoherent(false);
   pUAV->SetHasCounter(false);
-  pUAV->SetCompType(CompType::getI32());
-  pUAV->SetLowerBound(registerId);
+  pUAV->SetCompType(
+      CompType::getInvalid()); // This is what compiler generates for a raw UAV
+  pUAV->SetLowerBound(hlslBindIndex);
   pUAV->SetRangeSize(1);
+  pUAV->SetElementStride(1);
   pUAV->SetKind(DXIL::ResourceKind::RawBuffer);
+  auto HLSLType = DM.GetModule()->getTypeByName(RawUAVType());
+  if (HLSLType == nullptr) {
+    SmallVector<llvm::Type *, 1> Elements{Type::getInt32Ty(Ctx)};
+    HLSLType = llvm::StructType::create(Elements, RawUAVType());
+  }
+  pUAV->SetHLSLType(HLSLType->getPointerTo());
 
   auto pAnnotation = DM.GetTypeSystem().GetStructAnnotation(UAVStructTy);
   if (pAnnotation == nullptr) {
@@ -330,9 +336,18 @@ llvm::CallInst *CreateUAV(DxilModule &DM, IRBuilder<> &Builder,
     pAnnotation->GetFieldAnnotation(0).SetFieldName("count");
   }
 
-  auto *handle = CreateHandleForResource(DM, Builder, pUAV.get(), name);
-
+  auto *ret = pUAV.get();
   DM.AddUAV(std::move(pUAV));
+  return ret;
+}
+
+// Set up a UAV with structure of a single int
+llvm::CallInst *CreateUAVOnceForModule(hlsl::DxilModule &DM,
+                                       llvm::IRBuilder<> &Builder,
+                                       unsigned int hlslBindIndex,
+                                       const char *name) {
+  auto uav = CreateGlobalUAVResource(DM, hlslBindIndex, name);
+  auto *handle = CreateHandleForResource(DM, Builder, uav, name);
 
   return handle;
 }

lib/DxilPIXPasses/PixPassHelpers.h

@@ -34,8 +34,14 @@ class ScopedInstruction {
 
 void FindRayQueryHandlesForFunction(
     llvm::Function *F, llvm::SmallPtrSetImpl<llvm::Value *> &RayQueryHandles);
-llvm::CallInst *CreateUAV(hlsl::DxilModule &DM, llvm::IRBuilder<> &Builder,
-                          unsigned int registerId, const char *name);
+enum class PixUAVHandleMode { NonLib, Lib };
+llvm::CallInst *CreateUAVOnceForModule(hlsl::DxilModule &DM,
+                                       llvm::IRBuilder<> &Builder,
+                                       unsigned int hlslBindIndex,
+                                       const char *name);
+hlsl::DxilResource *CreateGlobalUAVResource(hlsl::DxilModule &DM,
+                                            unsigned int hlslBindIndex,
+                                            const char *name);
 llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,
                                         llvm::IRBuilder<> &Builder,
                                         hlsl::DxilResourceBase *resource,

tools/clang/test/HLSLFileCheck/pix/AccessTracking.hlsl

@@ -1,7 +1,7 @@
 // RUN: %dxc -ECSMain -Tcs_6_0 %s | %opt -S -hlsl-dxil-pix-shader-access-instrumentation,config=S0:1:1i1;U0:2:10i0;.0;0;0. | %FileCheck %s
 
-// Check we added the UAV:
-// CHECK:  %PIX_CountUAV_Handle = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 1, i32 1, i32 0, i1 false)
+// Check we added the UAV:                                                                      v----metadata position: not important for this check
+// CHECK:  %PIX_ShaderAccessUAV_Handle = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 1, i32 [[S:[0-9]+]], i32 0, i1 false)
 
 // check for correct out-of-bounds calculation
 // CHECK: CompareWithSlotLimit = icmp uge i32
@@ -12,7 +12,7 @@
 // CHECK: slotIndex = mul i32
 
 // Check for udpate of UAV:
-// CHECK: call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle %PIX_CountUAV_Handle
+// CHECK: call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle %PIX_ShaderAccessUAV_Handle
 
 
 ByteAddressBuffer inBuffer : register(t0);