[DirectX] Set whole-module flags prior to evaluating per-function flags (#139967)

Fixes #139024 and #139954

- Refactor DXILShaderFlags to compute the flags that apply to a whole
module before computing flags that apply individually to each function
- Make DXILResourceMap const, since it is not modified in
DXILShaderFlags
- Per-function shader flag analysis now initially starts with the set of
flags that apply to the whole module instead of starting from no flags.
This change fixes the above linked issues
- Fix shader flag tests affected by the above change
diff --git a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
index 8bdaf68..bd3349d 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
@@ -32,7 +32,7 @@
 using namespace llvm;
 using namespace llvm::dxil;
 
-static bool hasUAVsAtEveryStage(DXILResourceMap &DRM,
+static bool hasUAVsAtEveryStage(const DXILResourceMap &DRM,
                                 const ModuleMetadataInfo &MMDI) {
   if (DRM.uavs().empty())
     return false;
@@ -143,7 +143,7 @@
   }
 
   if (CSF.LowPrecisionPresent) {
-    if (CanSetNativeLowPrecisionMode)
+    if (CSF.NativeLowPrecisionMode)
       CSF.NativeLowPrecision = true;
     else
       CSF.MinimumPrecision = true;
@@ -207,9 +207,62 @@
   }
 }
 
+/// Set shader flags that apply to all functions within the module
+ComputedShaderFlags
+ModuleShaderFlags::gatherGlobalModuleFlags(const Module &M,
+                                           const DXILResourceMap &DRM,
+                                           const ModuleMetadataInfo &MMDI) {
+
+  ComputedShaderFlags CSF;
+
+  // Set DisableOptimizations flag based on the presence of OptimizeNone
+  // attribute of entry functions.
+  if (MMDI.EntryPropertyVec.size() > 0) {
+    CSF.DisableOptimizations = MMDI.EntryPropertyVec[0].Entry->hasFnAttribute(
+        llvm::Attribute::OptimizeNone);
+    // Ensure all entry functions have the same optimization attribute
+    for (const auto &EntryFunProps : MMDI.EntryPropertyVec)
+      if (CSF.DisableOptimizations !=
+          EntryFunProps.Entry->hasFnAttribute(llvm::Attribute::OptimizeNone))
+        EntryFunProps.Entry->getContext().diagnose(DiagnosticInfoUnsupported(
+            *(EntryFunProps.Entry), "Inconsistent optnone attribute "));
+  }
+
+  CSF.UAVsAtEveryStage = hasUAVsAtEveryStage(DRM, MMDI);
+
+  // Set the Max64UAVs flag if the number of UAVs is > 8
+  uint32_t NumUAVs = 0;
+  for (auto &UAV : DRM.uavs())
+    if (MMDI.ValidatorVersion < VersionTuple(1, 6))
+      NumUAVs++;
+    else // MMDI.ValidatorVersion >= VersionTuple(1, 6)
+      NumUAVs += UAV.getBinding().Size;
+  if (NumUAVs > 8)
+    CSF.Max64UAVs = true;
+
+  // Set the module flag that enables native low-precision execution mode.
+  // NativeLowPrecisionMode can only be set when the command line option
+  // -enable-16bit-types is provided. This is indicated by the dx.nativelowprec
+  // module flag being set
+  // This flag is needed even if the module does not use 16-bit types because a
+  // corresponding debug module may include 16-bit types, and tools that use the
+  // debug module may expect it to have the same flags as the original
+  if (auto *NativeLowPrec = mdconst::extract_or_null<ConstantInt>(
+          M.getModuleFlag("dx.nativelowprec")))
+    if (MMDI.ShaderModelVersion >= VersionTuple(6, 2))
+      CSF.NativeLowPrecisionMode = NativeLowPrec->getValue().getBoolValue();
+
+  // Set ResMayNotAlias to true if DXIL validator version < 1.8 and there
+  // are UAVs present globally.
+  if (CanSetResMayNotAlias && MMDI.ValidatorVersion < VersionTuple(1, 8))
+    CSF.ResMayNotAlias = !DRM.uavs().empty();
+
+  return CSF;
+}
+
 /// Construct ModuleShaderFlags for module Module M
 void ModuleShaderFlags::initialize(Module &M, DXILResourceTypeMap &DRTM,
-                                   DXILResourceMap &DRM,
+                                   const DXILResourceMap &DRM,
                                    const ModuleMetadataInfo &MMDI) {
 
   CanSetResMayNotAlias = MMDI.DXILVersion >= VersionTuple(1, 7);
@@ -217,16 +270,10 @@
   // flag to 1, thereby disabling the ability to set the ResMayNotAlias flag
   if (auto *ResMayAlias = mdconst::extract_or_null<ConstantInt>(
           M.getModuleFlag("dx.resmayalias")))
-    CanSetResMayNotAlias = !ResMayAlias->getValue().getBoolValue();
+    if (ResMayAlias->getValue().getBoolValue())
+      CanSetResMayNotAlias = false;
 
-  // NativeLowPrecisionMode can only be set when the command line option
-  // -enable-16bit-types is provided. This is indicated by the dx.nativelowprec
-  // module flag being set
-  CanSetNativeLowPrecisionMode = false;
-  if (auto *NativeLowPrec = mdconst::extract_or_null<ConstantInt>(
-          M.getModuleFlag("dx.nativelowprec")))
-    if (MMDI.ShaderModelVersion >= VersionTuple(6, 2))
-      CanSetNativeLowPrecisionMode = NativeLowPrec->getValue().getBoolValue();
+  ComputedShaderFlags GlobalSFMask = gatherGlobalModuleFlags(M, DRM, MMDI);
 
   CallGraph CG(M);
 
@@ -252,7 +299,7 @@
         continue;
       }
 
-      ComputedShaderFlags CSF;
+      ComputedShaderFlags CSF = GlobalSFMask;
       for (const auto &BB : *F)
         for (const auto &I : BB)
           updateFunctionFlags(CSF, I, DRTM, MMDI);
@@ -273,43 +320,6 @@
       // Merge SCCSF with that of F
       FunctionFlags[F].merge(SCCSF);
   }
