[HLSL][SPIRV] Add CLI option -fspv-extension

clang/include/clang/Driver/Options.td

@@ -9213,6 +9213,11 @@ def metal : DXCFlag<"metal">, HelpText<"Generate Metal library">;
 def fspv_target_env_EQ : Joined<["-"], "fspv-target-env=">, Group<dxc_Group>,
   HelpText<"Specify the target environment">,
   Values<"vulkan1.2, vulkan1.3">;
+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.">;
 def no_wasm_opt : Flag<["--"], "no-wasm-opt">,
   Group<m_Group>,
   HelpText<"Disable the wasm-opt optimizer">,

clang/lib/Driver/ToolChains/HLSL.cpp

@@ -13,6 +13,7 @@
 #include "clang/Driver/Job.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/TargetParser/Triple.h"
+#include <regex>
 
 using namespace clang::driver;
 using namespace clang::driver::tools;
@@ -173,6 +174,39 @@ bool isLegalValidatorVersion(StringRef ValVersionStr, const Driver &D) {
   return true;
 }
 
+std::string getSpirvExtArg(ArrayRef<std::string> SpvExtensionArgs) {
+  if (SpvExtensionArgs.empty()) {
+    return "-spirv-ext=all";
+  }
+
+  std::string LlvmOption =
+      (Twine("-spirv-ext=+") + SpvExtensionArgs.front()).str();
+  SpvExtensionArgs = SpvExtensionArgs.slice(1);
+  for (auto Extension : SpvExtensionArgs) {
+    LlvmOption = (Twine(LlvmOption) + ",+" + Extension).str();
+  }
+  return LlvmOption;
+}
+
+bool isValidSPIRVExtensionName(const std::string &str) {
+  std::regex pattern("SPV_[a-zA-Z0-9_]+");
+  return std::regex_match(str, pattern);
+}
+
+// SPIRV extension names are of the form `SPV_[a-zA-Z0-9_]+`. We want to
+// disallow obviously invalid names to avoid issues when parsing `spirv-ext`.
+bool checkExtensionArgsAreValid(ArrayRef<std::string> SpvExtensionArgs,
+                                const Driver &Driver) {
+  bool AllValid = true;
+  for (auto Extension : SpvExtensionArgs) {
+    if (!isValidSPIRVExtensionName(Extension)) {
+      Driver.Diag(diag::err_drv_invalid_value)
+          << "-fspv_extension" << Extension;
+      AllValid = false;
+    }
+  }
+  return AllValid;
+}
 } // namespace
 
 void tools::hlsl::Validator::ConstructJob(Compilation &C, const JobAction &JA,
@@ -301,6 +335,17 @@ HLSLToolChain::TranslateArgs(const DerivedArgList &Args, StringRef BoundArch,
     DAL->append(A);
   }
 
+  if (getArch() == llvm::Triple::spirv) {
+    std::vector<std::string> SpvExtensionArgs =
+        Args.getAllArgValues(options::OPT_fspv_extension_EQ);
+    if (checkExtensionArgsAreValid(SpvExtensionArgs, getDriver())) {
+      std::string LlvmOption = getSpirvExtArg(SpvExtensionArgs);
+      DAL->AddSeparateArg(nullptr, Opts.getOption(options::OPT_mllvm),
+                          LlvmOption);
+    }
+    Args.claimAllArgs(options::OPT_fspv_extension_EQ);
+  }
+
   if (!DAL->hasArg(options::OPT_O_Group)) {
     DAL->AddJoinedArg(nullptr, Opts.getOption(options::OPT_O), "3");
   }

clang/test/Driver/dxc_fspv_extension.hlsl

@@ -0,0 +1,25 @@
+// If there is no `-fspv-extension`, enable all extensions.
+// RUN: %clang_dxc -spirv -Tlib_6_7 -### %s 2>&1 | FileCheck %s -check-prefix=ALL
+// ALL: "-spirv-ext=all"
+
+// Convert the `-fspv-extension` into `spirv-ext`.
+// RUN: %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=SPV_TEST1 2>&1 | FileCheck %s -check-prefix=TEST1
+// TEST1: "-spirv-ext=+SPV_TEST1"
+
+// Merge multiple extensions into a single `spirv-ext` option.
+// RUN: %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=SPV_TEST1 -fspv-extension=SPV_TEST2 2>&1 | FileCheck %s -check-prefix=TEST2
+// TEST2: "-spirv-ext=+SPV_TEST1,+SPV_TEST2"
+
+// Check for the error message if the extension name is not properly formed.
+// RUN: not %clang_dxc -spirv -Tlib_6_7 -### %s -fspv-extension=TEST1 -fspv-extension=SPV_GOOD -fspv-extension=TEST2 2>&1 | FileCheck %s -check-prefix=FAIL
+// FAIL: invalid value 'TEST1' in '-fspv_extension'
+// FAIL: invalid value 'TEST2' in '-fspv_extension'
+
+// If targeting DXIL, the `-spirv-ext` should not be passed to the backend.
+// RUN: %clang_dxc -Tlib_6_7 -### %s 2>&1 | FileCheck %s -check-prefix=DXIL
+// DXIL-NOT: spirv-ext
+
+// If targeting DXIL, the `-fspv-extension` option is meaningless, and there should be a warning.
+// RUN: %clang_dxc -Tlib_6_7 -### %s -fspv-extension=SPV_TEST 2>&1 | FileCheck %s -check-prefix=WARN
+// WARN: warning: argument unused during compilation: '-fspv-extension=SPV_TEST'
+// WARN-NOT: spirv-ext