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

[S. Bharadwaj Yadavalli via llvm-commits]

================
@@ -13,36 +13,88 @@
 
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/DiagnosticPrinter.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) {
+namespace {
+/// A simple Wrapper DiagnosticInfo that generates Module-level diagnostic
+/// for ShaderFlagsAnalysis pass
+class DiagnosticInfoShaderFlags : public DiagnosticInfo {
+private:
+  const Twine &Msg;
+  const Module &Mod;
+
+public:
+  /// \p M is the module for which the diagnostic is being emitted. \p Msg is
+  /// the message to show. Note that this class does not copy this message, so
+  /// this reference must be valid for the whole life time of the diagnostic.
+  DiagnosticInfoShaderFlags(const Module &M, const Twine &Msg,
+                            DiagnosticSeverity Severity = DS_Error)
+      : DiagnosticInfo(DK_Unsupported, Severity), Msg(Msg), Mod(M) {}
+
+  void print(DiagnosticPrinter &DP) const override {
+    DP << Mod.getName() << ": " << Msg << '\n';
+  }
+};
+} // namespace
+
+static void updateFlags(ComputedShaderFlags &CSF, const Instruction &I) {
   Type *Ty = I.getType();
-  if (Ty->isDoubleTy()) {
-    Flags.Doubles = true;
+  bool DoubleTyInUse = Ty->isDoubleTy();
+  for (Value *Op : I.operands()) {
+    DoubleTyInUse |= Op->getType()->isDoubleTy();
+  }
+
+  if (DoubleTyInUse) {
+    CSF.Doubles = true;
     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
+      CSF.DX11_1_DoubleExtensions = true;
       break;
     }
   }
 }
 
-ComputedShaderFlags ComputedShaderFlags::computeFlags(Module &M) {
-  ComputedShaderFlags Flags;
-  for (const auto &F : M)
+static bool compareFuncSFPairs(const FuncShaderFlagsMask &First,
+                               const FuncShaderFlagsMask &Second) {
+  return (First.first->getName().compare(Second.first->getName()) < 0);
+}
+
+static DXILModuleShaderFlagsInfo computeFlags(Module &M) {
+  DXILModuleShaderFlagsInfo MSFI;
+  for (auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+    // Each of the functions in a module are unique. Hence no prior shader flags
+    // mask of the function should be present.
+    if (MSFI.hasShaderFlagsMask(&F)) {
+      M.getContext().diagnose(
+          DiagnosticInfoShaderFlags(M, "Shader Flags mask for Function '" +
+                                           F.getName() + "' already exists"));
+    }
+    ComputedShaderFlags CSF{};
     for (const auto &BB : F)
       for (const auto &I : BB)
-        updateFlags(Flags, I);
-  return Flags;
+        updateFlags(CSF, I);
+    // Insert shader flag mask for function F
+    MSFI.FuncShaderFlagsVec.push_back({&F, CSF});
+  }
+  // Sort MSFI.FuncShaderFlagsVec for later lookup that uses binary search
+  llvm::sort(MSFI.FuncShaderFlagsVec, compareFuncSFPairs);
----------------
bharadwajy wrote:

> If `MapVector` isn't an option, you can use `lower_bound` to figure out where to insert any new elements and still keep the vector sorted. That way you don't need to a do a full sort each time a new element is added.
>

`llvm::sort` is run once per module on the vector of pairs after shader flags masks for all functions are inserted.

Collection of shader flags in the analysis pass and subsequent queries in later passes follows the "distinct pattern" outlined in [LLVM Programmer's Manual](https://llvm.org/docs/ProgrammersManual.html#a-sorted-vector). Hence the choice of `SmallVector` as opposed to `MapVector` or `DenseVector`.

https://github.com/llvm/llvm-project/pull/112967