-
-  // Set DisableOptimizations flag based on the presence of OptimizeNone
-  // attribute of entry functions.
-  if (MMDI.EntryPropertyVec.size() > 0) {
-    CombinedSFMask.DisableOptimizations =
-        MMDI.EntryPropertyVec[0].Entry->hasFnAttribute(
-            llvm::Attribute::OptimizeNone);
-    // Ensure all entry functions have the same optimization attribute
-    for (const auto &EntryFunProps : MMDI.EntryPropertyVec)
-      if (CombinedSFMask.DisableOptimizations !=
-          EntryFunProps.Entry->hasFnAttribute(llvm::Attribute::OptimizeNone))
-        EntryFunProps.Entry->getContext().diagnose(DiagnosticInfoUnsupported(
-            *(EntryFunProps.Entry), "Inconsistent optnone attribute "));
-  }
-
-  // Set ResMayNotAlias to true if DXIL validator version < 1.8 and there
-  // are UAVs present globally.
-  if (CanSetResMayNotAlias && MMDI.ValidatorVersion < VersionTuple(1, 8))
-    CombinedSFMask.ResMayNotAlias = !DRM.uavs().empty();
-
-  // Set the module flag that enables native low-precision execution mode. This
-  // is needed even if the module does not use 16-bit types because a
-  // corresponding debug module may include 16-bit types, and tools that use the
-  // debug module may expect it to have the same flags as the original
-  CombinedSFMask.NativeLowPrecisionMode = CanSetNativeLowPrecisionMode;
-
-  // Set the Max64UAVs flag if the number of UAVs is > 8
-  uint32_t NumUAVs = 0;
-  for (auto &UAV : DRM.uavs())
-    if (MMDI.ValidatorVersion < VersionTuple(1, 6))
-      NumUAVs++;
-    else // MMDI.ValidatorVersion >= VersionTuple(1, 6)
-      NumUAVs += UAV.getBinding().Size;
-  if (NumUAVs > 8)
-    CombinedSFMask.Max64UAVs = true;
-
-  CombinedSFMask.UAVsAtEveryStage = hasUAVsAtEveryStage(DRM, MMDI);
 }
 
 void ComputedShaderFlags::print(raw_ostream &OS) const {
diff --git a/llvm/lib/Target/DirectX/DXILShaderFlags.h b/llvm/lib/Target/DirectX/DXILShaderFlags.h
index c4eef4e7..f94f799 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.h
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.h
@@ -85,15 +85,15 @@
 };
 
 struct ModuleShaderFlags {
-  void initialize(Module &, DXILResourceTypeMap &DRTM, DXILResourceMap &DRM,
-                  const ModuleMetadataInfo &MMDI);
+  void initialize(Module &, DXILResourceTypeMap &DRTM,
+                  const DXILResourceMap &DRM, const ModuleMetadataInfo &MMDI);
   const ComputedShaderFlags &getFunctionFlags(const Function *) const;
   const ComputedShaderFlags &getCombinedFlags() const { return CombinedSFMask; }
 
 private:
-  // Booleans set by module flags
-  bool CanSetResMayNotAlias;         // dx.resmayalias
-  bool CanSetNativeLowPrecisionMode; // dx.nativelowprec
+  // This boolean is inversely set by the LLVM module flag dx.resmayalias to
+  // determine whether or not the ResMayNotAlias DXIL module flag can be set
+  bool CanSetResMayNotAlias;
 
   /// Map of Function-Shader Flag Mask pairs representing properties of each of
   /// the functions in the module. Shader Flags of each function represent both
@@ -101,6 +101,9 @@
   DenseMap<const Function *, ComputedShaderFlags> FunctionFlags;
   /// Combined Shader Flag Mask of all functions of the module
   ComputedShaderFlags CombinedSFMask{};
+  ComputedShaderFlags gatherGlobalModuleFlags(const Module &M,
+                                              const DXILResourceMap &,
+                                              const ModuleMetadataInfo &);
   void updateFunctionFlags(ComputedShaderFlags &, const Instruction &,
                            DXILResourceTypeMap &, const ModuleMetadataInfo &);
 };
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-cs.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-cs.ll
index 8bf242f..3f2ec9a 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-cs.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-cs.ll
@@ -8,7 +8,7 @@
 ; CHECK-NEXT: ;       Disable shader optimizations
 
 ; CHECK: ; Shader Flags for Module Functions
