[DXIL] Remove incompatible metadata types when preparing DXIL. (llvm#136386)

This PR introduces a Metadata Node Kind allowlist. The purpose is to
prevent newer Metadata Node Kinds to be used and inserted into the
outputted DXIL module. Only the metadata kinds that are accepted in the
DXIL Validator are on the allowlist. The Github DXC validator doesn't
support these newer Metadata Node Kinds, so we need to filter them out.

We introduce this restrictive allowlist into LLVM and strip all metadata
that isn't found in the list.

The accompanying test would add the `llvm.loop.mustprogress` metadata
node kind, but thanks to the allowlist, filters it out, and so the
whitelist is proven to work.
The test also has two separate metadata kinds that are on the allowlist,
and remain after the DXIL Prepare pass.

Author: bob80905

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}