Add support for SPV_INTEL_function_variants (KhronosGroup#3246)

This PR implements
[SPV_INTEL_function_variants](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/spirv-extensions/SPV_INTEL_function_variants.asciidoc).
It adds an optional SPIR-V to SPIR-V specialization pass that converts a
multitarget module into a targeted one.

The multitarget module does not have a LLVM IR representation, the
extension only describes the specialization algorithm that takes place
before converting the SPIR-V module into LLVM-IR. For this reason, it is
only implemented as a part of SPIRVReader and not SPIRVWriter.

The specialization is controlled by the user supplying the target device
category, family, architecture, target ISA, supported features and/or
supported capabilities via CLI flags. For example, to specialize for an
Intel x86_64 CPU with Lion Cove microarchitecture that supports SSE,
SSE2, SSE3, SSE4.1, SSE4.2, SSE4a, AVX, AVX2 and AVX512f features and
Addresses, Linkage, Kernel, Int64 and Int8 capabilities, the user needs
to provide the following flags:
```
llvm-spirv -r \
  --spirv-ext=+SPV_INTEL_function_variants \
  --fnvar-spec-enable \
  --fnvar-spv-out targeted.spv \
  --fnvar-category 1  --fnvar-family 1  --fnvar-arch 15 \
  --fnvar-target 4  --fnvar-features '4,5,6,7,8,9,10,11,12' \
  --fnvar-capabilities '4,5,6,11,39' \
  multitarget.spv -o targeted.bc
```
Omitting a flag means that the target device supports all values for the
flag. For example, in the above example, leaving out the
`--fnvar-features` flag means that that the target device supports all
features available for the x86_64 target.

The integer values passed to the CLI flags are taken from a proposed
[targets _registry_](intel/llvm#18822)
accompanying the extension. (Capabilities correspond directly to the
values defined in the SPIR-V specification). During the specialization
pass, the specialization pass compares these CLI-supplied integers with
the operands of `OpSpecConstantTargetINTEL`,
`OpSpecConstantArchitectureINTEL` and `OpSpecConstantCapabilitiesINTEL`
instructions in the input multitarget module, converts these
instructions to constant true/false and proceeds with the specialization
according to the rules described in the extension.

Providing the CLI values as raw integer is not the most user friendly,
and the translator does not validate the values in any way (eg.,
checking that feature X is allowed for target Y). This can be improved
after the _registry_ is merged and more mature (version >0).

Note: `--spirv-debug` can be used to print out details about what's
happening when evaluating the above spec constants. It's useful for
getting an insight into why a certain function variant got selected if
the selection does not match the expected outcome.

Author: jzadnik

include/LLVMSPIRVExtensions.inc

@@ -80,3 +80,4 @@ EXT(SPV_INTEL_subgroup_matrix_multiply_accumulate)
 EXT(SPV_KHR_bfloat16)
 EXT(SPV_INTEL_ternary_bitwise_function)
 EXT(SPV_INTEL_int4)
+EXT(SPV_INTEL_function_variants)

include/LLVMSPIRVOpts.h

@@ -252,6 +252,50 @@ class TranslatorOpts {
   void setUseLLVMTarget(bool Flag) noexcept { UseLLVMTarget = Flag; }
   bool getUseLLVMTarget() const noexcept { return UseLLVMTarget; }
 
+  void setFnVarCategory(uint32_t Category) noexcept {
+    FnVarCategory = Category;
+  }
+  std::optional<uint32_t> getFnVarCategory() const noexcept {
+    return FnVarCategory;
+  }
+
+  void setFnVarFamily(uint32_t Family) noexcept { FnVarFamily = Family; }
+  std::optional<uint32_t> getFnVarFamily() const noexcept {
+    return FnVarFamily;
+  }
+
+  void setFnVarArch(uint32_t Arch) noexcept { FnVarArch = Arch; }
+  std::optional<uint32_t> getFnVarArch() const noexcept { return FnVarArch; }
+
+  void setFnVarTarget(uint32_t Target) noexcept { FnVarTarget = Target; }
+  std::optional<uint32_t> getFnVarTarget() const noexcept {
+    return FnVarTarget;
+  }
+
+  void setFnVarFeatures(std::vector<uint32_t> Features) noexcept {
+    FnVarFeatures = Features;
+  }
+  std::vector<uint32_t> getFnVarFeatures() const noexcept {
+    return FnVarFeatures;
+  }
+
+  void setFnVarCapabilities(std::vector<uint32_t> Capabilities) noexcept {
+    FnVarCapabilities = Capabilities;
+  }
+  std::vector<uint32_t> getFnVarCapabilities() const noexcept {
+    return FnVarCapabilities;
+  }
+
+  void setFnVarSpecEnable(bool Val) noexcept { FnVarSpecEnable = Val; }
+  bool getFnVarSpecEnable() const noexcept { return FnVarSpecEnable; }
+
+  void setFnVarSpvOut(std::string Val) noexcept { FnVarSpvOut = Val; }
+  std::string getFnVarSpvOut() const noexcept { return FnVarSpvOut; }
+
+  // Check that options passed to --fnvar-xxx flags make sense. Return true on
+  // success, false on failure.
+  bool validateFnVarOpts() const;
+
 private:
   // Common translation options
   VersionNumber MaxVersion = VersionNumber::MaximumVersion;
@@ -301,6 +345,15 @@ class TranslatorOpts {
 
   bool PreserveAuxData = false;
 
+  std::optional<uint32_t> FnVarCategory = std::nullopt;
+  std::optional<uint32_t> FnVarFamily = std::nullopt;
+  std::optional<uint32_t> FnVarArch = std::nullopt;
+  std::optional<uint32_t> FnVarTarget = std::nullopt;
+  std::vector<uint32_t> FnVarFeatures = {};
+  std::vector<uint32_t> FnVarCapabilities = {};
+  std::string FnVarSpvOut = "";
+  bool FnVarSpecEnable = false;
+
   BuiltinFormat SPIRVBuiltinFormat = BuiltinFormat::Function;
 
   // Convert LLVM to SPIR-V using the LLVM SPIR-V Backend target

lib/SPIRV/CMakeLists.txt

@@ -39,6 +39,7 @@ set(SRC_LIST
   libSPIRV/SPIRVType.cpp
   libSPIRV/SPIRVValue.cpp
   libSPIRV/SPIRVError.cpp
+  libSPIRV/SPIRVFnVar.cpp
 )
 add_llvm_library(LLVMSPIRVLib
   ${SRC_LIST}

lib/SPIRV/LLVMSPIRVOpts.cpp

@@ -44,6 +44,7 @@
 #include <llvm/ADT/SmallVector.h>
 #include <llvm/ADT/StringRef.h>
 #include <llvm/IR/IntrinsicInst.h>
+#include <optional>
 
 using namespace llvm;
 using namespace SPIRV;
@@ -89,3 +90,25 @@ std::vector<std::string> TranslatorOpts::getAllowedSPIRVExtensionNames(
   }
   return AllowExtNames;
 }
+
+bool TranslatorOpts::validateFnVarOpts() const {
+  if (getFnVarCategory() == std::nullopt &&
+      (getFnVarFamily() != std::nullopt || getFnVarArch() != std::nullopt)) {
+    errs() << "FnVar: Device category must be specified if the family or "
+              "architecture are specified.";
+    return false;
+  }
+
+  if (getFnVarFamily() == std::nullopt && getFnVarArch() != std::nullopt) {
+    errs() << "FnVar: Device family must be specified if the architecture is "
+              "specified.";
+    return false;
+  }
+
+  if (getFnVarTarget() == std::nullopt && !getFnVarFeatures().empty()) {
+    errs() << "Device target must be specified if the features are specified.";
+    return false;
+  }
+
+  return true;
+}

lib/SPIRV/SPIRVReader.cpp

@@ -41,6 +41,7 @@
 #include "SPIRVAsm.h"
 #include "SPIRVBasicBlock.h"
 #include "SPIRVExtInst.h"
+#include "SPIRVFnVar.h"
 #include "SPIRVFunction.h"
 #include "SPIRVInstruction.h"
 #include "SPIRVInternal.h"
@@ -1756,6 +1757,20 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
   case OpLabel:
     return mapValue(BV, BasicBlock::Create(*Context, BV->getName(), F));
 
+  case OpSpecConstantArchitectureINTEL:
+    llvm_unreachable(
+        "Encountered non-specialized OpSpecConstantArchitectureINTEL");
+    return nullptr;
+
+  case OpSpecConstantTargetINTEL:
+    llvm_unreachable("Encountered non-specialized OpSpecConstantTargetINTEL");
+    return nullptr;
+
+  case OpSpecConstantCapabilitiesINTEL:
+    llvm_unreachable(
+        "Encountered non-specialized OpSpecConstantCapabilitiesINTEL");
+    return nullptr;
+
   default:
     // do nothing
     break;
@@ -5599,6 +5614,31 @@ bool llvm::readSpirv(LLVMContext &C, const SPIRV::TranslatorOpts &Opts,
   if (!BM)
     return false;
 
+  if (Opts.getFnVarSpecEnable()) {
+    if (!specializeFnVariants(BM.get(), ErrMsg)) {
+      return false;
+    }
+
+    // Write out the specialized/targeted module
+    if (!BM->getFnVarSpvOut().empty()) {
+      auto SaveOpt = SPIRVUseTextFormat;
+      auto OFSSpv = std::ofstream(BM->getFnVarSpvOut(), std::ios::binary);
+      SPIRVUseTextFormat = false;
+      OFSSpv << *BM;
+      if (BM->getError(ErrMsg) != SPIRVEC_Success) {
+        return false;
+      }
+      SPIRVUseTextFormat = SaveOpt;
+    }
+  }
+
+  if (BM->getExtension().find("SPV_INTEL_function_variants") !=
+      BM->getExtension().end()) {
+    ErrMsg = "Instructions from SPV_INTEL_function_variants are not "
+             "convertible to LLVM IR.";
+    return false;
+  }
+
   M = convertSpirvToLLVM(C, *BM, Opts, ErrMsg).release();
 
   if (!M)

lib/SPIRV/libSPIRV/SPIRVEntry.cpp

@@ -42,6 +42,7 @@
 #include "SPIRVBasicBlock.h"
 #include "SPIRVDebug.h"
 #include "SPIRVDecorate.h"
+#include "SPIRVFnVar.h"
 #include "SPIRVFunction.h"
 #include "SPIRVInstruction.h"
 #include "SPIRVMemAliasingINTEL.h"

lib/SPIRV/libSPIRV/SPIRVEntry.h

@@ -914,6 +914,9 @@ class SPIRVCapability : public SPIRVEntryNoId<OpCapability> {
       return ExtensionID::SPV_INTEL_subgroup_requirements;
     case CapabilityFPFastMathModeINTEL:
       return ExtensionID::SPV_INTEL_fp_fast_math_mode;
+    case CapabilityFunctionVariantsINTEL:
+    case CapabilitySpecConditionalINTEL:
+      return ExtensionID::SPV_INTEL_function_variants;
     default:
       return {};
     }

lib/SPIRV/libSPIRV/SPIRVErrorEnum.h

@@ -28,6 +28,8 @@ _SPIRV_OP(UnspecifiedMemoryModel, "Unspecified Memory Model.")
 _SPIRV_OP(RepeatedMemoryModel, "Expects a single OpMemoryModel instruction.")
 _SPIRV_OP(UnsupportedVarArgFunction,
           "Variadic functions other than 'printf' are not supported in SPIR-V.")
+_SPIRV_OP(InvalidNumberOfOperands,
+          "Number of operands does not match the expected count.")
 
 /* This is the last error code to have a maximum valid value to compare to */
 _SPIRV_OP(InternalMaxErrorCode, "Unknown error code")