Add DXILVersion for DXIL metadata abstractions

python3kgae · python3kgae · commit 6d0a795baacc · 2024-05-07T12:39:30.000-04:00
diff --git a/llvm/include/llvm/Transforms/Utils/DXIL.h b/llvm/include/llvm/Transforms/Utils/DXIL.h
@@ -0,0 +1,66 @@
+//===- DXIL.h - Abstractions for DXIL constructs ----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// \file This file defines various abstractions for transforming between DXIL's
+// and LLVM's representations of shader metadata.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_UTILS_DXIL_H
+#define LLVM_TRANSFORMS_UTILS_DXIL_H
+
+#include "llvm/Support/Error.h"
+#include "llvm/TargetParser/Triple.h"
+
+namespace llvm {
+class Module;
+
+namespace dxil {
+
+class DXILVersion {
+  unsigned Major = 0;
+  unsigned Minor = 0;
+
+public:
+  DXILVersion() = default;
+  DXILVersion(unsigned Major, unsigned Minor) : Major(Major), Minor(Minor) {}
+
+  /// Get the DXILVersion for \c M
+  static Expected<DXILVersion> get(Module &M);
+  /// Read the DXILVersion from the DXIL metadata in \c M
+  static Expected<DXILVersion> readDXIL(Module &M);
+
+  /// Returns true if no DXILVersion is set
+  bool empty() { return Major == 0 && Minor == 0; }
+
+  /// Remove any non-DXIL LLVM representations of the DXILVersion from \c M.
+  void strip(Module &M);
+  /// Embed the LLVM representation of the DXILVersion into \c M.
+  void embed(Module &M);
+  /// Remove any DXIL representation of the DXILVersion from \c M.
+  void stripDXIL(Module &M);
+  /// Embed a DXIL representation of the DXILVersion into \c M.
+  void embedDXIL(Module &M);
+
+  void print(raw_ostream &OS) const {
+    // Format like Triple ArchName.
+    OS << "dxilv" << Major << "." << Minor;
+  }
+  LLVM_DUMP_METHOD void dump() const { print(errs()); }
+};
+
+inline raw_ostream &operator<<(raw_ostream &OS, const DXILVersion &V) {
+  V.print(OS);
+  return OS;
+}
+
+} // namespace dxil
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_UTILS_DXIL_H
+
diff --git a/llvm/lib/Target/DirectX/DXILMetadata.cpp b/llvm/lib/Target/DirectX/DXILMetadata.cpp
@@ -81,52 +81,6 @@ void dxil::createShaderModelMD(Module &M) {
   Entry->addOperand(MDNode::get(Ctx, Vals));
 }
 
-void dxil::createDXILVersionMD(Module &M) {
-  Triple TT(Triple::normalize(M.getTargetTriple()));
-  VersionTuple Ver = VersionTuple(1, 0);
-  switch (TT.getSubArch()) {
-  case Triple::DXILSubArch_v1_0:
-    Ver = VersionTuple(1, 0);
-    break;
-  case Triple::DXILSubArch_v1_1:
-    Ver = VersionTuple(1, 1);
-    break;
-  case Triple::DXILSubArch_v1_2:
-    Ver = VersionTuple(1, 2);
-    break;
-  case Triple::DXILSubArch_v1_3:
-    Ver = VersionTuple(1, 3);
-    break;
-  case Triple::DXILSubArch_v1_4:
-    Ver = VersionTuple(1, 4);
-    break;
-  case Triple::DXILSubArch_v1_5:
-    Ver = VersionTuple(1, 5);
-    break;
-  case Triple::DXILSubArch_v1_6:
-    Ver = VersionTuple(1, 6);
-    break;
-  case Triple::DXILSubArch_v1_7:
-    Ver = VersionTuple(1, 7);
-    break;
-  case Triple::DXILSubArch_v1_8:
-    Ver = VersionTuple(1, 8);
-    break;
-  case Triple::NoSubArch:
-    break;
-  default:
-    llvm_unreachable("Unsupported subarch for DXIL generation.");
-    break;
-  }
-  LLVMContext &Ctx = M.getContext();
-  IRBuilder<> B(Ctx);
-  NamedMDNode *Entry = M.getOrInsertNamedMetadata("dx.version");
-  Metadata *Vals[2];
-  Vals[0] = ConstantAsMetadata::get(B.getInt32(Ver.getMajor()));
-  Vals[1] = ConstantAsMetadata::get(B.getInt32(Ver.getMinor().value_or(0)));
-  Entry->addOperand(MDNode::get(Ctx, Vals));
-}
-
 static uint32_t getShaderStage(Triple::EnvironmentType Env) {
   return (uint32_t)Env - (uint32_t)llvm::Triple::Pixel;
 }
diff --git a/llvm/lib/Target/DirectX/DXILMetadata.h b/llvm/lib/Target/DirectX/DXILMetadata.h
@@ -33,7 +33,6 @@ class ValidatorVersionMD {
 };
 
 void createShaderModelMD(Module &M);
-void createDXILVersionMD(Module &M);
 void createEntryMD(Module &M, const uint64_t ShaderFlags);
 
 } // namespace dxil
