[DirectX] Infrastructure to collect shader flags for each function (llvm#112967)

Currently, ShaderFlagsAnalysis pass represents various module-level
properties as well as function-level properties of a DXIL Module using a
single mask. However, one mask per function is needed for accurate
computation of shader flags mask, such as for entry function metadata
creation.

This change introduces a structure that wraps a sorted vector of
function-shader flag mask pairs that represent function properties
instead of a single shader flag mask that represents module properties
and properties of all functions. The result type of ShaderFlagsAnalysis
pass is changed to newly-defined structure type instead of a single
shader flags mask.

This allows accurate computation of shader flags of an entry function
(and all functions in a library shader) for use during its metadata
generation (DXILTranslateMetadata pass) and its feature flags in DX
container globals construction (DXContainerGlobals pass) based on the
shader flags mask of functions. However, note that the change to
implement propagation of such callee-based shader flags mask computation
is planned in a follow-on PR. Consequently, this PR changes shader flag
mask computation in DXILTranslateMetadata and DXContainerGlobals passes
to simply be a union of module flags and shader flags of all functions,
thereby retaining the existing effect of using a single shader flag
mask.

Author: bharadwajy

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -78,13 +78,13 @@ bool DXContainerGlobals::runOnModule(Module &M) {
 }
 
 GlobalVariable *DXContainerGlobals::getFeatureFlags(Module &M) {
-  const uint64_t FeatureFlags =
-      static_cast<uint64_t>(getAnalysis<ShaderFlagsAnalysisWrapper>()
-                                .getShaderFlags()
-                                .getFeatureFlags());
+  uint64_t CombinedFeatureFlags = getAnalysis<ShaderFlagsAnalysisWrapper>()
+                                      .getShaderFlags()
+                                      .getCombinedFlags()
+                                      .getFeatureFlags();
 
   Constant *FeatureFlagsConstant =
-      ConstantInt::get(M.getContext(), APInt(64, FeatureFlags));
+      ConstantInt::get(M.getContext(), APInt(64, CombinedFeatureFlags));
   return buildContainerGlobal(M, FeatureFlagsConstant, "dx.sfi0", "SFI0");
 }
 

llvm/lib/Target/DirectX/DXILShaderFlags.cpp

@@ -13,36 +13,54 @@
 
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 using namespace llvm::dxil;
 
-static void updateFlags(ComputedShaderFlags &Flags, const Instruction &I) {
-  Type *Ty = I.getType();
-  if (Ty->isDoubleTy()) {
-    Flags.Doubles = true;
+static void updateFunctionFlags(ComputedShaderFlags &CSF,
+                                const Instruction &I) {
+  if (!CSF.Doubles)
+    CSF.Doubles = I.getType()->isDoubleTy();
+
+  if (!CSF.Doubles) {
+    for (Value *Op : I.operands())
+      CSF.Doubles |= Op->getType()->isDoubleTy();
+  }
+  if (CSF.Doubles) {
     switch (I.getOpcode()) {
     case Instruction::FDiv:
     case Instruction::UIToFP:
     case Instruction::SIToFP:
     case Instruction::FPToUI:
     case Instruction::FPToSI:
-      Flags.DX11_1_DoubleExtensions = true;
+      // TODO: To be set if I is a call to DXIL intrinsic DXIL::Opcode::Fma
+      // https://github.com/llvm/llvm-project/issues/114554
+      CSF.DX11_1_DoubleExtensions = true;
       break;
     }
   }
 }
 
-ComputedShaderFlags ComputedShaderFlags::computeFlags(Module &M) {
-  ComputedShaderFlags Flags;
-  for (const auto &F : M)
+void ModuleShaderFlags::initialize(const Module &M) {
+  // Collect shader flags for each of the functions
+  for (const auto &F : M.getFunctionList()) {
+    if (F.isDeclaration())
+      continue;
+    ComputedShaderFlags CSF;
     for (const auto &BB : F)
       for (const auto &I : BB)
-        updateFlags(Flags, I);
-  return Flags;
+        updateFunctionFlags(CSF, I);
+    // Insert shader flag mask for function F
+    FunctionFlags.push_back({&F, CSF});
+    // Update combined shader flags mask
+    CombinedSFMask.merge(CSF);
+  }
+  llvm::sort(FunctionFlags);
 }
 
 void ComputedShaderFlags::print(raw_ostream &OS) const {
@@ -63,20 +81,58 @@ void ComputedShaderFlags::print(raw_ostream &OS) const {
   OS << ";\n";
 }
 
+/// Return the shader flags mask of the specified function Func.
+const ComputedShaderFlags &
+ModuleShaderFlags::getFunctionFlags(const Function *Func) const {
+  const auto Iter = llvm::lower_bound(
+      FunctionFlags, Func,
+      [](const std::pair<const Function *, ComputedShaderFlags> FSM,
+         const Function *FindFunc) { return (FSM.first < FindFunc); });
+  assert((Iter != FunctionFlags.end() && Iter->first == Func) &&
+         "No Shader Flags Mask exists for function");
+  return Iter->second;
+}
+
+//===----------------------------------------------------------------------===//
+// ShaderFlagsAnalysis and ShaderFlagsAnalysisPrinterPass
+
+// Provide an explicit template instantiation for the static ID.
 AnalysisKey ShaderFlagsAnalysis::Key;
 
-ComputedShaderFlags ShaderFlagsAnalysis::run(Module &M,
-                                             ModuleAnalysisManager &AM) {
-  return ComputedShaderFlags::computeFlags(M);
+ModuleShaderFlags ShaderFlagsAnalysis::run(Module &M,
+                                           ModuleAnalysisManager &AM) {
+  ModuleShaderFlags MSFI;
+  MSFI.initialize(M);
+  return MSFI;
 }
 
 PreservedAnalyses ShaderFlagsAnalysisPrinter::run(Module &M,
                                                   ModuleAnalysisManager &AM) {
-  ComputedShaderFlags Flags = AM.getResult<ShaderFlagsAnalysis>(M);
-  Flags.print(OS);
+  const ModuleShaderFlags &FlagsInfo = AM.getResult<ShaderFlagsAnalysis>(M);
+  // Print description of combined shader flags for all module functions
+  OS << "; Combined Shader Flags for Module\n";
+  FlagsInfo.getCombinedFlags().print(OS);
+  // Print shader flags mask for each of the module functions
+  OS << "; Shader Flags for Module Functions\n";
+  for (const auto &F : M.getFunctionList()) {
+    if (F.isDeclaration())
+      continue;
+    auto SFMask = FlagsInfo.getFunctionFlags(&F);
+    OS << formatv("; Function {0} : {1:x8}\n;\n", F.getName(),
+                  (uint64_t)(SFMask));
+  }
+
   return PreservedAnalyses::all();
 }
 
+//===----------------------------------------------------------------------===//
+// ShaderFlagsAnalysis and ShaderFlagsAnalysisPrinterPass
+
+bool ShaderFlagsAnalysisWrapper::runOnModule(Module &M) {
+  MSFI.initialize(M);
+  return false;
+}
+
 char ShaderFlagsAnalysisWrapper::ID = 0;
 
 INITIALIZE_PASS(ShaderFlagsAnalysisWrapper, "dx-shader-flag-analysis",

llvm/lib/Target/DirectX/DXILShaderFlags.h

@@ -14,12 +14,14 @@
 #ifndef LLVM_TARGET_DIRECTX_DXILSHADERFLAGS_H
 #define LLVM_TARGET_DIRECTX_DXILSHADERFLAGS_H
 
+#include "llvm/IR/Function.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
+#include <memory>
 
 namespace llvm {
 class Module;
@@ -43,15 +45,23 @@ struct ComputedShaderFlags {
   constexpr uint64_t getMask(int Bit) const {
     return Bit != -1 ? 1ull << Bit : 0;
   }
+
+  uint64_t getModuleFlags() const {
+    uint64_t ModuleFlags = 0;
+#define DXIL_MODULE_FLAG(DxilModuleBit, FlagName, Str)                         \
+  ModuleFlags |= FlagName ? getMask(DxilModuleBit) : 0ull;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return ModuleFlags;
+  }
+
   operator uint64_t() const {
-    uint64_t FlagValue = 0;
+    uint64_t FlagValue = getModuleFlags();
 #define SHADER_FEATURE_FLAG(FeatureBit, DxilModuleBit, FlagName, Str)          \
   FlagValue |= FlagName ? getMask(DxilModuleBit) : 0ull;
-#define DXIL_MODULE_FLAG(DxilModuleBit, FlagName, Str)                         \
-  FlagValue |= FlagName ? getMask(DxilModuleBit) : 0ull;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
     return FlagValue;
   }
+
   uint64_t getFeatureFlags() const {
     uint64_t FeatureFlags = 0;
 #define SHADER_FEATURE_FLAG(FeatureBit, DxilModuleBit, FlagName, Str)          \
@@ -60,21 +70,43 @@ struct ComputedShaderFlags {
     return FeatureFlags;
   }
 
-  static ComputedShaderFlags computeFlags(Module &M);
+  void merge(const uint64_t IVal) {
+#define SHADER_FEATURE_FLAG(FeatureBit, DxilModuleBit, FlagName, Str)          \
+  FlagName |= (IVal & getMask(DxilModuleBit));
+#define DXIL_MODULE_FLAG(DxilModuleBit, FlagName, Str)                         \
+  FlagName |= (IVal & getMask(DxilModuleBit));
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return;
+  }
+
   void print(raw_ostream &OS = dbgs()) const;
   LLVM_DUMP_METHOD void dump() const { print(); }
 };
 
+struct ModuleShaderFlags {
+  void initialize(const Module &);
+  const ComputedShaderFlags &getFunctionFlags(const Function *) const;
+  const ComputedShaderFlags &getCombinedFlags() const { return CombinedSFMask; }
+
+private:
+  /// Vector of sorted Function-Shader Flag mask pairs representing properties
+  /// of each of the functions in the module. Shader Flags of each function
+  /// represent both module-level and function-level flags
+  SmallVector<std::pair<Function const *, ComputedShaderFlags>> FunctionFlags;
+  /// Combined Shader Flag Mask of all functions of the module
+  ComputedShaderFlags CombinedSFMask{};
+};
+
 class ShaderFlagsAnalysis : public AnalysisInfoMixin<ShaderFlagsAnalysis> {
   friend AnalysisInfoMixin<ShaderFlagsAnalysis>;
   static AnalysisKey Key;
 
 public:
   ShaderFlagsAnalysis() = default;
 
-  using Result = ComputedShaderFlags;
+  using Result = ModuleShaderFlags;
 
-  ComputedShaderFlags run(Module &M, ModuleAnalysisManager &AM);
+  ModuleShaderFlags run(Module &M, ModuleAnalysisManager &AM);
 };
 
 /// Printer pass for ShaderFlagsAnalysis results.
@@ -92,19 +124,16 @@ class ShaderFlagsAnalysisPrinter
 /// This is required because the passes that will depend on this are codegen
 /// passes which run through the legacy pass manager.
 class ShaderFlagsAnalysisWrapper : public ModulePass {
-  ComputedShaderFlags Flags;
+  ModuleShaderFlags MSFI;
 
 public:
   static char ID;
 
   ShaderFlagsAnalysisWrapper() : ModulePass(ID) {}
 
-  const ComputedShaderFlags &getShaderFlags() { return Flags; }
+  const ModuleShaderFlags &getShaderFlags() { return MSFI; }
 
-  bool runOnModule(Module &M) override {
-    Flags = ComputedShaderFlags::computeFlags(M);
-    return false;
-  }
+  bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();

llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp

@@ -286,11 +286,6 @@ static MDTuple *emitTopLevelLibraryNode(Module &M, MDNode *RMD,
   MDTuple *Properties = nullptr;
   if (ShaderFlags != 0) {
     SmallVector<Metadata *> MDVals;
-    // FIXME: ShaderFlagsAnalysis pass needs to collect and provide
-    // ShaderFlags for each entry function. Currently, ShaderFlags value
-    // provided by ShaderFlagsAnalysis pass is created by walking *all* the
-    // function instructions of the module. Is it is correct to use this value
-    // for metadata of the empty library entry?
     MDVals.append(
         getTagValueAsMetadata(EntryPropsTag::ShaderFlags, ShaderFlags, Ctx));
     Properties = MDNode::get(Ctx, MDVals);
@@ -302,7 +297,7 @@ static MDTuple *emitTopLevelLibraryNode(Module &M, MDNode *RMD,
 
 static void translateMetadata(Module &M, const DXILResourceMap &DRM,
                               const Resources &MDResources,
-                              const ComputedShaderFlags &ShaderFlags,
+                              const ModuleShaderFlags &ShaderFlags,
                               const ModuleMetadataInfo &MMDI) {
   LLVMContext &Ctx = M.getContext();
   IRBuilder<> IRB(Ctx);
@@ -318,23 +313,27 @@ static void translateMetadata(Module &M, const DXILResourceMap &DRM,
   // See https://github.com/llvm/llvm-project/issues/57928
   MDTuple *Signatures = nullptr;
 
-  if (MMDI.ShaderProfile == Triple::EnvironmentType::Library)
+  if (MMDI.ShaderProfile == Triple::EnvironmentType::Library) {
+    // Get the combined shader flag mask of all functions in the library to be
+    // used as shader flags mask value associated with top-level library entry
+    // metadata.
+    uint64_t CombinedMask = ShaderFlags.getCombinedFlags();
     EntryFnMDNodes.emplace_back(
-        emitTopLevelLibraryNode(M, ResourceMD, ShaderFlags));
-  else if (MMDI.EntryPropertyVec.size() > 1) {
+        emitTopLevelLibraryNode(M, ResourceMD, CombinedMask));
+  } else if (MMDI.EntryPropertyVec.size() > 1) {
     M.getContext().diagnose(DiagnosticInfoTranslateMD(
         M, "Non-library shader: One and only one entry expected"));
   }
 
   for (const EntryProperties &EntryProp : MMDI.EntryPropertyVec) {
-    // FIXME: ShaderFlagsAnalysis pass needs to collect and provide
-    // ShaderFlags for each entry function. For now, assume shader flags value
-    // of entry functions being compiled for lib_* shader profile viz.,
-    // EntryPro.Entry is 0.
-    uint64_t EntryShaderFlags =
-        (MMDI.ShaderProfile == Triple::EnvironmentType::Library) ? 0
-                                                                 : ShaderFlags;
+    const ComputedShaderFlags &EntrySFMask =
+        ShaderFlags.getFunctionFlags(EntryProp.Entry);
+
+    // If ShaderProfile is Library, mask is already consolidated in the
+    // top-level library node. Hence it is not emitted.
+    uint64_t EntryShaderFlags = 0;
     if (MMDI.ShaderProfile != Triple::EnvironmentType::Library) {
+      EntryShaderFlags = EntrySFMask;
       if (EntryProp.ShaderStage != MMDI.ShaderProfile) {
         M.getContext().diagnose(DiagnosticInfoTranslateMD(
             M,
@@ -361,8 +360,7 @@ PreservedAnalyses DXILTranslateMetadata::run(Module &M,
                                              ModuleAnalysisManager &MAM) {
   const DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
   const dxil::Resources &MDResources = MAM.getResult<DXILResourceMDAnalysis>(M);
-  const ComputedShaderFlags &ShaderFlags =
-      MAM.getResult<ShaderFlagsAnalysis>(M);
+  const ModuleShaderFlags &ShaderFlags = MAM.getResult<ShaderFlagsAnalysis>(M);
   const dxil::ModuleMetadataInfo MMDI = MAM.getResult<DXILMetadataAnalysis>(M);
 
   translateMetadata(M, DRM, MDResources, ShaderFlags, MMDI);
@@ -393,7 +391,7 @@ class DXILTranslateMetadataLegacy : public ModulePass {
         getAnalysis<DXILResourceWrapperPass>().getResourceMap();
     const dxil::Resources &MDResources =
         getAnalysis<DXILResourceMDWrapper>().getDXILResource();
-    const ComputedShaderFlags &ShaderFlags =
+    const ModuleShaderFlags &ShaderFlags =
         getAnalysis<ShaderFlagsAnalysisWrapper>().getShaderFlags();
     dxil::ModuleMetadataInfo MMDI =
         getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();

llvm/test/CodeGen/DirectX/ShaderFlags/double-extensions-obj-test.ll

@@ -0,0 +1,16 @@
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s
+
+target triple = "dxil-pc-shadermodel6.7-library"
+define double @div(double %a, double %b) #0 {
+  %res = fdiv double %a, %b
+  ret double %res
+}
+
+attributes #0 = { convergent norecurse nounwind "hlsl.export"}
+
+; CHECK: - Name:            SFI0
+; CHECK-NEXT:     Size:            8
+; CHECK-NEXT:     Flags:
+; CHECK:       Doubles:         true
+; CHECK:       DX11_1_DoubleExtensions:         true
+