[DXIL] Remove incompatible metadata types when preparing DXIL.

[bob80905]

This PR introduces a Metadata Node Kind whitelist. The purpose is to prevent newer Metadata Node Kinds to be used and inserted into the outputted DXIL module. The Github DXC validator doesn't support these newer Metadata Node Kinds, so we need to filter them out.
In Github DXC, lib\HLSL\DxilPreparePasses.cpp:814, there is a loop:

for (unsigned I = LLVMContext::MD_fpmath + 1;
     I <= LLVMContext::MD_dereferenceable_or_null; ++I) {
  IllegalMDSet.insert(I);
}

that loops up to and including the last metadata node kind that should be disallowed in the current context. I believe this code shows that there are no other metadata nodes to be allowed, setting aside the explicit list shown in lib\IR\LLVMContext.cpp:34
So, given that this is the world of acceptable metadata node kinds, we introduce this restrictive whitelist 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 whitelist, filters it out, and so the whitelist is proven to work.

[llvmbot]

@llvm/pr-subscribers-hlsl

@llvm/pr-subscribers-backend-directx

Author: Joshua Batista (bob80905)

Changes

This PR introduces a Metadata Node Kind whitelist. The purpose is to prevent newer Metadata Node Kinds to be used and inserted into the outputted DXIL module. The Github DXC validator doesn't support these newer Metadata Node Kinds, so we need to filter them out.
In Github DXC, lib\HLSL\DxilPreparePasses.cpp:814, there is a loop:

for (unsigned I = LLVMContext::MD_fpmath + 1;
     I <= LLVMContext::MD_dereferenceable_or_null; ++I) {
  IllegalMDSet.insert(I);
}

that loops up to and including the last metadata node kind that should be disallowed in the current context. I believe this code shows that there are no other metadata nodes to be allowed, setting aside the explicit list shown in lib\IR\LLVMContext.cpp:34
So, given that this is the world of acceptable metadata node kinds, we introduce this restrictive whitelist into LLVM and strip all metadata that isn't found in the list.

The accompanying test would add the loop.llvm.mustprogress metadata node kind, but thanks to the whitelist, filters it out, and so the whitelist is proven to work.

---

Full diff: https://github.com/llvm/llvm-project/pull/136386.diff

2 Files Affected:

• (added) clang/test/CodeGenHLSL/metadata-stripping.hlsl (+13)
• (modified) llvm/lib/Target/DirectX/DXILPrepare.cpp (+20)