-; CHECK: ; Function main : 0x00000000
+; CHECK: ; Function main : 0x00000001
 ; The test source in this file generated from the following command:
 ; clang -cc1 -triple dxil-pc-shadermodel6.0-compute -x hlsl -emit-llvm -O0 -o - <<EOF
 ; [numthreads(1,1,1)]
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-lib.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-lib.ll
index 3bf582c..c6e3cc9 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-lib.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/disable-opt-lib.ll
@@ -8,7 +8,7 @@
 ; CHECK-NEXT: ;       Disable shader optimizations
 
 ; CHECK: ; Shader Flags for Module Functions
-; CHECK: ; Function main : 0x00000000
+; CHECK: ; Function main : 0x00000001
 ; The test source in this file generated from the following command:
 ; clang -cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -emit-llvm -O0 -o - <<EOF
 
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.5.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.5.ll
index af6001b..6079071 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.5.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.5.ll
@@ -13,7 +13,7 @@
 ; CHECK:    UAVs at every shader stage
 ; CHECK-NOT:    64 UAV slots
 
-; CHECK: Function test : 0x00000000
+; CHECK: Function test : 0x00010000
 define void @test() "hlsl.export" {
   ; RWBuffer<float> Buf : register(u0, space0)
   %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.6.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.6.ll
index 7e1d73b..4f1a1e7 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.6.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs-array-valver1.6.ll
@@ -14,7 +14,7 @@
 ; CHECK:        UAVs at every shader stage
 ; CHECK:        64 UAV slots
 
-; CHECK: Function test : 0x00000000
+; CHECK: Function test : 0x00018000
 define void @test() "hlsl.export" {
   ; RWBuffer<float> Buf : register(u0, space0)
   %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs.ll
index a97fe5d..e0d4ac7 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/max-64-uavs.ll
@@ -13,7 +13,7 @@
 ; CHECK:       64 UAV slots
 
 ; Note: 64 UAV slots does not get set per-function
-; CHECK: Function test : 0x00000000
+; CHECK: Function test : 0x00008000
 define void @test() "hlsl.export" {
   ; RWBuffer<float> Buf : register(u0, space0)
   %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.7.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.7.ll
index fc6560e..c0ea6b4 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.7.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/res-may-not-alias-sm6.7.ll
@@ -19,7 +19,7 @@
 ; CHECK:       Any UAV may not alias any other UAV
 ;
 
-; CHECK: Function loadUAV : 0x00000000
+; CHECK: Function loadUAV : 0x200010000
 define float @loadUAV() #0 {
   %res = call target("dx.TypedBuffer", float, 1, 0, 0)
       @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
@@ -29,7 +29,7 @@
   ret float %val
 }
 
-; CHECK: Function loadSRV : 0x00000010
+; CHECK: Function loadSRV : 0x200010010
 define float @loadSRV() #0 {
   %res = tail call target("dx.RawBuffer", float, 0, 0)
       @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.7.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.7.ll
index 552f513..bd1de58 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.7.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-lib-valver1.7.ll
@@ -12,7 +12,7 @@
 ; CHECK: Note: shader requires additional functionality:
 ; CHECK:        UAVs at every shader stage
 
-; CHECK: Function test : 0x00000000
+; CHECK: Function test : 0x00010000
 define void @test() "hlsl.export" {
   ; RWBuffer<float> Buf : register(u0, space0)
   %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs.ll
index d3f556b..bdb07b4 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/uavs-at-every-stage-vs.ll
@@ -13,7 +13,7 @@
 ; CHECK: Note: shader requires additional functionality:
 ; CHECK:        UAVs at every shader stage
 
-; CHECK: Function VSMain : 0x00000000
+; CHECK: Function VSMain : 0x00010000
 define void @VSMain() {
   ; RWBuffer<float> Buf : register(u0, space0)
   %buf0 = call target("dx.TypedBuffer", float, 1, 0, 1)
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/use-native-low-precision-1.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/use-native-low-precision-1.ll
index cb3b486..ba844b1 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/use-native-low-precision-1.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/use-native-low-precision-1.ll
@@ -18,19 +18,19 @@
 ;CHECK-NEXT: ;
 ;CHECK-NEXT: ; Shader Flags for Module Functions
 
-;CHECK-LABEL: ; Function add_i16 : 0x00000020
+;CHECK-LABEL: ; Function add_i16 : 0x00800020
 define i16 @add_i16(i16 %a, i16 %b) "hlsl.export" {
   %sum = add i16 %a, %b
   ret i16 %sum
 }
 
-;CHECK-LABEL: ; Function add_i32 : 0x00000000
+;CHECK-LABEL: ; Function add_i32 : 0x00800000
 define i32 @add_i32(i32 %a, i32 %b) "hlsl.export" {
   %sum = add i32 %a, %b
   ret i32 %sum
 }
 
-;CHECK-LABEL: ; Function add_half : 0x00000020
+;CHECK-LABEL: ; Function add_half : 0x00800020
 define half @add_half(half %a, half %b) "hlsl.export" {
   %sum = fadd half %a, %b
   ret half %sum