From fc840ef7710b6015ad6ed4862caee2025714ca28 Mon Sep 17 00:00:00 2001
From: Damyan Pepper <damyanp@microsoft.com>
Date: Mon, 24 Aug 2026 18:54:07 -0700
Subject: [PATCH] [PIX] Recompute shader flags after NURI instrumentation

The non-uniform resource index pass inserts WaveActiveAllEqual to test whether a dynamic index is uniform. It inserts these calls after the compiler computed the shader flags. The module can therefore declare no wave operations and still contain them, and the validator refuses that module.

The pass computes the flags again, but only when it changed the module.

Assisted-by: Copilot

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 40dc9de3-617e-4caf-ab0d-fba0a033ed93
---
 ...NonUniformResourceIndexInstrumentation.cpp |  4 ++
 tools/clang/unittests/HLSL/PixTest.cpp        | 44 +++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/lib/DxilPIXPasses/DxilNonUniformResourceIndexInstrumentation.cpp b/lib/DxilPIXPasses/DxilNonUniformResourceIndexInstrumentation.cpp
index d1e7fa0fd8..ef58abe355 100644
--- a/lib/DxilPIXPasses/DxilNonUniformResourceIndexInstrumentation.cpp
+++ b/lib/DxilPIXPasses/DxilNonUniformResourceIndexInstrumentation.cpp
@@ -161,6 +161,10 @@ bool DxilNonUniformResourceIndexInstrumentation::runOnModule(Module &M) {
   PIXPassHelpers::EraseIfUnused(DM, AtomicOpFunc);
 
   if (modified) {
+    // Recompute shader flags after inserting WaveActiveAllEqual so the
+    // declared flags match the module.
+    DM.CollectShaderFlagsForModule();
+
     DM.ReEmitDxilResources();
 
     if (OSOverride != nullptr) {
diff --git a/tools/clang/unittests/HLSL/PixTest.cpp b/tools/clang/unittests/HLSL/PixTest.cpp
index d455fcfea5..239b6900be 100644
--- a/tools/clang/unittests/HLSL/PixTest.cpp
+++ b/tools/clang/unittests/HLSL/PixTest.cpp
@@ -184,6 +184,7 @@ class PixTest : public ::testing::Test {
   TEST_METHOD(Validation_ControlValidModulePasses)
   TEST_METHOD(Validation_ControlInvalidModuleFails)
   TEST_METHOD(Validation_ControlBoilerplateOnlyFailureIsRejected)
+  TEST_METHOD(Validation_NonUniformResourceIndex_WaveOpsFlag)
 
   dxc::DxCompilerDllLoader m_dllSupport;
   VersionSupportInfo m_ver;
@@ -4350,3 +4351,46 @@ TEST_F(PixTest, Validation_ControlBoilerplateOnlyFailureIsRejected) {
   VERIFY_IS_FALSE(realDiagnostic.Significant.empty());
   VERIFY_IS_FALSE(IsPermittedValidationException(realDiagnostic));
 }
+
+TEST_F(PixTest, Validation_NonUniformResourceIndex_WaveOpsFlag) {
+  const char *source = R"x(
+Texture2D textures[]  : register(t0);
+SamplerState samp     : register(s0);
+
+cbuffer Constants : register(b0)
+{
+    uint index;
+};
+
+float4 main(float4 pos : SV_Position) : SV_Target
+{
+    return textures[index].Sample(samp, pos.xy);
+})x";
+
+  // This index is dynamic and unmarked, so the pass instruments it; an
+  // index already marked NonUniformResourceIndex would be skipped.
+  // Instrumentation inserts WaveActiveAllEqual, which requires the WaveOps
+  // shader flag.
+  auto compiled = Compile(m_dllSupport, source, L"ps_6_6", {L"-Od"});
+  CComPtr<IDxcBlob> dxil = FindModule(DFCC_ShaderDebugInfoDXIL, compiled);
+
+  CComPtr<IDxcOptimizer> pOptimizer;
+  VERIFY_SUCCEEDED(
+      m_dllSupport.CreateInstance(CLSID_DxcOptimizer, &pOptimizer));
+  std::array<LPCWSTR, 4> Options = {
+      L"-opt-mod-passes", L"-dxil-dbg-value-to-dbg-declare",
+      L"-dxil-annotate-with-virtual-regs",
+      L"-hlsl-dxil-non-uniform-resource-index-instrumentation"};
+
+  CComPtr<IDxcBlob> pOptimizedModule;
+  CComPtr<IDxcBlobEncoding> pText;
+  VERIFY_SUCCEEDED(pOptimizer->RunOptimizer(
+      dxil, Options.data(), Options.size(), &pOptimizedModule, &pText));
+
+  VerifyInstrumentedModuleIsValid(pOptimizedModule,
+                                  "non-uniform resource index instrumentation");
+
+  VERIFY_ARE_NOT_EQUAL(
+      std::string::npos,
+      Disassemble(pOptimizedModule).find("dx.op.waveActiveAllEqual"));
+}