[DirectX] Implement Shader Flag Analysis for UAVsAtEveryStage (llvm#137085)

Fixes llvm#112272

In addition to the implementation of the UAVsAtEveryStage shader flag
analysis, several unrelated tests have had the `dx.valver` module
metadata defined to avoid setting the UAVsAtEveryStage shader flag in
them.
Example:
```
!dx.valver = !{!0}
!0 = !{i32 1, i32 8}
```

---------

Co-authored-by: Justin Bogner <[email protected]>

Author: 2 people

llvm/lib/Target/DirectX/DXILShaderFlags.cpp

@@ -32,6 +32,35 @@
 using namespace llvm;
 using namespace llvm::dxil;
 
+static bool hasUAVsAtEveryStage(DXILResourceMap &DRM,
+                                const ModuleMetadataInfo &MMDI) {
+  if (DRM.uavs().empty())
+    return false;
+
+  switch (MMDI.ShaderProfile) {
+  default:
+    return false;
+  case Triple::EnvironmentType::Compute:
+  case Triple::EnvironmentType::Pixel:
+    return false;
+  case Triple::EnvironmentType::Vertex:
+  case Triple::EnvironmentType::Geometry:
+  case Triple::EnvironmentType::Hull:
+  case Triple::EnvironmentType::Domain:
+    return true;
+  case Triple::EnvironmentType::Library:
+  case Triple::EnvironmentType::RayGeneration:
+  case Triple::EnvironmentType::Intersection:
+  case Triple::EnvironmentType::AnyHit:
+  case Triple::EnvironmentType::ClosestHit:
+  case Triple::EnvironmentType::Miss:
+  case Triple::EnvironmentType::Callable:
+  case Triple::EnvironmentType::Mesh:
+  case Triple::EnvironmentType::Amplification:
+    return MMDI.ValidatorVersion < VersionTuple(1, 8);
+  }
+}
+
 static bool checkWaveOps(Intrinsic::ID IID) {
   // Currently unsupported intrinsics
   // case Intrinsic::dx_wave_getlanecount:
@@ -266,6 +295,8 @@ void ModuleShaderFlags::initialize(Module &M, DXILResourceTypeMap &DRTM,
       NumUAVs += UAV.getBinding().Size;
   if (NumUAVs > 8)
     CombinedSFMask.Max64UAVs = true;
+
+  CombinedSFMask.UAVsAtEveryStage = hasUAVsAtEveryStage(DRM, MMDI);
 }
 
 void ComputedShaderFlags::print(raw_ostream &OS) const {

llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.5.ll

@@ -7,9 +7,10 @@
 target triple = "dxil-pc-shadermodel6.7-library"
 
 ; CHECK:      Combined Shader Flags for Module
-; CHECK-NEXT: Shader Flags Value: 0x00000000
+; CHECK-NEXT: Shader Flags Value: 0x00010000
 
-; CHECK-NOT: Note: shader requires additional functionality:
+; CHECK: Note: shader requires additional functionality:
+; CHECK:    UAVs at every shader stage
 ; CHECK-NOT:    64 UAV slots
 
 ; CHECK: Function test : 0x00000000
@@ -26,10 +27,8 @@ define void @test() "hlsl.export" {
   ret void
 }
 
-; Set validator version to 1.5
 !dx.valver = !{!1}
 !1 = !{i32 1, i32 5}
 
-; Set this flag to 1 to prevent the ResMayNotAlias flag from being set
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"dx.resmayalias", i32 1}

llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.6.ll

@@ -8,9 +8,10 @@
 target triple = "dxil-pc-shadermodel6.7-library"
 
 ; CHECK:      Combined Shader Flags for Module
-; CHECK-NEXT: Shader Flags Value: 0x00008000
+; CHECK-NEXT: Shader Flags Value: 0x00018000
 
 ; CHECK: Note: shader requires additional functionality:
+; CHECK:        UAVs at every shader stage
 ; CHECK:        64 UAV slots
 
 ; CHECK: Function test : 0x00000000
@@ -27,19 +28,15 @@ define void @test() "hlsl.export" {
   ret void
 }
 
-; Set validator version to 1.6
 !dx.valver = !{!1}
 !1 = !{i32 1, i32 6}
 
-; Set this flag to 1 to prevent the ResMayNotAlias flag from being set
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"dx.resmayalias", i32 1}
 
 ; DXC: - Name:            SFI0
 ; DXC-NEXT:     Size:            8
 ; DXC-NEXT:     Flags:
-; DXC-NOT:   {{[A-Za-z]+: +true}}
-; DXC:              Max64UAVs:         true
-; DXC-NOT:   {{[A-Za-z]+: +true}}
+; DXC:       Max64UAVs:         true
 ; DXC:       NextUnusedBit:   false
 ; DXC: ...

llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs.ll

@@ -55,13 +55,13 @@ define void @test() "hlsl.export" {
 }
 
 !llvm.module.flags = !{!0}
+!dx.valver = !{!1}
 !0 = !{i32 1, !"dx.resmayalias", i32 1}
+!1 = !{i32 1, i32 8}
 
 ; DXC: - Name:            SFI0
 ; DXC-NEXT:     Size:            8
 ; DXC-NEXT:     Flags:
-; DXC-NOT:   {{[A-Za-z]+: +true}}
-; DXC:              Max64UAVs:         true
-; DXC-NOT:   {{[A-Za-z]+: +true}}
+; DXC:       Max64UAVs:         true
 ; DXC:       NextUnusedBit:   false
 ; DXC: ...

llvm/test/CodeGen/DirectX/ShaderFlags/res-may-alias-1.ll

@@ -35,7 +35,9 @@ define float @loadSRV() #0 {
 }
 
 !llvm.module.flags = !{!0}
