[SPIRV][HLSL] Add DXC compatibility option for extension (#151554)

The default behaviour in DXC is to allow all extesions the compiler
knows about. We did the same in clang: all extensions that clang knows
about. However, this causes the shader to use different extensions
because the two compilers have different sets of extensions.

To avoid using a new extension when moving from DXC to Clang, we add the
special DXC suboptions to `-fspv-extension`. If `-fspv-extension=DXC` is
used, then the available extensions will be those available in DXC.

---------

Co-authored-by: Chris B <[email protected]>

Author: s-perron

clang/include/clang/Driver/Options.td

@@ -9468,8 +9468,12 @@ def fspv_target_env_EQ : Joined<["-"], "fspv-target-env=">, Group<dxc_Group>,
 def fspv_extension_EQ
     : Joined<["-"], "fspv-extension=">,
       Group<dxc_Group>,
-      HelpText<"Specify the available SPIR-V extensions. If this option is not "
-               "specified, then all extensions are available.">;
+      HelpText<
+          "Specify the available SPIR-V extensions. If this option is not "
+          "specified, then all extensions are available. If KHR is specified, "
+          "then all KHR extensions will be available. If DXC is specifided, "
+          "then all extensions implemented by the DirectX Shader compiler will "
+          "be available. This option is useful for moving from DXC to Clang.">;
 def fvk_use_dx_layout
     : DXCFlag<"fvk-use-dx-layout">,
       HelpText<"Use DirectX memory layout for Vulkan resources.">;

clang/lib/Driver/ToolChains/HLSL.cpp

@@ -173,24 +173,72 @@ bool isLegalValidatorVersion(StringRef ValVersionStr, const Driver &D) {
   return true;
 }
 
-std::string getSpirvExtArg(ArrayRef<std::string> SpvExtensionArgs) {
+void getSpirvExtOperand(StringRef SpvExtensionArg, raw_ostream &out) {
+  // The extensions that are commented out are supported in DXC, but the SPIR-V
+  // backend does not know about them yet.
+  static const std::vector<StringRef> DxcSupportedExtensions = {
+      "SPV_KHR_16bit_storage", "SPV_KHR_device_group",
+      "SPV_KHR_fragment_shading_rate", "SPV_KHR_multiview",
+      "SPV_KHR_post_depth_coverage", "SPV_KHR_non_semantic_info",
+      "SPV_KHR_shader_draw_parameters", "SPV_KHR_ray_tracing",
+      "SPV_KHR_shader_clock", "SPV_EXT_demote_to_helper_invocation",
+      "SPV_EXT_descriptor_indexing", "SPV_EXT_fragment_fully_covered",
+      "SPV_EXT_fragment_invocation_density",
+      "SPV_EXT_fragment_shader_interlock", "SPV_EXT_mesh_shader",
+      "SPV_EXT_shader_stencil_export", "SPV_EXT_shader_viewport_index_layer",
+      // "SPV_AMD_shader_early_and_late_fragment_tests",
+      "SPV_GOOGLE_hlsl_functionality1", "SPV_GOOGLE_user_type",
+      "SPV_KHR_ray_query", "SPV_EXT_shader_image_int64",
+      "SPV_KHR_fragment_shader_barycentric", "SPV_KHR_physical_storage_buffer",
+      "SPV_KHR_vulkan_memory_model",
+      // "SPV_KHR_compute_shader_derivatives",
+      // "SPV_KHR_maximal_reconvergence",
+      "SPV_KHR_float_controls", "SPV_NV_shader_subgroup_partitioned",
+      // "SPV_KHR_quad_control"
+  };
+
+  if (SpvExtensionArg.starts_with("SPV_")) {
+    out << "+" << SpvExtensionArg;
+    return;
+  }
+
+  if (SpvExtensionArg.compare_insensitive("DXC") == 0) {
+    bool first = true;
+    std::string Operand;
+    for (StringRef E : DxcSupportedExtensions) {
+      if (!first)
+        out << ",";
+      else
+        first = false;
+      out << "+" << E;
+    }
+    return;
+  }
+  out << SpvExtensionArg;
+  return;
+}
+
+SmallString<1024> getSpirvExtArg(ArrayRef<std::string> SpvExtensionArgs) {
   if (SpvExtensionArgs.empty()) {
-    return "-spirv-ext=all";
+    return StringRef("-spirv-ext=all");
   }
 
-  std::string LlvmOption =
-      (Twine("-spirv-ext=+") + SpvExtensionArgs.front()).str();
+  llvm::SmallString<1024> LlvmOption;
+  raw_svector_ostream out(LlvmOption);
+
+  out << "-spirv-ext=";
+  getSpirvExtOperand(SpvExtensionArgs[0], out);
+
   SpvExtensionArgs = SpvExtensionArgs.slice(1);
-  for (auto Extension : SpvExtensionArgs) {
-    if (Extension != "KHR")
-      Extension = (Twine("+") + Extension).str();
-    LlvmOption = (Twine(LlvmOption) + "," + Extension).str();
+  for (StringRef Extension : SpvExtensionArgs) {
+    out << ",";
+    getSpirvExtOperand(Extension, out);
   }
   return LlvmOption;
 }
 
 bool isValidSPIRVExtensionName(const std::string &str) {
-  std::regex pattern("KHR|SPV_[a-zA-Z0-9_]+");
+  std::regex pattern("dxc|DXC|khr|KHR|SPV_[a-zA-Z0-9_]+");
   return std::regex_match(str, pattern);
 }
 
@@ -371,7 +419,7 @@ HLSLToolChain::TranslateArgs(const DerivedArgList &Args, StringRef BoundArch,
     std::vector<std::string> SpvExtensionArgs =
         Args.getAllArgValues(options::OPT_fspv_extension_EQ);
     if (checkExtensionArgsAreValid(SpvExtensionArgs, getDriver())) {
-      std::string LlvmOption = getSpirvExtArg(SpvExtensionArgs);
+      SmallString<1024> LlvmOption = getSpirvExtArg(SpvExtensionArgs);
       DAL->AddSeparateArg(nullptr, Opts.getOption(options::OPT_mllvm),
                           LlvmOption);
     }

clang/test/Driver/dxc_fspv_extension.hlsl

@@ -14,6 +14,14 @@
 // RUN: %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=SPV_TEST1 -fspv-extension=KHR -fspv-extension=SPV_TEST2 2>&1 | FileCheck %s -check-prefix=TEST3
 // TEST3: "-spirv-ext=+SPV_TEST1,KHR,+SPV_TEST2"
 
+// Merge KHR with other extensions.
+// RUN: %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=KHR -fspv-extension=SPV_TEST2 2>&1 | FileCheck %s -check-prefix=TEST4
+// TEST4: "-spirv-ext=KHR,+SPV_TEST2"
+
+// Convert DXC to a list of SPV extensions.
+// RUN: %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=DXC 2>&1 | FileCheck %s -check-prefix=TEST5
+// TEST5: "-spirv-ext={{(\+SPV_[a-zA-Z0-9_]+,?)+}}"
+
 // Check for the error message if the extension name is not properly formed.
 // RUN: not %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=KHR_BAD -fspv-extension=TEST1 -fspv-extension=SPV_GOOD -fspv-extension=TEST2 2>&1 | FileCheck %s -check-prefix=FAIL
 // FAIL: invalid value 'KHR_BAD' in '-fspv-extension'

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -33,6 +33,24 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::SPV_EXT_arithmetic_fence},
         {"SPV_EXT_demote_to_helper_invocation",
          SPIRV::Extension::Extension::SPV_EXT_demote_to_helper_invocation},
+        {"SPV_EXT_descriptor_indexing",
+         SPIRV::Extension::Extension::SPV_EXT_descriptor_indexing},
+        {"SPV_EXT_fragment_fully_covered",
+         SPIRV::Extension::Extension::SPV_EXT_fragment_fully_covered},
+        {"SPV_EXT_fragment_invocation_density",
+         SPIRV::Extension::Extension::SPV_EXT_fragment_invocation_density},
+        {"SPV_EXT_fragment_shader_interlock",
+         SPIRV::Extension::Extension::SPV_EXT_fragment_shader_interlock},
+        {"SPV_EXT_mesh_shader",
+         SPIRV::Extension::Extension::SPV_EXT_mesh_shader},
+        {"SPV_EXT_shader_stencil_export",
+         SPIRV::Extension::Extension::SPV_EXT_shader_stencil_export},
+        {"SPV_EXT_shader_viewport_index_layer",
+         SPIRV::Extension::Extension::SPV_EXT_shader_viewport_index_layer},
+        {"SPV_GOOGLE_hlsl_functionality1",
+         SPIRV::Extension::Extension::SPV_GOOGLE_hlsl_functionality1},
+        {"SPV_GOOGLE_user_type",
+         SPIRV::Extension::Extension::SPV_GOOGLE_user_type},
         {"SPV_INTEL_arbitrary_precision_integers",
          SPIRV::Extension::Extension::SPV_INTEL_arbitrary_precision_integers},
         {"SPV_INTEL_cache_controls",
@@ -58,6 +76,19 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::SPV_INTEL_memory_access_aliasing},
         {"SPV_INTEL_joint_matrix",
          SPIRV::Extension::Extension::SPV_INTEL_joint_matrix},
+        {"SPV_KHR_16bit_storage",
+         SPIRV::Extension::Extension::SPV_KHR_16bit_storage},
+        {"SPV_KHR_device_group",
+         SPIRV::Extension::Extension::SPV_KHR_device_group},
+        {"SPV_KHR_fragment_shading_rate",
+         SPIRV::Extension::Extension::SPV_KHR_fragment_shading_rate},
+        {"SPV_KHR_multiview", SPIRV::Extension::Extension::SPV_KHR_multiview},
+        {"SPV_KHR_post_depth_coverage",
+         SPIRV::Extension::Extension::SPV_KHR_post_depth_coverage},
+        {"SPV_KHR_shader_draw_parameters",
+         SPIRV::Extension::Extension::SPV_KHR_shader_draw_parameters},
+        {"SPV_KHR_ray_tracing",
+         SPIRV::Extension::Extension::SPV_KHR_ray_tracing},
         {"SPV_KHR_uniform_group_instructions",
          SPIRV::Extension::Extension::SPV_KHR_uniform_group_instructions},
         {"SPV_KHR_no_integer_wrap_decoration",
@@ -90,6 +121,17 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::SPV_KHR_cooperative_matrix},
         {"SPV_KHR_non_semantic_info",
          SPIRV::Extension::Extension::SPV_KHR_non_semantic_info},
+        {"SPV_KHR_ray_query", SPIRV::Extension::Extension::SPV_KHR_ray_query},
+        {"SPV_EXT_shader_image_int64",
+         SPIRV::Extension::Extension::SPV_EXT_shader_image_int64},
+        {"SPV_KHR_fragment_shader_barycentric",
+         SPIRV::Extension::Extension::SPV_KHR_fragment_shader_barycentric},
+        {"SPV_KHR_physical_storage_buffer",
+         SPIRV::Extension::Extension::SPV_KHR_physical_storage_buffer},
+        {"SPV_KHR_vulkan_memory_model",
+         SPIRV::Extension::Extension::SPV_KHR_vulkan_memory_model},
+        {"SPV_NV_shader_subgroup_partitioned",
+         SPIRV::Extension::Extension::SPV_NV_shader_subgroup_partitioned},
         {"SPV_INTEL_long_composites",
          SPIRV::Extension::Extension::SPV_INTEL_long_composites},
         {"SPV_INTEL_fp_max_error",