[DirectX] Validating Root flags are denying shader stage

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -205,7 +205,57 @@ getRootDescriptorsBindingInfo(const mcdxbc::RootSignatureDesc &RSD,
   return RDs;
 }
 
-static void validateRootSignatureBindings(Module &M,
+static void reportIfDeniedShaderStageAccess(Module &M, dxbc::RootFlags Flags,
+                                            dxbc::RootFlags Mask) {
+  if ((Flags & Mask) == Mask) {
+    SmallString<128> Message;
+    raw_svector_ostream OS(Message);
+    OS << "Shader has root bindings but root signature uses a DENY flag to "
+          "disallow root binding access to the shader stage.";
+    M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+  }
+}
+
+static void validateDeniedStagedNotInUse(Module &M,
+                                         const mcdxbc::RootSignatureDesc &RSD,
+                                         const dxil::ModuleMetadataInfo &MMI) {
+  dxbc::RootFlags Flags = dxbc::RootFlags(RSD.Flags);
+
+  switch (MMI.ShaderProfile) {
+  case Triple::Pixel:
+    reportIfDeniedShaderStageAccess(M, Flags,
+                                    dxbc::RootFlags::DenyPixelShaderRootAccess);
+    break;
+  case Triple::Vertex:
+    reportIfDeniedShaderStageAccess(
+        M, Flags, dxbc::RootFlags::DenyVertexShaderRootAccess);
+    break;
+  case Triple::Geometry:
+    reportIfDeniedShaderStageAccess(
+        M, Flags, dxbc::RootFlags::DenyGeometryShaderRootAccess);
+    break;
+  case Triple::Hull:
+    reportIfDeniedShaderStageAccess(M, Flags,
+                                    dxbc::RootFlags::DenyHullShaderRootAccess);
+    break;
+  case Triple::Domain:
+    reportIfDeniedShaderStageAccess(
+        M, Flags, dxbc::RootFlags::DenyDomainShaderRootAccess);
+    break;
+  case Triple::Mesh:
+    reportIfDeniedShaderStageAccess(M, Flags,
+                                    dxbc::RootFlags::DenyMeshShaderRootAccess);
+    break;
+  case Triple::Amplification:
+    reportIfDeniedShaderStageAccess(
+        M, Flags, dxbc::RootFlags::DenyAmplificationShaderRootAccess);
+    break;
+  default:
+    llvm_unreachable("Invalid triple to shader stage conversion");
+  }
+}
+
+static bool validateRootSignatureBindings(Module &M,
                                           const mcdxbc::RootSignatureDesc &RSD,
                                           dxil::ModuleMetadataInfo &MMI,
                                           DXILResourceMap &DRM) {
@@ -275,23 +325,27 @@ static void validateRootSignatureBindings(Module &M,
             Builder.findOverlapping(ReportedBinding);
         reportOverlappingRegisters(M, ReportedBinding, Overlaping);
       });
-    SmallVector<ResourceInfo::ResourceBinding> RDs =
-        getRootDescriptorsBindingInfo(RSD, Visibility);
-    for (const auto &ResList :
-         {std::make_pair(ResourceClass::SRV, DRM.srvs()),
-          std::make_pair(ResourceClass::UAV, DRM.uavs()),
-          std::make_pair(ResourceClass::CBuffer, DRM.cbuffers()),
-          std::make_pair(ResourceClass::Sampler, DRM.samplers())}) {
-      for (auto Res : ResList.second) {
-        llvm::dxil::ResourceInfo::ResourceBinding ResBinding = Res.getBinding();
-        llvm::hlsl::BindingInfo::BindingRange ResRange(
-            ResBinding.LowerBound, ResBinding.LowerBound + ResBinding.Size);
-
-        if (!Info.isBound(ResList.first, ResBinding.Space, ResRange))
-          reportRegNotBound(M, ResList.first, ResBinding);
-      }
-      checkInvalidHandleTy(M, RDs, ResList.second);
+  bool HasBindings = false;
+  SmallVector<ResourceInfo::ResourceBinding> RDs =
+      getRootDescriptorsBindingInfo(RSD, Visibility);
+  for (const auto &ResList :
+       {std::make_pair(ResourceClass::SRV, DRM.srvs()),
+        std::make_pair(ResourceClass::UAV, DRM.uavs()),
+        std::make_pair(ResourceClass::CBuffer, DRM.cbuffers()),
+        std::make_pair(ResourceClass::Sampler, DRM.samplers())}) {
+    for (auto Res : ResList.second) {
+      llvm::dxil::ResourceInfo::ResourceBinding ResBinding = Res.getBinding();
+      llvm::hlsl::BindingInfo::BindingRange ResRange(
+          ResBinding.LowerBound, ResBinding.LowerBound + ResBinding.Size);
+
+      if (!Info.isBound(ResList.first, ResBinding.Space, ResRange))
+        reportRegNotBound(M, ResList.first, ResBinding);
+      else
+        HasBindings = true;
+    }
+    checkInvalidHandleTy(M, RDs, ResList.second);
   }
+  return HasBindings;
 }
 
 static mcdxbc::RootSignatureDesc *
@@ -316,7 +370,9 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
                                        "DXILResourceImplicitBinding pass");
 
   if (mcdxbc::RootSignatureDesc *RSD = getRootSignature(RSBI, MMI)) {
-    validateRootSignatureBindings(M, *RSD, MMI, DRM);
+    bool HasBindings = validateRootSignatureBindings(M, *RSD, MMI, DRM);
+    if (HasBindings && MMI.ShaderProfile != Triple::Compute)
+      validateDeniedStagedNotInUse(M, *RSD, MMI);
   }
 }
 

