[DirectX] Validating Root flags are denying shader stage

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -160,6 +160,41 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   }
 }
 
+static void reportIfDeniedShaderStageAccess(Module &M,
+                                            const dxbc::RootFlags &Flags,
+                                            const 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 std::optional<dxbc::RootFlags>
+getEnvironmentDenyFlagMask(Triple::EnvironmentType ShaderProfile) {
+  switch (ShaderProfile) {
+  case Triple::Pixel:
+    return dxbc::RootFlags::DenyPixelShaderRootAccess;
+  case Triple::Vertex:
+    return dxbc::RootFlags::DenyVertexShaderRootAccess;
+  case Triple::Geometry:
+    return dxbc::RootFlags::DenyGeometryShaderRootAccess;
+  case Triple::Hull:
+    return dxbc::RootFlags::DenyHullShaderRootAccess;
+  case Triple::Domain:
+    return dxbc::RootFlags::DenyDomainShaderRootAccess;
+  case Triple::Mesh:
+    return dxbc::RootFlags::DenyMeshShaderRootAccess;
+  case Triple::Amplification:
+    return dxbc::RootFlags::DenyAmplificationShaderRootAccess;
+  default:
+    return std::nullopt;
+  }
+}
+
 static void validateRootSignature(Module &M,
                                   const mcdxbc::RootSignatureDesc &RSD,
                                   dxil::ModuleMetadataInfo &MMI,
@@ -225,7 +260,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()];
@@ -236,22 +273,34 @@ static void validateRootSignature(Module &M,
         BoundRegs.findBoundReg(RC, Binding.Space, Binding.LowerBound,
                                Binding.LowerBound + Binding.Size - 1);
 
-    if (Reg != nullptr) {
-      const auto *ParamInfo =
-          static_cast<const mcdxbc::RootParameterInfo *>(Reg->Cookie);
-
-      if (RC != ResourceClass::SRV && RC != ResourceClass::UAV)
-        continue;
-
-      if (ParamInfo->Type == dxbc::RootParameterType::DescriptorTable)
-        continue;
-
-      if (RK != ResourceKind::RawBuffer && RK != ResourceKind::StructuredBuffer)
-        reportInvalidHandleTyError(M, RC, Binding);
-    } else {
+    if (!Reg) {
       reportRegNotBound(M, RC, Binding);
+      continue;
+    }
+
+    const auto *ParamInfo =
+        static_cast<const mcdxbc::RootParameterInfo *>(Reg->Cookie);
+
+    const bool IsSRVOrUAV =
+        RC == ResourceClass::SRV || RC == ResourceClass::UAV;
+    const bool IsDescriptorTable =
+        ParamInfo->Type == dxbc::RootParameterType::DescriptorTable;
+    const bool IsRawOrStructuredBuffer =
+        RK != ResourceKind::RawBuffer && RK != ResourceKind::StructuredBuffer;
+    if (IsSRVOrUAV && !IsDescriptorTable && IsRawOrStructuredBuffer) {
+      reportInvalidHandleTyError(M, RC, Binding);
+      continue;
     }
+
+    HasBindings = true;
   }
+
+  if (!HasBindings)
+    return;
+
+  if (std::optional<dxbc::RootFlags> Mask =
+          getEnvironmentDenyFlagMask(MMI.ShaderProfile))
+    reportIfDeniedShaderStageAccess(M, dxbc::RootFlags(RSD.Flags), *Mask);
 }
 
 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-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-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

llvm/test/CodeGen/DirectX/rootsignature-validation-not-dening-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 }