[DirectX] Add shader flags for heap resources

llvm/lib/Target/DirectX/DXILShaderFlags.cpp

@@ -34,7 +34,10 @@ using namespace llvm::dxil;
 
 static bool hasUAVsAtEveryStage(const DXILResourceMap &DRM,
                                 const ModuleMetadataInfo &MMDI) {
-  if (DRM.uavs().empty())
+  // Heap resources do not count towards hasUAVsAtEveryStage.
+  bool HasUAVWithBinding = any_of(
+      DRM.uavs(), [](const ResourceInfo &RI) { return RI.hasBinding(); });
+  if (!HasUAVWithBinding)
     return false;
 
   switch (MMDI.ShaderProfile) {
@@ -274,6 +277,18 @@ void ModuleShaderFlags::updateFunctionFlags(ComputedShaderFlags &CSF,
       }
       break;
     }
+    case Intrinsic::dx_resource_handlefromheap: {
+      dxil::ResourceTypeInfo &RTI = DRTM[cast<TargetExtType>(II->getType())];
+      bool IsSamplerHeap = RTI.isSampler();
+      CSF.SamplerDescriptorHeapIndexing |= IsSamplerHeap;
+      CSF.ResourceDescriptorHeapIndexing |= !IsSamplerHeap;
+
+      if (!CSF.ResMayNotAlias && CanSetResMayNotAlias && RTI.isUAV() &&
+          MMDI.ValidatorVersion >= VersionTuple(1, 8)) {
+        CSF.ResMayNotAlias = true;
+      }
+      break;
+    }
     case Intrinsic::dx_resource_load_typedbuffer: {
       dxil::ResourceTypeInfo &RTI =
           DRTM[cast<TargetExtType>(II->getArgOperand(0)->getType())];
@@ -338,7 +353,10 @@ ModuleShaderFlags::gatherGlobalModuleFlags(const Module &M,
 
   // Set the Max64UAVs flag if the number of UAVs is > 8
   uint32_t NumUAVs = 0;
-  for (auto &UAV : DRM.uavs())
+  for (auto &UAV : DRM.uavs()) {
+    // Heap resources do not count towards Max64UAVs flag.
+    if (!UAV.hasBinding())
+      continue;
     if (MMDI.ValidatorVersion < VersionTuple(1, 6)) {
       NumUAVs++;
     } else { // MMDI.ValidatorVersion >= VersionTuple(1, 6)
@@ -348,6 +366,7 @@ ModuleShaderFlags::gatherGlobalModuleFlags(const Module &M,
         NewNum = ~0U;
       NumUAVs = NewNum;
     }
+  }
   if (NumUAVs > 8)
     CSF.Max64UAVs = true;
 

llvm/test/CodeGen/DirectX/ShaderFlags/heap-resources-max64uavs.ll

@@ -0,0 +1,42 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; RUN: llc %s -disable-dxil-remove-unused-resources --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
+
+; This test verifies that the Max64UAVs flag is *not* set if there
+; are 9 or more heap UAVs, since a heap UAVs do not count.
+
+target triple = "dxil-pc-shadermodel6.6-library"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NEXT: Shader Flags Value: 0x40000000
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK-NOT:  64 UAV slots
+; CHECK:      Resource descriptor heap indexing
+
+; CHECK: Function test : 0x40000000
+define void @test() "hlsl.export" {
+
+  %1 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 0)
+  %2 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 1)
+  %3 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 2)
+  %4 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 3)
+  %5 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 4)
+  %6 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 5)
+  %7 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 6)
+  %8 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 7)
+  %9 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefromheap(i32 8)
+
+  ret void
+}
+
+!dx.valver = !{!0}
+!0 = !{i32 1, i32 8}
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            8
+; DXC-NEXT:     Flags:
+; DXC:             Max64UAVs: false
+; DXC:             ResourceDescriptorHeapIndexing: true
+; DXC:             SamplerDescriptorHeapIndexing: false
+; DXC:       NextUnusedBit:   false
+

llvm/test/CodeGen/DirectX/ShaderFlags/heap-resources.ll

@@ -0,0 +1,52 @@
+; RUN: opt -S --passes="print-dx-shader-flags" -mtriple=dxil-pc-shadermodel6.6-library 2>&1 %s | FileCheck %s --check-prefixes=CHECK,CHECK-66
+; RUN: opt -S --passes="print-dx-shader-flags" -mtriple=dxil-pc-shadermodel6.7-library 2>&1 %s | FileCheck %s --check-prefixes=CHECK,CHECK-67
+; RUN: llc %s -disable-dxil-remove-unused-resources -mtriple=dxil-pc-shadermodel6.7-library --filetype=obj -o - | \
+; RUN:   obj2yaml | FileCheck %s --check-prefix=DXC
+
+; This test makes sure that the shader flag 'Resource descriptor heap indexing'
+; is set when the shader uses llvm.dx.resource.handlefromheap intrinsic to get
+; a resource from a resource descriptor heap, and that the shader flag 
+; `Sampler descriptor heap indexing` is set when the shader uses the same
+; intrinsic to get a sampler from a sampler descriptor heap.
+
+; It also checks that the flag "Any UAV may not alias any other UAV" is set for
+; a function that uses a heap UAV resource, but only for shader model 6.7 and higher. 
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-66:   Shader Flags Value: 0xc0000000
+; CHECK-67:   Shader Flags Value: 0x2c0000000
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK:        Resource descriptor heap indexing
+; CHECK:        Sampler descriptor heap indexing
+
+; CHECK: Note: extra DXIL module flags:
+; CHECK-67:     Any UAV may not alias any other UAV
+;
+; CHECK-66: Function test_1 : 0x40000000
+; CHECK-67: Function test_1 : 0x240000000
+define void @test_1() "hlsl.export" {
+  ; RWBuffer<float4> Buf = ResourceDescriptorHeap[3]
+  %typed = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
+              @llvm.dx.resource.handlefromheap.tdx.TypedBuffer_v4f32_1_0_0(i32 3)
+  ret void
+}
+
+; CHECK: Function test_2 : 0x80000000
+define void @test_2() "hlsl.export" {
+  ; SamplerState Samp = SamplerDescriptorHeap[100];
+  %samp = call target("dx.Sampler", 0)
+              @llvm.dx.resource.handlefromheap.tdx.Sampler_0(i32 100)
+  ret void
+}
+
+!dx.valver = !{!0}
+!0 = !{i32 1, i32 8}
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            8
+; DXC-NEXT:     Flags:
+; DXC:             ResourceDescriptorHeapIndexing: true
+; DXC:             SamplerDescriptorHeapIndexing: true
+; DXC:       NextUnusedBit:   false
+

llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs-heap.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 resource from heap 
+; does *not* set the module and shader feature flag UAVsAtEveryStage.
+
+target triple = "dxil-pc-shadermodel6.5-vertex"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NOT: Shader Flags Value: 0x400000000
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK-NOT:   UAVs at every shader stage
+; CHECK:       Resource descriptor heap indexing
+
+; CHECK: Function VSMain : 0x40000000
+define void @VSMain() {
+  ; RWBuffer<float> Buf : register(u0, space0)
+  %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
+       @llvm.dx.resource.handlefromheap.tdx.TypedBuffer_f32_1_0t(i32 7)
+  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:         false
+; DXC:       NextUnusedBit:   false
+; DXC: ...