llvm/test/CodeGen/DirectX/rootsignature-validation-deny-no-binding.ll

@@ -0,0 +1,17 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' %s 
+; This is a valid case where no resource is being used
+target triple = "dxil-pc-shadermodel6.6-pixel"
+
+define void @CSMain() #0 {
+entry:
+  ret void
+}
+attributes #0 = { noinline nounwind "exp-shader"="cs" "hlsl.numthreads"="1,2,1" "hlsl.shader"="geometry" }
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3, !4}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}
+!3 = !{ !"RootFlags", i32 294 } ; 294 = deny_pixel/hull/vertex/amplification_shader_root_access
+!4 = !{ !"RootSRV", i32 0, i32 1, i32 0, i32 0 }

llvm/test/CodeGen/DirectX/rootsignature-validation-deny-shader.ll

@@ -0,0 +1,16 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' %s 
+; Valid scenario where shader stage is not blocked from accessing root bindings
+target triple = "dxil-pc-shadermodel6.6-geometry"
+
+define void @CSMain() #0 {
+entry:
+  ret void
+}
+attributes #0 = { noinline nounwind "exp-shader"="cs" "hlsl.numthreads"="1,2,1" "hlsl.shader"="geometry" }
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{ !"RootFlags", i32 294 } ; 294 = deny_pixel/hull/vertex/amplification_shader_root_access
+!3 = !{ !"RootCBV", i32 0, i32 1, i32 0, i32 0 }

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-multiple-shader.ll

@@ -0,0 +1,20 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' %s 2>&1 | FileCheck %s
+; CHECK: error: Shader has root bindings but root signature uses a DENY flag to disallow root binding access to the shader stage.
+target triple = "dxil-pc-shadermodel6.6-pixel"
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 2, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}
+!3 = !{!"RootFlags", i32 294} ; 294 = deny_pixel/hull/vertex/amplification_shader_root_access

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-single-shader.ll

@@ -0,0 +1,21 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' %s 2>&1 | FileCheck %s
+
+; CHECK: error: Shader has root bindings but root signature uses a DENY flag to disallow root binding access to the shader stage.
+target triple = "dxil-pc-shadermodel6.6-pixel"
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 2, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}
+!3 = !{!"RootFlags", i32 32} ; 32 = deny_pixel_shader_root_access