diff --git a/clang/test/CodeGenHLSL/metadata-stripping.hlsl b/clang/test/CodeGenHLSL/metadata-stripping.hlsl
new file mode 100644
index 0000000000000..0178f5efe4c01
--- /dev/null
+++ b/clang/test/CodeGenHLSL/metadata-stripping.hlsl
@@ -0,0 +1,13 @@
+// RUN: %clang --driver-mode=dxc -T cs_6_0 -Fo x.dxil %s | FileCheck %s
+// CHECK-NOT: llvm.loop.mustprogress
+
+StructuredBuffer<uint4> X : register(t0);
+StructuredBuffer<float4> In : register(t1);
+RWStructuredBuffer<float4> Out : register(u0);
+
+[numthreads(1, 1, 1)]
+void main(uint3 dispatch_thread_id : SV_DispatchThreadID) {
+  for (uint I = 0; I < X[dispatch_thread_id].x; ++I) {
+    Out[dispatch_thread_id] = In[dispatch_thread_id];
+  }
+}
diff --git a/llvm/lib/Target/DirectX/DXILPrepare.cpp b/llvm/lib/Target/DirectX/DXILPrepare.cpp
index de97de209048b..c54a5ac5c1227 100644
--- a/llvm/lib/Target/DirectX/DXILPrepare.cpp
+++ b/llvm/lib/Target/DirectX/DXILPrepare.cpp
@@ -189,6 +189,26 @@ class DXILPrepareModule : public ModulePass {
       for (auto &BB : F) {
         IRBuilder<> Builder(&BB);
         for (auto &I : make_early_inc_range(BB)) {
+
+          // TODO: Audit this list - is it enough? Too much?
+          static unsigned DXILCompatibleMDs[] = {
+              LLVMContext::MD_dbg,
+              LLVMContext::MD_tbaa,
+              LLVMContext::MD_prof,
+              LLVMContext::MD_fpmath,
+              LLVMContext::MD_range,
+              LLVMContext::MD_tbaa_struct,
+              LLVMContext::MD_invariant_load,
+              LLVMContext::MD_alias_scope,
+              LLVMContext::MD_noalias,
+              LLVMContext::MD_nontemporal,
+              LLVMContext::MD_mem_parallel_loop_access,
+              LLVMContext::MD_nonnull,
+              LLVMContext::MD_dereferenceable,
+              LLVMContext::MD_dereferenceable_or_null,
+          };
+          I.dropUnknownNonDebugMetadata(DXILCompatibleMDs);
+
           if (I.getOpcode() == Instruction::FNeg) {
             Builder.SetInsertPoint(&I);
             Value *In = I.getOperand(0);

[llvmbot]

@llvm/pr-subscribers-clang

Author: Joshua Batista (bob80905)

Changes

This PR introduces a Metadata Node Kind whitelist. The purpose is to prevent newer Metadata Node Kinds to be used and inserted into the outputted DXIL module. The Github DXC validator doesn't support these newer Metadata Node Kinds, so we need to filter them out.
In Github DXC, lib\HLSL\DxilPreparePasses.cpp:814, there is a loop:

for (unsigned I = LLVMContext::MD_fpmath + 1;
     I <= LLVMContext::MD_dereferenceable_or_null; ++I) {
  IllegalMDSet.insert(I);
}

that loops up to and including the last metadata node kind that should be disallowed in the current context. I believe this code shows that there are no other metadata nodes to be allowed, setting aside the explicit list shown in lib\IR\LLVMContext.cpp:34
So, given that this is the world of acceptable metadata node kinds, we introduce this restrictive whitelist into LLVM and strip all metadata that isn't found in the list.

The accompanying test would add the loop.llvm.mustprogress metadata node kind, but thanks to the whitelist, filters it out, and so the whitelist is proven to work.

---

Full diff: https://github.com/llvm/llvm-project/pull/136386.diff

2 Files Affected:

• (added) clang/test/CodeGenHLSL/metadata-stripping.hlsl (+13)
• (modified) llvm/lib/Target/DirectX/DXILPrepare.cpp (+20)

diff --git a/clang/test/CodeGenHLSL/metadata-stripping.hlsl b/clang/test/CodeGenHLSL/metadata-stripping.hlsl
new file mode 100644
index 0000000000000..0178f5efe4c01
--- /dev/null
+++ b/clang/test/CodeGenHLSL/metadata-stripping.hlsl
@@ -0,0 +1,13 @@
+// RUN: %clang --driver-mode=dxc -T cs_6_0 -Fo x.dxil %s | FileCheck %s
+// CHECK-NOT: llvm.loop.mustprogress
+
+StructuredBuffer<uint4> X : register(t0);
+StructuredBuffer<float4> In : register(t1);
+RWStructuredBuffer<float4> Out : register(u0);
+
+[numthreads(1, 1, 1)]
+void main(uint3 dispatch_thread_id : SV_DispatchThreadID) {
+  for (uint I = 0; I < X[dispatch_thread_id].x; ++I) {
+    Out[dispatch_thread_id] = In[dispatch_thread_id];
+  }
+}
diff --git a/llvm/lib/Target/DirectX/DXILPrepare.cpp b/llvm/lib/Target/DirectX/DXILPrepare.cpp
index de97de209048b..c54a5ac5c1227 100644
--- a/llvm/lib/Target/DirectX/DXILPrepare.cpp
+++ b/llvm/lib/Target/DirectX/DXILPrepare.cpp
@@ -189,6 +189,26 @@ class DXILPrepareModule : public ModulePass {
       for (auto &BB : F) {
         IRBuilder<> Builder(&BB);
         for (auto &I : make_early_inc_range(BB)) {
+
+          // TODO: Audit this list - is it enough? Too much?
+          static unsigned DXILCompatibleMDs[] = {
+              LLVMContext::MD_dbg,
+              LLVMContext::MD_tbaa,
+              LLVMContext::MD_prof,
+              LLVMContext::MD_fpmath,
+              LLVMContext::MD_range,
+              LLVMContext::MD_tbaa_struct,
+              LLVMContext::MD_invariant_load,
+              LLVMContext::MD_alias_scope,
+              LLVMContext::MD_noalias,
+              LLVMContext::MD_nontemporal,
+              LLVMContext::MD_mem_parallel_loop_access,
+              LLVMContext::MD_nonnull,
+              LLVMContext::MD_dereferenceable,
+              LLVMContext::MD_dereferenceable_or_null,
+          };
+          I.dropUnknownNonDebugMetadata(DXILCompatibleMDs);
+
           if (I.getOpcode() == Instruction::FNeg) {
             Builder.SetInsertPoint(&I);
             Value *In = I.getOperand(0);

[bogner]

I don't think this list is correct. Looking at what we have here and considering what the DXIL validator handles in ValidateInstructionMetadata:

• MD_dbg probably doesn't do anything in this list, given the function is about "non-debug metadata"
• There are a couple of dxil-specific MD node types that we need to let through ("dx.controlflow.hints", "dx.precise", "dx.nonuniform")
• TBAA metadata is allowed but has further validations. We can probably just let it through for now though and file a follow up to make sure we're matching what the validator accepts.
• Loop metadata is allowed but has further validations. Unfortunately, we do need to do some work if we want to accept "llvm.loop.unroll.full" but not "llvm.loop.mustprogress"
• All other metadata types are rejected, except MD_range, MD_noalias, and MD_alias_scope

So I guess we want to break out a getCompatibleInstructionMDs() function that populates a list with the ones that are allowed (both simple enum ones like MD_noalias and ones we need to look up like Ctx.getMDKindID("dx.nonuniform"), and then call that outside of the loop and reuse it for the per-instruction calls to dropUnknownNonDebugMetadata

I think it makes sense to just do the above in this PR, and drop MD_tbaa and MD_loop for now, but file two issues to do the correct validation / downgrade for those two in particular.

[bogner]

This test can be simplified a lot. We can probably drop everything DXIL specific and reduce this down to just a function with a couple of blocks that are mostly empty. Notably only the br instruction with the !llvm.loop metadata is really essential.