[Vulkan] Broke out implicit device requirements into SPIRVSupport

src/target/spirv/build_vulkan.cc

@@ -75,7 +75,7 @@ runtime::Module BuildSPIRV(IRModule mod, Target target, bool webgpu_restriction)
 
   mod = tir::transform::PointerValueTypeRewrite()(std::move(mod));
 
-  CodeGenSPIRV cg;
+  CodeGenSPIRV cg(target);
 
   for (auto kv : mod->functions) {
     ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenSPIRV: Can only take PrimFunc";

src/target/spirv/codegen_spirv.cc

@@ -37,14 +37,17 @@
 namespace tvm {
 namespace codegen {
 
+CodeGenSPIRV::CodeGenSPIRV(Target target) : spirv_support_(target) {}
+
 runtime::VulkanShader CodeGenSPIRV::BuildFunction(const PrimFunc& f, const std::string& name) {
   this->InitFuncState();
   ICHECK(f->HasNonzeroAttr(tir::attr::kNoAlias)) << "SPIRV only takes restricted memory model";
   std::vector<Var> pod_args;
   uint32_t num_buffer = 0;
 
   // Currently, all storage and uniform buffer arguments are passed as
-  // a single descriptor set at index 0.
+  // a single descriptor set at index 0.  If ever non-zero, must
+  // ensure it is less than maxBoundDescriptorSets.
   const uint32_t descriptor_set = 0;
 
   for (Var arg : f->params) {
@@ -114,7 +117,7 @@ void CodeGenSPIRV::InitFuncState() {
   var_map_.clear();
   storage_info_.clear();
   analyzer_.reset(new arith::Analyzer());
-  builder_.reset(new spirv::IRBuilder());
+  builder_.reset(new spirv::IRBuilder(spirv_support_));
   builder_->InitHeader();
 }
 

src/target/spirv/codegen_spirv.h

@@ -25,6 +25,7 @@
 #define TVM_TARGET_SPIRV_CODEGEN_SPIRV_H_
 
 #include <tvm/arith/analyzer.h>
+#include <tvm/target/target.h>
 #include <tvm/tir/analysis.h>
 #include <tvm/tir/expr.h>
 #include <tvm/tir/function.h>
@@ -38,6 +39,7 @@
 #include "../../runtime/thread_storage_scope.h"
 #include "../../runtime/vulkan/vulkan_shader.h"
 #include "ir_builder.h"
+#include "spirv_support.h"
 
 namespace tvm {
 namespace codegen {
@@ -50,6 +52,14 @@ using namespace tir;
 class CodeGenSPIRV : public ExprFunctor<spirv::Value(const PrimExpr&)>,
                      public StmtFunctor<void(const Stmt&)> {
  public:
+  /*!
+   * \brief Initialize the codegen based on a specific target.
+   *
+   * \param target The target for which code should be generated.  The
+   * device_type for this target must be kDLVulkan.
+   */
+  CodeGenSPIRV(Target target);
+
   /*!
    * \brief Compile and add function f to the current module.
    * \param f The function to be added.
@@ -131,6 +141,8 @@ class CodeGenSPIRV : public ExprFunctor<spirv::Value(const PrimExpr&)>,
   spirv::Value GetThreadIndex(const IterVar& iv, const PrimExpr& extent);
   spirv::Value CreateStorageSync(const CallNode* op);
   void Scalarize(const PrimExpr& e, std::function<void(int i, spirv::Value v)> f);
+  // SPIRV-related capabilities of the target
+  SPIRVSupport spirv_support_;
   // The builder
   std::unique_ptr<spirv::IRBuilder> builder_;
   // Work group size of three

src/target/spirv/ir_builder.cc

@@ -23,33 +23,35 @@
  */
 #include "ir_builder.h"
 
+#include <spirv.hpp>
+
 namespace tvm {
 namespace codegen {
 namespace spirv {
 
 // implementations
 
+IRBuilder::IRBuilder(const SPIRVSupport& support) : spirv_support_(support) {}
+
 void IRBuilder::InitHeader() {
   ICHECK_EQ(header_.size(), 0U);
   header_.push_back(spv::MagicNumber);
 
-  // Use the spirv version as indicated in the SDK.
-#if SPV_VERSION >= 0x10300
-  header_.push_back(0x10300);
-#else
+  // Target SPIR-V version 1.0.  Additional functionality will be
+  // enabled through extensions.
   header_.push_back(0x10000);
-#endif
 
   // generator: set to 0, unknown
   header_.push_back(0U);
   // Bound: set during Finalize
   header_.push_back(0U);
   // Schema: reserved
   header_.push_back(0U);
-  // shader
-  ib_.Begin(spv::OpCapability).Add(spv::CapabilityShader).Commit(&header_);
-  // Declare int64 capability by default
-  ib_.Begin(spv::OpCapability).Add(spv::CapabilityInt64).Commit(&header_);
+
+  // Declare CapabilityShader by default.  All other capabilities are
+  // determined by the types declared.
+  capabilities_used_.insert(spv::CapabilityShader);
+
   // memory model
   ib_.Begin(spv::OpMemoryModel)
       .AddSeq(spv::AddressingModelLogical, spv::MemoryModelGLSL450)
@@ -71,6 +73,30 @@ void IRBuilder::InitPreDefs() {
   ib_.Begin(spv::OpTypeFunction).AddSeq(t_void_func_, t_void_).Commit(&global_);
 }
 
+std::vector<uint32_t> IRBuilder::Finalize() {
+  std::vector<uint32_t> data;
+  // Index for upper bound of id numbers.
+  const int kBoundLoc = 3;
+  header_[kBoundLoc] = id_counter_;
+  data.insert(data.end(), header_.begin(), header_.end());
+  for (const auto& capability : capabilities_used_) {
+    ib_.Begin(spv::OpCapability).Add(capability).Commit(&data);
+  }
+  for (const auto& ext_name : extensions_used_) {
+    ib_.Begin(spv::OpExtension).Add(ext_name).Commit(&data);
+  }
+  data.insert(data.end(), extended_instruction_section_.begin(),
+              extended_instruction_section_.end());
+  data.insert(data.end(), entry_.begin(), entry_.end());
+  data.insert(data.end(), exec_mode_.begin(), exec_mode_.end());
+  data.insert(data.end(), debug_.begin(), debug_.end());
+  data.insert(data.end(), decorate_.begin(), decorate_.end());
+  data.insert(data.end(), global_.begin(), global_.end());
+  data.insert(data.end(), func_header_.begin(), func_header_.end());
+  data.insert(data.end(), function_.begin(), function_.end());
+  return data;
+}
+
 SType IRBuilder::GetSType(const DataType& dtype) {
   if (dtype == DataType::Int(32)) {
     return t_int32_;
@@ -145,16 +171,19 @@ SType IRBuilder::GetStructArrayType(const SType& value_type, uint32_t num_elems)
       .AddSeq(struct_type, 0, spv::DecorationOffset, 0)
       .Commit(&decorate_);
 
-#if SPV_VERSION < 0x10300
-  // NOTE: BufferBlock was deprecated in SPIRV 1.3
-  // use StorageClassStorageBuffer instead.
-  // runtime array are always decorated as BufferBlock(shader storage buffer)
-  if (num_elems == 0) {
-    this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBufferBlock);
+  // Runtime array are always decorated as Block or BufferBlock
+  // (shader storage buffer)
+  if (spirv_support_.supports_storage_buffer_storage_class) {
+    // If SPIRV 1.3+, or with extension
+    // SPV_KHR_storage_buffer_storage_class, BufferBlock is
+    // deprecated.
+    extensions_used_.insert("SPV_KHR_storage_buffer_storage_class");
+    this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBlock);
+  } else {
+    if (num_elems == 0) {
+      this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBufferBlock);
+    }
   }
-#else
-  this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBlock);
-#endif
   struct_array_type_tbl_[key] = struct_type;
   return struct_type;
 }
@@ -186,13 +215,14 @@ Value IRBuilder::FloatImm(const SType& dtype, double value) {
 
 Value IRBuilder::BufferArgument(const SType& value_type, uint32_t descriptor_set,
                                 uint32_t binding) {
-  // NOTE: BufferBlock was deprecated in SPIRV 1.3
-  // use StorageClassStorageBuffer instead.
-#if SPV_VERSION >= 0x10300
-  spv::StorageClass storage_class = spv::StorageClassStorageBuffer;
-#else
-  spv::StorageClass storage_class = spv::StorageClassUniform;
-#endif
+  // If SPIRV 1.3+, or with extension SPV_KHR_storage_buffer_storage_class, BufferBlock is
+  // deprecated.
+  spv::StorageClass storage_class;
+  if (spirv_support_.supports_storage_buffer_storage_class) {
+    storage_class = spv::StorageClassStorageBuffer;
+  } else {
+    storage_class = spv::StorageClassUniform;
+  }
 
   SType sarr_type = GetStructArrayType(value_type, 0);
   SType ptr_type = GetPointerType(sarr_type, storage_class);
@@ -383,6 +413,8 @@ Value IRBuilder::GetConst_(const SType& dtype, const uint64_t* pvalue) {
 }
 
 SType IRBuilder::DeclareType(const DataType& dtype) {
+  AddCapabilityFor(dtype);
+
   if (dtype.lanes() == 1) {
     SType t;
     t.id = id_counter_++;
@@ -410,6 +442,60 @@ SType IRBuilder::DeclareType(const DataType& dtype) {
   }
 }
 
+void IRBuilder::AddCapabilityFor(const DataType& dtype) {
+  // Declare appropriate capabilities for int/float types
+  if (dtype.is_int() || dtype.is_uint()) {
+    if (dtype.bits() == 8) {
+      ICHECK(spirv_support_.supports_int8) << "Vulkan target does not support Int8 capability";
+      capabilities_used_.insert(spv::CapabilityInt8);
+    } else if (dtype.bits() == 16) {
+      ICHECK(spirv_support_.supports_int16) << "Vulkan target does not support Int16 capability";
+      capabilities_used_.insert(spv::CapabilityInt16);
+    } else if (dtype.bits() == 64) {
+      ICHECK(spirv_support_.supports_int64) << "Vulkan target does not support Int64 capability";
+      capabilities_used_.insert(spv::CapabilityInt64);
+    }
+
+  } else if (dtype.is_float()) {
+    if (dtype.bits() == 16) {
+      ICHECK(spirv_support_.supports_float16)
+          << "Vulkan target does not support Float16 capability";
+      capabilities_used_.insert(spv::CapabilityFloat16);
+    } else if (dtype.bits() == 64) {
+      ICHECK(spirv_support_.supports_float64)
+          << "Vulkan target does not support Float64 capability";
+      capabilities_used_.insert(spv::CapabilityFloat64);
+    }
+  }
+
+  // Declare ability to read type to/from storage buffers.  Doing so
+  // here is a little bit overzealous, should be relaxed in the
+  // future.  Requiring StorageBuffer8BitAccess in order to declare an
+  // Int8 prevents use of an 8-bit loop iterator on a device that
+  // supports Int8 but doesn't support 8-bit buffer access.
+  if (dtype.bits() == 8) {
+    ICHECK(spirv_support_.supports_storage_buffer_8bit_access)
+        << "Vulkan target does not support StorageBuffer8BitAccess";
+    capabilities_used_.insert(spv::CapabilityStorageBuffer8BitAccess);
+    extensions_used_.insert("SPV_KHR_8bit_storage");
+
+    ICHECK(spirv_support_.supports_storage_buffer_storage_class)
+        << "Illegal Vulkan target description.  "
+        << "Vulkan spec requires extension VK_KHR_storage_buffer_storage_class "
+        << "if VK_KHR_8bit_storage is supported";
+  } else if (dtype.bits() == 16) {
+    ICHECK(spirv_support_.supports_storage_buffer_8bit_access)
+        << "Vulkan target does not support StorageBuffer16BitAccess";
+
+    extensions_used_.insert("SPV_KHR_16bit_storage");
+    if (spirv_support_.supports_storage_buffer_storage_class) {
+      capabilities_used_.insert(spv::CapabilityStorageBuffer16BitAccess);
+    } else {
+      capabilities_used_.insert(spv::CapabilityStorageUniformBufferBlock16);
+    }
+  }
+}
+
 PhiValue IRBuilder::MakePhi(const SType& out_type, uint32_t num_incoming) {
   Value val = NewValue(out_type, kNormal);
   ib_.Begin(spv::OpPhi).AddSeq(out_type, val);

src/target/spirv/ir_builder.h

@@ -30,13 +30,16 @@
 // clang-format off
 #include <algorithm>
 #include <map>
+#include <set>
 #include <string>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 #include <spirv.hpp>
 // clang-format on
 
+#include "spirv_support.h"
+
 namespace tvm {
 namespace codegen {
 namespace spirv {
@@ -268,6 +271,14 @@ class InstrBuilder {
  */
 class IRBuilder {
  public:
+  /*!
+   * \brief Initialize the codegen based on a specific feature set.
+   *
+   * \param support The features in SPIRV that are supported by the
+   * target device.
+   */
+  explicit IRBuilder(const SPIRVSupport& support);
+
   /*! \brief Initialize header */
   void InitHeader();
   /*! \brief Initialize the predefined contents */
@@ -278,29 +289,21 @@ class IRBuilder {
    * \return The finalized binary instruction.
    */
   Value ExtInstImport(const std::string& name) {
+    auto it = ext_inst_tbl_.find(name);
+    if (it != ext_inst_tbl_.end()) {
+      return it->second;
+    }
     Value val = NewValue(SType(), kExtInst);
-    ib_.Begin(spv::OpExtInstImport).AddSeq(val, name).Commit(&header_);
+    ib_.Begin(spv::OpExtInstImport).AddSeq(val, name).Commit(&extended_instruction_section_);
+    ext_inst_tbl_[name] = val;
     return val;
   }
   /*!
    * \brief Get the final binary built from the builder
    * \return The finalized binary instruction.
    */
-  std::vector<uint32_t> Finalize() {
-    std::vector<uint32_t> data;
-    // set bound
-    const int kBoundLoc = 3;
-    header_[kBoundLoc] = id_counter_;
-    data.insert(data.end(), header_.begin(), header_.end());
-    data.insert(data.end(), entry_.begin(), entry_.end());
-    data.insert(data.end(), exec_mode_.begin(), exec_mode_.end());
-    data.insert(data.end(), debug_.begin(), debug_.end());
-    data.insert(data.end(), decorate_.begin(), decorate_.end());
-    data.insert(data.end(), global_.begin(), global_.end());
-    data.insert(data.end(), func_header_.begin(), func_header_.end());
-    data.insert(data.end(), function_.begin(), function_.end());
-    return data;
-  }
+  std::vector<uint32_t> Finalize();
+
   /*!
    * \brief Create new label
    * \return The created new label
@@ -599,6 +602,19 @@ class IRBuilder {
   Value GetConst_(const SType& dtype, const uint64_t* pvalue);
   // declare type
   SType DeclareType(const DataType& dtype);
+
+  // Declare the appropriate SPIR-V capabilities and extensions to use
+  // this data type.
+  void AddCapabilityFor(const DataType& dtype);
+
+  /*! \brief SPIRV-related capabilities of the target
+   *
+   * This SPIRVSupport object is owned by the same CodeGenSPIRV
+   * object that owns the IRBuilder.  Therefore, safe to use a
+   * reference as the CodeGenSPIRV will live longer.
+   */
+  const SPIRVSupport& spirv_support_;
+
   /*! \brief internal instruction builder  */
   InstrBuilder ib_;
   /*! \brief Current label */
@@ -623,9 +639,22 @@ class IRBuilder {
   std::map<std::pair<uint32_t, spv::StorageClass>, SType> pointer_type_tbl_;
   /*! \brief map from constant int to its value */
   std::map<std::pair<uint32_t, uint64_t>, Value> const_tbl_;
-  /*! \brief Header segment, include import */
+  /*! \brief map from name of a ExtInstImport to its value */
+  std::map<std::string, Value> ext_inst_tbl_;
+
+  /*! \brief Header segment
+   *
+   * 5 words long, described in "First Words of Physical Layout"
+   * section of SPIR-V documentation.
+   */
   std::vector<uint32_t> header_;
-  /*! \brief engtry point segment */
+  /*! \brief SPIR-V capabilities used by this module. */
+  std::set<spv::Capability> capabilities_used_;
+  /*! \brief SPIR-V extensions used by this module. */
+  std::set<std::string> extensions_used_;
+  /*! \brief entry point segment */
+  std::vector<uint32_t> extended_instruction_section_;
+  /*! \brief entry point segment */
   std::vector<uint32_t> entry_;
   /*! \brief Header segment */
   std::vector<uint32_t> exec_mode_;