-
 !0 = !{i32 1, !"dx.resmayalias", i32 1}
 
+!dx.valver = !{!1}
+!1 = !{i32 1, i32 8}
+
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}

llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.6.ll

@@ -34,4 +34,7 @@ define float @loadSRV() #0 {
   ret float %val
 }
 
+!dx.valver = !{!0}
+!0 = !{i32 1, i32 8}
+
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}

llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.7.ll

@@ -1,14 +1,18 @@
 ; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
 
 ; This test checks to ensure the behavior of the DXIL shader flag analysis
-; for the flag ResMayNotAlias is correct when the DXIL Version is >= 1.7. The
-; ResMayNotAlias flag (0x20000000) should be set on all functions if there are
-; one or more UAVs present globally in the module.
+; for the flag ResMayNotAlias is correct when the DXIL Version is >= 1.7 and the
+; DXIL Validator Version < 1.8. The ResMayNotAlias flag (0x20000000) should be
+; set on all functions if there are one or more UAVs present globally in the
+; module.
 
 target triple = "dxil-pc-shadermodel6.7-library"
 
 ; CHECK:      Combined Shader Flags for Module
-; CHECK-NEXT: Shader Flags Value: 0x200000010
+; CHECK-NEXT: Shader Flags Value: 0x200010010
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK:       UAVs at every shader stage
 
 ; CHECK: Note: extra DXIL module flags:
 ; CHECK:       Raw and structured buffers
@@ -35,4 +39,7 @@ define float @loadSRV() #0 {
   ret float %val
 }
 
+!dx.valver = !{!0}
+!0 = !{i32 1, i32 7}
+
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}

llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-valver1.8.ll

@@ -35,7 +35,7 @@ define float @loadSRV() #0 {
   ret float %val
 }
 
-!dx.valver = !{!1}
-!1 = !{i32 1, i32 8}
+!dx.valver = !{!0}
+!0 = !{i32 1, i32 8}
 
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}

llvm/test/CodeGen/DirectX/ShaderFlags/typed-uav-load-additional-formats.ll

@@ -44,7 +44,9 @@ define void @noload(<4 x float> %val) #0 {
 }
 
 !llvm.module.flags = !{!0}
+!dx.valver = !{!1}
 !0 = !{i32 1, !"dx.resmayalias", i32 1}
+!1 = !{i32 1, i32 8}
 
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}
 

llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.7.ll

@@ -0,0 +1,35 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
+
+; This test ensures that a library shader with a UAV gets the module and
+; shader feature flag UAVsAtEveryStage when the DXIL validator version is < 1.8
+
+target triple = "dxil-pc-shadermodel6.5-library"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NEXT: Shader Flags Value: 0x00010000
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK:        UAVs at every shader stage
+
+; CHECK: Function test : 0x00000000
+define void @test() "hlsl.export" {
+  ; RWBuffer<float> Buf : register(u0, space0)
+  %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
+       @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0t(
+           i32 0, i32 0, i32 1, i32 0, i1 false)
+  ret void
+}
+
+!dx.valver = !{!1}
+!1 = !{i32 1, i32 7}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"dx.resmayalias", i32 1}
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            8
+; DXC-NEXT:     Flags:
+; DXC:       UAVsAtEveryStage:         true
+; DXC:       NextUnusedBit:   false
+; DXC: ...

llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.8.ll

@@ -0,0 +1,27 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+
+; This test ensures that a library shader with a UAV does not get the module and
+; shader feature flag UAVsAtEveryStage when the DXIL validator version is >= 1.8
+
+target triple = "dxil-pc-shadermodel6.5-library"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NEXT: Shader Flags Value: 0x00000000
+
+; CHECK-NOT: Note: shader requires additional functionality:
+; CHECK-NOT:        UAVs at every shader stage
+
+; CHECK: Function test : 0x00000000
+define void @test() "hlsl.export" {
+  ; RWBuffer<float> Buf : register(u0, space0)
+  %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
+       @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0t(
+           i32 0, i32 0, i32 1, i32 0, i1 false)
+  ret void
+}
+
+!dx.valver = !{!1}
+!1 = !{i32 1, i32 8}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"dx.resmayalias", i32 1}

llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs.ll

@@ -0,0 +1,36 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; TODO: Remove this comment and add 'RUN' to the line below once vertex shaders are supported by llc
+; llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
+
+; This test ensures that a Vertex shader with a UAV gets the module and
+; shader feature flag UAVsAtEveryStage
+
+target triple = "dxil-pc-shadermodel6.5-vertex"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NEXT: Shader Flags Value: 0x00010000
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK:        UAVs at every shader stage
+
+; CHECK: Function VSMain : 0x00000000
+define void @VSMain() {
+  ; RWBuffer<float> Buf : register(u0, space0)
+  %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
+       @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0t(
+           i32 0, i32 0, i32 1, i32 0, i1 false)
+  ret void
+}
+
+!dx.valver = !{!1}
+!1 = !{i32 1, i32 8}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"dx.resmayalias", i32 1}
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            8
+; DXC-NEXT:     Flags:
+; DXC:       UAVsAtEveryStage:         true
+; DXC:       NextUnusedBit:   false
+; DXC: ...