[spirv] Specify "KHR" to enable all KHR extensions (#1195)

docs/SPIR-V.rst

@@ -311,7 +311,8 @@ in the CodeGen using the **explicit** mode, turned on by the
 ``-fspv-extension=<extension>`` command-line option. Only extensions supplied
 via ``-fspv-extension=`` will be used. If that does not suffice, errors will
 be emitted explaining what additional extensions are required to translate what
-specific feature in the source code.
+specific feature in the source code. If you want to allow all KHR extensions, 
+you can use ``-fspv-extension=KHR``.
 
 Legalization, optimization, validation
 --------------------------------------
@@ -2651,7 +2652,8 @@ codegen for Vulkan:
   `HLSL semantic and Vulkan Location`_ for more details.
 - ``-fspv-reflect``: Emits additional SPIR-V instructions to aid reflection.
 - ``-fspv-extension=<extension>``: Only allows using ``<extension>`` in CodeGen.
-  If you want to allow multiple extensions, provide more than one such option.
+  If you want to allow multiple extensions, provide more than one such option. If you
+  want to allow *all* KHR extensions, use ``-fspv-extension=KHR``.
 - ``-fspv-target-env=<env>``: Specifies the target environment for this compilation.
   The current valid options are ``vulkan1.0`` and ``vulkan1.1``. If no target
   environment is provided, ``vulkan1.0`` is used as default.

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -30,6 +30,7 @@ namespace spirv {
 
 /// A list of SPIR-V extensions known to our CodeGen.
 enum class Extension {
+  KHR = 0,
   KHR_device_group,
   KHR_multiview,
   KHR_shader_draw_parameters,
@@ -60,6 +61,8 @@ public:
   static Extension getExtensionSymbol(llvm::StringRef name);
   /// Translates extension symbol to name.
   static const char *getExtensionName(Extension symbol);
+  /// Returns true if the given extension is a KHR extension.
+  static bool isKHRExtension(llvm::StringRef name);
 
   /// Returns the names of all known extensions as a string.
   std::string getKnownExtensions(const char *delimiter, const char *prefix = "",

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -40,6 +40,18 @@ FeatureManager::FeatureManager(DiagnosticsEngine &de,
 }
 
 bool FeatureManager::allowExtension(llvm::StringRef name) {
+  // Special case: If we are asked to allow "SPV_KHR" extension, it indicates
+  // that we should allow using *all* KHR extensions.
+  if (getExtensionSymbol(name) == Extension::KHR) {
+    bool result = true;
+    for (uint32_t i = 0; i < static_cast<uint32_t>(Extension::Unknown); ++i) {
+      llvm::StringRef extName(getExtensionName(static_cast<Extension>(i)));
+      if (isKHRExtension(extName))
+        result = result && allowExtension(extName);
+    }
+    return result;
+  }
+
   const auto symbol = getExtensionSymbol(name);
   if (symbol == Extension::Unknown) {
     emitError("unknown SPIR-V extension '%0'", {}) << name;
@@ -84,6 +96,7 @@ bool FeatureManager::requestTargetEnv(spv_target_env requestedEnv,
 
 Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
   return llvm::StringSwitch<Extension>(name)
+      .Case("KHR", Extension::KHR)
       .Case("SPV_KHR_device_group", Extension::KHR_device_group)
       .Case("SPV_KHR_multiview", Extension::KHR_multiview)
       .Case("SPV_KHR_shader_draw_parameters",
@@ -104,6 +117,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
 
 const char *FeatureManager::getExtensionName(Extension symbol) {
   switch (symbol) {
+  case Extension::KHR:
+    return "KHR";
   case Extension::KHR_device_group:
     return "SPV_KHR_device_group";
   case Extension::KHR_multiview:
@@ -128,6 +143,10 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
   return "<unknown extension>";
 }
 
+bool FeatureManager::isKHRExtension(llvm::StringRef name) {
+  return name.startswith_lower("spv_khr_");
+}
+
 std::string FeatureManager::getKnownExtensions(const char *delimiter,
                                                const char *prefix,
                                                const char *postfix) {

tools/clang/test/CodeGenSPIRV/spirv.ext.allow-all-khr.hlsl

@@ -0,0 +1,12 @@
+// Run: %dxc -T vs_6_1 -E main -fspv-extension=KHR
+
+// CHECK: OpExtension "SPV_KHR_shader_draw_parameters"
+// CHECK: OpExtension "SPV_KHR_device_group"
+// CHECK: OpExtension "SPV_KHR_multiview"
+
+float4 main(
+    [[vk::builtin("BaseVertex")]]    int baseVertex  : A,
+   [[vk::builtin("DeviceIndex")]]    int deviceIndex : C,
+                                     uint viewid     : SV_ViewID) : B {
+    return baseVertex + viewid;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -1185,6 +1185,9 @@ TEST_F(FileTest, SpirvExtensionCLForbid) {
 TEST_F(FileTest, SpirvExtensionCLUnknown) {
   runFileTest("spirv.ext.cl.unknown.hlsl", Expect::Failure);
 }
+TEST_F(FileTest, SpirvExtensionAllowAllKHR) {
+  runFileTest("spirv.ext.allow-all-khr.hlsl");
+}
 
 // For shader stage input/output interface
 // For semantic SV_Position, SV_ClipDistance, SV_CullDistance