[DXIL] Remove incompatible metadata types when preparing DXIL.

llvm/lib/Target/DirectX/DXILPrepare.cpp

@@ -161,6 +161,15 @@ class DXILPrepareModule : public ModulePass {
                          Builder.getPtrTy(PtrTy->getAddressSpace())));
   }
 
+  static std::array<unsigned, 6> 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};
+  }
+
 public:
   bool runOnModule(Module &M) override {
     PointerTypeMap PointerTypes = PointerTypeAnalysis::run(M);
@@ -176,6 +185,9 @@ class DXILPrepareModule : public ModulePass {
     VersionTuple ValVer = MetadataInfo.ValidatorVersion;
     bool SkipValidation = ValVer.getMajor() == 0 && ValVer.getMinor() == 0;
 
+    // construct allowlist of valid metadata node kinds
+    std::array<unsigned, 6> DXILCompatibleMDs = getCompatibleInstructionMDs(M);
+
     for (auto &F : M.functions()) {
       F.removeFnAttrs(AttrMask);
       F.removeRetAttrs(AttrMask);
@@ -189,6 +201,9 @@ class DXILPrepareModule : public ModulePass {
       for (auto &BB : F) {
         IRBuilder<> Builder(&BB);
         for (auto &I : make_early_inc_range(BB)) {
+
+          I.dropUnknownNonDebugMetadata(DXILCompatibleMDs);
+
           if (I.getOpcode() == Instruction::FNeg) {
             Builder.SetInsertPoint(&I);
             Value *In = I.getOperand(0);

llvm/test/CodeGen/DirectX/metadata-stripping.ll

@@ -0,0 +1,32 @@
+; RUN: opt -S --dxil-prepare %s | FileCheck %s
+
+; Test that only metadata nodes that are valid in DXIL are allowed through
+
+target triple = "dxilv1.0-unknown-shadermodel6.0-compute"
+
+; Function Attrs: noinline nounwind memory(readwrite, inaccessiblemem: none)
+define void @main(i32* %ptr) {
+entry:  
+  ; metadata ID changes to 0 once the current !0 and !1 are removed
+  ; since they aren't in the allowlist. range needs a payload.
+  ; CHECK: %val = load i32, ptr %ptr, align 4, !range [[RANGEMD:![0-9]+]]
+  %val = load i32, ptr %ptr, align 4, !range !2
+
+  %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{{$}}
+  br i1 %cmp.i, label %_Z4mainDv3_j.exit, label %_Z4mainDv3_j.exit, !llvm.loop !0
+
+_Z4mainDv3_j.exit:                                ; preds = %for.body.i, %entry
+  ret void
+}
+
+; These next check lines check that only the range metadata remains
+; 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-NOT: !{!"llvm.loop.mustprogress"}
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.mustprogress"}
+!2 = !{i32 1, i32 5}