[PIX] Recompute shader flags after NURI instrumentation

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) {

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"));
+}