[SPIRV] Filter disallowed extensions for env (#150051)

Not all SPIR-V extensions are allows in every environment. When we use
the `-spirv-ext=all` option, the backend currently believes that all
extensions can be used.

This commit filters out the extensions on the command line to remove
those that are not known to be allowed for the current environment.

Alternatives considered: I considered modifying the
SPIRVExtensionsParser::parse to use a different list of extensions for
"all" depending on the target triple. However that does not work because
the target triple is not available, and cannot be made available in a
reasonable way.

Fixes #147717

---------

Co-authored-by: Victor Lomuller <[email protected]>

Author: s-perron

llvm/lib/Target/SPIRV/MCTargetDesc/SPIRVBaseInfo.cpp

@@ -38,8 +38,15 @@ struct CapabilityEntry {
   Capability::Capability ReqCapability;
 };
 
+struct EnvironmentEntry {
+  OperandCategory::OperandCategory Category;
+  uint32_t Value;
+  Environment::Environment AllowedEnvironment;
+};
+
 using namespace OperandCategory;
 using namespace Extension;
+using namespace Environment;
 using namespace Capability;
 using namespace InstructionSet;
 #define GET_SymbolicOperands_DECL
@@ -48,6 +55,8 @@ using namespace InstructionSet;
 #define GET_ExtensionEntries_IMPL
 #define GET_CapabilityEntries_DECL
 #define GET_CapabilityEntries_IMPL
+#define GET_EnvironmentEntries_DECL
+#define GET_EnvironmentEntries_IMPL
 #define GET_ExtendedBuiltins_DECL
 #define GET_ExtendedBuiltins_IMPL
 #include "SPIRVGenTables.inc"
@@ -133,6 +142,23 @@ getSymbolicOperandCapabilities(SPIRV::OperandCategory::OperandCategory Category,
   return Capabilities;
 }
 
+EnvironmentList getSymbolicOperandAllowedEnvironments(
+    SPIRV::OperandCategory::OperandCategory Category, uint32_t Value) {
+  EnvironmentList Environments;
+  const SPIRV::EnvironmentEntry *Environment =
+      SPIRV::lookupEnvironmentByCategoryAndValue(Category, Value);
+  auto TableEnd = ArrayRef(SPIRV::EnvironmentEntries).end();
+  while (Environment && Environment->Category == Category &&
+         Environment->Value == Value) {
+    Environments.push_back(static_cast<SPIRV::Environment::Environment>(
+        Environment->AllowedEnvironment));
+    if (++Environment == TableEnd)
+      break;
+  }
+
+  return Environments;
+}
+
 CapabilityList
 getCapabilitiesEnabledByExtension(SPIRV::Extension::Extension Extension) {
   const SPIRV::ExtensionEntry *Entry =

llvm/lib/Target/SPIRV/MCTargetDesc/SPIRVBaseInfo.h

@@ -37,6 +37,11 @@ namespace Capability {
 #include "SPIRVGenTables.inc"
 } // namespace Capability
 
+namespace Environment {
+#define GET_Environment_DECL
+#include "SPIRVGenTables.inc"
+} // namespace Environment
+
 namespace SourceLanguage {
 #define GET_SourceLanguage_DECL
 #include "SPIRVGenTables.inc"
@@ -241,6 +246,7 @@ enum InstFlags {
 
 using CapabilityList = SmallVector<SPIRV::Capability::Capability, 8>;
 using ExtensionList = SmallVector<SPIRV::Extension::Extension, 8>;
+using EnvironmentList = SmallVector<SPIRV::Environment::Environment, 8>;
 
 std::string
 getSymbolicOperandMnemonic(SPIRV::OperandCategory::OperandCategory Category,
@@ -254,6 +260,8 @@ getSymbolicOperandMaxVersion(SPIRV::OperandCategory::OperandCategory Category,
 CapabilityList
 getSymbolicOperandCapabilities(SPIRV::OperandCategory::OperandCategory Category,
                                uint32_t Value);
+EnvironmentList getSymbolicOperandAllowedEnvironments(
+    SPIRV::OperandCategory::OperandCategory Category, uint32_t Value);
 CapabilityList
 getCapabilitiesEnabledByExtension(SPIRV::Extension::Extension Extension);
 ExtensionList

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -12,7 +12,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "SPIRVCommandLine.h"
-#include "llvm/ADT/StringRef.h"
+#include "MCTargetDesc/SPIRVBaseInfo.h"
+#include "llvm/TargetParser/Triple.h"
 #include <algorithm>
 #include <map>
 
@@ -171,3 +172,23 @@ StringRef SPIRVExtensionsParser::checkExtensions(
   }
   return StringRef();
 }
+
+std::set<SPIRV::Extension::Extension>
+SPIRVExtensionsParser::getValidExtensions(const Triple &TT) {
+  std::set<SPIRV::Extension::Extension> R;
+  SPIRV::Environment::Environment CurrentEnvironment =
+      SPIRV::Environment::Environment::EnvOpenCL;
+  if (TT.getOS() == Triple::Vulkan)
+    CurrentEnvironment = SPIRV::Environment::Environment::EnvVulkan;
+
+  for (const auto &[ExtensionName, ExtensionEnum] : SPIRVExtensionMap) {
+    EnvironmentList AllowedEnv = getSymbolicOperandAllowedEnvironments(
+        SPIRV::OperandCategory::OperandCategory::ExtensionOperand,
+        ExtensionEnum);
+
+    if (std::count(AllowedEnv.begin(), AllowedEnv.end(), CurrentEnvironment))
+      R.insert(ExtensionEnum);
+  }
+
+  return R;
+}

llvm/lib/Target/SPIRV/SPIRVCommandLine.h

@@ -21,6 +21,7 @@
 
 namespace llvm {
 class StringRef;
+class Triple;
 
 /// Command line parser for toggling SPIR-V extensions.
 struct SPIRVExtensionsParser
@@ -42,6 +43,11 @@ struct SPIRVExtensionsParser
   static StringRef
   checkExtensions(const std::vector<std::string> &ExtNames,
                   std::set<SPIRV::Extension::Extension> &AllowedExtensions);
+
+  /// Returns the list of extensions that are valid for a particular
+  /// target environment (i.e., OpenCL or Vulkan).
+  static std::set<SPIRV::Extension::Extension>
+  getValidExtensions(const Triple &TT);
 };
 
 } // namespace llvm

llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp

@@ -166,7 +166,13 @@ void SPIRVSubtarget::initAvailableExtInstSets() {
 void SPIRVSubtarget::initAvailableExtensions(
     const std::set<SPIRV::Extension::Extension> &AllowedExtIds) {
   AvailableExtensions.clear();
-  AvailableExtensions.insert_range(AllowedExtIds);
+  const std::set<SPIRV::Extension::Extension> &ValidExtensions =
+      SPIRVExtensionsParser::getValidExtensions(TargetTriple);
+
+  for (const auto &Ext : AllowedExtIds) {
+    if (ValidExtensions.count(Ext))
+      AvailableExtensions.insert(Ext);
+  }
 
   accountForAMDShaderTrinaryMinmax();
 }