[DirectX] Validating Root flags are denying shader stage

[joaosaffran]

Root Signature Flags, allow flags to block compilation of certain shader stages. This PR implements a validation and notify the user if they compile a root signature that is denying such shader stage.
Closes: #153062

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

Changes

Root Signature Flags, allow flags to block compilation of certain shader stages. This PR implements a validation and notify the user if they compile a root signature that is denying such shader stage.
Closes: #153062

---

Full diff: https://github.com/llvm/llvm-project/pull/153287.diff

4 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp (+71-24)
• (added) llvm/test/CodeGen/DirectX/rootsignature-validation-deny-shader.ll (+16)
• (added) llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-multiple-shader.ll (+17)
• (added) llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-single-shader.ll (+17)

diff --git a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
index 3721b5f539b8c..3897056d5081a 100644
--- a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
+++ b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
@@ -21,7 +21,6 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/Support/DXILABI.h"
-#include <cstdint>
 
 #define DEBUG_TYPE "dxil-post-optimization-validation"
 
@@ -169,15 +168,16 @@ reportDescriptorTableMixingTypes(Module &M, uint32_t Location,
   M.getContext().diagnose(DiagnosticInfoGeneric(Message));
 }
 
-static void reportOverlowingRange(Module &M, const dxbc::RTS0::v2::DescriptorRange &Range) {
+static void
+reportOverlowingRange(Module &M, const dxbc::RTS0::v2::DescriptorRange &Range) {
   SmallString<128> Message;
   raw_svector_ostream OS(Message);
-  OS << "Cannot append range with implicit lower " 
-      << "bound after an unbounded range "
-      << getResourceClassName(toResourceClass(static_cast<dxbc::DescriptorRangeType>(Range.RangeType)))
-      << "(register=" << Range.BaseShaderRegister << ", space=" << 
-      Range.RegisterSpace
-      << ") exceeds maximum allowed value.";
+  OS << "Cannot append range with implicit lower "
+     << "bound after an unbounded range "
+     << getResourceClassName(toResourceClass(
+            static_cast<dxbc::DescriptorRangeType>(Range.RangeType)))
+     << "(register=" << Range.BaseShaderRegister
+     << ", space=" << Range.RegisterSpace << ") exceeds maximum allowed value.";
   M.getContext().diagnose(DiagnosticInfoGeneric(Message));
 }
 
@@ -262,12 +262,57 @@ getRootDescriptorsBindingInfo(const mcdxbc::RootSignatureDesc &RSD,
   return RDs;
 }
 
+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 validateRootFlags(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:
+    break;
+  }
+}
 
 static void validateDescriptorTables(Module &M,
-                                     const mcdxbc::RootSignatureDesc &RSD,
-                                     dxil::ModuleMetadataInfo &MMI,
-                                     DXILResourceMap &DRM) {
+                                     const mcdxbc::RootSignatureDesc &RSD) {
   for (const mcdxbc::RootParameterInfo &ParamInfo : RSD.ParametersContainer) {
     if (static_cast<dxbc::RootParameterType>(ParamInfo.Header.ParameterType) !=
         dxbc::RootParameterType::DescriptorTable)
@@ -283,30 +328,31 @@ static void validateDescriptorTables(Module &M,
 
     uint64_t AppendingOffset = 0;
 
-
     for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
       dxbc::DescriptorRangeType RangeType =
           static_cast<dxbc::DescriptorRangeType>(Range.RangeType);
-      
+
       uint64_t Offset = AppendingOffset;
-      if(Range.OffsetInDescriptorsFromTableStart != ~0U)
+      if (Range.OffsetInDescriptorsFromTableStart != ~0U)
         Offset = Range.OffsetInDescriptorsFromTableStart;
-      
-      if(Offset > ~0U)
+
+      if (Offset > ~0U)
         reportOverlowingRange(M, Range);
-      if(Range.NumDescriptors == ~0U) {
+      if (Range.NumDescriptors == ~0U) {
         AppendingOffset = (uint64_t)~0U + (uint64_t)1ULL;
-      } else { 
-        uint64_t UpperBound = (uint64_t)Range.BaseShaderRegister + (uint64_t)Range.NumDescriptors - (uint64_t)1U;
-        if(UpperBound > ~0U)
+      } else {
+        uint64_t UpperBound = (uint64_t)Range.BaseShaderRegister +
+                              (uint64_t)Range.NumDescriptors - (uint64_t)1U;
+        if (UpperBound > ~0U)
           reportOverlowingRange(M, Range);
 
-        uint64_t AppendingUpperBound = (uint64_t)Offset + (uint64_t)Range.NumDescriptors - (uint64_t)1U;
-        if(AppendingUpperBound > ~0U)
+        uint64_t AppendingUpperBound =
+            (uint64_t)Offset + (uint64_t)Range.NumDescriptors - (uint64_t)1U;
+        if (AppendingUpperBound > ~0U)
           reportOverlowingRange(M, Range);
         AppendingOffset = Offset + Range.NumDescriptors;
       }
-      
+
       if (RangeType == dxbc::DescriptorRangeType::Sampler) {
         HasSampler = true;
       } else {
@@ -441,7 +487,8 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   if (mcdxbc::RootSignatureDesc *RSD = getRootSignature(RSBI, MMI)) {
     validateRootSignatureBindings(M, *RSD, MMI, DRM);
-    validateDescriptorTables(M, *RSD, MMI, DRM);
+    validateDescriptorTables(M, *RSD);
+    validateRootFlags(M, *RSD, MMI);
   }
 }
 
diff --git a/llvm/test/CodeGen/DirectX/rootsignature-validation-deny-shader.ll b/llvm/test/CodeGen/DirectX/rootsignature-validation-deny-shader.ll
new file mode 100644
index 0000000000000..234909e82b792
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/rootsignature-validation-deny-shader.ll
@@ -0,0 +1,16 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' %s 2>&1 
+; expected-no-diagnostics
+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}
+!2 = !{ !"RootFlags", i32 294 } ; 32 = deny_pixel/hull/vertex/amplification_shader_root_access
+
diff --git a/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-multiple-shader.ll b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-multiple-shader.ll
new file mode 100644
index 0000000000000..9286c31db2de0
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-multiple-shader.ll
@@ -0,0 +1,17 @@
+; 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-hull"
+
+define void @CSMain() #0 {
+entry:
+  ret void
+}
+attributes #0 = { noinline nounwind "exp-shader"="cs" "hlsl.numthreads"="1,2,1" "hlsl.shader"="hull" }
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{ !"RootFlags", i32 294 } ; 32 = deny_pixel/hull/vertex/amplification_shader_root_access
+
diff --git a/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-single-shader.ll b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-single-shader.ll
new file mode 100644
index 0000000000000..7294346900415
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-deny-single-shader.ll
@@ -0,0 +1,17 @@
+; 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"
+
+define void @CSMain() #0 {
+entry:
+  ret void
+}
+attributes #0 = { noinline nounwind "exp-shader"="cs" "hlsl.numthreads"="1,2,1" "hlsl.shader"="pixel" }
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{ !"RootFlags", i32 32 } ; 32 = deny_pixel_shader_root_access
+

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.