[DirectX] Add DXIL validation of `llvm.loop` metadata (#164292)

This pr adds the equivalent validation of `llvm.loop` metadata that is
[done in
DXC](https://github.com/microsoft/DirectXShaderCompiler/blob/8f21027f2ad5dcfa63a275cbd278691f2c8fad33/lib/DxilValidation/DxilValidation.cpp#L3010).

This is done as follows:
- Add `llvm.loop` to the metadata allow-list in `DXILTranslateMetadata`
- Iterate through all `llvm.loop` metadata nodes and strip all
incompatible ones
- Raise an error for ill-formed nodes that are compatible with DXIL

Resolves: https://github.com/llvm/llvm-project/issues/137387

7 files changed, 443 insertions, 70 deletions

diff --git a/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp b/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp
index 1e4797b..e345bda 100644
--- a/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp
+++ b/llvm/lib/Target/DirectX/DXILTranslateMetadata.cpp
@@ -36,9 +36,10 @@ using namespace llvm;
 using namespace llvm::dxil;
 
 namespace {
-/// A simple Wrapper DiagnosticInfo that generates Module-level diagnostic
-/// for TranslateMetadata pass
-class DiagnosticInfoTranslateMD : public DiagnosticInfo {
+
+/// A simple wrapper of DiagnosticInfo that generates module-level diagnostic
+/// for the DXILValidateMetadata pass
+class DiagnosticInfoValidateMD : public DiagnosticInfo {
 private:
   const Twine &Msg;
   const Module &Mod;
@@ -47,9 +48,9 @@ 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.
-  DiagnosticInfoTranslateMD(const Module &M,
-                            const Twine &Msg LLVM_LIFETIME_BOUND,
-                            DiagnosticSeverity Severity = DS_Error)
+  DiagnosticInfoValidateMD(const Module &M,
+                           const Twine &Msg LLVM_LIFETIME_BOUND,
+                           DiagnosticSeverity Severity = DS_Error)
       : DiagnosticInfo(DK_Unsupported, Severity), Msg(Msg), Mod(M) {}
 
   void print(DiagnosticPrinter &DP) const override {
@@ -57,6 +58,16 @@ public:
   }
 };
 
+static void reportError(Module &M, Twine Message,
+                        DiagnosticSeverity Severity = DS_Error) {
+  M.getContext().diagnose(DiagnosticInfoValidateMD(M, Message, Severity));
+}
+
+static void reportLoopError(Module &M, Twine Message,
+                            DiagnosticSeverity Severity = DS_Error) {
+  reportError(M, Twine("Invalid \"llvm.loop\" metadata: ") + Message, Severity);
+}
+
 enum class EntryPropsTag {
   ShaderFlags = 0,
   GSState,
@@ -314,25 +325,122 @@ static void translateBranchMetadata(Module &M, Instruction *BBTerminatorInst) {
   BBTerminatorInst->setMetadata("hlsl.controlflow.hint", nullptr);
 }
 
-static std::array<unsigned, 6> getCompatibleInstructionMDs(llvm::Module &M) {
+// Determines if the metadata node will be compatible with DXIL's loop metadata
+// representation.
+//
+// Reports an error for compatible metadata that is ill-formed.
+static bool isLoopMDCompatible(Module &M, Metadata *MD) {
+  // DXIL only accepts the following loop hints:
+  std::array<StringLiteral, 3> ValidHintNames = {"llvm.loop.unroll.count",
+                                                 "llvm.loop.unroll.disable",
+                                                 "llvm.loop.unroll.full"};
+
+  MDNode *HintMD = dyn_cast<MDNode>(MD);
+  if (!HintMD || HintMD->getNumOperands() == 0)
+    return false;
+
+  auto *HintStr = dyn_cast<MDString>(HintMD->getOperand(0));
+  if (!HintStr)
+    return false;
+
+  if (!llvm::is_contained(ValidHintNames, HintStr->getString()))
+    return false;
+
+  auto ValidCountNode = [](MDNode *CountMD) -> bool {
+    if (CountMD->getNumOperands() == 2)
+      if (auto *Count = dyn_cast<ConstantAsMetadata>(CountMD->getOperand(1)))
+        if (isa<ConstantInt>(Count->getValue()))
+          return true;
+    return false;
+  };
+
+  if (HintStr->getString() == "llvm.loop.unroll.count") {
+    if (!ValidCountNode(HintMD)) {
+      reportLoopError(M, "\"llvm.loop.unroll.count\" must have 2 operands and "
+                         "the second must be a constant integer");
+      return false;
+    }
+  } else if (HintMD->getNumOperands() != 1) {
+    reportLoopError(
+        M, "\"llvm.loop.unroll.disable\" and \"llvm.loop.unroll.full\" "
+           "must be provided as a single operand");
+    return false;
+  }
+
+  return true;
+}
+
+static void translateLoopMetadata(Module &M, Instruction *I, MDNode *BaseMD) {
+  // A distinct node has the self-referential form: !0 = !{ !0, ... }
+  auto IsDistinctNode = [](MDNode *Node) -> bool {
+    return Node && Node->getNumOperands() != 0 && Node == Node->getOperand(0);
+  };
+
+  // Set metadata to null to remove empty/ill-formed metadata from instruction
+  if (BaseMD->getNumOperands() == 0 || !IsDistinctNode(BaseMD))
+    return I->setMetadata("llvm.loop", nullptr);
+
+  // It is valid to have a chain of self-refential loop metadata nodes, as
+  // below. We will collapse these into just one when we reconstruct the
+  // metadata.
+  //
+  // Eg:
+  // !0 = !{!0, !1}
+  // !1 = !{!1, !2}
+  // !2 = !{!"llvm.loop.unroll.disable"}
+  //
+  // So, traverse down a potential self-referential chain
+  while (1 < BaseMD->getNumOperands() &&
+         IsDistinctNode(dyn_cast<MDNode>(BaseMD->getOperand(1))))
+    BaseMD = dyn_cast<MDNode>(BaseMD->getOperand(1));
+
+  // To reconstruct a distinct node we create a temporary node that we will
+  // then update to create a self-reference.
+  llvm::TempMDTuple TempNode = llvm::MDNode::getTemporary(M.getContext(), {});
+  SmallVector<Metadata *> CompatibleOperands = {TempNode.get()};
+
+  // Iterate and reconstruct the metadata nodes that contains any hints,
+  // stripping any unrecognized metadata.
+  ArrayRef<MDOperand> Operands = BaseMD->operands();
+  for (auto &Op : Operands.drop_front())
+    if (isLoopMDCompatible(M, Op.get()))
+      CompatibleOperands.push_back(Op.get());
+
+  if (2 < CompatibleOperands.size())
+    reportLoopError(M, "Provided conflicting hints");
+
+  MDNode *CompatibleLoopMD = MDNode::get(M.getContext(), CompatibleOperands);
+  TempNode->replaceAllUsesWith(CompatibleLoopMD);
+
+  I->setMetadata("llvm.loop", CompatibleLoopMD);
+}
+
+using InstructionMDList = std::array<unsigned, 7>;
+
+static InstructionMDList getCompatibleInstructionMDs(llvm::Module &M) {
   return {
       M.getMDKindID("dx.nonuniform"),    M.getMDKindID("dx.controlflow.hints"),
       M.getMDKindID("dx.precise"),       llvm::LLVMContext::MD_range,
-      llvm::LLVMContext::MD_alias_scope, llvm::LLVMContext::MD_noalias};
+      llvm::LLVMContext::MD_alias_scope, llvm::LLVMContext::MD_noalias,
+      M.getMDKindID("llvm.loop")};
 }
 
 static void translateInstructionMetadata(Module &M) {
   // construct allowlist of valid metadata node kinds
-  std::array<unsigned, 6> DXILCompatibleMDs = getCompatibleInstructionMDs(M);
+  InstructionMDList DXILCompatibleMDs = getCompatibleInstructionMDs(M);
+  unsigned char MDLoopKind = M.getContext().getMDKindID("llvm.loop");
 
   for (Function &F : M) {
     for (BasicBlock &BB : F) {
       // This needs to be done first so that "hlsl.controlflow.hints" isn't
-      // removed in the whitelist below
+      // removed in the allow-list below
       if (auto *I = BB.getTerminator())
         translateBranchMetadata(M, I);
 
       for (auto &I : make_early_inc_range(BB)) {
+        if (isa<BranchInst>(I))
+          if (MDNode *LoopMD = I.getMetadata(MDLoopKind))
+            translateLoopMetadata(M, &I, LoopMD);
         I.dropUnknownNonDebugMetadata(DXILCompatibleMDs);
       }
     }
@@ -389,31 +497,23 @@ static void translateGlobalMetadata(Module &M, DXILResourceMap &DRM,
     uint64_t CombinedMask = ShaderFlags.getCombinedFlags();
     EntryFnMDNodes.emplace_back(
         emitTopLevelLibraryNode(M, ResourceMD, CombinedMask));
-  } else if (MMDI.EntryPropertyVec.size() > 1) {
-    M.getContext().diagnose(DiagnosticInfoTranslateMD(
-        M, "Non-library shader: One and only one entry expected"));
-  }
+  } else if (1 < MMDI.EntryPropertyVec.size())
+    reportError(M, "Non-library shader: One and only one entry expected");
 
   for (const EntryProperties &EntryProp : MMDI.EntryPropertyVec) {
-    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,
-            "Shader stage '" +
-                Twine(getShortShaderStage(EntryProp.ShaderStage) +
-                      "' for entry '" + Twine(EntryProp.Entry->getName()) +
-                      "' different from specified target profile '" +
-                      Twine(Triple::getEnvironmentTypeName(MMDI.ShaderProfile) +
-                            "'"))));
-      }
+      EntryShaderFlags = ShaderFlags.getFunctionFlags(EntryProp.Entry);
+      if (EntryProp.ShaderStage != MMDI.ShaderProfile)
+        reportError(
+            M, "Shader stage '" +
+                   Twine(getShortShaderStage(EntryProp.ShaderStage)) +
+                   "' for entry '" + Twine(EntryProp.Entry->getName()) +
+                   "' different from specified target profile '" +
+                   Twine(Triple::getEnvironmentTypeName(MMDI.ShaderProfile) +
+                         "'"));
     }
+
     EntryFnMDNodes.emplace_back(emitEntryMD(EntryProp, Signatures, ResourceMD,
                                             EntryShaderFlags,
                                             MMDI.ShaderProfile));
@@ -454,45 +554,34 @@ PreservedAnalyses DXILTranslateMetadata::run(Module &M,
   return PreservedAnalyses::all();
 }
 
-namespace {
-class DXILTranslateMetadataLegacy : public ModulePass {
-public:
-  static char ID; // Pass identification, replacement for typeid
-  explicit DXILTranslateMetadataLegacy() : ModulePass(ID) {}
-
-  StringRef getPassName() const override { return "DXIL Translate Metadata"; }
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<DXILResourceTypeWrapperPass>();
-    AU.addRequired<DXILResourceWrapperPass>();
-    AU.addRequired<ShaderFlagsAnalysisWrapper>();
-    AU.addRequired<DXILMetadataAnalysisWrapperPass>();
-    AU.addRequired<RootSignatureAnalysisWrapper>();
-
-    AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
-    AU.addPreserved<DXILResourceBindingWrapperPass>();
-    AU.addPreserved<DXILResourceWrapperPass>();
-    AU.addPreserved<RootSignatureAnalysisWrapper>();
-    AU.addPreserved<ShaderFlagsAnalysisWrapper>();
-  }
+void DXILTranslateMetadataLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<DXILResourceTypeWrapperPass>();
+  AU.addRequired<DXILResourceWrapperPass>();
+  AU.addRequired<ShaderFlagsAnalysisWrapper>();
+  AU.addRequired<DXILMetadataAnalysisWrapperPass>();
+  AU.addRequired<RootSignatureAnalysisWrapper>();
+
+  AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
+  AU.addPreserved<DXILResourceBindingWrapperPass>();
+  AU.addPreserved<DXILResourceWrapperPass>();
+  AU.addPreserved<RootSignatureAnalysisWrapper>();
+  AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+}
 
-  bool runOnModule(Module &M) override {
-    DXILResourceMap &DRM =
-        getAnalysis<DXILResourceWrapperPass>().getResourceMap();
-    DXILResourceTypeMap &DRTM =
-        getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
-    const ModuleShaderFlags &ShaderFlags =
-        getAnalysis<ShaderFlagsAnalysisWrapper>().getShaderFlags();
-    dxil::ModuleMetadataInfo MMDI =
-        getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
-
-    translateGlobalMetadata(M, DRM, DRTM, ShaderFlags, MMDI);
-    translateInstructionMetadata(M);
-    return true;
-  }
-};
+bool DXILTranslateMetadataLegacy::runOnModule(Module &M) {
+  DXILResourceMap &DRM =
+      getAnalysis<DXILResourceWrapperPass>().getResourceMap();
+  DXILResourceTypeMap &DRTM =
+      getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
+  const ModuleShaderFlags &ShaderFlags =
+      getAnalysis<ShaderFlagsAnalysisWrapper>().getShaderFlags();
+  dxil::ModuleMetadataInfo MMDI =
+      getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
 
-} // namespace
+  translateGlobalMetadata(M, DRM, DRTM, ShaderFlags, MMDI);
+  translateInstructionMetadata(M);
+  return true;
+}
 
 char DXILTranslateMetadataLegacy::ID = 0;
 
diff --git a/llvm/lib/Target/DirectX/DXILTranslateMetadata.h b/llvm/lib/Target/DirectX/DXILTranslateMetadata.h
index 4c1ffac..cfb8aaa 100644
--- a/llvm/lib/Target/DirectX/DXILTranslateMetadata.h
+++ b/llvm/lib/Target/DirectX/DXILTranslateMetadata.h
@@ -10,6 +10,7 @@
 #define LLVM_TARGET_DIRECTX_DXILTRANSLATEMETADATA_H
 
 #include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
 
 namespace llvm {
 
@@ -20,6 +21,22 @@ public:
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
 };
 
+/// Wrapper pass for the legacy pass manager.
+///
+/// This is required because the passes that will depend on this are codegen
+/// passes which run through the legacy pass manager.
+class DXILTranslateMetadataLegacy : public ModulePass {
+public:
+  static char ID; // Pass identification, replacement for typeid
+  explicit DXILTranslateMetadataLegacy() : ModulePass(ID) {}
+
+  StringRef getPassName() const override { return "DXIL Translate Metadata"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  bool runOnModule(Module &M) override;
+};
+
 } // namespace llvm
 
 #endif // LLVM_TARGET_DIRECTX_DXILTRANSLATEMETADATA_H
diff --git a/llvm/test/CodeGen/DirectX/Metadata/loop-md-errs.ll b/llvm/test/CodeGen/DirectX/Metadata/loop-md-errs.ll
new file mode 100644
index 0000000..fbe4653
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/Metadata/loop-md-errs.ll
@@ -0,0 +1,113 @@
+; RUN: split-file %s %t
+; RUN: not opt -S --dxil-translate-metadata %t/args.ll 2>&1 | FileCheck %t/args.ll
+; RUN: not opt -S --dxil-translate-metadata %t/bad-count.ll 2>&1 | FileCheck %t/bad-count.ll
+; RUN: not opt -S --dxil-translate-metadata %t/invalid-disable.ll 2>&1 | FileCheck %t/invalid-disable.ll
+; RUN: not opt -S --dxil-translate-metadata %t/invalid-full.ll 2>&1 | FileCheck %t/invalid-full.ll
+
+; Test that loop metadata is validated as with the DXIL validator
+
+;--- args.ll
+
+; CHECK: Invalid "llvm.loop" metadata: Provided conflicting hints
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+!1 = !{!1, !2, !3} ; conflicting args
+!2 = !{!"llvm.loop.unroll.full"}
+!3 = !{!"llvm.loop.unroll.disable"}
+
+;--- bad-count.ll
+
+; CHECK: "llvm.loop.unroll.count" must have 2 operands and the second must be a constant integer
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+!1 = !{!1, !2}
+!2 = !{!"llvm.loop.unroll.count", !"not an int"} ; invalid count parameters
+
+;--- invalid-disable.ll
+
+; CHECK: Invalid "llvm.loop" metadata: "llvm.loop.unroll.disable" and "llvm.loop.unroll.full" must be provided as a single operand
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+!1 = !{!1, !2}
+!2 = !{!"llvm.loop.unroll.disable", i32 0} ; invalid second operand
+
+
+;--- invalid-full.ll
+
+; CHECK: Invalid "llvm.loop" metadata: "llvm.loop.unroll.disable" and "llvm.loop.unroll.full" must be provided as a single operand
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+!1 = !{!1, !2}
+!2 = !{!"llvm.loop.unroll.full", i32 0} ; invalid second operand
diff --git a/llvm/test/CodeGen/DirectX/Metadata/loop-md-stripped.ll b/llvm/test/CodeGen/DirectX/Metadata/loop-md-stripped.ll
new file mode 100644
index 0000000..09d8aec
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/Metadata/loop-md-stripped.ll
@@ -0,0 +1,58 @@
+; RUN: split-file %s %t
+; RUN: opt -S --dxil-translate-metadata %t/not-distinct.ll 2>&1 | FileCheck %t/not-distinct.ll
+; RUN: opt -S --dxil-translate-metadata %t/not-md.ll 2>&1 | FileCheck %t/not-md.ll
+
+; Test that DXIL incompatible loop metadata is stripped
+
+;--- not-distinct.ll
+
+; Ensure it is stripped because it is not provided a distinct loop parent
+; CHECK-NOT: {!"llvm.loop.unroll.disable"}
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+!1 = !{!"llvm.loop.unroll.disable"} ; first node must be a distinct self-reference
+
+
+;--- not-md.ll
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  ; CHECK: br label %loop.header, !llvm.loop ![[#LOOP_MD:]]
+  br label %loop.header, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+; CHECK: ![[#LOOP_MD:]] = distinct !{![[#LOOP_MD]]}
+
+!1 = !{!1, i32 0} ; second operand is not a metadata node
diff --git a/llvm/test/CodeGen/DirectX/Metadata/loop-md-valid.ll b/llvm/test/CodeGen/DirectX/Metadata/loop-md-valid.ll
new file mode 100644
index 0000000..a189c0e
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/Metadata/loop-md-valid.ll
@@ -0,0 +1,95 @@
+; RUN: split-file %s %t
+; RUN: opt -S --dxil-translate-metadata %t/count.ll | FileCheck %t/count.ll
+; RUN: opt -S --dxil-translate-metadata %t/disable.ll | FileCheck %t/disable.ll
+; RUN: opt -S --dxil-translate-metadata %t/full.ll | FileCheck %t/full.ll
+
+;--- count.ll
+
+; Test that we collapse a self-referential chain and allow a unroll.count hint
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  ; CHECK: br label %loop.header, !llvm.loop ![[#LOOP_MD:]]
+  br label %loop.header, !llvm.loop !0
+
+exit:
+  ret void
+}
+
+; CHECK: ![[#LOOP_MD]] = distinct !{![[#LOOP_MD]], ![[#COUNT:]]}
+; CHECK: ![[#COUNT]] = !{!"llvm.loop.unroll.count", i6 4}
+
+!0 = !{!0, !1}
+!1 = !{!1, !2}
+!2 = !{!"llvm.loop.unroll.count", i6 4}
+
+;--- disable.ll
+
+; Test that we allow a disable hint
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  ; CHECK: br label %loop.header, !llvm.loop ![[#LOOP_MD:]]
+  br label %loop.header, !llvm.loop !0
+
+exit:
+  ret void
+}
+
+; CHECK: ![[#LOOP_MD]] = distinct !{![[#LOOP_MD]], ![[#DISABLE:]]}
+; CHECK: ![[#DISABLE]] = !{!"llvm.loop.unroll.disable"}
+
+!0 = !{!0, !1}
+!1 = !{!"llvm.loop.unroll.disable"}
+
+;--- full.ll
+
+; Test that we allow a full hint
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-library"
+
+define void @example_loop(i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  %i.next = add nsw i32 %i, 1
+  ; CHECK: br label %loop.header, !llvm.loop ![[#LOOP_MD:]]
+  br label %loop.header, !llvm.loop !0
+
+exit:
+  ret void
+}
+
+; CHECK: ![[#LOOP_MD]] = distinct !{![[#LOOP_MD]], ![[#FULL:]]}
+; CHECK: ![[#FULL]] = !{!"llvm.loop.unroll.full"}
+
+!0 = !{!0, !1}
+!1 = !{!"llvm.loop.unroll.full"}
diff --git a/llvm/test/CodeGen/DirectX/Metadata/multiple-entries-cs-error.ll b/llvm/test/CodeGen/DirectX/Metadata/multiple-entries-cs-error.ll
index 9697d438..5740ee1 100644
--- a/llvm/test/CodeGen/DirectX/Metadata/multiple-entries-cs-error.ll
+++ b/llvm/test/CodeGen/DirectX/Metadata/multiple-entries-cs-error.ll
@@ -1,4 +1,4 @@
-; RUN: not opt -S  -S -dxil-translate-metadata %s 2>&1 | FileCheck %s
+; RUN: not opt -S -dxil-translate-metadata %s 2>&1 | FileCheck %s
 target triple = "dxil-pc-shadermodel6.8-compute"
 
 ; CHECK: Non-library shader: One and only one entry expected
diff --git a/llvm/test/CodeGen/DirectX/metadata-stripping.ll b/llvm/test/CodeGen/DirectX/metadata-stripping.ll
index 531ab6c..53716ff 100644
--- a/llvm/test/CodeGen/DirectX/metadata-stripping.ll
+++ b/llvm/test/CodeGen/DirectX/metadata-stripping.ll
@@ -14,7 +14,7 @@ entry:
 
   %cmp.i = icmp ult i32 1, 2
   ; Ensure that the !llvm.loop metadata node gets dropped.
-  ; CHECK: br i1 %cmp.i, label %_Z4mainDv3_j.exit, label %_Z4mainDv3_j.exit{{$}}
+  ; CHECK: br i1 %cmp.i, label %_Z4mainDv3_j.exit, label %_Z4mainDv3_j.exit, !llvm.loop [[LOOPMD:![0-9]+]]
   br i1 %cmp.i, label %_Z4mainDv3_j.exit, label %_Z4mainDv3_j.exit, !llvm.loop !0
 
 _Z4mainDv3_j.exit:                                ; preds = %for.body.i, %entry
@@ -25,7 +25,8 @@ _Z4mainDv3_j.exit:                                ; preds = %for.body.i, %entry
 ; No more metadata should be necessary, the rest (the current 0 and 1)
 ; should be removed.
 ; CHECK-NOT: !{!"llvm.loop.mustprogress"}
-; CHECK: [[RANGEMD]] = !{i32 1, i32 5}
+; CHECK-DAG: [[RANGEMD]] = !{i32 1, i32 5}
+; CHECK-DAG: [[LOOPMD]] = distinct !{[[LOOPMD]]}
 ; CHECK-NOT: !{!"llvm.loop.mustprogress"}
 !0 = distinct !{!0, !1}
 !1 = !{!"llvm.loop.mustprogress"}