[llvm] [SPIR-V] Fix removeCapabilityIf not pruning MinimalCaps (PR #206478)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/206478

removeCapabilityIf erased from AllCaps but not MinimalCaps, which is what the AsmPrinter iterates to emit OpCapability, so pruned capabilities were still emitted

>From 79960f9f53cfe1d508ee5ce4404821a28efcaeb4 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Mon, 29 Jun 2026 15:04:01 +0200
Subject: [PATCH] [SPIR-V] Fix removeCapabilityIf not pruning MinimalCaps

removeCapabilityIf erased from AllCaps but not MinimalCaps, which is what the AsmPrinter iterates to emit OpCapability, so pruned capabilities were still emitted
---
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp |  4 +-
 llvm/unittests/Target/SPIRV/CMakeLists.txt    |  1 +
 .../Target/SPIRV/SPIRVModuleAnalysisTests.cpp | 42 +++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 llvm/unittests/Target/SPIRV/SPIRVModuleAnalysisTests.cpp

diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index aed16fd785af8..be6e5c26b7bc7 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -939,8 +939,10 @@ void SPIRV::RequirementHandler::addAvailableCaps(const CapabilityList &ToAdd) {
 void SPIRV::RequirementHandler::removeCapabilityIf(
     const Capability::Capability ToRemove,
     const Capability::Capability IfPresent) {
-  if (AllCaps.contains(IfPresent))
+  if (AllCaps.contains(IfPresent)) {
     AllCaps.erase(ToRemove);
+    llvm::erase(MinimalCaps, ToRemove);
+  }
 }
 
 namespace llvm {
diff --git a/llvm/unittests/Target/SPIRV/CMakeLists.txt b/llvm/unittests/Target/SPIRV/CMakeLists.txt
index e20adc886106c..1c73c5c2ffce1 100644
--- a/llvm/unittests/Target/SPIRV/CMakeLists.txt
+++ b/llvm/unittests/Target/SPIRV/CMakeLists.txt
@@ -22,6 +22,7 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_target_unittest(SPIRVTests
   SPIRVConvergenceRegionAnalysisTests.cpp
   SPIRVGlobalRegistryTests.cpp
+  SPIRVModuleAnalysisTests.cpp
   SPIRVSortBlocksTests.cpp
   SPIRVPartialOrderingVisitorTests.cpp
   SPIRVAPITest.cpp
diff --git a/llvm/unittests/Target/SPIRV/SPIRVModuleAnalysisTests.cpp b/llvm/unittests/Target/SPIRV/SPIRVModuleAnalysisTests.cpp
new file mode 100644
index 0000000000000..f0af1cc21e027
--- /dev/null
+++ b/llvm/unittests/Target/SPIRV/SPIRVModuleAnalysisTests.cpp
@@ -0,0 +1,42 @@
+//===- SPIRVModuleAnalysisTests.cpp ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRVModuleAnalysis.h"
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+// A capability pruned by removeCapabilityIf must be dropped from the minimal
+// capabilities too, since those are the ones the AsmPrinter emits as
+// OpCapability.
+TEST(SPIRVModuleAnalysisTest, RemoveCapabilityIfPrunesMinimalCaps) {
+  SPIRV::RequirementHandler Reqs;
+  Reqs.addCapability(SPIRV::Capability::BitInstructions);
+  Reqs.addCapability(SPIRV::Capability::Shader);
+
+  Reqs.removeCapabilityIf(SPIRV::Capability::BitInstructions,
+                          SPIRV::Capability::Shader);
+
+  EXPECT_FALSE(llvm::is_contained(Reqs.getMinimalCapabilities(),
+                                  SPIRV::Capability::BitInstructions));
+  EXPECT_TRUE(llvm::is_contained(Reqs.getMinimalCapabilities(),
+                                 SPIRV::Capability::Shader));
+}
+
+// When the guarding capability is absent, nothing is removed.
+TEST(SPIRVModuleAnalysisTest, RemoveCapabilityIfNoopWhenGuardAbsent) {
+  SPIRV::RequirementHandler Reqs;
+  Reqs.addCapability(SPIRV::Capability::BitInstructions);
+
+  Reqs.removeCapabilityIf(SPIRV::Capability::BitInstructions,
+                          SPIRV::Capability::Shader);
+
+  EXPECT_TRUE(llvm::is_contained(Reqs.getMinimalCapabilities(),
+                                 SPIRV::Capability::BitInstructions));
+}