[DirectX] Validating Root flags are denying shader stage

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -160,6 +160,18 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   }
 }
 
+static void reportIfDeniedShaderStageAccess(Module &M, dxbc::RootFlags Flags,
+                                            dxbc::RootFlags Mask) {
+  if ((Flags & Mask) != Mask)
+    return;
+
+  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 validateRootSignature(Module &M,
                                   const mcdxbc::RootSignatureDesc &RSD,
                                   dxil::ModuleMetadataInfo &MMI,
@@ -225,7 +237,9 @@ static void validateRootSignature(Module &M,
             Builder.findOverlapping(ReportedBinding);
         reportOverlappingRegisters(M, ReportedBinding, Overlaping);
       });
+
   const hlsl::BoundRegs &BoundRegs = Builder.takeBoundRegs();
+  bool HasBindings = false;
   for (const ResourceInfo &RI : DRM) {
     const ResourceInfo::ResourceBinding &Binding = RI.getBinding();
     const dxil::ResourceTypeInfo &RTI = DRTM[RI.getHandleTy()];
@@ -240,18 +254,63 @@ static void validateRootSignature(Module &M,
       const auto *ParamInfo =
           static_cast<const mcdxbc::RootParameterInfo *>(Reg->Cookie);
 
-      if (RC != ResourceClass::SRV && RC != ResourceClass::UAV)
+      if (RC != ResourceClass::SRV && RC != ResourceClass::UAV) {
+        HasBindings = true;
         continue;
+      }
 
-      if (ParamInfo->Type == dxbc::RootParameterType::DescriptorTable)
+      if (ParamInfo->Type == dxbc::RootParameterType::DescriptorTable) {
+        HasBindings = true;
         continue;
+      }
 
-      if (RK != ResourceKind::RawBuffer && RK != ResourceKind::StructuredBuffer)
+      if (RK != ResourceKind::RawBuffer &&
+          RK != ResourceKind::StructuredBuffer) {
         reportInvalidHandleTyError(M, RC, Binding);
+        continue;
+      }
+      HasBindings = true;
+
     } else {
       reportRegNotBound(M, RC, Binding);
     }
   }
+
+  if (HasBindings && MMI.ShaderProfile != Triple::Compute) {
+    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 mcdxbc::RootSignatureDesc *

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,21 @@
+; 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"
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="geometry" {
+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, ptr nonnull @CB.str)
+  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 2, 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, 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-root-descriptor.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"
+
+@SB.str = private unnamed_addr constant [3 x i8] c"SB\00", align 1
+
+define void @CSMain() "hlsl.shader"="pixel" {
+entry:
+  %SB = tail call target("dx.RawBuffer", i32, 0, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr nonnull @SB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"DescriptorTable", i32 0, !4}
+!4 = !{!"SRV", i32 1, i32 0, i32 0, i32 -1, i32 4}
+!3 = !{!"RootFlags", i32 32} ; 32 = deny_pixel_shader_root_access

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

@@ -0,0 +1,19 @@
+; 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"
+
+@SB.str = private unnamed_addr constant [3 x i8] c"SB\00", align 1
+
+define void @CSMain() "hlsl.shader"="pixel" {
+entry:
+  %SB = tail call target("dx.RawBuffer", i32, 0, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr nonnull @SB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"RootSRV", i32 0, i32 0, i32 0, i32 4}
+!3 = !{!"RootFlags", i32 32} ; 32 = deny_pixel_shader_root_access