[llvm] [DXIL] Consume Metadata Analysis information in passes (PR #108034)

[S. Bharadwaj Yadavalli via llvm-commits]

================
@@ -65,18 +69,269 @@ static void emitResourceMetadata(Module &M, const DXILResourceMap &DRM,
       MDNode::get(M.getContext(), {SRVMD, UAVMD, CBufMD, SmpMD}));
 }
 
+static StringRef getShortShaderStage(Triple::EnvironmentType Env) {
+  switch (Env) {
+  case Triple::Pixel:
+    return "ps";
+  case Triple::Vertex:
+    return "vs";
+  case Triple::Geometry:
+    return "gs";
+  case Triple::Hull:
+    return "hs";
+  case Triple::Domain:
+    return "ds";
+  case Triple::Compute:
+    return "cs";
+  case Triple::Library:
+    return "lib";
+  case Triple::Mesh:
+    return "ms";
+  case Triple::Amplification:
+    return "as";
+  default:
+    break;
+  }
+  llvm_unreachable("Unsupported environment for DXIL generation.");
+  return "";
+}
+
+static uint32_t getShaderStage(Triple::EnvironmentType Env) {
+  return (uint32_t)Env - (uint32_t)llvm::Triple::Pixel;
+}
+
+struct ShaderEntryMDInfo : EntryProperties {
+
+  enum EntryPropsTag {
+    ShaderFlagsTag = 0,
+    GSStateTag,
+    DSStateTag,
+    HSStateTag,
+    NumThreadsTag,
+    AutoBindingSpaceTag,
+    RayPayloadSizeTag,
+    RayAttribSizeTag,
+    ShaderKindTag,
+    MSStateTag,
+    ASStateTag,
+    WaveSizeTag,
+    EntryRootSigTag,
+  };
+
+  ShaderEntryMDInfo(EntryProperties &EP, LLVMContext &C,
+                    Triple::EnvironmentType SP, MDTuple *MDR = nullptr,
+                    uint64_t ShaderFlags = 0)
+      : EntryProperties(EP), Ctx(C), EntryShaderFlags(ShaderFlags),
+        MDResources(MDR), ShaderProfile(SP) {};
+
+  MDTuple *getAsMetadata() {
+    MDTuple *Properties = constructEntryPropMetadata();
+    // FIXME: Add support to construct Signatures
+    // See https://github.com/llvm/llvm-project/issues/57928
+    MDTuple *Signatures = nullptr;
+    return constructEntryMetadata(Signatures, MDResources, Properties);
+  }
+
+private:
+  LLVMContext &Ctx;
+  // Shader Flags for the Entry - from ShadeFLagsAnalysis pass
+  uint64_t EntryShaderFlags{0};
+  MDTuple *MDResources{nullptr};
+  Triple::EnvironmentType ShaderProfile{
+      Triple::EnvironmentType::UnknownEnvironment};
+  // Each entry point metadata record specifies:
+  //  * reference to the entry point function global symbol
+  //  * unmangled name
+  //  * list of signatures
+  //  * list of resources
+  //  * list of tag-value pairs of shader capabilities and other properties
+
+  MDTuple *constructEntryMetadata(MDTuple *Signatures, MDTuple *Resources,
+                                  MDTuple *Properties) {
+    Metadata *MDVals[5];
+    MDVals[0] =
+        Entry ? ValueAsMetadata::get(const_cast<Function *>(Entry)) : nullptr;
+    MDVals[1] = MDString::get(Ctx, Entry ? Entry->getName() : "");
+    MDVals[2] = Signatures;
+    MDVals[3] = Resources;
+    MDVals[4] = Properties;
+    return MDNode::get(Ctx, MDVals);
+  }
+
+  SmallVector<Metadata *> getTagValueAsMetadata(EntryPropsTag Tag,
+                                                uint64_t Value) {
+    SmallVector<Metadata *> MDVals;
+    MDVals.emplace_back(
+        ConstantAsMetadata::get(ConstantInt::get(Type::getInt32Ty(Ctx), Tag)));
+    switch (Tag) {
+    case ShaderFlagsTag:
+      MDVals.emplace_back(ConstantAsMetadata::get(
+          ConstantInt::get(Type::getInt64Ty(Ctx), Value)));
+      break;
+    case ShaderKindTag:
+      MDVals.emplace_back(ConstantAsMetadata::get(
+          ConstantInt::get(Type::getInt32Ty(Ctx), Value)));
+      break;
+    default:
+      assert(false && "NYI: Unhandled entry property tag");
+    }
+    return MDVals;
+  }
+
+  MDTuple *constructEntryPropMetadata() {
+    SmallVector<Metadata *> MDVals;
+    if (EntryShaderFlags != 0)
+      MDVals.append(getTagValueAsMetadata(ShaderFlagsTag, EntryShaderFlags));
+
+    if (Entry != nullptr) {
+      // FIXME: support more props.
+      // See https://github.com/llvm/llvm-project/issues/57948.
+      // Add shader kind for lib entries.
+      if (ShaderProfile == Triple::EnvironmentType::Library &&
+          ShaderStage != Triple::EnvironmentType::Library)
+        MDVals.append(
+            getTagValueAsMetadata(ShaderKindTag, getShaderStage(ShaderStage)));
+
+      if (ShaderStage == Triple::EnvironmentType::Compute) {
+        MDVals.emplace_back(ConstantAsMetadata::get(
+            ConstantInt::get(Type::getInt32Ty(Ctx), NumThreadsTag)));
+        std::vector<Metadata *> NumThreadVals;
+        NumThreadVals.emplace_back(ConstantAsMetadata::get(
+            ConstantInt::get(Type::getInt32Ty(Ctx), NumThreadsX)));
+        NumThreadVals.emplace_back(ConstantAsMetadata::get(
+            ConstantInt::get(Type::getInt32Ty(Ctx), NumThreadsY)));
+        NumThreadVals.emplace_back(ConstantAsMetadata::get(
+            ConstantInt::get(Type::getInt32Ty(Ctx), NumThreadsZ)));
+        MDVals.emplace_back(MDNode::get(Ctx, NumThreadVals));
+      }
+    }
+    if (MDVals.empty())
+      return nullptr;
+    return MDNode::get(Ctx, MDVals);
+  }
+};
+
+static void createEntryMD(Module &M, const uint64_t ShaderFlags,
+                          const dxil::ModuleMetadataInfo &MDAnalysisInfo) {
+  auto &Ctx = M.getContext();
+  // FIXME: generate metadata for resource.
+  MDTuple *MDResources = nullptr;
+  if (auto *NamedResources = M.getNamedMetadata("dx.resources"))
+    MDResources = dyn_cast<MDTuple>(NamedResources->getOperand(0));
----------------
bharadwajy wrote:

> This FIXME seems wrong - we're generating the resource metadata in emitResourceMetadata. Why not make that function return the `NamedMDNode` rather than looking it up again?

Changed per suggestion.

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