Extend TranslatorOpts to handle SPIR-V extensions

Added new command line option to specify which particular extensions are
enabled/disabled.

Signed-off-by: Alexey Sachkov <alexey.sachkov@intel.com>

README.md

@@ -114,6 +114,27 @@ the version of the SPIR-V file which is being generated/consumed.
 
 Allowed values are `1.0`/`1.1`.
 
+### Handling SPIR-V extensions generated by the translator
+
+By default, during SPIR-V generation, the translator doesn't use any extensions.
+However, during SPIR-V consumption, the translator accepts input files that use
+any known extensions.
+
+If certain extensions are required to be enabled or disabled, the following
+command line option can be used:
+
+* ``--spirv-ext=`` - this options allows controlling which extensions are
+  allowed/disallowed
+
+Valid value for this option is comma-separated list of extension names prefixed
+with ``+`` or ``-`` - plus means allow to use extension, minus means disallow
+to use extension. There is one more special value which can be used as extension
+name in this option: ``all`` - it affects all extension which are known to the
+translator.
+
+If ``--spirv-ext`` contains the name of an extension which is not known for the
+translator, it will emit an error.
+
 ## Branching strategy
 
 Code on the master branch in this repository is intended to be compatible with master/trunk branch of the [llvm](https://github.com/llvm-mirror/llvm) project. That is, for an OpenCL kernel compiled to llvm bitcode by the latest version(built with the latest git commit or svn revision) of Clang it should be possible to translate it to SPIR-V with the llvm-spirv tool.

include/LLVMSPIRVExtensions.inc

@@ -0,0 +1,15 @@
+
+#ifndef EXT
+  #error "EXT macro must be defined"
+#endif
+
+EXT(SPV_KHR_no_integer_wrap_decoration)
+EXT(SPV_INTEL_subgroups)
+EXT(SPV_INTEL_media_block_io)
+EXT(SPV_INTEL_device_side_avc_motion_estimation)
+EXT(SPV_INTEL_fpga_loop_controls)
+EXT(SPV_INTEL_fpga_memory_attributes)
+EXT(SPV_INTEL_unstructured_loop_controls)
+EXT(SPV_INTEL_fpga_reg)
+EXT(SPV_INTEL_blocking_pipes)
+

include/LLVMSPIRVOpts.h

@@ -33,14 +33,15 @@
 //===----------------------------------------------------------------------===//
 /// \file LLVMSPIRVOpts.h
 ///
-/// This files declares helper classes to handle SPIR-V versions.
+/// This files declares helper classes to handle SPIR-V versions and extensions.
 ///
 //===----------------------------------------------------------------------===//
 #ifndef SPIRV_LLVMSPIRVOPTS_H
 #define SPIRV_LLVMSPIRVOPTS_H
 
 #include <cassert>
 #include <cstdint>
+#include <map>
 
 namespace SPIRV {
 
@@ -55,21 +56,47 @@ enum class VersionNumber : uint32_t {
   MaximumVersion = SPIRV_1_1
 };
 
+enum class ExtensionID : uint32_t {
+  First,
+#define EXT(X) X,
+#include "LLVMSPIRVExtensions.inc"
+#undef EXT
+  Last,
+};
+
 /// \brief Helper class to manage SPIR-V translation
 class TranslatorOpts {
 public:
+  using ExtensionsStatusMap = std::map<ExtensionID, bool>;
+
   TranslatorOpts() = default;
 
-  TranslatorOpts(VersionNumber Max) : MaxVersion(Max) {}
+  TranslatorOpts(VersionNumber Max, const ExtensionsStatusMap &Map = {})
+      : MaxVersion(Max), ExtStatusMap(Map) {}
 
   bool isAllowedToUseVersion(VersionNumber RequestedVersion) const {
     return RequestedVersion <= MaxVersion;
   }
 
+  bool isAllowedToUseExtension(ExtensionID Extension) const {
+    auto I = ExtStatusMap.find(Extension);
+    if (ExtStatusMap.end() == I)
+      return false;
+
+    return I->second;
+  }
+
   VersionNumber getMaxVersion() const { return MaxVersion; }
 
+  void enableAllExtensions() {
+#define EXT(X) ExtStatusMap[ExtensionID::X] = true;
+#include "LLVMSPIRVExtensions.inc"
+#undef EXT
+  }
+
 private:
   VersionNumber MaxVersion = VersionNumber::MaximumVersion;
+  ExtensionsStatusMap ExtStatusMap;
 };
 
 } // namespace SPIRV

lib/SPIRV/SPIRVReader.cpp

@@ -3053,6 +3053,9 @@ llvm::convertSpirvToLLVM(LLVMContext &C, SPIRVModule &BM, std::string &ErrMsg) {
 bool llvm::readSpirv(LLVMContext &C, std::istream &IS, Module *&M,
                      std::string &ErrMsg) {
   SPIRV::TranslatorOpts DefaultOpts;
+  // As it is stated in the documentation, the translator accepts all SPIR-V
+  // extensions by default
+  DefaultOpts.enableAllExtensions();
   return llvm::readSpirv(C, DefaultOpts, IS, M, ErrMsg);
 }
 

lib/SPIRV/SPIRVWriter.cpp

@@ -1978,6 +1978,9 @@ bool isValidLLVMModule(Module *M, SPIRVErrorLog &ErrorLog) {
 
 bool llvm::writeSpirv(Module *M, std::ostream &OS, std::string &ErrMsg) {
   SPIRV::TranslatorOpts DefaultOpts;
+  // To preserve old behavior of the translator, let's enable all extensions
+  // by default in this API
+  DefaultOpts.enableAllExtensions();
   return llvm::writeSpirv(M, DefaultOpts, OS, ErrMsg);
 }
 

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -1756,7 +1756,13 @@ bool convertSpirv(std::istream &IS, std::ostream &OS, std::string &ErrMsg,
                   bool FromText, bool ToText) {
   auto SaveOpt = SPIRVUseTextFormat;
   SPIRVUseTextFormat = FromText;
+  // Conversion from/to SPIR-V text representation is a side feature of the
+  // translator which is mostly intended for debug usage. So, this step cannot
+  // be customized to enable/disable particular extensions or restrict/allow
+  // particular SPIR-V versions: all known SPIR-V versions are allowed, all
+  // known SPIR-V extensions are enabled during this conversion
   SPIRV::TranslatorOpts DefaultOpts;
+  DefaultOpts.enableAllExtensions();
   SPIRVModuleImpl M(DefaultOpts);
   IS >> M;
   if (M.getError(ErrMsg) != SPIRVEC_Success) {

lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -402,6 +402,11 @@ class SPIRVModule {
     return TranslationOpts.getMaxVersion();
   }
 
+  virtual bool
+  isAllowedToUseExtension(ExtensionID RequestedExtension) const final {
+    return TranslationOpts.isAllowedToUseExtension(RequestedExtension);
+  }
+
   // I/O functions
   friend spv_ostream &operator<<(spv_ostream &O, SPIRVModule &M);
   friend std::istream &operator>>(std::istream &I, SPIRVModule &M);

test/spirv-extensions-control.ll

@@ -0,0 +1,53 @@
+; Bunch of negative tests:
+;
+; RUN: not llvm-spirv --spirv-ext=EXT 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-FORMAT
+; RUN: not llvm-spirv --spirv-ext=+EXT 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-EXT
+; RUN: not llvm-spirv --spirv-ext=-EXT 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-EXT
+; RUN: not llvm-spirv --spirv-ext=- 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-FORMAT
+; RUN: not llvm-spirv --spirv-ext=+ 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-FORMAT
+; RUN: not llvm-spirv --spirv-ext=, 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-FORMAT
+; RUN: not llvm-spirv --spirv-ext= 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-FORMAT
+;
+; Bunch of positive tests:
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+all -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=-all -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_subgroups -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=-SPV_INTEL_subgroups -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=+all,-SPV_INTEL_subgroups -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=-all,+SPV_INTEL_subgroups -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+; RUN: llvm-spirv %t.bc --spirv-ext=-SPV_INTEL_subgroups,+SPV_INTEL_subgroups -o - 2>&1 | FileCheck %s --check-prefix=CHECK-VALID
+;
+; CHECK-INVALID-FORMAT: Invalid value of --spirv-ext
+; CHECK-INVALID-EXT: Unknown extension 'EXT' was specified
+;
+; CHECK-VALID-NOT: Unknown extension '{{.*}}' was specified
+
+target datalayout = "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir-unknown-unknown"
+
+; Function Attrs: nounwind
+define spir_kernel void @foo(i32 addrspace(1)* %a) #0 !kernel_arg_addr_space !1 !kernel_arg_access_qual !2 !kernel_arg_type !3 !kernel_arg_base_type !4 !kernel_arg_type_qual !5 {
+entry:
+  %a.addr = alloca i32 addrspace(1)*, align 4
+  store i32 addrspace(1)* %a, i32 addrspace(1)** %a.addr, align 4
+  ret void
+}
+
+attributes #0 = { nounwind "less-precise-fpmad"="false" "no-frame-pointer-elim"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-realign-stack" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+!opencl.enable.FP_CONTRACT = !{}
+!opencl.spir.version = !{!6}
+!opencl.ocl.version = !{!6}
+!opencl.used.extensions = !{!7}
+!opencl.used.optional.core.features = !{!7}
+!opencl.compiler.options = !{!7}
+
+!1 = !{i32 1}
+!2 = !{!"none"}
+!3 = !{!"int*"}
+!4 = !{!"int*"}
+!5 = !{!""}
+!6 = !{i32 1, i32 2}
+!7 = !{}
+

tools/llvm-spirv/llvm-spirv.cpp

@@ -67,7 +67,10 @@
 
 #include <fstream>
 #include <iostream>
+#include <map>
 #include <memory>
+#include <set>
+#include <string>
 
 #define DEBUG_TYPE "spirv"
 
@@ -102,6 +105,14 @@ static cl::opt<VersionNumber> MaxSPIRVVersion(
                clEnumValN(VersionNumber::SPIRV_1_1, "1.1", "SPIR-V 1.1")),
     cl::init(VersionNumber::MaximumVersion));
 
+static cl::list<std::string>
+    SPVExt("spirv-ext", cl::CommaSeparated,
+           cl::desc("Specify list of allowed/disallowed extensions"),
+           cl::value_desc("+SPV_extenstion1_name,-SPV_extension2_name"),
+           cl::ValueRequired);
+
+using SPIRV::ExtensionID;
+
 #ifdef _SPIRV_SUPPORT_TEXT_FMT
 namespace SPIRV {
 // Use textual format for SPIRV.
@@ -268,14 +279,81 @@ static int regularizeLLVM() {
   return 0;
 }
 
+static int parseSPVExtOption(
+    SPIRV::TranslatorOpts::ExtensionsStatusMap &ExtensionsStatus) {
+  // Map name -> id for known extensions
+  std::map<std::string, ExtensionID> ExtensionNamesMap;
+#define _STRINGIFY(X) #X
+#define STRINGIFY(X) _STRINGIFY(X)
+#define EXT(X) ExtensionNamesMap[STRINGIFY(X)] = ExtensionID::X;
+#include "LLVMSPIRVExtensions.inc"
+#undef EXT
+#undef STRINGIFY
+#undef _STRINGIFY
+
+  // Set the initial state:
+  //  - during SPIR-V consumption, assume that any known extension is allowed.
+  //  - during SPIR-V generation, assume that any known extension is disallowed.
+  //  - during conversion to/from SPIR-V text representation, assume that any
+  //    known extension is allowed.
+  for (const auto &It : ExtensionNamesMap)
+    ExtensionsStatus[It.second] = IsReverse;
+
+  if (SPVExt.empty())
+    return 0; // Nothing to do
+
+  for (unsigned i = 0; i < SPVExt.size(); ++i) {
+    const std::string &ExtString = SPVExt[i];
+    if ('+' != ExtString.front() && '-' != ExtString.front()) {
+      errs() << "Invalid value of --spirv-ext, expected format is:\n"
+             << "\t--spirv-ext=+EXT_NAME,-EXT_NAME\n";
+      return -1;
+    }
+
+    auto ExtName = ExtString.substr(1);
+
+    if (ExtName.empty()) {
+      errs() << "Invalid value of --spirv-ext, expected format is:\n"
+             << "\t--spirv-ext=+EXT_NAME,-EXT_NAME\n";
+      return -1;
+    }
+
+    bool ExtStatus = ('+' == ExtString.front());
+    if ("all" == ExtName) {
+      // Update status for all known extensions
+      for (const auto &It : ExtensionNamesMap)
+        ExtensionsStatus[It.second] = ExtStatus;
+    } else {
+      // Reject unknown extensions
+      const auto &It = ExtensionNamesMap.find(ExtName);
+      if (ExtensionNamesMap.end() == It) {
+        errs() << "Unknown extension '" << ExtName << "' was specified via "
+               << "--spirv-ext option\n";
+        return -1;
+      }
+
+      ExtensionsStatus[It->second] = ExtStatus;
+    }
+  }
+
+  return 0;
+}
+
 int main(int Ac, char **Av) {
   EnablePrettyStackTrace();
   sys::PrintStackTraceOnErrorSignal(Av[0]);
   PrettyStackTraceProgram X(Ac, Av);
 
   cl::ParseCommandLineOptions(Ac, Av, "LLVM/SPIR-V translator");
 
-  SPIRV::TranslatorOpts Opts(MaxSPIRVVersion);
+  SPIRV::TranslatorOpts::ExtensionsStatusMap ExtensionsStatus;
+  // ExtensionsStatus will be properly initialized and update according to
+  // values passed via --spirv-ext option in parseSPVExtOption function.
+  int Ret = parseSPVExtOption(ExtensionsStatus);
+  if (0 != Ret)
+    return Ret;
+
+  SPIRV::TranslatorOpts Opts(MaxSPIRVVersion, ExtensionsStatus);
 
 #ifdef _SPIRV_SUPPORT_TEXT_FMT
   if (ToText && (ToBinary || IsReverse || IsRegularization)) {