diff --git a/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp b/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp
@@ -19,6 +19,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/Utils/DXIL.h"
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -48,7 +49,13 @@ bool DXILTranslateMetadata::runOnModule(Module &M) {
   if (ValVerMD.isEmpty())
     ValVerMD.update(VersionTuple(1, 0));
   dxil::createShaderModelMD(M);
-  dxil::createDXILVersionMD(M);
+  Expected<dxil::DXILVersion> DXILVer = dxil::DXILVersion::get(M);
+  if (auto E = DXILVer.takeError()) {
+    errs() << "Fail to get DXIL version " << toString(std::move(E)) << "\n";
+    return false;
+  }
+
+  DXILVer->embedDXIL(M);
 
   const dxil::Resources &Res =
       getAnalysis<DXILResourceWrapper>().getDXILResource();
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -20,6 +20,7 @@ add_llvm_component_library(LLVMTransformUtils
   CountVisits.cpp
   Debugify.cpp
   DemoteRegToStack.cpp
+  DXIL.cpp
   DXILUpgrade.cpp
   EntryExitInstrumenter.cpp
   EscapeEnumerator.cpp
diff --git a/llvm/lib/Transforms/Utils/DXIL.cpp b/llvm/lib/Transforms/Utils/DXIL.cpp
@@ -0,0 +1,119 @@
+//===- DXIL.cpp - Abstractions for DXIL constructs ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/DXIL.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+static Error errInvalid(const char *Msg) {
+  return createStringError(std::errc::invalid_argument, Msg);
+}
+
+template <typename... Ts>
+static Error errInvalid(const char *Fmt, const Ts &... Vals) {
+  return createStringError(std::errc::invalid_argument, Fmt, Vals...);
+}
+
+Expected<DXILVersion> DXILVersion::get(Module &M) {
+  Triple TT(Triple::normalize(M.getTargetTriple()));
+
+  if (!TT.isDXIL())
+    return errInvalid("Cannot get DXIL version for arch '%s'",
+                      TT.getArchName().str().c_str());
+
+  switch (TT.getSubArch()) {
+    case Triple::NoSubArch:
+    case Triple::DXILSubArch_v1_0:
+      return DXILVersion(1, 0);
+    case Triple::DXILSubArch_v1_1:
+      return DXILVersion(1, 1);
+    case Triple::DXILSubArch_v1_2:
+      return DXILVersion(1, 2);
+    case Triple::DXILSubArch_v1_3:
+      return DXILVersion(1, 3);
+    case Triple::DXILSubArch_v1_4:
+      return DXILVersion(1, 4);
+    case Triple::DXILSubArch_v1_5:
+      return DXILVersion(1, 5);
+    case Triple::DXILSubArch_v1_6:
+      return DXILVersion(1, 6);
+    case Triple::DXILSubArch_v1_7:
+      return DXILVersion(1, 7);
+    case Triple::DXILSubArch_v1_8:
+      return DXILVersion(1, 8);
+    default:
+    return errInvalid("Cannot get DXIL version for arch '%s'",
+                      TT.getArchName().str().c_str());
+  }
+}
+
+Expected<DXILVersion> DXILVersion::readDXIL(Module &M) {
+  NamedMDNode *DXILVersionMD = M.getNamedMetadata("dx.version");
+  if (!DXILVersionMD)
+    return DXILVersion();
+
+  if (DXILVersionMD->getNumOperands() != 1)
+    return errInvalid("dx.version must have one operand");
+
+  MDNode *N = DXILVersionMD->getOperand(0);
+  if (N->getNumOperands() != 2)
+    return errInvalid("dx.version must have 2 components, not %d",
+                      N->getNumOperands());
+
+  const auto *MajorOp = mdconst::dyn_extract<ConstantInt>(N->getOperand(0));
+  const auto *MinorOp = mdconst::dyn_extract<ConstantInt>(N->getOperand(1));
+  if (!MajorOp)
+    return errInvalid("dx.version major version must be an integer");
+  if (!MinorOp)
+    return errInvalid("dx.version minor version must be an integer");
+
+  return DXILVersion(MajorOp->getZExtValue(), MinorOp->getZExtValue());
+}
+
+void DXILVersion::strip(Module &M) {
+  M.setTargetTriple("dxil-ms-dx");
+}
+
+void DXILVersion::embed(Module &M) {
+  Triple TT(Triple::normalize(M.getTargetTriple()));
+  SmallString<64> Triple;
+  raw_svector_ostream OS(Triple);
+  print(OS);
+  OS << "-" << TT.getVendorName() << "-" << TT.getOSAndEnvironmentName();
+  M.setTargetTriple(OS.str());
+}
+
+void DXILVersion::stripDXIL(Module &M) {
+  if (NamedMDNode *V = M.getNamedMetadata("dx.version")) {
+    V->dropAllReferences();
+    V->eraseFromParent();
+  }
+}
+
+void DXILVersion::embedDXIL(Module &M) {
+  LLVMContext &Ctx = M.getContext();
+  IRBuilder<> B(Ctx);
+
+  Metadata *Vals[2];
+  Vals[0] = ConstantAsMetadata::get(B.getInt32(Major));
+  Vals[1] = ConstantAsMetadata::get(B.getInt32(Minor));
+  MDNode *MD = MDNode::get(Ctx, Vals);
+
+  NamedMDNode *V = M.getOrInsertNamedMetadata("dx.version");
+  if (V->getNumOperands())
+    V->setOperand(0, MD);
+  else
+    V->addOperand(MD);
+}
diff --git a/llvm/lib/Transforms/Utils/DXILUpgrade.cpp b/llvm/lib/Transforms/Utils/DXILUpgrade.cpp
@@ -11,6 +11,7 @@
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/DXIL.h"
 
 using namespace llvm;
 
@@ -33,6 +34,20 @@ static bool handleValVerMetadata(Module &M) {
   return true;
 }
 
+static bool handleDXILVerMetadata(Module &M) {
+  auto V = dxil::DXILVersion::readDXIL(M);
+  if (Error E = V.takeError()) {
+    report_fatal_error(std::move(E), /*gen_crash_diag=*/false);
+  }
+  if (V->empty())
+    return false;
+
+  LLVM_DEBUG(dbgs() << "DXIL: DXIL Version " << *V << "\n");
+  V->embed(M);
+  V->stripDXIL(M);
+  return true;
+}
+
 PreservedAnalyses DXILUpgradePass::run(Module &M, ModuleAnalysisManager &AM) {
   PreservedAnalyses PA;
   // We never add, remove, or change functions here.
@@ -41,6 +56,7 @@ PreservedAnalyses DXILUpgradePass::run(Module &M, ModuleAnalysisManager &AM) {
 
   bool Changed = false;
   Changed |= handleValVerMetadata(M);
+  Changed |= handleDXILVerMetadata(M);
 
   if (!Changed)
     return PreservedAnalyses::all();
diff --git a/llvm/test/CodeGen/DirectX/Metadata/dxilVer-default.ll b/llvm/test/CodeGen/DirectX/Metadata/dxilVer-default.ll
@@ -0,0 +1,12 @@
+; RUN: opt -S -dxil-metadata-emit %s | FileCheck %s
+target triple = "dxil-pc-shadermodel-vertex"
+
+; CHECK: !dx.version = !{![[DXVER:[0-9]+]]}
+; CHECK: ![[DXVER]] = !{i32 1, i32 0}
+
+define void @entry() #0 {
+entry:
+  ret void
+}
+
+attributes #0 = { noinline nounwind "hlsl.shader"="vertex" }
diff --git a/llvm/test/Transforms/DXILUpgrade/strip-dxilver.ll b/llvm/test/Transforms/DXILUpgrade/strip-dxilver.ll
@@ -0,0 +1,18 @@
+; RUN: opt -passes=dxil-upgrade -S < %s | FileCheck %s
+
+; Ensure that both the dxil version metadata and its operand are removed.
+; CHECK: !unrelated_md1 = !{!0}
+; CHECK-NOT: !dx.version
+; CHECK: !unrelated_md2 = !{!1}
+;
+; CHECK: !0 = !{i32 1234}
+; CHECK-NOT: !{i32 1, i32 7}
+; CHECK: !1 = !{i32 4321}
+
+!unrelated_md1 = !{!0}
+!dx.version = !{!1}
+!unrelated_md2 = !{!2}
+
+!0 = !{i32 1234}
+!1 = !{i32 1, i32 7}
+!2 = !{i32 4321}
